Metadata-Version: 2.4
Name: aegis-firewall
Version: 2.0.0
Summary: Python SDK for Aegis LLM Inference Firewall
Author-email: Aegis Team <team@aegis-firewall.dev>
License: MIT
Project-URL: Homepage, https://github.com/xaseai/aegis
Project-URL: Documentation, https://aegis-firewall.dev/docs
Project-URL: Repository, https://github.com/xaseai/aegis
Project-URL: Issues, https://github.com/xaseai/aegis/issues
Keywords: aegis,llm,firewall,security,ai-security,prompt-injection
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Security
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: httpx>=0.25.0
Requires-Dist: pydantic>=2.5.0
Requires-Dist: cbor2>=5.5.0
Requires-Dist: cryptography>=41.0.0
Provides-Extra: websocket
Requires-Dist: websockets>=12.0; extra == "websocket"
Provides-Extra: all
Requires-Dist: websockets>=12.0; extra == "all"
Provides-Extra: dev
Requires-Dist: pytest>=7.4.0; extra == "dev"
Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
Requires-Dist: black>=23.7.0; extra == "dev"
Requires-Dist: ruff>=0.0.285; extra == "dev"
Requires-Dist: mypy>=1.5.0; extra == "dev"

# 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
