Skip to content

Commit c6bf04b

Browse files
committed
Use the shared LiteLLM helper for Librarian and OIE eval.
Shortlist, edition remap, metadata enrich, Module B, and the B2 judge all go through litellm_router instead of importing litellm directly.
1 parent 3edb217 commit c6bf04b

10 files changed

Lines changed: 357 additions & 203 deletions

File tree

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
"""Shared LiteLLM wrapper for chat, embeddings, Module B, and GSoC/OIE.
2+
3+
PromptHandler stays DB-coupled (embedding contract + RAG). Everything else
4+
that only needs a completion or embedding should call this module instead of
5+
``import litellm`` so retry, rate-limit detection, and response parsing stay
6+
in one place.
7+
"""
8+
9+
from __future__ import annotations
10+
11+
import logging
12+
import os
13+
import time
14+
from typing import Any, Callable, List, Optional, Sequence
15+
16+
from application.prompt_client.llm_error_utils import is_rate_limit_error
17+
18+
logger = logging.getLogger(__name__)
19+
20+
LlmFn = Callable[[str, str], str]
21+
22+
23+
def retry_policy() -> tuple[int, int]:
24+
"""Return (max_retries, sleep_seconds) from CRE_LLM_* env vars."""
25+
return (
26+
int(os.environ.get("CRE_LLM_MAX_RETRIES", "2")),
27+
int(os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS", "15")),
28+
)
29+
30+
31+
def get_litellm() -> Any:
32+
"""Import LiteLLM or raise the same RuntimeError PromptHandler uses."""
33+
try:
34+
import litellm # type: ignore
35+
except ImportError as exc:
36+
raise RuntimeError(
37+
"litellm package is required for PromptHandler LLM calls"
38+
) from exc
39+
return litellm
40+
41+
42+
def with_rate_limit_retry(
43+
fn: Callable[[], Any],
44+
*,
45+
context: str,
46+
max_retries: Optional[int] = None,
47+
retry_sleep_seconds: Optional[int] = None,
48+
) -> Any:
49+
default_retries, default_sleep = retry_policy()
50+
retries = default_retries if max_retries is None else max_retries
51+
sleep_s = default_sleep if retry_sleep_seconds is None else retry_sleep_seconds
52+
for attempt in range(retries + 1):
53+
try:
54+
return fn()
55+
except Exception as err:
56+
if not is_rate_limit_error(err) or attempt >= retries:
57+
raise
58+
logger.info(
59+
"rate/quota limited during %s; sleeping %ss (attempt %s/%s)",
60+
context,
61+
sleep_s,
62+
attempt + 1,
63+
retries + 1,
64+
)
65+
time.sleep(sleep_s)
66+
raise RuntimeError("unreachable: retry loop exited unexpectedly")
67+
68+
69+
def extract_content_text(response: Any, *, strict: bool = True) -> str:
70+
"""Pull assistant text from a LiteLLM completion (object or dict)."""
71+
choices = getattr(response, "choices", None)
72+
if not choices and isinstance(response, dict):
73+
choices = response.get("choices")
74+
if not choices:
75+
if strict:
76+
raise ValueError("LLM response did not contain choices")
77+
return ""
78+
first = choices[0]
79+
msg = getattr(first, "message", None)
80+
if msg is None and isinstance(first, dict):
81+
msg = first.get("message")
82+
if msg is None:
83+
if strict:
84+
raise ValueError("LLM response did not contain message content")
85+
return ""
86+
content = getattr(msg, "content", None)
87+
if content is None and isinstance(msg, dict):
88+
content = msg.get("content")
89+
if content is None:
90+
if strict:
91+
raise ValueError("LLM response did not contain message content")
92+
return ""
93+
if isinstance(content, list):
94+
return "".join(
95+
part.get("text", "") if isinstance(part, dict) else str(part)
96+
for part in content
97+
).strip()
98+
return str(content).strip()
99+
100+
101+
def extract_embeddings(response: Any) -> List[List[float]]:
102+
data = getattr(response, "data", None)
103+
if data is None and isinstance(response, dict):
104+
data = response.get("data")
105+
if not isinstance(data, list):
106+
raise ValueError("Embedding response missing data list")
107+
vectors: List[List[float]] = []
108+
for item in data:
109+
emb = getattr(item, "embedding", None)
110+
if emb is None and isinstance(item, dict):
111+
emb = item.get("embedding")
112+
if not isinstance(emb, list):
113+
raise ValueError("Embedding item missing vector")
114+
vectors.append([float(x) for x in emb])
115+
return vectors
116+
117+
118+
def completion(
119+
*,
120+
model: str,
121+
messages: Sequence[dict[str, Any]],
122+
context: str = "LiteLLM completion",
123+
client: Any = None,
124+
max_retries: Optional[int] = None,
125+
retry_sleep_seconds: Optional[int] = None,
126+
**kwargs: Any,
127+
) -> Any:
128+
"""LiteLLM chat completion with shared rate-limit retry."""
129+
llm = client if client is not None else get_litellm()
130+
131+
def _call() -> Any:
132+
return llm.completion(model=model, messages=list(messages), **kwargs)
133+
134+
return with_rate_limit_retry(
135+
_call,
136+
context=context,
137+
max_retries=max_retries,
138+
retry_sleep_seconds=retry_sleep_seconds,
139+
)
140+
141+
142+
def embedding(
143+
*,
144+
model: str,
145+
input: str | List[str],
146+
context: str = "LiteLLM embeddings",
147+
client: Any = None,
148+
max_retries: Optional[int] = None,
149+
retry_sleep_seconds: Optional[int] = None,
150+
**kwargs: Any,
151+
) -> Any:
152+
llm = client if client is not None else get_litellm()
153+
154+
def _call() -> Any:
155+
return llm.embedding(model=model, input=input, **kwargs)
156+
157+
return with_rate_limit_retry(
158+
_call,
159+
context=context,
160+
max_retries=max_retries,
161+
retry_sleep_seconds=retry_sleep_seconds,
162+
)
163+
164+
165+
def completion_text(
166+
*,
167+
model: str,
168+
messages: Sequence[dict[str, Any]],
169+
strict: bool = True,
170+
**kwargs: Any,
171+
) -> str:
172+
return extract_content_text(
173+
completion(model=model, messages=messages, **kwargs),
174+
strict=strict,
175+
)
176+
177+
178+
def system_user_fn(
179+
model: str,
180+
*,
181+
temperature: float = 0.0,
182+
extra_try_kwargs: Optional[dict[str, Any]] = None,
183+
**kwargs: Any,
184+
) -> LlmFn:
185+
"""Build ``(system, user) -> text`` for Librarian / OIE call sites."""
186+
187+
def _call(system: str, user: str) -> str:
188+
messages = [
189+
{"role": "system", "content": system},
190+
{"role": "user", "content": user},
191+
]
192+
if extra_try_kwargs:
193+
try:
194+
return completion_text(
195+
model=model,
196+
messages=messages,
197+
temperature=temperature,
198+
**extra_try_kwargs,
199+
**kwargs,
200+
)
201+
except Exception:
202+
logger.debug(
203+
"LiteLLM extra kwargs failed for model=%s; retrying without them",
204+
model,
205+
exc_info=True,
206+
)
207+
return completion_text(
208+
model=model,
209+
messages=messages,
210+
temperature=temperature,
211+
**kwargs,
212+
)
213+
214+
return _call

0 commit comments

Comments
 (0)