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
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 _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 settings.llm_provider, settings.llm_model
    if role == "hyde_llm":
        return (
            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"

    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`."""
    provider, model = _resolve_role_provider_model(role)
    client = get_llm_client(provider=provider, model=model)
    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]}