agent_lab/backend/deepseek_mcp_client.py

"""
DeepSeek MCP Client
Enables DeepSeek R1 to use MCP tools via prompt-based invocation.
"""

import json
import requests
from typing import Dict, List, Optional, Any
from dataclasses import dataclass

@dataclass
class MCPTool:
    name: str
    description: str
    parameters: Dict[str, Any]

class DeepSeekMCPClient:
    """
    Client for integrating DeepSeek with MCP server.
    Handles tool schema injection and JSON-based tool calling.
    """
    
    def __init__(self, mcp_server_url: str = "http://localhost:8001"):
        self.mcp_url = mcp_server_url
        self.tools: List[MCPTool] = []
        self.load_tools()
    
    def load_tools(self):
        """Fetch available tools from MCP server."""
        try:
            response = requests.get(f"{self.mcp_url}/tools", timeout=5)
            if response.status_code == 200:
                tools_data = response.json()
                self.tools = [
                    MCPTool(
                        name=t['name'],
                        description=t['description'],
                        parameters=t['input_schema']
                    )
                    for t in tools_data
                ]
                print(f"[MCP Client] Loaded {len(self.tools)} tools")
            else:
                print(f"[MCP Client] Failed to load tools: {response.status_code}")
        except Exception as e:
            print(f"[MCP Client] Error connecting to MCP server: {e}")
    
    def get_tool_schema_prompt(self, include_tools: Optional[List[str]] = None) -> str:
        """
        Generate system prompt with tool schemas.
        If include_tools is provided, only those tools will be included.
        """
        if not self.tools:
            return ""
        
        # Filter tools if requested
        tools_to_show = self.tools
        if include_tools:
            tools_to_show = [t for t in self.tools if t.name in include_tools]
            # If no matches found with filtering, fall back to first 10 for safety
            if not tools_to_show:
                tools_to_show = self.tools[:10]
        else:
            tools_to_show = self.tools[:40] # Default limit for token safety
            
        tools_desc = "You have access to the following specialized tools:\n\n"
        for tool in tools_to_show:
            tools_desc += f"### {tool.name}\n"
            tools_desc += f"{tool.description}\n"
            tools_desc += f"JSON Schema: {json.dumps(tool.parameters)}\n\n"
        
        tools_desc += """
[PROTOCOL]: To execute a protocol, response MUST include a JSON block:
```json
{
  "tool_call": {
    "name": "tool_name",
    "arguments": {"arg1": "value1"}
  }
}
```
"""
        return tools_desc
    
    def extract_tool_calls(self, response_text: str) -> List[Dict]:
        """Extract tool calls from DeepSeek response."""
        tool_calls = []
        
        # Look for JSON blocks
        if "```json" in response_text:
            start = response_text.find("```json") + 7
            end = response_text.find("```", start)
            json_str = response_text[start:end].strip()
            
            try:
                data = json.loads(json_str)
                if "tool_call" in data:
                    tool_calls.append(data["tool_call"])
            except json.JSONDecodeError as e:
                print(f"[MCP Client] JSON parse error: {e}")
        
        return tool_calls
    
    def invoke_tool(self, tool_name: str, arguments: Dict) -> Dict:
        """Invoke a tool on the MCP server."""
        try:
            response = requests.post(
                f"{self.mcp_url}/invoke",
                json={"tool": tool_name, "args": arguments},
                timeout=30
            )
            if response.status_code == 200:
                return response.json()
            else:
                return {"error": f"HTTP {response.status_code}", "detail": response.text}
        except Exception as e:
            return {"error": str(e)}
    
    def process_response_with_tools(self, deepseek_response: str) -> tuple[str, List[Dict]]:
        """
        Process DeepSeek response, execute any tool calls, and return results.
        Returns: (final_text, tool_results)
        """
        tool_calls = self.extract_tool_calls(deepseek_response)
        tool_results = []
        
        if tool_calls:
            print(f"[MCP Client] Executing {len(tool_calls)} tool calls...")
            for call in tool_calls:
                result = self.invoke_tool(call['name'], call.get('arguments', {}))
                tool_results.append({
                    "tool": call['name'],
                    "result": result
                })
        
        # Remove JSON blocks from response for cleaner output
        clean_text = deepseek_response
        if "```json" in clean_text:
            start = clean_text.find("```json")
            end = clean_text.find("```", start + 7) + 3
            clean_text = clean_text[:start] + clean_text[end:]
        
        return clean_text.strip(), tool_results

# Global instance
mcp_client = DeepSeekMCPClient()