88 lines
2.4 KiB
Python
88 lines
2.4 KiB
Python
"""OpenAI-compatible LLM client for multi-server support."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from openai import AsyncOpenAI
|
|
|
|
from backend.models import LLMConfig
|
|
|
|
|
|
# ── Server Config Loading ─────────────────────────────
|
|
|
|
def _load_llm_configs() -> list[LLMConfig]:
|
|
config_path = os.getenv(
|
|
"LLMS_CONFIG_PATH", str(Path(__file__).parent.parent / "config" / "llms.json")
|
|
)
|
|
import json
|
|
with open(config_path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
return [LLMConfig(**item) for item in data]
|
|
|
|
|
|
_llm_configs: list[LLMConfig] = []
|
|
_clients: dict[str, AsyncOpenAI] = {} # alias -> client instance
|
|
|
|
|
|
def get_llm_configs() -> list[LLMConfig]:
|
|
"""Return all LLM configurations."""
|
|
return _llm_configs
|
|
|
|
|
|
def get_client_for_alias(alias: str) -> tuple[AsyncOpenAI, LLMConfig]:
|
|
"""Get the AsyncOpenAI client and config for a given alias.
|
|
|
|
Creates a lazily cached client per alias.
|
|
"""
|
|
if alias not in _clients:
|
|
cfg = None
|
|
for c in _llm_configs:
|
|
if c.alias == alias:
|
|
cfg = c
|
|
break
|
|
if cfg is None:
|
|
raise ValueError(f"Unknown LLM alias: {alias}")
|
|
|
|
client = AsyncOpenAI(
|
|
api_key=cfg.api_key,
|
|
base_url=cfg.base_url,
|
|
timeout=120.0,
|
|
)
|
|
_clients[alias] = client
|
|
|
|
return _clients[alias], next(c for c in _llm_configs if c.alias == alias)
|
|
|
|
|
|
async def chat_complete(
|
|
alias: str,
|
|
system_prompt: str,
|
|
user_content: str,
|
|
temperature: float = 0.3,
|
|
) -> str:
|
|
"""Send a single-turn chat completion and return the assistant message."""
|
|
client, cfg = get_client_for_alias(alias)
|
|
|
|
response = await client.chat.completions.create(
|
|
model=cfg.model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_content},
|
|
],
|
|
temperature=temperature,
|
|
)
|
|
|
|
content = response.choices[0].message.content or ""
|
|
return content.strip()
|
|
|
|
|
|
def get_context_size(alias: str) -> int | None:
|
|
"""Get the context window size for a given alias."""
|
|
for c in _llm_configs:
|
|
if c.alias == alias:
|
|
return c.context_size
|
|
return None
|
|
|
|
|
|
# ── Initialize on import ──────────────────────────────
|
|
|
|
_llm_configs = _load_llm_configs()
|