# aegis-firewall

Python SDK for Aegis LLM Inference Firewall.

## Installation

```bash
pip install aegis-firewall
```

## Quick Start

### Synchronous Client

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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

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

# Run tests
pytest

# Format code
black aegis/

# Type check
mypy aegis/
```

## License

MIT
