apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-04-28 18:36:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

chore: linting

Browse source
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-01-15 00:05:53 -08:00
parent 3375aeb9bc
commit 7ae68455b3
20 changed files with 128 additions and 103 deletions
Download patch file Download diff file Expand all files Collapse all files

49
surfsense_backend/app/agents/new_chat/tools/mcp_client.py
View file

@ -18,7 +18,9 @@ logger = logging.getLogger(__name__)
class MCPClient:
"""Client for communicating with an MCP server."""
def __init__(self, command: str, args: list[str], env: dict[str, str] | None = None):
def __init__(
self, command: str, args: list[str], env: dict[str, str] | None = None
):
"""Initialize MCP client.
Args:
@ -44,18 +46,16 @@ class MCPClient:
# Merge env vars with current environment
server_env = os.environ.copy()
server_env.update(self.env)
# Create server parameters with env
server_params = StdioServerParameters(
command=self.command,
args=self.args,
env=server_env
command=self.command, args=self.args, env=server_env
)
# Spawn server process and create session
# Note: Cannot combine these context managers because ClientSession
# needs the read/write streams from stdio_client
async with stdio_client(server=server_params) as (read, write):
async with stdio_client(server=server_params) as (read, write): # noqa: SIM117
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
@ -85,7 +85,9 @@ class MCPClient:
"""
if not self.session:
raise RuntimeError("Not connected to MCP server. Use 'async with client.connect():'")
raise RuntimeError(
"Not connected to MCP server. Use 'async with client.connect():'"
)
try:
# Call tools/list RPC method
@ -93,11 +95,15 @@ class MCPClient:
tools = []
for tool in response.tools:
tools.append({
"name": tool.name,
"description": tool.description or "",
"input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
})
tools.append(
{
"name": tool.name,
"description": tool.description or "",
"input_schema": tool.inputSchema
if hasattr(tool, "inputSchema")
else {},
}
)
logger.info("Listed %d tools from MCP server", len(tools))
return tools
@ -121,10 +127,14 @@ class MCPClient:
"""
if not self.session:
raise RuntimeError("Not connected to MCP server. Use 'async with client.connect():'")
raise RuntimeError(
"Not connected to MCP server. Use 'async with client.connect():'"
)
try:
logger.info("Calling MCP tool '%s' with arguments: %s", tool_name, arguments)
logger.info(
"Calling MCP tool '%s' with arguments: %s", tool_name, arguments
)
# Call tools/call RPC method
response = await self.session.call_tool(tool_name, arguments=arguments)
@ -147,12 +157,17 @@ class MCPClient:
# Handle validation errors from MCP server responses
# Some MCP servers (like server-memory) return extra fields not in their schema
if "Invalid structured content" in str(e):
logger.warning("MCP server returned data not matching its schema, but continuing: %s", e)
logger.warning(
"MCP server returned data not matching its schema, but continuing: %s",
e,
)
# Try to extract result from error message or return a success message
return "Operation completed (server returned unexpected format)"
raise
except (ValueError, TypeError, AttributeError, KeyError) as e:
logger.error("Failed to call MCP tool '%s': %s", tool_name, e, exc_info=True)
logger.error(
"Failed to call MCP tool '%s': %s", tool_name, e, exc_info=True
)
return f"Error calling tool: {e!s}"

25
surfsense_backend/app/agents/new_chat/tools/mcp_tool.py
View file

@ -21,7 +21,8 @@ logger = logging.getLogger(__name__)
def _create_dynamic_input_model_from_schema(
tool_name: str, input_schema: dict[str, Any],
tool_name: str,
input_schema: dict[str, Any],
) -> type[BaseModel]:
"""Create a Pydantic model from MCP tool's JSON schema.
@ -41,15 +42,18 @@ def _create_dynamic_input_model_from_schema(
for param_name, param_schema in properties.items():
param_description = param_schema.get("description", "")
is_required = param_name in required_fields
# Use Any type for complex schemas to preserve structure
# This allows the MCP server to do its own validation
from typing import Any as AnyType
from pydantic import Field
if is_required:
field_definitions[param_name] = (AnyType, Field(..., description=param_description))
field_definitions[param_name] = (
AnyType,
Field(..., description=param_description),
)
else:
field_definitions[param_name] = (
AnyType | None,
@ -88,7 +92,7 @@ async def _create_mcp_tool_from_definition(
async def mcp_tool_call(**kwargs) -> str:
"""Execute the MCP tool call via the client."""
logger.info(f"MCP tool '{tool_name}' called with params: {kwargs}")
try:
# Connect to server and call tool
async with mcp_client.connect():
@ -114,7 +118,8 @@ async def _create_mcp_tool_from_definition(
async def load_mcp_tools(
session: AsyncSession, search_space_id: int,
session: AsyncSession,
search_space_id: int,
) -> list[StructuredTool]:
"""Load all MCP tools from user's active MCP server connectors.
@ -150,7 +155,9 @@ async def load_mcp_tools(
env = server_config.get("env", {})
if not command:
logger.warning(f"MCP connector {connector.id} missing command, skipping")
logger.warning(
f"MCP connector {connector.id} missing command, skipping"
)
continue
# Create MCP client
@ -168,7 +175,9 @@ async def load_mcp_tools(
# Create LangChain tools from definitions
for tool_def in tool_definitions:
try:
tool = await _create_mcp_tool_from_definition(tool_def, mcp_client)
tool = await _create_mcp_tool_from_definition(
tool_def, mcp_client
)
tools.append(tool)
except Exception as e:
logger.exception(

3
surfsense_backend/app/agents/new_chat/tools/registry.py
View file

@ -283,7 +283,8 @@ async def build_tools_async(
):
try:
mcp_tools = await load_mcp_tools(
dependencies["db_session"], dependencies["search_space_id"],
dependencies["db_session"],
dependencies["search_space_id"],
)
tools.extend(mcp_tools)
logging.info(