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.model_gateway import ModelGatewayAdapter
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, adapter: ModelGatewayAdapter) -> ProbeResult:
    # API identity probes query embedding. Document embedding remains worker-only.
    result = await adapter.embed_query("用于能力探测的虚构地质问题。")
    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,
    )


async def probe_rerank(_: Settings, adapter: ModelGatewayAdapter) -> ProbeResult:
    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,
    )


async def probe_chat(_: Settings, adapter: ModelGatewayAdapter) -> ProbeResult:
    model: str | None = None
    request_id: str | None = None
    elapsed_ms = 0.0
    content_seen = False
    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,
    )


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, ModelGatewayAdapter], Awaitable[ProbeResult]],
    settings: Settings,
    adapter: ModelGatewayAdapter,
) -> ProbeResult:
    try:
        return await operation(settings, adapter)
    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:
    adapter: ModelGatewayAdapter | None = None
    try:
        settings = Settings()
        adapter = ModelGatewayAdapter.from_settings(settings)
    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),
    )
    try:
        results = []
        for capability, operation in probes:
            result = await run_probe(capability, operation, settings, adapter)
            results.append(result)
            write_json_line(asdict(result))
        return 0 if all(result.status == "ok" for result in results) else 1
    finally:
        await adapter.aclose()


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


if __name__ == "__main__":
    main()