from pydantic import BaseModel, Field
from typing import List
from core.llm_router import get_llm

class TrapScan(BaseModel):
    detected_traps: List[str] = Field(default_factory=list)
    overall_risk: str = Field(description="low/medium/high/critical")
    critical_issues: List[str] = Field(default_factory=list)
    citation_grounding_note: str = ""

llm = get_llm(task_type="critical").with_structured_output(TrapScan)
llm.invoke("Some prompt")
print("Done")