soc_ai/model_probe.py
Rama main · réplica de lectura del repositorio original.
"""Structured-output contract probe for candidate analyst backends.
``soc-ai model-probe`` answers, in minutes, the question that previously took a
debugging session: *can this model actually drive soc-ai's triage contract, and
under which settings?* It runs the REAL synthesizer agent (same builders and
prompt the pipeline uses) against a canned scenario N times and tallies the
outcomes into the failure classes this project has recorded from live serving
models — schema-retry exhaustion, no tool call, HTTP failures, timeouts.
Run it before pointing prod at a new analyst backend, and again for each
candidate setting::
soc-ai model-probe --model qwen3.6-35b-cpu
soc-ai model-probe --model qwen3.6-35b-cpu --output-mode native
soc-ai model-probe --model qwen3.6-35b-cpu --tool-choice required
soc-ai model-probe --model new-thing -n 12 --min-ok 10 # CI-gateable
The probe reports ``served_backend`` from the gateway's own attribution
headers, so an aliased or fallback-routed model name cannot misattribute the
result (the exact trap that derailed the 2026-08-03 investigation).
"""
from __future__ import annotations
import time
from collections.abc import Callable, Coroutine
from typing import Any
from soc_ai.config import Settings
# A realistic, self-contained benign-DNS triage scenario. Deliberately easy:
# the probe measures CONTRACT compliance (schema, tool call, JSON), not verdict
# quality — an easy scenario keeps a failed probe unambiguous about what broke.
PROBE_SCENARIO = """Alert: ET POLICY Observed DNS Query to .top TLD
src=10.0.0.42 dest=8.8.8.8 host=WKSTN-114 time=2026-08-03T14:02:11Z
Evidence gathered:
- dns.question.name = cdn-assets-01.top
- 3 queries in 60s, all NOERROR, resolving to 104.21.x.x (Cloudflare)
- Destination reputation: no hits on any blocklist
- Process: chrome.exe, user browsing session active
- No subsequent connection to the resolved IP beyond TLS 443
- Host has no other alerts in the last 30 days
Write the TriageReport now."""
def classify_probe_exception(exc: BaseException) -> str:
"""Map an exception from a probe attempt to a failure-class label.
The labels line up with the failure taxonomy in the lesser-model runbook so
a probe report reads as a diagnosis: ``schema_retry_exhausted`` points at
output-shape wobble (try ``--output-mode native``), ``http_5xx`` at the
serving stack, ``timeout`` at generation speed vs the configured budgets.
"""
from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior # noqa: PLC0415
if isinstance(exc, ModelHTTPError):
return f"http_{exc.status_code}"
if isinstance(exc, UnexpectedModelBehavior):
msg = str(exc).lower()
if "retries" in msg:
return "schema_retry_exhausted"
return "unexpected_model_behavior"
if isinstance(exc, TimeoutError):
return "timeout"
return type(exc).__name__
def summarize_probe(
*,
model: str,
output_mode: str,
tool_choice_required: bool,
outcomes: list[tuple[str, str]],
served_backend: dict[str, Any] | None,
elapsed_s: float,
) -> dict[str, Any]:
"""Fold raw (label, detail) outcomes into the probe report dict."""
tally: dict[str, int] = {}
failures: list[str] = []
for label, detail in outcomes:
tally[label] = tally.get(label, 0) + 1
if label != "OK" and len(failures) < 3:
failures.append(detail)
ok = tally.get("OK", 0)
return {
"model": model,
"output_mode": output_mode,
"tool_choice_required": tool_choice_required,
"n": len(outcomes),
"ok": ok,
"usable_rate": (ok / len(outcomes)) if outcomes else 0.0,
"tally": tally,
"failures": failures,
"served_backend": served_backend,
"elapsed_s": round(elapsed_s, 1),
}
async def _run_once(agent: Any, prompt: str) -> tuple[str, str]:
"""One probe attempt: run the agent, return (label, detail)."""
try:
result = await agent.run(prompt)
out = result.output
return ("OK", f"{out.verdict} conf={out.confidence}")
except Exception as exc:
detail = " ".join(str(exc).split())
return (classify_probe_exception(exc), detail)
async def probe_model(
settings: Settings,
*,
n: int = 6,
output_mode: str = "tool",
run_once: Callable[[Any, str], Coroutine[Any, Any, tuple[str, str]]] | None = None,
) -> dict[str, Any]:
"""Probe ``settings.analyst_model`` with the real synth agent, N attempts.
Attempts run SEQUENTIALLY on purpose: a lesser backend (single llama.cpp
slot, small vLLM) queues concurrent requests, which would fold queue wait
into the timing and can flip timeout outcomes — sequential keeps each
attempt's result attributable to the model alone.
``run_once`` is injectable for tests; the default runs the agent for real.
"""
from soc_ai.agent._gateway_retry import capture_backend_attribution # noqa: PLC0415
from soc_ai.agent.models import build_synthesizer_model # noqa: PLC0415
from soc_ai.agent.orchestrator import build_synth_first_agent # noqa: PLC0415
runner = run_once if run_once is not None else _run_once
agent = build_synth_first_agent(
build_synthesizer_model(settings, temperature=settings.synthesizer_temperature),
output_mode=output_mode,
)
outcomes: list[tuple[str, str]] = []
started = time.monotonic()
with capture_backend_attribution() as attribution:
for _ in range(n):
outcomes.append(await runner(agent, PROBE_SCENARIO))
return summarize_probe(
model=settings.analyst_model,
output_mode=output_mode,
tool_choice_required=settings.analyst_tool_choice_required,
outcomes=outcomes,
served_backend=dict(attribution) if attribution else None,
elapsed_s=time.monotonic() - started,
)
def format_probe_report(rep: dict[str, Any]) -> str:
"""Human-readable probe report for the CLI."""
lines = [
f"model : {rep['model']}",
f"output mode : {rep['output_mode']}"
+ (" (tool_choice=required)" if rep["tool_choice_required"] else ""),
f"usable : {rep['ok']}/{rep['n']} ({rep['usable_rate']:.0%})",
f"tally : {rep['tally']}",
f"elapsed : {rep['elapsed_s']}s",
]
backend = rep.get("served_backend") or {}
if backend.get("api_base"):
served = backend["api_base"]
fallbacks = backend.get("attempted_fallbacks")
lines.append(
f"served by : {served}"
+ (f" (fallbacks attempted: {fallbacks})" if fallbacks not in (None, "0") else "")
)
for f in rep.get("failures", []):
lines.append(f" failure: {f}")
return "\n".join(lines)
# ── Battery: the four knob configurations + a deterministic recommendation ───
# Fixed probe order. Baseline first — it is the comparison anchor and, on a
# healthy backend, also the fastest way to confirm the battery itself works.
BATTERY_CONFIGS: tuple[tuple[str, bool], ...] = (
("tool", False), # baseline: today's prod behavior
("native", False),
("prompted", False),
("tool", True),
)
# Equal-rate speed bar: recommend on speed alone only when the candidate is at
# least this much faster than baseline (elapsed <= baseline * 0.75).
_SPEEDUP_BAR = 0.75
def _config_label(mode: str, required: bool) -> str:
return f"{mode}+required" if required else mode
def recommend(configs: list[dict[str, Any]]) -> dict[str, Any] | None:
"""Pick the configuration to recommend over the baseline, or None.
Deterministic and explainable by design: rank by usable rate desc then
elapsed asc; recommend only a STRICT improvement over the baseline (higher
rate, or equal rate and >= 25% faster). ``prompted`` — the
weakest-guarantees mode — is recommended only when it is the sole
configuration at the top usable rate, so it never displaces a tied
``native``/``tool`` result on speed. Returns the two knob values plus a
human-readable reason, or None when the baseline stands (including when
everything failed — "configure nothing" is the honest answer to that).
"""
baseline = next(
(c for c in configs if c["output_mode"] == "tool" and not c["tool_choice_required"]),
None,
)
if baseline is None or not configs:
return None
candidates = [c for c in configs if c is not baseline and c["usable_rate"] > 0]
if not candidates:
return None
top_rate = max(c["usable_rate"] for c in candidates)
top = [c for c in candidates if c["usable_rate"] == top_rate]
# Demote prompted unless it is the sole top-rate configuration.
non_prompted = [c for c in top if c["output_mode"] != "prompted"]
if non_prompted:
top = non_prompted
best = min(top, key=lambda c: c["elapsed_s"])
label = _config_label(best["output_mode"], best["tool_choice_required"])
rate_str = f"{best['ok']}/{best['n']}"
if best["usable_rate"] > baseline["usable_rate"]:
reason = f"{label}: {rate_str} usable vs baseline {baseline['ok']}/{baseline['n']}"
elif (
best["usable_rate"] == baseline["usable_rate"]
and baseline["elapsed_s"] > 0
and best["elapsed_s"] <= baseline["elapsed_s"] * _SPEEDUP_BAR
):
speedup = baseline["elapsed_s"] / max(best["elapsed_s"], 0.1)
reason = f"{label}: {rate_str} usable, {speedup:.1f}x faster than tool mode"
else:
return None
return {
"synthesizer_output_mode": best["output_mode"],
"analyst_tool_choice_required": best["tool_choice_required"],
"config": label,
"reason": reason,
}
async def run_battery(
settings: Settings,
*,
model: str | None = None,
n: int = 2,
run_once: Callable[[Any, str], Coroutine[Any, Any, tuple[str, str]]] | None = None,
per_attempt_timeout_s: float = 120.0,
on_progress: Callable[[str, int, int], None] | None = None,
) -> dict[str, Any]:
"""Probe *model* under every battery configuration; attach a recommendation.
Sequential across configurations for the same reason probe attempts are
sequential within one. ``per_attempt_timeout_s`` bounds each attempt so one
hung backend read cannot wedge the whole battery (a timed-out attempt lands
as a ``timeout`` outcome, it does not abort the run). ``on_progress`` is
called with (config label, 1-based index, total) before each configuration —
the API's status object hooks in there.
"""
import asyncio # noqa: PLC0415
base_runner = run_once if run_once is not None else _run_once
async def bounded_runner(agent: Any, prompt: str) -> tuple[str, str]:
try:
return await asyncio.wait_for(base_runner(agent, prompt), per_attempt_timeout_s)
except TimeoutError:
return ("timeout", f"attempt exceeded {per_attempt_timeout_s:.0f}s")
probe_settings = settings
if model:
probe_settings = settings.model_copy(update={"analyst_model": model})
started = time.monotonic()
configs: list[dict[str, Any]] = []
total = len(BATTERY_CONFIGS)
for i, (mode, required) in enumerate(BATTERY_CONFIGS, start=1):
if on_progress is not None:
on_progress(_config_label(mode, required), i, total)
cfg_settings = probe_settings.model_copy(update={"analyst_tool_choice_required": required})
configs.append(
await probe_model(cfg_settings, n=n, output_mode=mode, run_once=bounded_runner)
)
return {
"model": probe_settings.analyst_model,
"n_per_config": n,
"configs": configs,
"recommendation": recommend(configs),
"elapsed_s": round(time.monotonic() - started, 1),
}