aegis-firewall

Python SDK for Aegis LLM Inference Firewall.

Installation

pip install aegis-firewall

Quick Start

Synchronous Client

from aegis import AegisClient, AegisFirewall

# Create client
client = AegisClient(
    base_url="http://localhost:8080",
    api_key="your-api-key",
)

# Create firewall wrapper
firewall = AegisFirewall(client)

# Add user message
user_span = firewall.add_user_message("Send email to john@example.com")

# Execute tool call with taint tracking
result = firewall.safe_execute_tool_call(
    "tool.send_email",
    {"to": "john@example.com", "subject": "Hello"},
    [user_span.id]
)

if result["success"]:
    print("Tool call executed successfully")
else:
    print(f"Tool call failed: {result['error']}")

Asynchronous Client

import asyncio
from aegis import AegisAsyncClient, AegisAsyncFirewall

async def main():
    async with AegisAsyncClient(
        base_url="http://localhost:8080",
        api_key="your-api-key",
    ) as client:
        firewall = AegisAsyncFirewall(client)
        
        # Add user message
        user_span = await firewall.add_user_message("Send email")
        
        # Execute tool call
        result = await firewall.safe_execute_tool_call(
            "tool.send_email",
            {"to": "john@example.com"},
            [user_span.id]
        )
        
        print(f"Success: {result['success']}")

asyncio.run(main())

Features

• ✅ Type-safe API - Full type hints with Pydantic
• ✅ Async support - Both sync and async clients
• ✅ Taint tracking - Automatic span management
• ✅ Policy builder - Programmatic policy creation
• ✅ Approval workflows - Human-in-the-loop support
• ✅ Error handling - Rich error types

API Reference

AegisClient

client = AegisClient(
    base_url="http://localhost:8080",
    api_key="optional-api-key",
    timeout=30.0,
    retries=3,
)

Methods

• validate_tool_call(tool_call) - Validate tool call
• create_span(span_data) - Create new span
• get_span(span_id) - Get span by ID
• list_decisions(session_id=None) - List decisions
• get_approval(approval_id) - Get approval
• list_approvals() - List pending approvals
• approve(approval_id) - Approve request
• deny(approval_id, reason) - Deny request
• get_trace(session_id) - Get trace
• verify_trace(bundle_data) - Verify trace
• health() - Health check

AegisFirewall

firewall = AegisFirewall(client)

Methods

• add_user_message(content) - Add user message
• add_document(content, sensitivity=...) - Add RAG document
• add_system_prompt(content) - Add system prompt
• execute_tool_call(capability, args, source_span_ids) - Execute tool call
• safe_execute_tool_call(...) - Safe execution with error handling
• get_session_spans() - Get all spans
• clear_session() - Clear session

PolicyBuilder

from aegis import PolicyBuilder, TrustLevel

policy = (
    PolicyBuilder()
    .allow("tool.read_calendar", TrustLevel.UNTRUSTED_USER)
    .allow("tool.send_email", TrustLevel.UNTRUSTED_USER, requires_approval=True)
    .deny("payments.*", TrustLevel.UNTRUSTED_DOCUMENT)
    .block_untrusted_documents()
    .block_secret_data()
    .require_approval_for_high_risk()
    .auto_deny_after(30)
    .build()
)

# Export as .aegis format
aegis_format = PolicyBuilder().allow("tool.send_email").to_aegis_format()

Examples

RAG with Document Isolation

from aegis import AegisClient, AegisFirewall, SensitivityLevel

client = AegisClient(base_url="http://localhost:8080")
firewall = AegisFirewall(client)

# Add user query
user_span = firewall.add_user_message("What is our pricing?")

# Add retrieved documents (automatically marked as untrusted)
doc = firewall.add_document("Pricing: $99/month", SensitivityLevel.PUBLIC)

# Tool call sourced only from documents will be BLOCKED
result = firewall.safe_execute_tool_call(
    "tool.send_email",
    {"to": "attacker@evil.com"},
    [doc.id]  # Only document source - blocked by taint tracking!
)

assert not result["success"], "Should be blocked"

Approval Workflow

def handle_approval(approval_id: str):
    print(f"Waiting for approval: {approval_id}")
    # Notify admin, poll for status, etc.

result = firewall.safe_execute_tool_call(
    "tool.execute_code",
    {"code": "print('hello')"},
    [user_span.id],
    on_approval_required=handle_approval
)

Async Example

import asyncio
from aegis import AegisAsyncClient, AegisAsyncFirewall

async def main():
    async with AegisAsyncClient("http://localhost:8080") as client:
        firewall = AegisAsyncFirewall(client)
        
        user_span = await firewall.add_user_message("Hello")
        result = await firewall.execute_tool_call(
            "tool.search_web",
            {"query": "weather"},
            [user_span.id]
        )
        
        print(f"Allowed: {result.allowed}")

asyncio.run(main())

Error Handling

from aegis import (
    CapabilityBlockedError,
    ApprovalRequiredError,
    NetworkError,
)

try:
    result = client.validate_tool_call(tool_call)
except CapabilityBlockedError as e:
    print(f"Blocked: {e.capability} - {e.reason}")
except ApprovalRequiredError as e:
    print(f"Approval needed: {e.approval_id}")
    await handle_approval(e.approval_id)
except NetworkError as e:
    print(f"Network error: {e.status_code}")

Development

# Install dev dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Format code
black aegis/

# Type check
mypy aegis/

License

MIT