135 lines
5.2 KiB
Python
135 lines
5.2 KiB
Python
"""Startup model connectivity probe helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from copy import deepcopy
|
|
|
|
from ..logging import get_logger
|
|
from ...core.database import AsyncSessionLocal
|
|
from ...services.model.model_config import build_tier_runtime
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def startup_log(message: str, **fields) -> None:
|
|
logger.info(f"[后端] {message}", **fields)
|
|
|
|
|
|
async def probe_model_connectivity_on_startup() -> None:
|
|
from ...services.llm_service import llm_service
|
|
|
|
low_runtime = build_tier_runtime("low")
|
|
high_runtime = build_tier_runtime("high")
|
|
low_provider = str(low_runtime.get("provider") or "glm").strip() or "glm"
|
|
low_model = str(low_runtime.get("model") or "qwen-plus").strip() or "qwen-plus"
|
|
high_provider = str(high_runtime.get("provider") or "openai").strip() or "openai"
|
|
high_model = str(high_runtime.get("model") or "GLM-5").strip() or "GLM-5"
|
|
|
|
probe_targets: list[tuple[str, str, str]] = [
|
|
("LOW", low_provider, low_model),
|
|
("HIGH", high_provider, high_model),
|
|
]
|
|
deduped_targets: list[tuple[str, str, str]] = []
|
|
seen: set[tuple[str, str]] = set()
|
|
for level, provider, model in probe_targets:
|
|
key = (provider.lower(), model)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
deduped_targets.append((level, provider, model))
|
|
|
|
lines: list[str] = [
|
|
"┌──────────────────────────────┐",
|
|
"│ 模型连通性检测 │",
|
|
]
|
|
|
|
for level, provider, model in deduped_targets:
|
|
try:
|
|
runtime = None
|
|
async with AsyncSessionLocal() as db:
|
|
try:
|
|
resolved_function_name = (
|
|
"classify_intent_ai_pmo_kernel" if level == "LOW" else "classify_intent_complex"
|
|
)
|
|
runtime = await llm_service._resolve_engine_for_call(
|
|
db,
|
|
tenant_id=1,
|
|
function_name=resolved_function_name,
|
|
)
|
|
except Exception as resolve_exc:
|
|
logger.warning(
|
|
"startup_model_probe_runtime_resolve_failed",
|
|
level=level,
|
|
provider=provider,
|
|
model=model,
|
|
error=str(resolve_exc),
|
|
)
|
|
if not isinstance(runtime, dict) or not runtime:
|
|
runtime = llm_service._env_provider_runtime(provider)
|
|
if level == "LOW":
|
|
runtime["model"] = low_model
|
|
runtime["base_url"] = str(low_runtime.get("base_url") or runtime.get("base_url") or "")
|
|
elif level == "HIGH":
|
|
runtime["model"] = high_model
|
|
runtime["base_url"] = str(high_runtime.get("base_url") or runtime.get("base_url") or "")
|
|
except Exception as exc:
|
|
logger.error(
|
|
"startup_model_probe_runtime_invalid",
|
|
level=level,
|
|
provider=provider,
|
|
model=model,
|
|
error=str(exc),
|
|
)
|
|
lines.append(f"│ {model:<12} ✗ 失败 (0.0s) │")
|
|
continue
|
|
|
|
started_at = time.perf_counter()
|
|
ok = False
|
|
error_text = ""
|
|
circuit_cls = type(llm_service)
|
|
saved_circuit_state = deepcopy(getattr(circuit_cls, "_provider_circuit_state", {}))
|
|
try:
|
|
result = await asyncio.wait_for(
|
|
llm_service._chat_completion(
|
|
provider=str(runtime.get("provider") or provider),
|
|
prompt="ping",
|
|
system_prompt="你是连通性检测助手，只回复 OK。",
|
|
api_key=str(runtime.get("api_key") or ""),
|
|
base_url=str(runtime.get("base_url") or ""),
|
|
model=str(runtime.get("model") or model),
|
|
temperature=0.0,
|
|
max_tokens=1,
|
|
function_name="startup_model_probe",
|
|
response_mode="fast",
|
|
),
|
|
timeout=8.0,
|
|
)
|
|
ok = bool(str(result or "").strip())
|
|
if not ok:
|
|
error_text = "empty_response"
|
|
except Exception as exc:
|
|
error_text = str(exc)
|
|
logger.error(
|
|
"startup_model_probe_failed",
|
|
level=level,
|
|
provider=provider,
|
|
model=model,
|
|
error=error_text,
|
|
)
|
|
finally:
|
|
circuit_cls._provider_circuit_state = saved_circuit_state
|
|
|
|
elapsed = time.perf_counter() - started_at
|
|
if ok:
|
|
lines.append(f"│ {model:<12} ✓ 正常 ({elapsed:.1f}s) │")
|
|
else:
|
|
lines.append(f"│ {model:<12} ✗ 失败 ({elapsed:.1f}s) │")
|
|
if error_text:
|
|
startup_log("模型连通性失败（不阻断启动）", model=model, provider=provider, error=error_text)
|
|
|
|
lines.append("└──────────────────────────────┘")
|
|
for line in lines:
|
|
logger.info(line)
|