diff --git a/application/prompt_client/litellm_router.py b/application/prompt_client/litellm_router.py new file mode 100644 index 000000000..e9a8b5114 --- /dev/null +++ b/application/prompt_client/litellm_router.py @@ -0,0 +1,215 @@ +"""Shared LiteLLM wrapper for chat, embeddings, Module B, and GSoC/OIE. + +PromptHandler stays DB-coupled (embedding contract + RAG). Everything else +that only needs a completion or embedding should call this module instead of +``import litellm`` so retry, rate-limit detection, and response parsing stay +in one place. +""" + +from __future__ import annotations + +from cre_logging import get_logger + +logger = get_logger(__name__) + +import os +import time +from typing import Any, Callable, List, Optional, Sequence + +from application.prompt_client.llm_error_utils import is_rate_limit_error + +LlmFn = Callable[[str, str], str] + + +def retry_policy() -> tuple[int, int]: + """Return (max_retries, sleep_seconds) from CRE_LLM_* env vars.""" + return ( + int(os.environ.get("CRE_LLM_MAX_RETRIES", "2")), + int(os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS", "15")), + ) + + +def get_litellm() -> Any: + """Import LiteLLM or raise the same RuntimeError PromptHandler uses.""" + try: + import litellm # type: ignore + except ImportError as exc: + raise RuntimeError( + "litellm package is required for PromptHandler LLM calls" + ) from exc + return litellm + + +def with_rate_limit_retry( + fn: Callable[[], Any], + *, + context: str, + max_retries: Optional[int] = None, + retry_sleep_seconds: Optional[int] = None, +) -> Any: + default_retries, default_sleep = retry_policy() + retries = default_retries if max_retries is None else max_retries + sleep_s = default_sleep if retry_sleep_seconds is None else retry_sleep_seconds + for attempt in range(retries + 1): + try: + return fn() + except Exception as err: + if not is_rate_limit_error(err) or attempt >= retries: + raise + logger.info( + "rate/quota limited during %s; sleeping %ss (attempt %s/%s)", + context, + sleep_s, + attempt + 1, + retries + 1, + ) + time.sleep(sleep_s) + raise RuntimeError("unreachable: retry loop exited unexpectedly") + + +def extract_content_text(response: Any, *, strict: bool = True) -> str: + """Pull assistant text from a LiteLLM completion (object or dict).""" + choices = getattr(response, "choices", None) + if not choices and isinstance(response, dict): + choices = response.get("choices") + if not choices: + if strict: + raise ValueError("LLM response did not contain choices") + return "" + first = choices[0] + msg = getattr(first, "message", None) + if msg is None and isinstance(first, dict): + msg = first.get("message") + if msg is None: + if strict: + raise ValueError("LLM response did not contain message content") + return "" + content = getattr(msg, "content", None) + if content is None and isinstance(msg, dict): + content = msg.get("content") + if content is None: + if strict: + raise ValueError("LLM response did not contain message content") + return "" + if isinstance(content, list): + return "".join( + part.get("text", "") if isinstance(part, dict) else str(part) + for part in content + ).strip() + return str(content).strip() + + +def extract_embeddings(response: Any) -> List[List[float]]: + data = getattr(response, "data", None) + if data is None and isinstance(response, dict): + data = response.get("data") + if not isinstance(data, list): + raise ValueError("Embedding response missing data list") + vectors: List[List[float]] = [] + for item in data: + emb = getattr(item, "embedding", None) + if emb is None and isinstance(item, dict): + emb = item.get("embedding") + if not isinstance(emb, list): + raise ValueError("Embedding item missing vector") + vectors.append([float(x) for x in emb]) + return vectors + + +def completion( + *, + model: str, + messages: Sequence[dict[str, Any]], + context: str = "LiteLLM completion", + client: Any = None, + max_retries: Optional[int] = None, + retry_sleep_seconds: Optional[int] = None, + **kwargs: Any, +) -> Any: + """LiteLLM chat completion with shared rate-limit retry.""" + llm = client if client is not None else get_litellm() + + def _call() -> Any: + return llm.completion(model=model, messages=list(messages), **kwargs) + + return with_rate_limit_retry( + _call, + context=context, + max_retries=max_retries, + retry_sleep_seconds=retry_sleep_seconds, + ) + + +def embedding( + *, + model: str, + input: str | List[str], + context: str = "LiteLLM embeddings", + client: Any = None, + max_retries: Optional[int] = None, + retry_sleep_seconds: Optional[int] = None, + **kwargs: Any, +) -> Any: + llm = client if client is not None else get_litellm() + + def _call() -> Any: + return llm.embedding(model=model, input=input, **kwargs) + + return with_rate_limit_retry( + _call, + context=context, + max_retries=max_retries, + retry_sleep_seconds=retry_sleep_seconds, + ) + + +def completion_text( + *, + model: str, + messages: Sequence[dict[str, Any]], + strict: bool = True, + **kwargs: Any, +) -> str: + return extract_content_text( + completion(model=model, messages=messages, **kwargs), + strict=strict, + ) + + +def system_user_fn( + model: str, + *, + temperature: float = 0.0, + extra_try_kwargs: Optional[dict[str, Any]] = None, + **kwargs: Any, +) -> LlmFn: + """Build ``(system, user) -> text`` for Librarian / OIE call sites.""" + + def _call(system: str, user: str) -> str: + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ] + if extra_try_kwargs: + try: + return completion_text( + model=model, + messages=messages, + temperature=temperature, + **extra_try_kwargs, + **kwargs, + ) + except Exception: + logger.debug( + "LiteLLM extra kwargs failed for model=%s; retrying without them", + model, + exc_info=True, + ) + return completion_text( + model=model, + messages=messages, + temperature=temperature, + **kwargs, + ) + + return _call diff --git a/application/prompt_client/prompt_client.py b/application/prompt_client/prompt_client.py index 81a7a284f..81dddcd95 100644 --- a/application/prompt_client/prompt_client.py +++ b/application/prompt_client/prompt_client.py @@ -9,7 +9,7 @@ from io import BytesIO from urllib.parse import urlparse -from application.prompt_client import embed_alignment +from application.prompt_client import embed_alignment, litellm_router from scipy import sparse from sklearn.metrics.pairwise import cosine_similarity @@ -26,7 +26,6 @@ import json import re import requests -import time SIMILARITY_THRESHOLD = float(os.environ.get("CHATBOT_SIMILARITY_THRESHOLD", "0.7")) @@ -48,54 +47,11 @@ def _safe_truncate_for_log(text: str, limit: int = 600) -> str: def _extract_content_text(response: Any) -> str: - choices = getattr(response, "choices", None) - if not choices and isinstance(response, dict): - choices = response.get("choices") - if not choices: - raise ValueError("LLM response did not contain choices") - msg = choices[0].message - content = getattr(msg, "content", None) - if content is None and isinstance(msg, dict): - content = msg.get("content") - if content is None: - raise ValueError("LLM response did not contain message content") - if isinstance(content, list): - return "".join( - c.get("text", "") if isinstance(c, dict) else str(c) for c in content - ).strip() - return str(content).strip() + return litellm_router.extract_content_text(response) def _extract_embeddings(response: Any) -> List[List[float]]: - data = getattr(response, "data", None) - if data is None and isinstance(response, dict): - data = response.get("data") - if not isinstance(data, list): - raise ValueError("Embedding response missing data list") - vectors: List[List[float]] = [] - for item in data: - emb = getattr(item, "embedding", None) - if emb is None and isinstance(item, dict): - emb = item.get("embedding") - if not isinstance(emb, list): - raise ValueError("Embedding item missing vector") - vectors.append([float(x) for x in emb]) - return vectors - - -def _is_llm_rate_limit_error(err: Exception) -> bool: - msg = str(err).lower() - if "rate limit" in msg or "too many requests" in msg: - return True - if "resource exhausted" in msg or "quota" in msg or "exceeded quota" in msg: - return True - status = ( - getattr(err, "status_code", None) - or getattr(err, "status", None) - or getattr(err, "http_status", None) - or getattr(err, "code", None) - ) - return status == 429 + return litellm_router.extract_embeddings(response) def _render_chat_prompt(*, question: str, retrieved_knowledge: Optional[str]) -> str: @@ -715,13 +671,7 @@ class PromptHandler: embeddings_instance = None # instance of our in_memory_embeddings singletton def __init__(self, database: db.Node_collection, load_all_embeddings=False) -> None: - try: - import litellm # type: ignore - except ImportError as e: - raise RuntimeError( - "litellm package is required for PromptHandler LLM calls" - ) from e - self._litellm = litellm + self._litellm = litellm_router.get_litellm() self.chat_model = os.environ.get( "CRE_LLM_CHAT_MODEL", "gemini/gemini-2.5-flash" ) @@ -729,9 +679,8 @@ def __init__(self, database: db.Node_collection, load_all_embeddings=False) -> N "CRE_EMBED_MODEL", "gemini/gemini-embedding-001" ) self.align_model = os.environ.get("CRE_EMBED_ALIGN_MODEL", self.chat_model) - self._llm_max_retries = int(os.environ.get("CRE_LLM_MAX_RETRIES", "2")) - self._llm_retry_sleep_seconds = int( - os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS", "15") + self._llm_max_retries, self._llm_retry_sleep_seconds = ( + litellm_router.retry_policy() ) expected_dim_raw = os.environ.get("CRE_EMBED_EXPECTED_DIM", "").strip() self._expected_embed_dim = int(expected_dim_raw) if expected_dim_raw else None @@ -773,21 +722,12 @@ def __init__(self, database: db.Node_collection, load_all_embeddings=False) -> N ) def _with_llm_rate_limit_retry(self, fn: Any, *, context: str) -> Any: - for attempt in range(self._llm_max_retries + 1): - try: - return fn() - except Exception as e: - if not _is_llm_rate_limit_error(e) or attempt >= self._llm_max_retries: - raise - logger.info( - "rate/quota limited during %s; sleeping %ss (attempt %s/%s)", - context, - self._llm_retry_sleep_seconds, - attempt + 1, - self._llm_max_retries + 1, - ) - time.sleep(self._llm_retry_sleep_seconds) - raise RuntimeError("unreachable: retry loop exited unexpectedly") + return litellm_router.with_rate_limit_retry( + fn, + context=context, + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, + ) def get_model_name(self) -> str: return self.chat_model @@ -811,11 +751,15 @@ def _litellm_get_text_embeddings( else self._truncate_one(text) ) - def _call() -> Any: - return self._litellm.embedding(model=self.embed_model, input=payload) - vectors = _extract_embeddings( - self._with_llm_rate_limit_retry(_call, context="LiteLLM embeddings") + litellm_router.embedding( + model=self.embed_model, + input=payload, + client=self._litellm, + context="LiteLLM embeddings", + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, + ) ) if self._expected_embed_dim is not None: for v in vectors: @@ -843,10 +787,14 @@ def create_chat_completion(self, prompt: str, closest_object_str: str) -> str: {"role": "user", "content": rag_instruction}, ] - def _call() -> Any: - return self._litellm.completion(model=self.chat_model, messages=messages) - - resp = self._with_llm_rate_limit_retry(_call, context="LiteLLM chat completion") + resp = litellm_router.completion( + model=self.chat_model, + messages=messages, + client=self._litellm, + context="LiteLLM chat completion", + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, + ) return _extract_content_text(resp) def align_embedding_span_json( @@ -865,35 +813,32 @@ def align_embedding_span_json( }, } - def _call_with_json_schema() -> Any: - return self._litellm.completion( + try: + resp = litellm_router.completion( model=self.align_model, messages=messages, + client=self._litellm, + context="LiteLLM align_embedding_span_json", + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, response_format=strict_format, temperature=0.2, ) - - def _call_json_object_fallback() -> Any: - return self._litellm.completion( - model=self.align_model, - messages=messages, - response_format={"type": "json_object"}, - temperature=0.2, - ) - - try: - resp = self._with_llm_rate_limit_retry( - _call_with_json_schema, context="LiteLLM align_embedding_span_json" - ) except Exception as e: logger.warning( "strict json_schema mode failed for model=%s: %s; retrying json_object", self.align_model, e, ) - resp = self._with_llm_rate_limit_retry( - _call_json_object_fallback, + resp = litellm_router.completion( + model=self.align_model, + messages=messages, + client=self._litellm, context="LiteLLM align_embedding_span_json fallback", + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, + response_format={"type": "json_object"}, + temperature=0.2, ) text = _extract_content_text(resp) @@ -929,10 +874,14 @@ def query_llm(self, raw_question: str) -> str: {"role": "user", "content": direct_instruction}, ] - def _call() -> Any: - return self._litellm.completion(model=self.chat_model, messages=messages) - - resp = self._with_llm_rate_limit_retry(_call, context="LiteLLM query_llm") + resp = litellm_router.completion( + model=self.chat_model, + messages=messages, + client=self._litellm, + context="LiteLLM query_llm", + max_retries=self._llm_max_retries, + retry_sleep_seconds=self._llm_retry_sleep_seconds, + ) return _extract_content_text(resp) def generate_embeddings_for(self, item_name: str): diff --git a/application/tests/litellm_router_test.py b/application/tests/litellm_router_test.py index a3f4016e0..dcced3048 100644 --- a/application/tests/litellm_router_test.py +++ b/application/tests/litellm_router_test.py @@ -4,9 +4,10 @@ import os import unittest +from types import SimpleNamespace from unittest.mock import Mock, patch -from application.prompt_client import llm_error_utils, prompt_client +from application.prompt_client import litellm_router, llm_error_utils, prompt_client class _FakeEmbeddingsSingleton: @@ -15,11 +16,17 @@ def with_ai_client(self, ai_client): return self +def _chat_resp(text: str) -> SimpleNamespace: + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=text))] + ) + + class TestLiteLLMRouter(unittest.TestCase): def tearDown(self) -> None: os.environ.pop("CRE_VALIDATE_EMBED_DIM_ON_INIT", None) - def test_prompt_handler_uses_litellm_directly(self) -> None: + def test_prompt_handler_uses_shared_router(self) -> None: os.environ["CRE_VALIDATE_EMBED_DIM_ON_INIT"] = "0" fake_embed_singleton = _FakeEmbeddingsSingleton() fake_db = Mock() @@ -31,6 +38,7 @@ def test_prompt_handler_uses_litellm_directly(self) -> None: with patch("application.prompt_client.prompt_client.logger.info"): ph = prompt_client.PromptHandler(fake_db) self.assertIs(ph.ai_client, ph) + self.assertIsNotNone(ph._litellm) def test_rate_limit_error_helper_detects_429(self) -> None: err = Exception("HTTP 429 too many requests") @@ -40,6 +48,41 @@ def test_rate_limit_error_helper_detects_quota_message(self) -> None: err = Exception("resource exhausted due to quota") self.assertTrue(llm_error_utils.is_rate_limit_error(err)) + def test_completion_retries_rate_limit_then_succeeds(self) -> None: + client = Mock( + completion=Mock( + side_effect=[ + Exception("HTTP 429 too many requests"), + _chat_resp("ok"), + ] + ) + ) + with patch("application.prompt_client.litellm_router.time.sleep"): + resp = litellm_router.completion( + model="gemini/gemini-2.5-flash", + messages=[{"role": "user", "content": "hi"}], + client=client, + max_retries=2, + retry_sleep_seconds=0, + ) + self.assertEqual(litellm_router.extract_content_text(resp), "ok") + self.assertEqual(client.completion.call_count, 2) + + def test_system_user_fn_returns_assistant_text(self) -> None: + client = Mock(completion=Mock(return_value=_chat_resp("cre-id"))) + with patch( + "application.prompt_client.litellm_router.get_litellm", + return_value=client, + ): + fn = litellm_router.system_user_fn("gemini/gemini-2.5-flash") + self.assertEqual(fn("sys", "user"), "cre-id") + + def test_extract_content_text_lenient_on_empty(self) -> None: + self.assertEqual( + litellm_router.extract_content_text({}, strict=False), + "", + ) + if __name__ == "__main__": unittest.main() diff --git a/application/tests/noise_filter/llm_classifier_test.py b/application/tests/noise_filter/llm_classifier_test.py index bd33808f8..cb4e4b2ac 100644 --- a/application/tests/noise_filter/llm_classifier_test.py +++ b/application/tests/noise_filter/llm_classifier_test.py @@ -3,6 +3,7 @@ Uses unittest (project-wide discovery pattern). The LLM is fully mocked -- no network calls. We swap LLMClassifier._litellm with a Mock and assert on the messages it receives and how responses are parsed back into verdicts. +Completions go through ``litellm_router.completion`` with that client. """ from __future__ import annotations @@ -277,7 +278,7 @@ def test_rate_limit_retried_then_succeeds(self) -> None: clf._litellm = Mock( completion=Mock(side_effect=[Exception("HTTP 429 too many requests"), ok]) ) - with patch("application.utils.noise_filter.llm_classifier.time.sleep"): + with patch("application.prompt_client.litellm_router.time.sleep"): out = clf.classify_batch([_record()]) self.assertEqual(out[0].label, "NOISE") self.assertEqual(clf._litellm.completion.call_count, 2) @@ -288,7 +289,7 @@ def test_rate_limit_exhausted_marks_batch_failed(self) -> None: clf._litellm = Mock( completion=Mock(side_effect=Exception("HTTP 429 too many requests")) ) - with patch("application.utils.noise_filter.llm_classifier.time.sleep"): + with patch("application.prompt_client.litellm_router.time.sleep"): out = clf.classify_batch([_record(), _record()]) self.assertEqual([v.label for v in out], ["UNCERTAIN", "UNCERTAIN"]) self.assertEqual( diff --git a/application/tests/test_smart_embeddings_e2e_llm.py b/application/tests/test_smart_embeddings_e2e_llm.py index 05d5d318a..972fce0b3 100644 --- a/application/tests/test_smart_embeddings_e2e_llm.py +++ b/application/tests/test_smart_embeddings_e2e_llm.py @@ -19,11 +19,10 @@ import pytest import requests -import litellm from pydantic import ValidationError from application.defs import cre_defs -from application.prompt_client import embed_alignment +from application.prompt_client import embed_alignment, litellm_router pytestmark = pytest.mark.llm_e2e @@ -58,20 +57,20 @@ def align_embedding_span_json( }, } try: - resp = litellm.completion( + resp = litellm_router.completion( model=self.model, messages=messages, response_format=strict_format, temperature=0.2, ) except Exception: - resp = litellm.completion( + resp = litellm_router.completion( model=self.model, messages=messages, response_format={"type": "json_object"}, temperature=0.2, ) - text = (resp.choices[0].message.content or "").strip() + text = litellm_router.extract_content_text(resp) try: payload = embed_alignment.AlignmentPayload.model_validate_json(text) return payload.model_dump() diff --git a/application/utils/noise_filter/llm_classifier.py b/application/utils/noise_filter/llm_classifier.py index cbf114dfe..21796b19e 100644 --- a/application/utils/noise_filter/llm_classifier.py +++ b/application/utils/noise_filter/llm_classifier.py @@ -1,15 +1,12 @@ """Module B Stage 2: LLM relevance classifier (recall-first). -Self-contained by design (decided 2026-06-18, Option B): this module talks to -LiteLLM directly rather than wrapping PromptHandler, whose constructor is -DB-coupled and whose retry/litellm members are private. We reuse the one -shared, public piece -- llm_error_utils.is_rate_limit_error -- inside a small -retry loop over the upstream CRE_LLM_MAX_RETRIES / CRE_LLM_RETRY_SLEEP_SECONDS -vars, so noise filtering and the chatbot share one retry policy. - -The classifier uses a dedicated cheap model (config.llm_model, default -gemini/gemini-2.5-flash-lite) and never falls back to CRE_LLM_CHAT_MODEL: -Module B is the cheap gate and must stay decoupled from the chatbot's model. +Self-contained by design (decided 2026-06-18, Option B): this module does not +construct PromptHandler (DB-coupled). Completions go through +``application.prompt_client.litellm_router`` so retry policy and parsing match +chat/embeddings. The classifier uses a dedicated cheap model +(config.llm_model, default gemini/gemini-2.5-flash-lite) and never falls back +to CRE_LLM_CHAT_MODEL: Module B is the cheap gate and must stay decoupled +from the chatbot's model. """ from __future__ import annotations @@ -19,12 +16,16 @@ logger = get_logger(__name__) import json -import os -import time from typing import Any, Iterator from pydantic import ValidationError +from application.prompt_client.litellm_router import ( + completion as litellm_completion, + extract_content_text, + get_litellm, + retry_policy, +) from application.prompt_client.llm_error_utils import is_rate_limit_error from application.utils.noise_filter.config_loader import NoiseFilterConfig from application.utils.noise_filter.prompts import ( @@ -109,14 +110,7 @@ def _batches(seq: list[Any], size: int) -> Iterator[list[Any]]: def _extract_text(resp: Any) -> str: """Pull message content from a LiteLLM response (object or dict shaped).""" - try: - choices = resp.choices if hasattr(resp, "choices") else resp["choices"] - first = choices[0] - msg = first.message if hasattr(first, "message") else first["message"] - content = msg.content if hasattr(msg, "content") else msg["content"] - return content or "" - except (AttributeError, KeyError, IndexError, TypeError): - return "" + return extract_content_text(resp, strict=False) def _is_schema_unsupported_error(err: Exception) -> bool: @@ -160,17 +154,8 @@ class LLMClassifier: def __init__(self, config: NoiseFilterConfig) -> None: self.config = config - try: - import litellm # type: ignore - except ImportError as e: - raise RuntimeError( - "litellm is required for the Module B Stage 2 classifier" - ) from e - self._litellm = litellm - self._max_retries = int(os.environ.get("CRE_LLM_MAX_RETRIES", "2")) - self._retry_sleep_seconds = int( - os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS", "15") - ) + self._litellm = get_litellm() + self._max_retries, self._retry_sleep_seconds = retry_policy() def classify_batch(self, records: list[ChangeRecord]) -> list[ClassifyResult]: """Classify records, one verdict per record in input order. @@ -239,22 +224,14 @@ def _call_llm(self, messages: list[dict]) -> str: return _extract_text(resp) def _completion_with_retry(self, **kwargs: Any) -> Any: - for attempt in range(self._max_retries + 1): - try: - return self._litellm.completion( - model=self.config.llm_model, temperature=0.0, **kwargs - ) - except Exception as e: - if not is_rate_limit_error(e) or attempt >= self._max_retries: - raise - logger.info( - "rate/quota limited; sleeping %ss (attempt %s/%s)", - self._retry_sleep_seconds, - attempt + 1, - self._max_retries + 1, - ) - time.sleep(self._retry_sleep_seconds) - raise RuntimeError("unreachable: retry loop exited unexpectedly") + return litellm_completion( + model=self.config.llm_model, + client=self._litellm, + max_retries=self._max_retries, + retry_sleep_seconds=self._retry_sleep_seconds, + temperature=0.0, + **kwargs, + ) def _parse(self, text: str, n: int) -> list[ClassifyResult]: verdicts = [_uncertain(MALFORMED_OUTPUT) for _ in range(n)]