🛡️ agent_safety_box

Battle-hardened safety wrapper for AI agents — zero external dependencies, production-grade, fully tested.

pip install agent-safety-box

##The Problem AI agents are powerful, but they fail unpredictably in production. They hallucinate, get stuck in loops, blow through budgets, and crash silently. As one developer put it, "we're spending 40% of our budget on garbage output." Debugging becomes guesswork. agent_safety_box is the answer .

Features

Quick Start

from agent_safety_box import AgentWrapper, CircuitBreaker, RateLimiter

cb = CircuitBreaker(failure_threshold=3, reset_timeout=60)
rl = RateLimiter(rate=5.0)

wrapper = AgentWrapper(
    agent=my_agent,
    budget=100.0,
    max_runtime=10.0,
    max_cost_per_run=5.0,
    max_retries=2,
    retry_delay=1.0,
    circuit_breaker=cb,
    rate_limiter=rl,
    output_validator=lambda o: isinstance(o, str) and len(o) > 0,
    tags={"env": "prod", "model": "gpt-4o"},
    log_path="logs/audit.jsonl",
)

with wrapper:
    result = wrapper.run("Summarise this document", cost=1.5)

AgentWrapperV2 — Advanced Pipeline

from agent_safety_box.wrapper_v2 import AgentWrapperV2
from agent_safety_box.middleware import HookManager, Bulkhead, AdaptiveTimeout

hooks = HookManager()
hooks.on_before(lambda ctx: print(f"→ {ctx['task']}"))
hooks.on_after(lambda ctx: print(f"✓ done"))
hooks.on_error(lambda ctx, exc: print(f"✗ {exc}"))

w = AgentWrapperV2(
    agent=my_agent,
    budget=500.0,
    hooks=hooks,
    bulkhead=Bulkhead(max_concurrent=10),
    adaptive_timeout=AdaptiveTimeout(initial=5.0),
    on_budget_warning=lambda rem, bud: print(f"⚠ Only {rem:.2f} left!"),
    budget_warning_pct=0.20,
    plugins={"logger": lambda ctx: my_logger.info(ctx)},
)

result = w.run("task", cost=1.0)
print(w.metrics)
# {
#   'total_runs': 1, 'successes': 1, 'errors': 0,
#   'spent': 1.0, 'remaining': 499.0,
#   'rolling': {'p50_s': 0.003, 'p95_s': 0.003, 'p99_s': 0.003,
#               'error_rate': 0.0, 'throughput_rps': 312.5, 'sample_size': 1},
#   'avg_cost_forecast': 1.0,
#   'bulkhead_active': 0, 'bulkhead_max': 10,
#   'adaptive_timeout_s': 5.0,
# }

Fallback Chain

from agent_safety_box.middleware import FallbackChain

chain = FallbackChain(
    primary=gpt4_agent,
    fallbacks=[gpt35_agent, local_llm_agent],
)
result = chain.run("task")
print(f"Used: {chain.last_used}")
print(chain.stats)

Async Support

from agent_safety_box import AsyncAgentWrapper

wrapper = AsyncAgentWrapper(agent=my_async_agent, budget=50.0)

async def main():
    result = await wrapper.run("Classify this text", cost=0.5)

asyncio.run(main())

Exceptions

All exceptions inherit from AgentSafetyBoxError:

Safety Pipeline (per run)

1. max_cost_per_run check
2. budget check  (thread/async-safe lock)
3. rate_limiter.acquire()
4. circuit_breaker.allow_request()
5. agent.run(task)  [with timeout]
6. output_validator(result)
7. actual_cost adjustment (if agent returns {"actual_cost": X})
8. budget commit
9. circuit_breaker.record_success/failure()
10. audit log write

Running Tests

# All 263 tests
python -m unittest discover -s agent_safety_box/tests -p "test_*.py" -v

# Single module
python -m unittest agent_safety_box.tests.test_ultimate -v

Architecture

agent_safety_box/
├── __init__.py          # Public API
├── wrapper.py           # AgentWrapper + AsyncAgentWrapper (v1)
├── wrapper_v2.py        # AgentWrapperV2 + AsyncAgentWrapperV2 (v2)
├── middleware.py        # HookManager, RollingMetrics, Bulkhead,
│                        # AdaptiveTimeout, CostForecaster, FallbackChain
├── safety.py            # CircuitBreaker + RateLimiter
├── exceptions.py        # All custom exceptions
├── logger.py            # AuditLogger (JSONL + rotation)
└── tests/
    ├── conftest.py      # Shared stubs and base class
    ├── test_budget.py   # Budget enforcement (31 tests)
    ├── test_execution.py# Timeout + retry (30 tests)
    ├── test_logger.py   # Audit logging (21 tests)
    ├── test_safety.py   # Circuit + rate limiter (31 tests)
    ├── test_stress.py   # Concurrency + chaos (18 tests)
    └── test_ultimate.py # Deep edge cases (132 tests)

License

MIT