AIRegulation-DocAnalysis/backend/app/api/routes/status.py

"""Define API routes for status."""

import asyncio
import time
from typing import Any

from fastapi import APIRouter

from app.config.settings import settings
from app.domain.retrieval import RetrievedChunk
from app.services.llm.llm_factory import get_llm_client, get_llm_factory
from app.shared.bootstrap import (
    get_bm25_retriever,
    get_binary_store,
    get_conversation_store,
    get_document_query_service,
    get_embedding_provider,
    get_reranker,
    get_vector_index,
)
from app.shared.model_usage_tracker import get_model_usage_tracker

router = APIRouter(prefix="/status", tags=["系统状态"])

# ---------------------------------------------------------------------------
# Simple TTL cache for /stats (avoids O(N) doc scan on every request)
# ---------------------------------------------------------------------------
_stats_cache: dict[str, Any] = {}
_stats_cache_time: float = 0.0
_STATS_TTL_SECONDS: float = 10.0

# ---------------------------------------------------------------------------
# AI model roles surfaced on the Status page (Task: System Status AI models)
# ---------------------------------------------------------------------------
_MODEL_ROLES: dict[str, str] = {
    "main_llm": "主问答 LLM",
    "hyde_llm": "HyDE 查询增强",
    "embedding": "Embedding",
    "reranker": "Reranker",
}


@router.get("/stats")
async def get_stats():
    """Return document statistics (cached for 10 s)."""
    global _stats_cache, _stats_cache_time
    now = time.time()
    if _stats_cache and (now - _stats_cache_time) < _STATS_TTL_SECONDS:
        return _stats_cache

    documents = get_document_query_service().list_documents()
    indexed = sum(1 for d in documents if d.status.value == "indexed")
    failed = sum(1 for d in documents if d.status.value == "failed")
    _stats_cache = {
        "documents_total": len(documents),
        "documents_indexed": indexed,
        "documents_failed": failed,
        "chunks_total": sum(d.chunk_count for d in documents),
    }
    _stats_cache_time = now
    return _stats_cache


@router.get("/config")
async def get_config():
    """Return system configuration."""
    return {
        "embedding_model": settings.embedding_model,
        "embedding_dim": settings.embedding_dim,
        "embedding_base_url": settings.embedding_base_url,
        "milvus_collection": settings.milvus_collection,
        "parser_backend": settings.parser_backend,
        "chunk_backend": settings.chunk_backend,
        "artifact_prefix": settings.document_parse_artifact_prefix,
        "parser_failure_mode": settings.parser_failure_mode,
        "llm_provider": settings.llm_provider,
        "llm_model": settings.llm_model,
        "document_metadata_path": settings.document_metadata_path,
    }


@router.get("/milvus/health")
async def milvus_health():
    """Return Milvus health (kept for backwards compat)."""
    return get_vector_index().health()


@router.get("/health")
async def get_health():
    """Return aggregate health of all backend services."""
    # --- Milvus ---
    try:
        milvus_info = get_vector_index().health()
        milvus_status = "ok" if milvus_info.get("connected") else "error"
    except Exception as exc:  # noqa: BLE001
        milvus_info = {}
        milvus_status = "error"
        milvus_info["error"] = str(exc)

    # --- MinIO ---
    try:
        minio_connected = get_binary_store().client.connected
        minio_status = "ok" if minio_connected else "error"
    except Exception:  # noqa: BLE001
        minio_status = "error"
        minio_connected = False

    # --- BM25 ---
    bm25 = get_bm25_retriever()

    # --- Sessions ---
    try:
        session_count = len(get_conversation_store().list_sessions())
    except Exception:  # noqa: BLE001
        session_count = 0

    return {
        "milvus": {"status": milvus_status, **milvus_info},
        "minio": {"status": minio_status, "connected": minio_connected},
        "bm25": {"available": bm25 is not None},
        "reranker": {
            "enabled": settings.reranker_enabled,
            "model": settings.reranker_model if settings.reranker_enabled else None,
        },
        "sessions": {
            "active": session_count,
            "max": settings.session_max_sessions,
        },
    }


def _normalize_llm_provider(raw_provider: str) -> str:
    """Normalize a raw LLM_PROVIDER/HYDE_LLM_PROVIDER settings string to the
    canonical LLMProvider enum value, the SAME way LLMFactory.create() does.

    TrackedLLMClient.chat() (tracked_client.py) always records usage under
    `self._inner.config.provider.value` — the NORMALIZED enum value produced by
    LLMFactory._parse_provider() — never the raw string a caller passed to
    get_llm_client(). Reusing that same normalization here (instead of
    duplicating the alias table) guarantees the tracker key this route reads
    always agrees with the key TrackedLLMClient wrote, even when the raw
    settings value is a non-canonical alias (e.g. "deepseek-v3") or different
    casing. Falls back to the raw string, unchanged, if it does not match any
    known provider/alias, so this passive status endpoint still renders
    (as "never_called") instead of raising on a misconfigured provider string.
    """
    try:
        return get_llm_factory()._parse_provider(raw_provider).value
    except ValueError:
        return raw_provider


