feat: LLM response validation with defaults, security headers middleware, cleaner stats endpoint
This commit is contained in:
@@ -14,6 +14,22 @@ from ..logging_setup import log
|
||||
|
||||
_llm_semaphore = asyncio.Semaphore(config.llm_max_concurrent)
|
||||
|
||||
_REPORT_DEFAULTS = {
|
||||
"verdict": "unknown",
|
||||
"summary": "No summary provided",
|
||||
"analysis": "No analysis provided",
|
||||
"severity_rating": "unknown",
|
||||
}
|
||||
|
||||
|
||||
def _validate_report(report: dict) -> dict:
|
||||
for field, default in _REPORT_DEFAULTS.items():
|
||||
if not report.get(field):
|
||||
report[field] = default
|
||||
if report["verdict"] not in ("safe", "suspicious", "malicious", "unknown"):
|
||||
report["verdict"] = "unknown"
|
||||
return report
|
||||
|
||||
|
||||
def _build_user_message(finding: dict) -> str:
|
||||
"""Build a concise prompt from a finding's data."""
|
||||
@@ -78,7 +94,8 @@ async def _attempt_llm_call(finding_data: dict) -> dict | None:
|
||||
content = message.get("content", "")
|
||||
if not content:
|
||||
raise ValueError("Empty message content")
|
||||
return json.loads(content)
|
||||
parsed = json.loads(content)
|
||||
return _validate_report(parsed)
|
||||
except (ValueError, json.JSONDecodeError) as e:
|
||||
raw = ""
|
||||
try:
|
||||
@@ -94,7 +111,7 @@ async def _attempt_llm_call(finding_data: dict) -> dict | None:
|
||||
stripped = raw.strip().strip("`").strip()
|
||||
if stripped.startswith("json\n"):
|
||||
stripped = stripped[5:]
|
||||
return json.loads(stripped)
|
||||
return _validate_report(json.loads(stripped))
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
log.warning(
|
||||
|
||||
Reference in New Issue
Block a user