ss-tools/backend/src/agent/_llm_health.py

# backend/src/agent/_llm_health.py
# #region AgentChat.LlmHealth [C:3] [TYPE Module] [SEMANTICS agent-chat,llm,health,status]
# @ingroup AgentChat
# @BRIEF LLM provider health check with in-memory cache (30s TTL).
#   Extracted from agent/app.py to avoid importing gradio in backend container.
# @LAYER Infrastructure
# @RELATION CALLED_BY -> [api/routes/agent_status.py]
# @RELATION CALLED_BY -> [agent/app.py]
# @RATIONALE agent/app.py imports gradio at module level (line 28). The backend
#   container does not have gradio installed (only the agent container does).
#   The /api/agent/llm-status endpoint needs _check_llm_provider_health but must
#   not trigger gradio import. This module isolates the health-check logic.
# @POST Returns status string: 'ok' | 'unavailable' | 'timeout' | 'auth_error'.
# @SIDE_EFFECT Makes a probe request to the LLM provider; caches result in module memory.

import time
from typing import Any

import httpx

from src.core.logger import logger

# ── LLM provider health cache ─────────────────────────────────────────
_llm_status: dict[str, Any] = {
    "status": "ok",
    "last_error": "",
    "last_check_ts": 0.0,
}
_LLM_CHECK_CACHE_TTL = 30  # seconds between health checks
_LLM_LAST_ERROR_KEY = "last_llm_error"
_LLM_LAST_ERROR_TS_KEY = "last_llm_error_ts"


# #region AgentChat.LlmHealth.Check [C:2] [TYPE Function] [SEMANTICS agent-chat,llm,health,check]
# @ingroup AgentChat
# @BRIEF Check LLM provider connectivity with in-memory cache (30s TTL).
# @POST Returns status string: 'ok' | 'unavailable' | 'timeout' | 'auth_error'.
# @SIDE_EFFECT Makes a probe request to the LLM provider; caches result in module memory.
# @RATIONALE Prevents sending every user request into a dead LLM backend.
# @REJECTED Module-level imports of langchain_core/openai were rejected — these
#   packages are only installed in the agent container, not the backend container.
#   The /api/agent/llm-status endpoint runs in the backend container and must not
#   fail with ModuleNotFoundError. All langchain/openai imports are lazy (inside
#   the function body) with ImportError handled gracefully.
async def _check_llm_provider_health() -> str:
    """Check LLM provider connectivity. Cached for _LLM_CHECK_CACHE_TTL seconds."""
    now = time.time()
    if now - _llm_status["last_check_ts"] < _LLM_CHECK_CACHE_TTL:
        return _llm_status["status"]

    # Lazy imports: langchain_core, langchain_openai, and openai are only
    # installed in the agent container, not the backend container.
    try:
        from langchain_core.messages import HumanMessage
        from langchain_openai import ChatOpenAI
        from openai import (
            APIConnectionError,
            APITimeoutError,
            AuthenticationError,
            RateLimitError,
        )
        from src.agent._llm_params import chat_openai_kwargs
    except ImportError as exc:
        _llm_status["status"] = "unavailable"
        _llm_status["last_error"] = f"LLM dependencies not installed: {exc}"
        _llm_status["last_check_ts"] = time.time()
        return "unavailable"

    try:
        from src.agent.langgraph_setup import _fetch_llm_config as _get_config

        config = await _get_config()
        if not config or not config.get("configured"):
            return _llm_status["status"]

        llm = ChatOpenAI(
            **chat_openai_kwargs(
                model=config.get("default_model", "gpt-4o-mini"),
                base_url=config.get("base_url", "https://api.openai.com/v1"),
                api_key=config.get("api_key", ""),
                max_tokens=1,
            )
        )
        await llm.ainvoke([HumanMessage(content="ping")])
        _llm_status["status"] = "ok"
        _llm_status["last_error"] = ""
        _llm_status["last_check_ts"] = time.time()
        return "ok"
    except (APIConnectionError, httpx.ConnectError):
        _llm_status["status"] = "unavailable"
        _llm_status["last_error"] = "Connection refused"
        _llm_status["last_check_ts"] = time.time()
        return "unavailable"
    except (APITimeoutError, httpx.ReadTimeout):
        _llm_status["status"] = "timeout"
        _llm_status["last_error"] = "Request timed out"
        _llm_status["last_check_ts"] = time.time()
        return "timeout"
    except AuthenticationError:
        _llm_status["status"] = "auth_error"
        _llm_status["last_error"] = "Invalid API key"
        _llm_status["last_check_ts"] = time.time()
        return "auth_error"
    except Exception as exc:
        if "usage_metadata.total_tokens" in str(exc):
            _llm_status["status"] = "ok"
            _llm_status["last_error"] = ""
            _llm_status["last_check_ts"] = time.time()
            return "ok"
        logger.explore("LLM health check failed", error=str(exc), extra={"src": "AgentChat.LlmHealth.Check"})
        return _llm_status["status"]  # return cached status on unexpected errors


# #endregion AgentChat.LlmHealth.Check
# #endregion AgentChat.LlmHealth