76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
"""LLM triage stage — offline tests of validation, acceptance gate, and skip."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from recon_triage.orchestrator import build_report_from_fixtures
|
|
from recon_triage.triage import llm
|
|
from recon_triage.triage.ranking import collect_identifiers
|
|
|
|
FIXTURES = Path(__file__).parent / "fixtures"
|
|
|
|
|
|
def _report(fixture_search_fn):
|
|
return build_report_from_fixtures(
|
|
FIXTURES, run_id="t", search_fn=fixture_search_fn, enable_nuclei=True
|
|
)
|
|
|
|
|
|
def test_skips_when_unconfigured(monkeypatch, fixture_search_fn):
|
|
monkeypatch.delenv("LLM_BASE_URL", raising=False)
|
|
assert not llm.is_configured()
|
|
triage = llm.run_triage(_report(fixture_search_fn))
|
|
assert triage.generated_by == "deterministic"
|
|
|
|
|
|
def test_extract_json_handles_fences():
|
|
assert llm._extract_json('```json\n{"a": 1}\n```') == {"a": 1}
|
|
assert llm._extract_json('prose {"a": 2} more') == {"a": 2}
|
|
assert llm._extract_json("not json") is None
|
|
|
|
|
|
def test_acceptance_gate_strips_invented_refs(fixture_search_fn):
|
|
from recon_triage.schema import TriageItem, TriageReport
|
|
|
|
report = _report(fixture_search_fn)
|
|
valid = collect_identifiers(report)
|
|
real_id = next(iter(valid))
|
|
triage = TriageReport(
|
|
prioritized_findings=[
|
|
TriageItem(summary="real", rationale="r", evidence_refs=[real_id], suggested_next_step=f"look at {real_id}"),
|
|
TriageItem(summary="invented", rationale="r", evidence_refs=["EDB-99999999"], suggested_next_step="run made-up-exploit"),
|
|
TriageItem(summary="mixed", rationale="r", evidence_refs=[real_id, "CVE-FAKE"], suggested_next_step="do CVE-FAKE"),
|
|
]
|
|
)
|
|
cleaned = llm._enforce_acceptance(triage, valid)
|
|
summaries = {i.summary for i in cleaned.prioritized_findings}
|
|
# invented item (no valid ref at all) is dropped entirely
|
|
assert "invented" not in summaries
|
|
assert "real" in summaries
|
|
# mixed item kept but bad ref stripped and bad next-step cleared
|
|
mixed = next(i for i in cleaned.prioritized_findings if i.summary == "mixed")
|
|
assert "CVE-FAKE" not in mixed.evidence_refs
|
|
assert mixed.suggested_next_step is None
|
|
|
|
|
|
def test_llm_path_with_mocked_transport(monkeypatch, fixture_search_fn):
|
|
"""Exercise the real run_triage LLM branch with a stubbed _call_llm (no network)."""
|
|
report = _report(fixture_search_fn)
|
|
valid = collect_identifiers(report)
|
|
real_id = next(r for r in valid if ":" in r) # a host:port id
|
|
|
|
def fake_call(report_json, **kwargs):
|
|
return (
|
|
'{"prioritized_findings": [{"summary": "test", "rationale": "r", '
|
|
f'"severity": "high", "evidence_refs": ["{real_id}"], '
|
|
f'"suggested_next_step": "review {real_id}", "confidence": 0.9}}]}}'
|
|
)
|
|
|
|
monkeypatch.setenv("LLM_BASE_URL", "http://fake:11434/v1")
|
|
monkeypatch.setattr(llm, "_call_llm", fake_call)
|
|
triage = llm.run_triage(report)
|
|
assert triage.generated_by == "llm"
|
|
assert triage.prioritized_findings[0].summary == "test"
|
|
assert triage.prioritized_findings[0].evidence_refs == [real_id]
|