141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
"""LLM client — single `chat()` entry point that dispatches to the active provider.
|
|
|
|
Provider selection is via `settings.llm_provider`. Supported:
|
|
- "groq" — hits Groq via the openai SDK (groq_base_url + groq_api_key).
|
|
- "anthropic" — native Anthropic SDK.
|
|
- "openai" — openai SDK with the standard base_url (also handles vLLM /
|
|
any OpenAI-compatible endpoint).
|
|
|
|
Every call opens a Langfuse `generation` span automatically, tagged with the
|
|
model name and populated with token usage from the provider's response.
|
|
Spans nest under whatever observation the caller has open, so the trace
|
|
hierarchy shows the wrapping step (e.g. `pick_for_question`) with the
|
|
`llm.chat` generation as a child carrying the usage.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
from typing import Callable
|
|
|
|
from anthropic import Anthropic
|
|
from openai import OpenAI
|
|
|
|
from api import langfuse_client as lf
|
|
from api.config import get_settings
|
|
|
|
|
|
# ── Provider clients (cached) ──
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _anthropic() -> Anthropic:
|
|
return Anthropic(api_key=get_settings().anthropic_api_key)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _groq() -> OpenAI:
|
|
s = get_settings()
|
|
return OpenAI(api_key=s.groq_api_key, base_url=s.groq_base_url)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _openai() -> OpenAI:
|
|
s = get_settings()
|
|
return OpenAI(api_key=s.openai_api_key, base_url=s.openai_base_url)
|
|
|
|
|
|
# ── Per-provider chat impls ──
|
|
|
|
@dataclass
|
|
class _Call:
|
|
text: str
|
|
usage: dict[str, int] # {"input": N, "output": M, "total": N+M}
|
|
|
|
|
|
def _chat_anthropic(system: str, user: str, max_tokens: int) -> _Call:
|
|
s = get_settings()
|
|
msg = _anthropic().messages.create(
|
|
model=s.anthropic_model,
|
|
max_tokens=max_tokens,
|
|
system=system,
|
|
messages=[{"role": "user", "content": user}],
|
|
)
|
|
text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
|
|
return _Call(
|
|
text=text,
|
|
usage={
|
|
"input": int(msg.usage.input_tokens),
|
|
"output": int(msg.usage.output_tokens),
|
|
"total": int(msg.usage.input_tokens + msg.usage.output_tokens),
|
|
},
|
|
)
|
|
|
|
|
|
def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> _Call:
|
|
resp = client.chat.completions.create(
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
messages=[
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
],
|
|
)
|
|
text = (resp.choices[0].message.content or "").strip()
|
|
u = resp.usage
|
|
usage = {
|
|
"input": int(u.prompt_tokens) if u else 0,
|
|
"output": int(u.completion_tokens) if u else 0,
|
|
"total": int(u.total_tokens) if u else 0,
|
|
}
|
|
return _Call(text=text, usage=usage)
|
|
|
|
|
|
def _chat_groq(system: str, user: str, max_tokens: int) -> _Call:
|
|
return _chat_openai_compat(_groq(), get_settings().groq_model, system, user, max_tokens)
|
|
|
|
|
|
def _chat_openai(system: str, user: str, max_tokens: int) -> _Call:
|
|
return _chat_openai_compat(_openai(), get_settings().openai_model, system, user, max_tokens)
|
|
|
|
|
|
_PROVIDERS: dict[str, Callable[[str, str, int], _Call]] = {
|
|
"groq": _chat_groq,
|
|
"anthropic": _chat_anthropic,
|
|
"openai": _chat_openai,
|
|
}
|
|
|
|
|
|
def _active_model() -> str:
|
|
s = get_settings()
|
|
return {
|
|
"groq": s.groq_model,
|
|
"anthropic": s.anthropic_model,
|
|
"openai": s.openai_model,
|
|
}.get(s.llm_provider.lower(), s.llm_provider)
|
|
|
|
|
|
# ── Public surface ──
|
|
|
|
def chat(*, system: str, user: str, max_tokens: int = 1024, span_name: str = "llm.chat") -> str:
|
|
"""Run a single chat completion. Opens a Langfuse generation span around
|
|
the call with the model name + token usage. `span_name` is the label
|
|
shown in Langfuse — pass something descriptive (e.g. "direct_answer.pick")
|
|
so traces are scannable."""
|
|
provider = get_settings().llm_provider.lower()
|
|
impl = _PROVIDERS.get(provider)
|
|
if impl is None:
|
|
raise ValueError(
|
|
f"unknown llm_provider {provider!r}; supported: {sorted(_PROVIDERS)}"
|
|
)
|
|
|
|
with lf.span(
|
|
span_name,
|
|
as_type="generation",
|
|
model=_active_model(),
|
|
input={"system": system, "user": user},
|
|
) as gen:
|
|
call = impl(system, user, max_tokens)
|
|
gen.update(output=call.text, usage_details=call.usage)
|
|
return call.text
|