MemoryAgents/backend/app/llm.py

from __future__ import annotations

import json
from typing import Any

import httpx

from .config import settings


class LLMUnavailable(RuntimeError):
    pass


class OpenAICompatibleClient:
    @property
    def configured(self) -> bool:
        return bool(settings.llm_api_key and settings.llm_base_url and settings.llm_model)

    async def chat(self, system: str, user: str, temperature: float = 0.2) -> str:
        if not self.configured:
            raise LLMUnavailable("未配置 LLM_API_KEY / LLM_BASE_URL / LLM_MODEL")
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                f"{settings.llm_base_url}/chat/completions",
                headers={"Authorization": f"Bearer {settings.llm_api_key}"},
                json={
                    "model": settings.llm_model,
                    "temperature": temperature,
                    "messages": [
                        {"role": "system", "content": system},
                        {"role": "user", "content": user},
                    ],
                },
            )
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"]

    async def json(self, system: str, user: str) -> dict[str, Any]:
        raw = await self.chat(system, user, temperature=0.1)
        cleaned = raw.strip()
        if "```" in cleaned:
            cleaned = cleaned.replace("```json", "").replace("```", "").strip()
        start, end = cleaned.find("{"), cleaned.rfind("}")
        if start >= 0 and end > start:
            cleaned = cleaned[start : end + 1]
        return json.loads(cleaned)


llm = OpenAICompatibleClient()