From c3719b99c25313d9614547c82312c4715b22de80 Mon Sep 17 00:00:00 2001 From: xm Date: Fri, 14 Aug 2026 12:04:44 +0200 Subject: [PATCH] BUG fix : LLmLite not used/called from embedded and webui --- serviette/server/llm.py | 83 ++++++++++++++++++++++++++++++++++-- serviette/server/reranker.py | 12 +++--- tests/test_server.py | 44 +++++++++++++++++++ 3 files changed, 131 insertions(+), 8 deletions(-) diff --git a/serviette/server/llm.py b/serviette/server/llm.py index ac067d3..d410a8b 100644 --- a/serviette/server/llm.py +++ b/serviette/server/llm.py @@ -149,15 +149,92 @@ async def close(self) -> None: self._client = None -_OPENAI_COMPATIBLE = {"openai", "litellm"} +class LiteLLMChat: + """LLM backend routed through LiteLLM (``type: litellm``). + + Unlike ``OpenAIChat`` this calls ``litellm.acompletion``, so the provider + prefix in ``model`` (``openrouter/...``, ``anthropic/...``, …) selects the + endpoint and ``api_key`` is forwarded to *that* provider. Routing a + litellm-typed config through the plain OpenAI client instead would send + the key to api.openai.com and 401 — which is exactly the bug this class + exists to prevent. + """ + + def __init__(self, config) -> None: + if not config.model: + raise ValueError( + "llm type 'litellm' requires a 'model' with a provider prefix, " + "e.g. 'openrouter/moonshotai/kimi-k3'." + ) + self._model = config.model + self._api_key = config.api_key + self._temperature = getattr(config, "temperature", None) + self._reasoning_effort = getattr(config, "reasoning_effort", None) + self._system_prompt = getattr(config, "system_prompt", None) + extra = config.model_dump( + exclude={ + "type", "model", "api_key", "temperature", "system_prompt", + "reasoning_effort", + } + ) + # Extra keys are per-call kwargs for litellm.acompletion (there is no + # persistent client object to construct). + self._call_kwargs = {k: v for k, v in extra.items() if v is not None} + + def _request_kwargs(self) -> dict: + kwargs: dict = {} + if self._temperature is not None: + kwargs["temperature"] = self._temperature + if self._reasoning_effort is not None: + kwargs["reasoning_effort"] = self._reasoning_effort + return kwargs + + async def _call(self, messages: list[dict]) -> str: + import litellm + + resp = await litellm.acompletion( + model=self._model, + api_key=self._api_key, + messages=messages, + **self._call_kwargs, + **self._request_kwargs(), + ) + return resp.choices[0].message.content or "" + + async def complete( + self, query: str, context: list[str], *, system_prompt: str | None = None + ) -> str: + return await self._call( + [ + { + "role": "system", + "content": system_prompt + or self._system_prompt + or _DEFAULT_SYSTEM_PROMPT, + }, + {"role": "user", "content": _build_prompt(query, context)}, + ] + ) + + async def raw(self, prompt: str) -> str: + return await self._call([{"role": "user", "content": prompt}]) + + async def close(self) -> None: + # No persistent client; nothing to release. + return None + + +_SUPPORTED = {"mock", "openai", "litellm"} def build_llm(config) -> AsyncLLM: if config.type == "mock": return MockLLM() - if config.type in _OPENAI_COMPATIBLE: + if config.type == "litellm": + return LiteLLMChat(config) + if config.type == "openai": return OpenAIChat(config) raise ValueError( f"Server-side LLM for type {config.type!r} is not implemented. " - "Supported: " + ", ".join(sorted(_OPENAI_COMPATIBLE)) + "Supported: " + ", ".join(sorted(_SUPPORTED)) ) diff --git a/serviette/server/reranker.py b/serviette/server/reranker.py index e7608a4..894602f 100644 --- a/serviette/server/reranker.py +++ b/serviette/server/reranker.py @@ -93,24 +93,26 @@ class LLMReranker: def __init__(self, config, llm_config=None, *, chat=None) -> None: # The scorer is any AsyncLLM (only ``raw`` is used). Tests inject a - # fake via ``chat``; production lazily builds an OpenAIChat from the - # reranker section, falling back to the top-level llm section. + # fake via ``chat``; production lazily builds the configured backend + # (build_llm) from the reranker section, falling back to the + # top-level llm section. self._chat = chat self._config = config self._llm_config = llm_config def _ensure_chat(self): if self._chat is None: - from serviette.server.llm import OpenAIChat - merged = dict(self._llm_config.model_dump() if self._llm_config else {}) overrides = self._config.model_dump(exclude={"type", "candidates"}) merged.update({k: v for k, v in overrides.items() if v is not None}) merged.setdefault("type", "openai") from serviette.config.schema import LLMConfig + from serviette.server.llm import build_llm - self._chat = OpenAIChat(LLMConfig(**merged)) + # build_llm, not OpenAIChat directly: a litellm-typed llm section + # must route through LiteLLM, not the plain OpenAI client. + self._chat = build_llm(LLMConfig(**merged)) return self._chat async def rerank( diff --git a/tests/test_server.py b/tests/test_server.py index df1eaf3..a4737e6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -73,6 +73,50 @@ def test_rag_disabled_without_llm(store_path, mock_server_embedder): assert resp.status_code == 501 +def test_litellm_routes_through_litellm_not_openai_client(store_path, mock_server_embedder): + """A litellm-typed llm section must call litellm.acompletion, not the OpenAI client. + + Regression: build_llm mapped ``type: litellm`` onto ``OpenAIChat``, which + built ``AsyncOpenAI`` with no base_url and sent the OpenRouter key to + api.openai.com (401). The /rag call in the traceback came from exactly + this path. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock, patch + + from serviette.config.schema import LLMConfig + from serviette.server.llm import LiteLLMChat, build_llm + + llm = build_llm( + LLMConfig( + type="litellm", + model="openrouter/moonshotai/kimi-k3", + api_key="sk-test-openrouter-key", + ) + ) + assert isinstance(llm, LiteLLMChat), "type: litellm must not build OpenAIChat" + + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="answer"))] + ) + + with _client(store_path, mock_server_embedder, llm=llm) as client, patch( + "litellm.acompletion", new=AsyncMock(side_effect=fake_acompletion) + ): + resp = client.post("/api/v1/rag", json={"query": "tell me about cats", "k": 2}) + + assert resp.status_code == 200 + # Provider prefix + key forwarded to litellm; no base_url (no OpenAI client). + assert captured["model"] == "openrouter/moonshotai/kimi-k3" + assert captured["api_key"] == "sk-test-openrouter-key" + assert "base_url" not in captured + assert resp.json()["answer"] == "answer" + + def test_legacy_unversioned_aliases_still_work(store_path, mock_server_embedder): """Pre-versioning routes are kept as deprecated aliases of /api/v1."""