1 line
19 KiB
JSON
1 line
19 KiB
JSON
{"repo": "angrysky56/project-synapse-mcp", "files": {"src/synapse_mcp/llm/providers.py": {"c56f9cc7b5949fdeda9f3c1ef37affde4f6d7ec6": {"t": 1785674009.880813, "n": 18403, "h": "540e7f337f56126f", "c": "\"\"\"LLM provider implementations for chat completions with strict JSON output.\n\nThree providers ship today:\n\n - :class:`OllamaProvider` — local Ollama server. Uses ``/api/chat`` with\n ``format: \"json\"`` (or a full JSON schema dict).\n - :class:`MinimaxProvider` — MiniMax via its OpenAI-compatible endpoint at\n ``https://api.minimax.io/v1``. Auth: ``Authorization: Bearer sk-cp-...``.\n - :class:`OpenRouterProvider` — OpenRouter chat completions at\n ``https://openrouter.ai/api/v1``. Proxies many model vendors behind a\n single OpenAI-style API; useful when you want to A/B different models\n without coding a new provider.\n\nAll providers expose the same async ``chat_json(system, user, *, ...)`` method\nwhich returns a parsed ``dict``. They also handle the well-known failure mode\nwhere models leak chain-of-thought or markdown around the JSON: see\n:func:`_salvage_json` for the cleanup pipeline.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport abc\nimport asyncio\nimport json\nimport os\nimport re\nfrom typing import Any\n\nimport aiohttp\n\nfrom ..utils.logging_config import get_logger\n\nlogger = get_logger(__name__)\n\n\nclass LlmResponseError(RuntimeError):\n \"\"\"Raised when an LLM call produces an unrecoverable response.\"\"\"\n\n\n# Reasoning models (DeepSeek-R1, QwQ, Gemma-think, etc.) emit chain-of-thought\n# inside <think>/<thought>/<reasoning> blocks even when asked for JSON. Strip\n# them before json.loads().\n_THINK_BLOCK_RE = re.compile(\n r\"<(?:think|thought|reasoning)>[\\s\\S]*?</(?:think|thought|reasoning)>\",\n re.IGNORECASE,\n)\n_FENCE_RE = re.compile(r\"```(?:json|JSON)?\\s*([\\s\\S]*?)\\s*```\")\n\n\ndef _salvage_json(content: str) -> dict[str, Any]:\n \"\"\"Extract a JSON object from an LLM response that may contain prose.\n\n Strategy, in order:\n\n 1. Strip ``<think>`` / ``<thought>`` / ``<reasoning>`` blocks.\n 2. Unwrap markdown ``` ``` ``` code fences (with or without ``json`` lang tag).\n 3. Slice between the first ``{`` and last ``}`` to skip preamble/postamble.\n 4. Repair common LLM JSON malformation: unescaped single quotes (``\\'`` →\n ``'``), bare single quotes inside strings, trailing comma before ``}``.\n 5. ``json.loads`` the result.\n\n Raises:\n LlmResponseError: if no valid JSON object can be extracted.\n \"\"\"\n if not content or not content.strip():\n raise LlmResponseError(\"LLM returned empty content\")\n\n cleaned = _THINK_BLOCK_RE.sub(\"\", content).strip()\n\n fence_match = _FENCE_RE.search(cleaned)\n if fence_match:\n cleaned = fence_match.group(1)\n elif \"```\" in cleaned:\n cleaned = cleaned.replace(\"```json\", \"\").replace(\"```\", \"\").strip()\n\n start = cleaned.find(\"{\")\n end = cleaned.rfind(\"}\")\n if start == -1 or end == -1 or end <= start:\n raise LlmResponseError(\n f\"No JSON object found in LLM response (preview: {content[:200]!r})\"\n )\n\n cleaned = cleaned[start : end + 1]\n\n # Repair step: some models emit \\' inside string values, or leave\n # trailing commas before } or ] — both are invalid JSON.\n repaired = (\n cleaned.replace(\"\\\\'\", \"'\")\n .replace(\",}\", \"}\")\n .replace(\",]\", \"]\")\n )\n\n try:\n return json.loads(repaired)\n except json.JSONDecodeError:\n raise LlmResponseError(\n f\"LLM response is not valid JSON: {repaired[:200]!r}\"\n ) from None\n\n\nclass LlmProvider(abc.ABC):\n \"\"\"Abstract base for chat-completion providers with JSON output.\"\"\"\n\n #: Short slug used in logs and env-var prefixes (``ollama``, ``minimax``…).\n name: str = \"abstract\"\n #: The model identifier this instance will call.\n model: str = \"\"\n #: The base URL this instance will hit.\n base_url: str = \"\"\n\n def __init__(\n self,\n base_url: str | None = None,\n model: str | None = None,\n ) -> None:\n self.base_url = base_url or \"\"\n self.model = model or \"\"\n\n @abc.abstractmethod\n async def chat_json(\n self,\n system: str,\n user: str,\n *,\n temperature: float = 0.0,\n max_tokens: int = 16384,\n schema: dict[str, Any] | None = None,\n timeout: int = 120,\n ) -> dict[str, Any]:\n \"\"\"Send a system+user prompt and return parsed JSON.\n\n Args:\n system: System prompt — hard rules, output schema description.\n user: User message — the actual content to process.\n temperature: Sampling temperature. Default 0 for determinism.\n max_tokens: Max output tokens.\n schema: Optional JSON schema. Providers that support structured\n outputs (OpenAI ``response_format=json_schema``, Ollama\n ``format=<schema>``) will use it; others fall back to a plain\n ``json_object`` mode.\n timeout: Per-request timeout in seconds.\n\n Returns:\n Parsed JSON object as a dict.\n\n Raises:\n LlmResponseError: on network error, non-2xx, empty body, or\n unparseable JSON.\n \"\"\"\n\n @abc.abstractmethod\n async def check_available(self) -> tuple[bool, str]:\n \"\"\"Probe the provider; return ``(ok, message)``.\n\n Should be cheap and fast — used at startup to fail fast with an\n actionable error instead of crashing on the first real call.\n \"\"\"\n\n\nclass OllamaProvider(LlmProvider):\n \"\"\"Local Ollama server. Uses ``/api/chat`` with ``format=json``.\"\"\"\n\n name = \"ollama\"\n\n def __init__(\n self,\n base_url: str | None = None,\n model: str | None = None,\n ) -> None:\n self.base_url = (\n base_url or os.getenv(\"OLLAMA_BASE_URL\") or \"http://localhost:11434\"\n ).rstrip(\"/\")\n self.model = model or os.getenv(\"OLLAMA_EXTRACTION_MODEL\") or \"gemma4:latest\"\n\n async def check_available(self) -> tuple[bool, str]:\n url = f\"{self.base_url}/api/tags\"\n try:\n timeout = aiohttp.ClientTimeout(total=5)\n async with aiohttp.ClientSession(timeout=timeout) as session:\n async with session.get(url) as resp:\n if resp.status >= 400:\n return False, f\"Ollama returned HTTP {resp.status}\"\n data = await resp.json()\n except Exception as e: # pylint: disable=broad-exception-caught\n return False, f\"Cannot reach Ollama at {self.base_url}: {e}\"\n\n # Ollama tag names may or may not include ':latest' — accept either.\n names = {m.get(\"name\", \"\") for m in data.get(\"models\", []) or []}\n wanted = {self.model, f\"{self.model}:latest\", self.model.split(\":\")[0]}\n if not names & wanted:\n return (\n False,\n f\"Model '{self.model}' not loaded in Ollama. \"\n f\"Run: ollama pull {self.model}\",\n )\n return True, \"ok\"\n\n async def chat_json(\n self,\n system: str,\n user: str,\n *,\n temperature: float = 0.0,\n max_tokens: int = 16384,\n schema: dict[str, Any] | None = None,\n timeout: int = 120,\n ) -> dict[str, Any]:\n # Ollama accepts either the string \"json\" or a full JSON schema dict\n # for `format`. Schema gives stronger guarantees on newer models.\n fmt: str | dict[str, Any] = schema if schema else \"json\"\n payload: dict[str, Any] = {\n \"model\": self.model,\n \"system\": system,\n \"messages\": [{\"role\": \"user\", \"content\": user}],\n \"stream\": False,\n \"format\": fmt,\n \"options\": {\n \"temperature\": temperature,\n \"num_ctx\": int(os.getenv(\"OLLAMA_NUM_CTX\", \"131072\")),\n \"num_predict\": max_tokens,\n \"num_keep\": -1,\n \"top_k\": 1,\n \"top_p\": 0.1,\n },\n # Allow reasoning models to think — segregated from `content` by\n # Ollama itself, so it doesn't pollute the JSON we parse.\n \"think\": os.getenv(\"OLLAMA_THINK\", \"true\").lower() == \"true\",\n }\n\n api_url = f\"{self.base_url}/api/chat\"\n client_timeout = aiohttp.ClientTimeout(total=timeout)\n try:\n async with aiohttp.ClientSession(timeout=client_timeout) as session:\n async with session.post(api_url, json=payload) as resp:\n if resp.status >= 400:\n body = await resp.text()\n raise LlmResponseError(\n f\"Ollama HTTP {resp.status}: {body[:300]}\"\n )\n result = await resp.json()\n except aiohttp.ClientError as e:\n raise LlmResponseError(f\"Ollama request failed: {e}\") from e\n except asyncio.TimeoutError as e:\n raise LlmResponseError(f\"Ollama request timed out after {timeout}s\") from e\n\n content = (result.get(\"message\") or {}).get(\"content\", \"\").strip()\n return _salvage_json(content)\n\n\nclass _OpenAICompatibleProvider(LlmProvider):\n \"\"\"Shared implementation for OpenAI-compatible chat APIs.\n\n Both MiniMax and OpenRouter implement the OpenAI ``/chat/completions``\n contract, so all the wire-format logic lives here. Subclasses just set\n the defaults (``default_base_url``, ``default_model``, ``api_key_env``).\n \"\"\"\n\n name = \"openai-compat\"\n api_key_env: str = \"\"\n default_base_url: str = \"\"\n default_model: str = \"\"\n\n def __init__(\n self,\n base_url: str | None = None,\n model: str | None = None,\n api_key: str | None = None,\n ) -> None:\n self.api_key = api_key or os.getenv(self.api_key_env, \"\")\n prefix = self.name.upper()\n self.base_url = (\n base_url or os.getenv(f\"{prefix}_BASE_URL\") or self.default_base_url\n ).rstrip(\"/\")\n self.model = model or os.getenv(f\"{prefix}_MODEL\") or self.default_model\n # Subclasses can append vendor-specific headers (e.g. OpenRouter\n # wants HTTP-Referer / X-Title for analytics).\n self.extra_headers: dict[str, str] = {}\n\n async def check_available(self) -> tuple[bool, str]:\n if not self.api_key:\n return False, f\"Missing API key (env: {self.api_key_env})\"\n url = f\"{self.base_url}/models\"\n try:\n headers = {\n \"Authorization\": f\"Bearer {self.api_key}\",\n **self.extra_headers,\n }\n timeout = aiohttp.ClientTimeout(total=10)\n async with aiohttp.ClientSession(timeout=timeout) as session:\n async with session.get(url, headers=headers) as resp:\n if resp.status == 401:\n return False, f\"{self.name}: invalid API key (HTTP 401)\"\n if resp.status >= 500:\n return False, f\"{self.name}: server error HTTP {resp.status}\"\n except Exception as e: # pylint: disable=broad-exception-caught\n return False, f\"{self.name}: cannot reach {self.base_url}: {e}\"\n return True, \"ok\"\n\n async def chat_json(\n self,\n system: str,\n user: str,\n *,\n temperature: float = 0.0,\n max_tokens: int = 16384,\n schema: dict[str, Any] | None = None,\n timeout: int = 120,\n ) -> dict[str, Any]:\n if not self.api_key:\n raise LlmResponseError(\n f\"{self.name}: API key missing (set {self.api_key_env})\"\n )\n\n # Try strict json_schema first if a schema was given; some servers\n # reject the json_schema variant — fall back to plain json_object on\n # 4xx. Most newer servers (OpenAI, OpenRouter for many models,\n # MiniMax-M2.7) accept json_schema; older/lighter ones don't.\n if schema:\n response_format: dict[str, Any] = {\n \"type\": \"json_schema\",\n \"json_schema\": {\n \"name\": \"synapse_response\",\n \"schema\": schema,\n \"strict\": True,\n },\n }\n else:\n response_format = {\"type\": \"json_object\"}\n\n payload: dict[str, Any] = {\n \"model\": self.model,\n \"messages\": [\n {\"role\": \"system\", \"content\": system},\n {\"role\": \"user\", \"content\": user},\n ],\n \"temperature\": temperature,\n \"max_tokens\": max_tokens,\n \"response_format\": response_format,\n \"stream\": False,\n }\n headers = {\n \"Authorization\": f\"Bearer {self.api_key}\",\n \"Content-Type\": \"application/json\",\n **self.extra_headers,\n }\n url = f\"{self.base_url}/chat/completions\"\n client_timeout = aiohttp.ClientTimeout(total=timeout)\n\n last_err: Exception | None = None\n result_json: dict[str, Any] | None = None\n for attempt in range(2):\n try:\n async with aiohttp.ClientSession(timeout=client_timeout) as session:\n async with session.post(url, json=payload, headers=headers) as resp:\n body_text = await resp.text()\n if resp.status >= 400:\n # Downgrade json_schema -> json_object on first 4xx\n # before giving up.\n if (\n attempt == 0\n and 400 <= resp.status < 500\n and payload[\"response_format\"].get(\"type\")\n == \"json_schema\"\n ):\n logger.warning(\n \"%s rejected json_schema (HTTP %d); \"\n \"retrying with json_object. Body: %s\",\n self.name,\n resp.status,\n body_text[:200],\n )\n payload[\"response_format\"] = {\"type\": \"json_object\"}\n continue\n raise LlmResponseError(\n f\"{self.name} HTTP {resp.status}: {body_text[:300]}\"\n )\n try:\n result_json = json.loads(body_text)\n except json.JSONDecodeError as e:\n raise LlmResponseError(\n f\"{self.name}: response body is not JSON: {e}\"\n ) from e\n break\n except (aiohttp.ClientError, asyncio.TimeoutError) as e:\n last_err = e\n if attempt == 1:\n raise LlmResponseError(f\"{self.name} request failed: {e}\") from e\n logger.warning(\"%s attempt %d failed: %s\", self.name, attempt + 1, e)\n await asyncio.sleep(0.5)\n\n if result_json is None:\n raise LlmResponseError(f\"{self.name}: exhausted retries ({last_err})\")\n\n choices = result_json.get(\"choices\") or []\n if not choices:\n raise LlmResponseError(f\"{self.name}: no choices in response\")\n content = (choices[0].get(\"message\") or {}).get(\"content\", \"\") or \"\"\n return _salvage_json(content.strip())\n\n\nclass MinimaxProvider(_OpenAICompatibleProvider):\n \"\"\"MiniMax via its OpenAI-compatible endpoint.\n\n Docs: https://platform.minimax.io/docs/token-plan/other-tools\n\n Set ``MINIMAX_API_KEY`` (the ``sk-cp-...`` Token Plan key). Optionally\n override ``MINIMAX_BASE_URL`` and ``MINIMAX_MODEL``.\n \"\"\"\n\n name = \"minimax\"\n api_key_env = \"MINIMAX_API_KEY\"\n default_base_url = \"https://api.minimax.io/v1\"\n # MiniMax-M2.7-highspeed is the latency-optimized variant — preferred for\n # insight synthesis where we make many small calls.\n default_model = \"MiniMax-M2.7-highspeed\"\n\n\nclass OpenRouterProvider(_OpenAICompatibleProvider):\n \"\"\"OpenRouter chat completions (proxies many model vendors).\n\n Docs: https://openrouter.ai/docs/api-reference/chat-completion\n\n Set ``OPENROUTER_API_KEY``. Optionally override ``OPENROUTER_MODEL``\n (defaults to a fast JSON-strong model) and ``OPENROUTER_BASE_URL``.\n\n OpenRouter recommends sending ``HTTP-Referer`` and ``X-Title`` headers\n for analytics and rate-limit allocation — set ``OPENROUTER_HTTP_REFERER``\n and ``OPENROUTER_X_TITLE`` if you have a public app URL.\n \"\"\"\n\n name = \"openrouter\"\n api_key_env = \"OPENROUTER_API_KEY\"\n default_base_url = \"https://openrouter.ai/api/v1\"\n # Sensible default: small, fast, strong on structured output.\n default_model = \"openai/gpt-4o-mini\"\n\n def __init__(\n self,\n base_url: str | None = None,\n model: str | None = None,\n api_key: str | None = None,\n ) -> None:\n super().__init__(base_url=base_url, model=model, api_key=api_key)\n referer = os.getenv(\"OPENROUTER_HTTP_REFERER\", \"\")\n title = os.getenv(\"OPENROUTER_X_TITLE\", \"project-synapse-mcp\")\n if referer:\n self.extra_headers[\"HTTP-Referer\"] = referer\n if title:\n self.extra_headers[\"X-Title\"] = title\n\n\n# Registry — lowercase keys, used by get_provider().\n_PROVIDERS: dict[str, type[LlmProvider]] = {\n \"ollama\": OllamaProvider,\n \"minimax\": MinimaxProvider,\n \"openrouter\": OpenRouterProvider,\n}\n\n\ndef get_provider(\n name: str | None = None,\n *,\n model: str | None = None,\n base_url: str | None = None,\n) -> LlmProvider:\n \"\"\"Resolve a provider by name.\n\n Args:\n name: Provider slug (``ollama``, ``minimax``, ``openrouter``).\n If None, reads ``LLM_PROVIDER`` env var, then defaults to ``ollama``.\n model: Override the provider's default model.\n base_url: Override the provider's default base URL.\n\n Returns:\n An instantiated :class:`LlmProvider` ready for ``chat_json`` calls.\n\n Raises:\n ValueError: if ``name`` is not a registered provider.\n \"\"\"\n resolved = (name or os.getenv(\"LLM_PROVIDER\") or \"ollama\").strip().lower()\n cls = _PROVIDERS.get(resolved)\n if cls is None:\n raise ValueError(\n f\"Unknown LLM provider: {resolved!r}. \"\n f\"Available: {sorted(_PROVIDERS.keys())}\"\n )\n return cls(base_url=base_url, model=model)\n"}}}} |