memory-system/llm_client.py

"""LLM Client with multi-provider failover and JSON mode support."""
import json
import logging
import time
import requests
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

COOLDOWN_SECONDS = 300  # 5 min before retrying primary after failure


class LLMClient:
    """LLM client supporting provider list with automatic failover."""

    def __init__(self, providers_config: List[Dict[str, Any]] = None):
        self.providers = []
        if providers_config:
            self.providers = sorted(providers_config, key=lambda p: p.get("priority", 999))

        # Normalize API URLs
        for p in self.providers:
            url = p.get("api_url", "").rstrip("/")
            if url and not url.endswith("/chat/completions"):
                url = url + "/chat/completions"
            p["api_url"] = url

        self._fallback_mode = False
        self._last_primary_failure = 0.0

    @property
    def available(self) -> bool:
        return len(self.providers) > 0

    def chat(self, system_prompt: str, user_prompt: str,
             temperature: float = 0.3, max_tokens: int = 800,
             json_mode: bool = False, thinking: Optional[dict] = None,
             prefill: str = None) -> Optional[str]:
        if not self.available:
            logger.warning("No LLM providers configured")
            return None

        for i, provider in enumerate(self.providers):
            is_primary = (i == 0)

            if is_primary and self._fallback_mode:
                if time.time() - self._last_primary_failure < COOLDOWN_SECONDS:
                    continue
                logger.info("Primary cooldown expired, retrying...")

            try:
                result = self._call(provider, system_prompt, user_prompt,
                                    temperature, max_tokens, json_mode, thinking, prefill)
                if is_primary and self._fallback_mode:
                    logger.info("Primary provider recovered, switched back")
                    self._fallback_mode = False
                return result
            except requests.exceptions.RequestException as e:
                logger.warning("Provider %s (%s) failed: %s",
                               provider.get('model', '?'),
                               provider.get('api_url', '?'), e)
                if is_primary:
                    self._fallback_mode = True
                    self._last_primary_failure = time.time()
                    logger.info("Switched to fallback provider (cooldown: %ds)", COOLDOWN_SECONDS)
                continue
            except ValueError as e:
                # json_mode returned non-JSON content (reasoning text from proxy)
                logger.warning("Provider %s returned non-JSON for json_mode, trying next: %s",
                               provider.get('model', '?'), str(e)[:100])
                if is_primary:
                    self._fallback_mode = True
                    self._last_primary_failure = time.time()
                continue

        logger.error("All LLM providers failed")
        return None

    def _call(self, provider: dict, system_prompt: str, user_prompt: str,
              temperature: float, max_tokens: int, json_mode: bool,
              thinking: Optional[dict] = None, prefill: str = None) -> str:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {provider['api_key']}",
        }
        payload = {
            "model": provider["model"],
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if json_mode:
            payload["response_format"] = {"type": "json_object"}
        if thinking is not None:
            payload["thinking"] = thinking
        if prefill:
            payload["messages"].append({"role": "assistant", "content": prefill})

        resp = requests.post(
            provider["api_url"], headers=headers, json=payload, timeout=60
        )
        resp.raise_for_status()
        data = resp.json()
        msg = data["choices"][0]["message"]

        reasoning = msg.get("reasoning_content", "") or ""
        content = msg.get("content") or ""

        # If content is empty but reasoning has content, use reasoning as fallback
        if not content and reasoning:
            content = reasoning

        if reasoning and content:
            logger.debug("LLM: reasoning=%d chars, content=%d chars", len(reasoning), len(content))

        # Detect when json_mode was requested but response is not JSON.
        # This catches two cases:
        # 1. Proxy merged reasoning into content field (reasoning is empty, content starts with analysis text)
        # 2. Model spent all tokens on reasoning (content was empty, reasoning was used as fallback)
        if json_mode and content:
            stripped = content.lstrip()
            if stripped and stripped[0] not in ('[', '{', '"'):
                raise ValueError(
                    "json_mode response is not JSON (starts with: %.50r)" % stripped
                )

        if json_mode and not content:
            logger.warning("JSON mode returned empty content for %s", provider.get("model", "?"))
            raise ValueError("JSON mode returned empty content")

        return content

    def get_status(self) -> dict:
        status = []
        for p in self.providers:
            status.append({
                "model": p["model"],
                "priority": p.get("priority"),
                "api_url": p["api_url"],
            })
        return {
            "fallback_mode": self._fallback_mode,
            "last_primary_failure_ago":
                int(time.time() - self._last_primary_failure) if self._fallback_mode else 0,
            "providers": status,
        }

    @property
    def current_provider(self) -> str:
        if not self.providers:
            return "none"
        idx = 1 if (self._fallback_mode and len(self.providers) > 1) else 0
        return self.providers[idx].get("model", "unknown")