1 line
16 KiB
JSON
1 line
16 KiB
JSON
{"repo": "alberthakobyan070/advanced-obsidian-rag", "files": {"src/llm/llm_client.py": {"90f083d83659cc377ec0a31dd34326f7c27a9480": {"t": 1785674012.3646433, "n": 15383, "h": "454cb9c2cc7f7b40", "c": "\"\"\"\nllm_client.py — One interface, any OpenAI- or Anthropic-compatible backend.\n\nThe rest of the pipeline NEVER imports a vendor SDK directly, so the whole\ngeneration layer swaps by editing one line in config.yaml.\n\n from src.llm.llm_client import LLMClient\n llm = LLMClient.from_config(cfg) # role=\"generation\"\n judge = LLMClient.from_config(cfg, role=\"eval.judge\")\n\nTwo ways to name a backend, and they coexist:\n\n1. PROVIDER REGISTRY (preferred). `providers:` in config.yaml holds named\n backends; a role points at one by name:\n\n providers:\n minimax:\n kind: anthropic # wire protocol, not vendor\n base_url: \"https://api.minimax.io/anthropic\"\n model: \"MiniMax-M3\"\n api_key_env: MINIMAX_API_KEY # the NAME of an env var\n api_key_prefix: \"sk-cp-\" # Token Plan key type\n generation:\n provider: minimax\n\n Keys are NEVER stored in config.yaml — a provider names the environment\n variable holding its key, and .env (gitignored) is loaded at startup. Add a\n provider for any service exposing an OpenAI-compatible /v1/chat/completions\n or an Anthropic-compatible endpoint: MiniMax, Moonshot/Kimi, xAI, OpenAI,\n a local KoboldCPP, an office proxy. `kind` selects the wire protocol only.\n\n2. LEGACY (unchanged, still fully supported). The three names `openai`,\n `anthropic` and `local` are RESERVED and always take the original code\n path with their original env vars (OPENAI_API_KEY / ANTHROPIC_API_KEY).\n Never add a registry entry under those names — existing configs, including\n the Docker bundle, depend on them resolving the old way.\n\nPer-role: any config section may carry `provider:` / `temperature:` /\n`max_tokens:`, so generation, HyDE and the eval judge can each run on a\ndifferent backend. `<role>.model` overrides the provider's default model.\n\nLocal note (KoboldCPP serving Gemma-4-E4B-it-Q8):\n KoboldCPP exposes an OpenAI-compatible API at /v1, so we reuse the OpenAI\n SDK and just point base_url at it with a dummy key. Context length must be\n set when LAUNCHING KoboldCPP (--contextsize N); it cannot be set per-request.\n\"\"\"\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass\nfrom typing import Any\n\nfrom src.utils.config_loader import Config\nfrom src.utils.logger import get_logger\n\nlog = get_logger(__name__)\n\n\n@dataclass\nclass LLMResponse:\n text: str\n model: str\n provider: str\n usage: dict[str, Any] | None = None\n\n\nclass LLMClient:\n \"\"\"Provider-agnostic chat completion wrapper.\"\"\"\n\n def __init__(\n self,\n provider: str,\n model: str,\n api_key: str | None = None,\n base_url: str | None = None,\n default_temperature: float = 0.1,\n default_max_tokens: int = 1500,\n backend: str | None = None,\n ):\n # `provider` is the wire protocol used by complete(); `backend` is the\n # configured registry alias (minimax, freellmapi, ...). Keeping both\n # prevents per-request backend comparisons from collapsing every\n # OpenAI-compatible service into the label \"openai\".\n self.provider = provider\n self.backend = backend or provider\n self.model = model\n self.default_temperature = default_temperature\n self.default_max_tokens = default_max_tokens\n self._client = self._build_client(provider, api_key, base_url)\n log.info(\n \"LLMClient ready: backend=%s provider=%s model=%s\",\n self.backend, provider, model,\n )\n\n # ---- construction -------------------------------------------------\n\n # Names that always take the legacy path below. A registry entry under one\n # of these would silently redirect every pre-registry config.\n RESERVED_PROVIDERS = (\"openai\", \"anthropic\", \"local\")\n\n @classmethod\n def from_config(cls, cfg: Config, role: str = \"generation\") -> \"LLMClient\":\n \"\"\"\n Build a client for a config section: \"generation\", \"eval.judge\", ...\n\n Resolution order: a `providers:` registry entry whose name matches\n `<role>.provider` wins; otherwise the reserved legacy names apply.\n \"\"\"\n provider = cfg.get(f\"{role}.provider\", \"anthropic\")\n temperature = cfg.get(f\"{role}.temperature\", 0.1)\n max_tokens = cfg.get(f\"{role}.max_tokens\", 1500)\n\n registry = cfg.get(\"providers\", {}) or {}\n if provider in registry and provider not in cls.RESERVED_PROVIDERS:\n return cls._from_registry(\n cfg, role, provider, registry[provider], temperature, max_tokens\n )\n if provider not in cls.RESERVED_PROVIDERS and registry:\n raise ValueError(\n f\"{role}.provider = {provider!r} is not in the `providers:` \"\n f\"registry and is not one of {cls.RESERVED_PROVIDERS}. \"\n f\"Known providers: {sorted(registry)}\"\n )\n\n if provider == \"anthropic\":\n return cls(\n provider=\"anthropic\",\n model=cfg.get(f\"{role}.model\", \"claude-sonnet-4-20250514\"),\n api_key=cfg.require_secret(\"ANTHROPIC_API_KEY\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n )\n if provider == \"openai\":\n return cls(\n provider=\"openai\",\n model=cfg.get(f\"{role}.model\", \"gpt-4o-mini\"),\n api_key=cfg.require_secret(\"OPENAI_API_KEY\"),\n base_url=cfg.get(f\"{role}.base_url\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n )\n if provider == \"local\":\n return cls(\n provider=\"local\",\n model=cfg.get(f\"{role}.local.model\", \"local-model\"),\n api_key=cfg.get(f\"{role}.local.api_key\", \"local\"),\n base_url=cfg.get(f\"{role}.local.base_url\", \"http://localhost:5001/v1\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n )\n raise ValueError(f\"Unknown LLM provider: {provider!r}\")\n\n @classmethod\n def from_provider_override(\n cls,\n cfg: Config,\n provider: str,\n model: str | None = None,\n role: str = \"generation\",\n ) -> \"LLMClient\":\n \"\"\"Build a request-local client for a configured backend.\n\n Registry backends use their own default model, not ``<role>.model``:\n that role value belongs to the process-wide default backend and may be\n invalid for the requested one. ``model`` is the only per-request model\n override. The supplied ``Config`` is read only and is never copied or\n mutated.\n\n Reserved legacy providers remain available and retain their established\n construction rules. For them, ``model`` overrides the configured role\n model when supplied.\n \"\"\"\n temperature = cfg.get(f\"{role}.temperature\", 0.1)\n max_tokens = cfg.get(f\"{role}.max_tokens\", 1500)\n registry = cfg.get(\"providers\", {}) or {}\n\n if provider in registry and provider not in cls.RESERVED_PROVIDERS:\n return cls._from_registry(\n cfg,\n role,\n provider,\n registry[provider],\n temperature,\n max_tokens,\n model_override=model,\n inherit_role_model=False,\n )\n if provider not in cls.RESERVED_PROVIDERS and registry:\n raise ValueError(\n f\"{role}.provider = {provider!r} is not in the `providers:` \"\n f\"registry and is not one of {cls.RESERVED_PROVIDERS}. \"\n f\"Known providers: {sorted(registry)}\"\n )\n\n if provider == \"anthropic\":\n return cls(\n provider=\"anthropic\",\n model=model or cfg.get(\n f\"{role}.model\", \"claude-sonnet-4-20250514\"\n ),\n api_key=cfg.require_secret(\"ANTHROPIC_API_KEY\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n backend=provider,\n )\n if provider == \"openai\":\n return cls(\n provider=\"openai\",\n model=model or cfg.get(f\"{role}.model\", \"gpt-4o-mini\"),\n api_key=cfg.require_secret(\"OPENAI_API_KEY\"),\n base_url=cfg.get(f\"{role}.base_url\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n backend=provider,\n )\n if provider == \"local\":\n return cls(\n provider=\"local\",\n model=model or cfg.get(f\"{role}.local.model\", \"local-model\"),\n api_key=cfg.get(f\"{role}.local.api_key\", \"local\"),\n base_url=cfg.get(\n f\"{role}.local.base_url\", \"http://localhost:5001/v1\"\n ),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n backend=provider,\n )\n raise ValueError(f\"Unknown LLM provider: {provider!r}\")\n\n @classmethod\n def _from_registry(\n cls,\n cfg: Config,\n role: str,\n name: str,\n spec: dict,\n temperature: float,\n max_tokens: int,\n model_override: str | None = None,\n inherit_role_model: bool = True,\n ) -> \"LLMClient\":\n \"\"\"Build a client from a `providers:` entry.\n\n The provider owns the backend facts (endpoint, model, which env var\n holds its key); the role may override `model` to run a different model\n on the same backend. That override is logged, because pointing a role\n at a new provider while leaving a stale `<role>.model` behind is the\n one way to misconfigure this quietly — the model name simply gets\n rejected by the new endpoint.\n \"\"\"\n spec = spec or {}\n kind = str(spec.get(\"kind\", \"openai\")).lower()\n if kind not in (\"openai\", \"anthropic\"):\n raise ValueError(\n f\"providers.{name}.kind = {kind!r}; expected 'openai' or \"\n f\"'anthropic' (the wire protocol, not the vendor)\"\n )\n\n spec_model = spec.get(\"model\")\n role_model = (\n cfg.get(f\"{role}.model\") if inherit_role_model else model_override\n )\n model = role_model or spec_model\n if not model:\n raise ValueError(\n f\"No model for role {role!r}: set providers.{name}.model \"\n f\"or {role}.model\"\n )\n\n # Keys come from the environment, never from config.yaml. A provider\n # with no api_key_env is a keyless local endpoint (KoboldCPP, Ollama,\n # a localhost proxy) — the OpenAI SDK still wants a non-empty string.\n key_env = spec.get(\"api_key_env\")\n if key_env:\n raw_key = (cfg.secret(key_env) if spec.get(\"api_key_optional\")\n else cfg.require_secret(key_env))\n expected_prefix = str(spec.get(\"api_key_prefix\") or \"\")\n if (raw_key and expected_prefix\n and not raw_key.startswith(expected_prefix)):\n raise RuntimeError(\n f\"{key_env} has the wrong credential type for \"\n f\"providers.{name}; expected a key beginning with \"\n f\"{expected_prefix!r}\"\n )\n api_key = raw_key or \"not-needed\"\n else:\n api_key = \"not-needed\"\n\n if role_model and spec_model and role_model != spec_model:\n log.warning(\n \"%s.model=%r overrides providers.%s.model=%r — make sure %s \"\n \"actually serves that model\",\n role, role_model, name, spec_model, name,\n )\n log.info(\"LLMClient[%s]: provider=%s kind=%s model=%s\", role, name, kind, model)\n\n return cls(\n provider=kind,\n model=model,\n api_key=api_key,\n base_url=spec.get(\"base_url\"),\n default_temperature=temperature,\n default_max_tokens=max_tokens,\n backend=name,\n )\n\n def _build_client(self, provider: str, api_key, base_url):\n if provider == \"anthropic\":\n try:\n import anthropic\n except ImportError as e:\n raise ImportError(\"pip install anthropic\") from e\n # base_url matters here too: several providers (MiniMax among them)\n # expose an Anthropic-compatible endpoint alongside their /v1 one.\n kwargs: dict[str, Any] = {\"api_key\": api_key}\n if base_url:\n kwargs[\"base_url\"] = base_url\n return anthropic.Anthropic(**kwargs)\n # openai + local both use the OpenAI SDK\n try:\n from openai import OpenAI\n except ImportError as e:\n raise ImportError(\"pip install openai\") from e\n kwargs: dict[str, Any] = {\"api_key\": api_key}\n if base_url:\n kwargs[\"base_url\"] = base_url\n return OpenAI(**kwargs)\n\n # ---- completion ---------------------------------------------------\n\n def complete(\n self,\n system: str,\n user: str,\n temperature: float | None = None,\n max_tokens: int | None = None,\n ) -> LLMResponse:\n \"\"\"Single-turn system+user completion. Synchronous.\"\"\"\n temp = self.default_temperature if temperature is None else temperature\n maxtok = self.default_max_tokens if max_tokens is None else max_tokens\n\n if self.provider == \"anthropic\":\n return self._complete_anthropic(system, user, temp, maxtok)\n return self._complete_openai_compatible(system, user, temp, maxtok)\n\n def _complete_anthropic(self, system, user, temp, maxtok) -> LLMResponse:\n resp = self._client.messages.create(\n model=self.model,\n system=system,\n max_tokens=maxtok,\n temperature=temp,\n messages=[{\"role\": \"user\", \"content\": user}],\n )\n text = \"\".join(\n block.text for block in resp.content if getattr(block, \"type\", \"\") == \"text\"\n )\n usage = None\n if getattr(resp, \"usage\", None):\n usage = {\n \"input_tokens\": resp.usage.input_tokens,\n \"output_tokens\": resp.usage.output_tokens,\n }\n return LLMResponse(text=text, model=self.model, provider=self.provider, usage=usage)\n\n def _complete_openai_compatible(self, system, user, temp, maxtok) -> LLMResponse:\n resp = self._client.chat.completions.create(\n model=self.model,\n temperature=temp,\n max_tokens=maxtok,\n messages=[\n {\"role\": \"system\", \"content\": system},\n {\"role\": \"user\", \"content\": user},\n ],\n )\n text = resp.choices[0].message.content or \"\"\n usage = None\n if getattr(resp, \"usage\", None):\n usage = {\n \"input_tokens\": resp.usage.prompt_tokens,\n \"output_tokens\": resp.usage.completion_tokens,\n }\n return LLMResponse(text=text, model=self.model, provider=self.provider, usage=usage)\n"}}}} |