92 lines
2.7 KiB
Python
92 lines
2.7 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.
|
|
|
|
The Anthropic and OpenAI/Groq clients are cached so HTTP connection pools
|
|
are reused across calls.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from functools import lru_cache
|
|
from typing import Callable
|
|
|
|
from anthropic import Anthropic
|
|
from openai import OpenAI
|
|
|
|
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 ──
|
|
|
|
def _chat_anthropic(system: str, user: str, max_tokens: int) -> str:
|
|
s = get_settings()
|
|
msg = _anthropic().messages.create(
|
|
model=s.anthropic_model,
|
|
max_tokens=max_tokens,
|
|
system=system,
|
|
messages=[{"role": "user", "content": user}],
|
|
)
|
|
return "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
|
|
|
|
|
|
def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> str:
|
|
resp = client.chat.completions.create(
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
messages=[
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
],
|
|
)
|
|
return (resp.choices[0].message.content or "").strip()
|
|
|
|
|
|
def _chat_groq(system: str, user: str, max_tokens: int) -> str:
|
|
return _chat_openai_compat(_groq(), get_settings().groq_model, system, user, max_tokens)
|
|
|
|
|
|
def _chat_openai(system: str, user: str, max_tokens: int) -> str:
|
|
return _chat_openai_compat(_openai(), get_settings().openai_model, system, user, max_tokens)
|
|
|
|
|
|
_PROVIDERS: dict[str, Callable[[str, str, int], str]] = {
|
|
"groq": _chat_groq,
|
|
"anthropic": _chat_anthropic,
|
|
"openai": _chat_openai,
|
|
}
|
|
|
|
|
|
# ── Public surface ──
|
|
|
|
def chat(*, system: str, user: str, max_tokens: int = 1024) -> str:
|
|
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)}"
|
|
)
|
|
return impl(system, user, max_tokens)
|