BudgetGuard that always denies the first LLM call and records everything the framework hands it for post-hoc assertions.
| 29 | |
| 30 | |
| 31 | class DenyingGuard: |
| 32 | """BudgetGuard that always denies the first LLM call and records |
| 33 | everything the framework hands it for post-hoc assertions.""" |
| 34 | |
| 35 | def __init__(self) -> None: |
| 36 | self.llm_checks: list[tuple[str, int]] = [] |
| 37 | self.tool_checks: list[tuple[str, str]] = [] |
| 38 | self.llm_records: list[tuple[str, dict]] = [] |
| 39 | |
| 40 | def check_before_llm(self, session_id: str, estimated_tokens: int) -> dict: |
| 41 | self.llm_checks.append((session_id, estimated_tokens)) |
| 42 | return { |
| 43 | "decision": "deny", |
| 44 | "resource": "llm_tokens", |
| 45 | "reason": "test cap exceeded", |
| 46 | } |
| 47 | |
| 48 | def check_before_tool(self, session_id: str, tool_name: str) -> dict | None: |
| 49 | self.tool_checks.append((session_id, tool_name)) |
| 50 | return None # allow |
| 51 | |
| 52 | def record_after_llm(self, session_id: str, usage: dict) -> None: |
| 53 | self.llm_records.append((session_id, usage)) |
| 54 | |
| 55 | |
| 56 | class AllowingGuard: |