didi-lot1-ai/ai_platform/modules/didi_brain/lint/detector.py

114 lines
3.2 KiB
Python

"""Single-pair NLI classification for the Lint pass.
Takes two claim texts and asks Qwen 397B whether they are EQUIVALENT,
CONTRADICTORY, or INCOMPARABLE. Returns a dataclass with label +
confidence + optional error. Never raises — all failure paths fall
through to INCOMPARABLE with an error string.
"""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from pathlib import Path
from shared.config import LlmRole
from shared.llm_client import LlmClient, LlmError
from shared.logging import get_logger
log = get_logger(__name__)
PROMPT_VERSION = "v1"
_PROMPT_PATH = (
Path(__file__).resolve().parent / "prompts" / f"pair_nli_{PROMPT_VERSION}.md"
)
ALLOWED_LABELS = {"EQUIVALENT", "CONTRADICTORY", "INCOMPARABLE"}
MAX_CLAIM_CHARS = 1500 # truncate long evidence before sending to LLM
PER_CALL_TIMEOUT_S = 30.0
_PROMPT: str | None = None
def _load_prompt() -> str:
global _PROMPT
if _PROMPT is None:
_PROMPT = _PROMPT_PATH.read_text(encoding="utf-8")
return _PROMPT
@dataclass(slots=True, frozen=True)
class PairClassification:
label: str
confidence: float
error: str | None = None
async def classify_pair(
llm: LlmClient,
*,
claim_a: str,
claim_b: str,
timeout_s: float = PER_CALL_TIMEOUT_S,
) -> PairClassification:
if not claim_a.strip() or not claim_b.strip():
return PairClassification(
label="INCOMPARABLE", confidence=0.0, error="empty_claim"
)
prompt = (
_load_prompt()
.replace("{claim_a}", claim_a[:MAX_CLAIM_CHARS])
.replace("{claim_b}", claim_b[:MAX_CLAIM_CHARS])
)
try:
result, _usage = await asyncio.wait_for(
llm.chat_json(
role=LlmRole.REASONING,
system=(
"You are a precise claim comparison classifier. "
"Respond with strictly valid JSON only, no commentary."
),
user=prompt,
max_tokens=120,
temperature=0.0,
),
timeout=timeout_s,
)
except asyncio.TimeoutError:
return PairClassification(
label="INCOMPARABLE", confidence=0.0, error="timeout"
)
except LlmError as e:
return PairClassification(
label="INCOMPARABLE", confidence=0.0, error=f"llm:{e}"
)
except Exception as e: # noqa: BLE001
return PairClassification(
label="INCOMPARABLE",
confidence=0.0,
error=f"{type(e).__name__}:{e}",
)
if not isinstance(result, dict):
return PairClassification(
label="INCOMPARABLE", confidence=0.0, error="non_dict_response"
)
raw_label = (result.get("label") or "").strip().upper()
try:
conf = float(result.get("confidence", 0))
except (TypeError, ValueError):
conf = 0.0
conf = max(0.0, min(1.0, conf))
if raw_label not in ALLOWED_LABELS:
return PairClassification(
label="INCOMPARABLE",
confidence=0.0,
error=f"bad_label:{raw_label[:40]}",
)
return PairClassification(label=raw_label, confidence=conf)