def _resolve_role_provider_model(role: str) -> tuple[str, str]:
    """Return the (provider, model) pair currently configured for one AI model role.

    For "hyde_llm" this mirrors the exact fallback logic already used in
    hyde_expander.py (settings.hyde_llm_provider or settings.llm_provider, same
    for model) so tracker lookups here always match what TrackedLLMClient
    recorded when HyDE actually ran.
    """
    if role == "main_llm":
        return _normalize_llm_provider(settings.llm_provider), settings.llm_model
    if role == "hyde_llm":
        return (
            _normalize_llm_provider(settings.hyde_llm_provider or settings.llm_provider),
            settings.hyde_llm_model or settings.llm_model,
        )
    if role == "embedding":
        return "embedding", settings.embedding_model
    if role == "reranker":
        return "reranker", settings.reranker_model
    raise ValueError(f"unknown model role: {role}")  # pragma: no cover - internal roles are fixed


def _build_model_status(role: str) -> dict[str, Any]:
    """Build one /status/models row for the given role from tracker data + live settings."""
    provider, model = _resolve_role_provider_model(role)
    entry = get_model_usage_tracker().get(provider, model)

    main_provider, main_model = _resolve_role_provider_model("main_llm")
    shares_usage_with = (
        "main_llm" if role != "main_llm" and (provider, model) == (main_provider, main_model) else None
    )

    enabled = True
    status = entry.status if entry else "never_called"
    if role == "reranker":
        enabled = settings.reranker_enabled
        if not enabled:
            # Config always wins: report "disabled" even if the reranker was
            # enabled and called successfully earlier in this process's life.
            status = "disabled"
    elif role == "hyde_llm":
        enabled = settings.hyde_enabled
        if not enabled:
            # Same "config always wins" override as the reranker branch above:
            # report "disabled" even if HyDE ran successfully before being
            # turned off in settings during this process's life.
            status = "disabled"

    return {
        "role": role,
        "role_label": _MODEL_ROLES[role],
        "provider": provider,
        "model": model,
        "enabled": enabled,
        "status": status,
        "total_tokens": entry.total_tokens if entry else 0,
        "call_count_ok": entry.call_count_ok if entry else 0,
        "call_count_error": entry.call_count_error if entry else 0,
        "last_called_at": entry.last_called_at.isoformat() if entry and entry.last_called_at else None,
        "last_latency_ms": entry.last_latency_ms if entry else None,
        "last_error": entry.last_error if entry else None,
        "shares_usage_with": shares_usage_with,
    }


@router.get("/models")
async def get_model_statuses():
    """Return connection status + cumulative token usage for all 4 tracked AI model roles.

    Passive: reads tracker state + settings only, makes no outbound network calls.
    """
    return {"models": [_build_model_status(role) for role in _MODEL_ROLES]}


async def _ping_main_or_hyde(role: str) -> None:
    """Send one minimal chat completion to the LLM configured for `role`.

    Skipped entirely for "hyde_llm" when settings.hyde_enabled is False,
    mirroring _ping_reranker()'s disabled-skip pattern: when HyDE is turned
    off (or reuses the main LLM, the default), issuing this ping would just be
    a redundant duplicate chat call against the same model for no benefit.
    "main_llm" is always pinged regardless of this check.
    """
    if role == "hyde_llm" and not settings.hyde_enabled:
        return
    provider, model = _resolve_role_provider_model(role)
    try:
        client = get_llm_client(provider=provider, model=model)
    except Exception as exc:  # noqa: BLE001 - record, then re-raise so gather() still isolates this ping
        # get_llm_client() can fail before any TrackedLLMClient exists to
        # record the outcome itself (e.g. missing API key, unsupported
        # provider string), so record the failure here directly, otherwise it
        # would be invisible on the /status/models page afterward.
        get_model_usage_tracker().record(provider=provider, model=model, success=False, error=str(exc))
        raise
    await asyncio.to_thread(client.chat, [{"role": "user", "content": "ping"}], max_tokens=1)


async def _ping_embedding() -> None:
    """Send one minimal embedding request."""
    await asyncio.to_thread(get_embedding_provider().embed_query, "ping")


async def _ping_reranker() -> None:
    """Send one minimal rerank request, only when the reranker is enabled."""
    reranker = get_reranker()
    if reranker is None:
        return
    # Minimal single-chunk probe — real content doesn't matter, only round-trip success.
    placeholder = RetrievedChunk(chunk_id="ping", doc_id="ping", doc_title="ping", text="ping", score=0.0)
    await asyncio.to_thread(reranker.rerank, "ping", [placeholder], 1)


@router.post("/models/ping")
async def ping_model_connections():
    """Actively test each configured model with a minimal request, then return fresh statuses.

    Each ping is isolated with return_exceptions=True so one model timing out
    or erroring does not prevent the other three from completing and being
    reported. Failures are still visible afterwards via _build_model_status()
    because the underlying clients record their own outcome into the tracker.
    """
    tasks = [
        _ping_main_or_hyde("main_llm"),
        _ping_main_or_hyde("hyde_llm"),
        _ping_embedding(),
        _ping_reranker(),
    ]
    await asyncio.gather(*tasks, return_exceptions=True)
    return {"models": [_build_model_status(role) for role in _MODEL_ROLES]}