src/utils/mcp_client.py

import json
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Dict, List, Optional

from src.utils.rag_service import retrieve_documents


class MCPClient:
    def __init__(
        self,
        server_url: Optional[str] = None,
        timeout_seconds: float = 5.0,
        enable_mock_tools: bool = True,
        cache_ttl_seconds: int = 300,
    ):
        self.server_url = (server_url or "").rstrip("/")
        self.timeout_seconds = float(timeout_seconds)
        self.enable_mock_tools = enable_mock_tools
        self.cache_ttl_seconds = max(1, int(cache_ttl_seconds))
        self._tool_cache: Dict[str, Dict[str, Any]] = {}

    def _cache_key(self, tool_name: str, arguments: Dict[str, Any]) -> str:
        return f"{tool_name}:{json.dumps(arguments, sort_keys=True, default=str)}"

    def _read_cache(self, key: str) -> Optional[Dict[str, Any]]:
        item = self._tool_cache.get(key)
        if not item:
            return None
        if (time.time() - item["timestamp"]) > self.cache_ttl_seconds:
            self._tool_cache.pop(key, None)
            return None
        return dict(item["value"])

    def _write_cache(self, key: str, value: Dict[str, Any]):
        self._tool_cache[key] = {"value": dict(value), "timestamp": time.time()}

    def _post_json(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        url = f"{self.server_url}{path}"
        body = json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(
            url=url,
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=self.timeout_seconds) as response:
            raw = response.read().decode("utf-8")
            return json.loads(raw) if raw else {}

    def _get_json(self, path: str) -> Dict[str, Any]:
        url = f"{self.server_url}{path}"
        request = urllib.request.Request(url=url, method="GET")
        with urllib.request.urlopen(request, timeout=self.timeout_seconds) as response:
            raw = response.read().decode("utf-8")
            return json.loads(raw) if raw else {}

    def discover_tools(self) -> List[Dict[str, Any]]:
        if self.server_url:
            try:
                remote = self._get_json("/tools")
                tools = remote.get("tools", []) if isinstance(remote, dict) else []
                if tools:
                    return tools
            except Exception:
                pass
        return self._mock_tools() if self.enable_mock_tools else []

    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        key = self._cache_key(tool_name, arguments)
        cached = self._read_cache(key)
        if cached is not None:
            cached["cache_hit"] = True
            return cached

        if self.server_url:
            try:
                payload = {"tool_name": tool_name, "arguments": arguments}
                result = self._post_json("/call", payload)
                response = {
                    "tool_name": tool_name,
                    "arguments": arguments,
                    "result": result,
                    "provider": "mcp_server",
                    "cache_hit": False,
                }
                self._write_cache(key, response)
                return response
            except (urllib.error.URLError, TimeoutError, ValueError):
                pass
            except Exception:
                pass

        if not self.enable_mock_tools:
            return {
                "tool_name": tool_name,
                "arguments": arguments,
                "error": "MCP server unavailable and mock tools disabled",
                "provider": "none",
                "cache_hit": False,
            }

        result = self._call_mock_tool(tool_name, arguments)
        response = {
            "tool_name": tool_name,
            "arguments": arguments,
            "result": result,
            "provider": "mock_mcp",
            "cache_hit": False,
        }
        self._write_cache(key, response)
        return response

    def call_tools_parallel(self, tool_requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        if not tool_requests:
            return []

        results: List[Dict[str, Any]] = []
        with ThreadPoolExecutor(max_workers=min(4, len(tool_requests))) as executor:
            future_to_request = {
                executor.submit(
                    self.call_tool,
                    request.get("tool_name", ""),
                    request.get("arguments", {}),
                ): request
                for request in tool_requests
            }
            for future in as_completed(future_to_request):
                request = future_to_request[future]
                try:
                    results.append(future.result())
                except Exception as exc:
                    results.append(
                        {
                            "tool_name": request.get("tool_name"),
                            "arguments": request.get("arguments", {}),
                            "error": str(exc),
                            "provider": "none",
                            "cache_hit": False,
                        }
                    )
        return results

    def _mock_tools(self) -> List[Dict[str, Any]]:
        return [
            {
                "name": "get_stock_price",
                "description": "Return recent stock price for a ticker symbol",
                "input_schema": {"symbol": "string"},
            },
            {
                "name": "fetch_sec_filing",
                "description": "Retrieve filing excerpts for company + year + report type",
                "input_schema": {
                    "company": "string",
                    "year": "string",
                    "report_type": "string",
                },
            },
            {
                "name": "calculate_ratio",
                "description": "Compute ratio from numerator and denominator",
                "input_schema": {
                    "numerator": "number",
                    "denominator": "number",
                    "metric_name": "string",
                },
            },
        ]

    def _call_mock_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        if tool_name == "get_stock_price":
            symbol = str(arguments.get("symbol", "AAPL")).upper()
            baseline = {
                "AAPL": 192.45,
                "MSFT": 414.88,
                "GOOGL": 176.12,
                "AMZN": 184.37,
                "META": 498.03,
            }
            price = baseline.get(symbol, 150.0)
            return {
                "symbol": symbol,
                "price": price,
                "currency": "USD",
                "as_of": time.strftime("%Y-%m-%d"),
            }

        if tool_name == "fetch_sec_filing":
            company = str(arguments.get("company", "")).strip().lower()
            year = str(arguments.get("year", "")).strip()
            report_type = str(arguments.get("report_type", "10-k")).strip().lower()
            query = f"{company} {report_type} {year}".strip()
            docs = retrieve_documents(query=query, mode="hybrid", limit=3, rerank=True)
            return {
                "company": company,
                "year": year,
                "report_type": report_type,
                "snippets": [
                    {
                        "file_name": doc.get("file_name"),
                        "year": doc.get("year"),
                        "quarter": doc.get("quarter"),
                        "snippet": doc.get("snippet"),
                    }
                    for doc in docs
                ],
            }

        if tool_name == "calculate_ratio":
            numerator = float(arguments.get("numerator", 0.0))
            denominator = float(arguments.get("denominator", 1.0))
            metric_name = str(arguments.get("metric_name", "ratio"))
            value = None if denominator == 0.0 else (numerator / denominator)
            return {
                "metric_name": metric_name,
                "numerator": numerator,
                "denominator": denominator,
                "value": value,
            }

        raise ValueError(f"Unsupported MCP tool: {tool_name}")