test_models.py

import sys
import yaml
from llm_triage.harness import triage_alert
from llm_triage.backends import GeminiBackend, OllamaBackend
from evaluation.run_baseline import load_alerts

def main():
    cfg = yaml.safe_load(open("config/experiment_config.yaml", "r"))
    alerts = load_alerts(r"data\ground_truth\alerts.jsonl")
    alert = alerts[0]
    
    print(f"Testing 1 alert across all {len(cfg['models'])} models...\n")
    
    for model_cfg in cfg["models"]:
        model_id = model_cfg["id"]
        btype = model_cfg["backend"]
        
        if btype == "gemini":
            backend = GeminiBackend(model_name=model_id)
        else:
            backend = OllamaBackend(model_name=model_id)
            
        print(f"--- MODEL: {model_id} ---")
        out, raw = triage_alert(backend, alert)
        
        if out is None:
            print(f"PARSE FAILURE! Raw output was:\n{raw}\n")
        else:
            print(f"SUCCESS! Parsed Output:\nSeverity: {out.severity}\nTechnique: {out.attack_technique}\nAction: {out.recommended_action}\nRationale: {out.rationale}\n")

if __name__ == "__main__":
    main()