RAG/backend/app/tools/provider_smoke.py

"""Minimal live probes for the three configured Bailian capabilities."""

from __future__ import annotations

import asyncio
import json
import sys
from collections.abc import Awaitable, Callable
from dataclasses import asdict, dataclass
from typing import Any

from app.adapters.bailian import (
    BailianChatAdapter,
    BailianEmbeddingAdapter,
    BailianRerankerAdapter,
)
from app.core.config import Settings
from app.core.secrets import SecretFileError
from app.ports.model_providers import ChatMessage, ModelProviderError


@dataclass(frozen=True, slots=True)
class ProbeResult:
    capability: str
    status: str
    model: str | None = None
    elapsed_ms: float | None = None
    request_id: str | None = None
    error_kind: str | None = None
    status_code: int | None = None


async def probe_embedding(settings: Settings, api_key: str) -> ProbeResult:
    adapter = BailianEmbeddingAdapter(
        api_key=api_key,
        base_url=settings.bailian_openai_base_url,
        model=settings.embedding_model,
        dimensions=settings.embedding_dimension,
        timeout_seconds=settings.model_timeout_seconds,
        max_retries=settings.model_max_retries,
    )
    try:
        result = await adapter.embed_documents(["用于能力探测的虚构地质文本。"])
        if len(result.vectors) != 1 or len(result.vectors[0]) != settings.embedding_dimension:
            raise RuntimeError("embedding contract mismatch")
        return ProbeResult(
            capability="embedding",
            status="ok",
            model=result.model,
            elapsed_ms=round(result.elapsed_ms, 2),
            request_id=result.request_id,
        )
    finally:
        await adapter.aclose()


async def probe_rerank(settings: Settings, api_key: str) -> ProbeResult:
    adapter = BailianRerankerAdapter(
        api_key=api_key,
        base_url=settings.bailian_rerank_base_url,
        model=settings.rerank_model,
        timeout_seconds=settings.model_timeout_seconds,
        max_retries=settings.model_max_retries,
    )
    try:
        result = await adapter.rerank(
            "哪段文本提到了斑岩铜矿？",
            ["虚构斑岩铜矿具有钾化带。", "虚构煤层采用测井曲线对比。"],
            top_n=1,
        )
        if len(result.items) != 1 or result.items[0].index not in (0, 1):
            raise RuntimeError("rerank contract mismatch")
        return ProbeResult(
            capability="rerank",
            status="ok",
            model=result.model,
            elapsed_ms=round(result.elapsed_ms, 2),
            request_id=result.request_id,
        )
    finally:
        await adapter.aclose()


async def probe_chat(settings: Settings, api_key: str) -> ProbeResult:
    adapter = BailianChatAdapter(
        api_key=api_key,
        base_url=settings.bailian_openai_base_url,
        model=settings.llm_model,
        timeout_seconds=settings.model_timeout_seconds,
        max_retries=settings.model_max_retries,
    )
    model: str | None = None
    request_id: str | None = None
    elapsed_ms = 0.0
    content_seen = False
    try:
        async for event in adapter.stream(
            [ChatMessage(role="user", content="只回复：能力正常")],
            max_tokens=16,
        ):
            model = event.model
            request_id = event.request_id or request_id
            elapsed_ms = max(elapsed_ms, event.elapsed_ms)
            content_seen = content_seen or bool(event.delta)
        if not content_seen:
            raise RuntimeError("chat stream contained no text")
        return ProbeResult(
            capability="chat",
            status="ok",
            model=model,
            elapsed_ms=round(elapsed_ms, 2),
            request_id=request_id,
        )
    finally:
        await adapter.aclose()


def failed_probe(capability: str, error: BaseException) -> ProbeResult:
    if isinstance(error, ModelProviderError):
        return ProbeResult(
            capability=capability,
            status="failed",
            request_id=error.request_id,
            error_kind=error.kind.value,
            status_code=error.status_code,
        )
    return ProbeResult(
        capability=capability,
        status="failed",
        error_kind="internal_contract_error",
    )


async def run_probe(
    capability: str,
    operation: Callable[[Settings, str], Awaitable[ProbeResult]],
    settings: Settings,
    api_key: str,
) -> ProbeResult:
    try:
        return await operation(settings, api_key)
    except Exception as exc:  # The output is deliberately reduced to a safe category.
        return failed_probe(capability, exc)


def write_json_line(payload: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")


async def async_main() -> int:
    try:
        settings = Settings()
        if any(
            "<workspace-id>" in url
            for url in (
                settings.bailian_openai_base_url,
                settings.bailian_rerank_base_url,
            )
        ):
            raise ValueError("workspace endpoint placeholders are not runnable")
        api_key = settings.bailian_api_key()
    except (SecretFileError, ValueError):
        write_json_line(
            {
                "capability": "configuration",
                "status": "failed",
                "error_kind": "invalid_local_configuration",
            }
        )
        return 2

    probes = (
        ("embedding", probe_embedding),
        ("rerank", probe_rerank),
        ("chat", probe_chat),
    )
    results = []
    for capability, operation in probes:
        result = await run_probe(capability, operation, settings, api_key)
        results.append(result)
        write_json_line(asdict(result))
    return 0 if all(result.status == "ok" for result in results) else 1


def main() -> None:
    raise SystemExit(asyncio.run(async_main()))


if __name__ == "__main__":
    main()