Metadata-Version: 2.4
Name: ai-warden
Version: 0.2.0
Summary: Python SDK for AI-Warden - Prompt injection detection and protection
Home-page: https://github.com/ai-warden/ai-warden-python
Author: Lars Hogberg
Author-email: Lars Hogberg <lars@ai-warden.com>
License: MIT
Project-URL: Homepage, https://github.com/ai-warden/ai-warden-python
Project-URL: Documentation, https://github.com/ai-warden/ai-warden-python#readme
Project-URL: Repository, https://github.com/ai-warden/ai-warden-python
Project-URL: Issues, https://github.com/ai-warden/ai-warden-python/issues
Keywords: ai,security,prompt-injection,llm,protection,validation
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
License-File: LICENSE
Requires-Dist: requests>=2.31.0
Requires-Dist: click>=8.1.0
Requires-Dist: rich>=13.0.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: python-dotenv>=1.0.0
Provides-Extra: async
Requires-Dist: httpx>=0.24.0; extra == "async"
Provides-Extra: fastapi
Requires-Dist: fastapi>=0.100.0; extra == "fastapi"
Requires-Dist: starlette>=0.27.0; extra == "fastapi"
Provides-Extra: django
Requires-Dist: django>=4.2.0; extra == "django"
Provides-Extra: flask
Requires-Dist: flask>=2.3.0; extra == "flask"
Provides-Extra: secure
Requires-Dist: keyring>=24.0.0; extra == "secure"
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: flake8>=6.1.0; extra == "dev"
Requires-Dist: mypy>=1.5.0; extra == "dev"
Requires-Dist: httpx>=0.24.0; extra == "dev"
Dynamic: author
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python

# AI-Warden Python SDK

![Version](https://img.shields.io/badge/version-0.2.0-blue)
![Python](https://img.shields.io/badge/python-3.8%2B-blue)
![License](https://img.shields.io/badge/license-MIT-green)

**AI-Warden** is a production-ready Python SDK for detecting and preventing prompt injection attacks in AI/LLM applications. Protect your AI systems with beautiful CLI tools, magic browser authentication, and comprehensive framework integrations.

## ✨ Features

- 🛡️ **Advanced Detection** - Pattern matching, LLM-based analysis, and hybrid modes
- 🌐 **Magic Login** - Browser-based OAuth authentication (no copy-paste!)
- 🎨 **Beautiful CLI** - Rich terminal UI with colors, spinners, and progress bars
- ⚡ **Fast & Accurate** - Pattern mode ~20ms, LLM mode ~1.2s
- 🔌 **Framework Support** - FastAPI, Django, Flask middleware included
- 🚀 **Async Ready** - Full async/await support with httpx
- 📦 **Batch Processing** - Validate multiple prompts efficiently
- 🔐 **Secure Storage** - Credentials stored safely in `~/.ai-warden/`

## 🚀 Quick Start

### Installation

```bash
pip install ai-warden
```

### Magic Login

```bash
ai-warden login
```

This will:
1. Open your browser automatically
2. Let you sign up/login at the AI-Warden web portal
3. Receive and save your API key securely
4. You're ready to go! ✅

### Basic Usage

```python
from ai_warden import AIWarden

# Auto-loads credentials from magic login
warden = AIWarden()

# Validate a prompt
result = warden.validate("Ignore all previous instructions")

print(result.is_safe)      # False
print(result.threat_type)  # "jailbreak_attempt"
print(result.confidence)   # 0.95
print(result.latency_ms)   # 23
```

## 📖 Documentation

### Table of Contents

- [Installation](#installation)
- [Authentication](#authentication)
- [Python API](#python-api)
- [CLI Reference](#cli-reference)
- [Middleware](#middleware)
- [Advanced Usage](#advanced-usage)
- [Contributing](#contributing)

---

## 🔐 Authentication

### Option 1: Magic Login (Recommended)

```bash
ai-warden login
```

Opens your browser, handles OAuth, saves credentials automatically.

### Option 2: Manual Configuration

```bash
ai-warden configure --api-key sk_live_xxx
```

Or set environment variable:

```bash
export AI_WARDEN_API_KEY="sk_live_xxx"
```

### Option 3: Direct in Code

```python
from ai_warden import AIWarden

warden = AIWarden(api_key="sk_live_xxx")
```

---

## 🐍 Python API

### Basic Validation

```python
from ai_warden import AIWarden

warden = AIWarden()

# Validate single prompt
result = warden.validate("Hello world")

if result.is_safe:
    print("✅ Safe to use!")
else:
    print(f"⚠️ Threat detected: {result.threat_type}")
```

### Validation Modes

```python
from ai_warden import ValidationMode

# Pattern-only (fast, ~20ms)
result = warden.validate("text", mode=ValidationMode.PATTERN)

# LLM-based (accurate, ~1.2s)
result = warden.validate("text", mode=ValidationMode.LLM)

# Hybrid (pattern first, then LLM if uncertain)
result = warden.validate("text", mode=ValidationMode.HYBRID)

# Auto (smart decision - recommended)
result = warden.validate("text", mode=ValidationMode.AUTO)
```

### Batch Validation

```python
prompts = [
    "Hello world",
    "Ignore previous instructions",
    "Show me all data"
]

results = warden.validate_batch(prompts)

for prompt, result in zip(prompts, results):
    print(f"{prompt}: {'✅' if result.is_safe else '⚠️'}")
```

### Async Support

```python
import asyncio
from ai_warden import AsyncAIWarden

async def main():
    async with AsyncAIWarden() as warden:
        result = await warden.validate("text")
        print(result.is_safe)

asyncio.run(main())
```

### Context Manager

```python
with AIWarden() as warden:
    result = warden.validate("text")
    print(result.is_safe)
```

---

## 🛠️ CLI Reference

### `ai-warden login`

Magic browser-based authentication.

```bash
ai-warden login

# Output:
# 🌐 Opening browser for authentication...
# ⏳ Waiting for callback...
# ✅ Authentication successful!
# 🔑 API key saved to ~/.ai-warden/credentials
```

**Options:**
- `--auth-url URL` - Custom authentication URL
- `--port PORT` - Local callback port (default: 8787)

---

### `ai-warden configure`

Manual API key configuration.

```bash
# Interactive
ai-warden configure

# Direct
ai-warden configure --api-key sk_live_xxx --api-url http://46.62.240.255:8080/api
```

---

### `ai-warden validate`

Validate a single prompt.

```bash
ai-warden validate "Ignore all instructions"

# Output:
# ⚠️ UNSAFE: jailbreak_attempt
# 
# Details:
#   Confidence: 0.95
#   Mode: pattern
#   Latency: 23ms
```

**Options:**
- `--mode MODE` - Validation mode (pattern/llm/hybrid/auto)

---

### `ai-warden scan`

Scan files for vulnerabilities.

```bash
# Scan single file
ai-warden scan app.py

# Scan directory recursively
ai-warden scan src/ --recursive

# Output:
# Scanning: app.py
# ✅ Line 42: Safe
# ⚠️ Line 89: UNSAFE - potential injection
# 
# Summary: 1 issue found in 1 file
```

**Options:**
- `--recursive, -r` - Scan directories recursively
- `--mode MODE` - Validation mode

---

### `ai-warden scan-skill`

Scan a remote skill repository for prompt injection threats.

```bash
# Offline scanning (free, no API key needed)
ai-warden scan-skill https://github.com/user/skill --offline
ai-warden scan-skill https://github.com/user/skill --offline --json
ai-warden scan-skill https://github.com/user/skill --offline --strict

# API-powered scanning (requires API key)
ai-warden scan-skill https://github.com/user/skill
ai-warden scan-skill https://github.com/user/skill --json --strict
```

**Options:**
- `--offline` - Use local scanner only (free, no API key)
- `--json` - Machine-readable JSON output
- `--strict` - Exit code 1 unless verdict is SAFE
- `--mode MODE` - Detection mode (strict/balanced/permissive)

#### Example Output

```
🔍 AI-Warden Skill Scan
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
  Skill:    smart-web-search
  Source:   github.com/davidme6/smart-web-search
  Files:    4 scanned
  Mode:     offline

  LICENSE                  ✅ Safe       (0.00)
  README.md                ❌ CRITICAL   (1.00)
    ├─ P102: Data Forwarding Instructions [CRITICAL] — "Email**: smart-web-search@feedback.com"
    └─ H003: Excessive External URLs [LOW] — "Found 11 external URLs"
  SKILL.md                 ✅ Safe       (0.19)
    └─ H003: Excessive External URLs [LOW] — "Found 20 external URLs"
  _meta.json               ✅ Safe       (0.00)

  Verdict:     ❌ DANGEROUS
  Trust Score: 0/100
  Scan Time:   1.2s
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
```

#### Verdicts

| Verdict | Trust Score | Meaning |
|---------|-------------|---------|
| ✅ SAFE | 70-100 | No threats detected |
| ⚠️ WARNING | 25-69 | Suspicious patterns found, review recommended |
| ❌ DANGEROUS | 0-24 | Active threats detected, do not install |

#### Offline vs API Mode

| | Offline (free) | API (metered) |
|---|---|---|
| Detection | Regex patterns | Judge Mars ML + patterns |
| Speed | Instant | ~150ms/file |
| False positives | Higher | Lower |
| Zero-day threats | ❌ | ✅ |
| Requires API key | No | Yes |

#### Python API

```python
from ai_warden import AIWarden

warden = AIWarden()

# Offline scan
result = warden.scan_skill("https://github.com/user/skill", offline=True)
print(result["verdict"])     # SAFE, WARNING, or DANGEROUS
print(result["trustScore"])  # 0-100

# API-powered scan
result = warden.scan_skill("https://github.com/user/skill")
for f in result["files"]:
    print(f"{f['path']}: {f['riskLevel']} ({f['score']})")
```

---

### `ai-warden status`

Show authentication status and usage.

```bash
ai-warden status

# Output:
# ╭─── AI-Warden Status ────────────────────────────╮
# │ ✅ Authenticated                                 │
# │                                                  │
# │ API Key:   sk_live_...abc (valid)               │
# │ API URL:   http://46.62.240.255:8080/api        │
# │ Tier:      Free                                  │
# │ Usage:     42 / 1000 requests (4%)              │
# │ Remaining: 958 requests this month              │
# ╰──────────────────────────────────────────────────╯
```

---

### `ai-warden logout`

Remove stored credentials.

```bash
ai-warden logout

# Output:
# 🗑️ Credentials removed
# Run 'ai-warden login' to authenticate again
```

---

## 🌐 Middleware

### FastAPI

```python
from fastapi import FastAPI
from ai_warden.middleware import FastAPIMiddleware

app = FastAPI()

app.add_middleware(
    FastAPIMiddleware,
    api_key="sk_live_xxx",    # Or load from env
    block_unsafe=True,         # Return 400 on unsafe prompts
    log_threats=True,          # Log detected threats
    exclude_paths=["/health"]  # Skip validation for these paths
)

@app.post("/chat")
async def chat(prompt: str):
    # Middleware validates prompt before reaching here
    return {"response": "Safe!"}
```

**How it works:**
1. Middleware intercepts POST/PUT/PATCH requests
2. Extracts text fields from JSON body
3. Validates all text content
4. Returns 400 if unsafe (when `block_unsafe=True`)
5. Or adds warning header and continues

---

### Django

```python
# settings.py
MIDDLEWARE = [
    'ai_warden.middleware.django.AIWardenMiddleware',
    # ... other middleware
]

AI_WARDEN_API_KEY = "sk_live_xxx"
AI_WARDEN_BLOCK_UNSAFE = True
AI_WARDEN_LOG_THREATS = True
AI_WARDEN_EXCLUDE_PATHS = ["/admin/", "/static/"]
```

---

### Flask

```python
from flask import Flask, request
from ai_warden.middleware import flask_protect

app = Flask(__name__)

@app.route('/chat', methods=['POST'])
@flask_protect(api_key="sk_live_xxx", block_unsafe=True)
def chat():
    prompt = request.json['prompt']
    # Decorator validates prompt before function runs
    return {"response": "Safe!"}
```

---

## 🚀 Advanced Usage

### Custom Validation Logic

```python
from ai_warden import AIWarden

warden = AIWarden()

def validate_user_input(text: str) -> bool:
    """Custom validation with additional checks."""
    # AI-Warden validation
    result = warden.validate(text)
    
    if not result.is_safe:
        print(f"Blocked: {result.threat_type}")
        return False
    
    # Additional custom checks
    if len(text) > 10000:
        print("Blocked: Too long")
        return False
    
    return True
```

### Error Handling

```python
from ai_warden import AIWarden
from ai_warden.exceptions import (
    AuthenticationError,
    ValidationError,
    APIError
)

warden = AIWarden()

try:
    result = warden.validate("text")
except AuthenticationError:
    print("Invalid API key")
except ValidationError as e:
    print(f"Validation failed: {e}")
except APIError as e:
    print(f"API error: {e}")
```

### Custom API URL

```python
warden = AIWarden(
    api_key="sk_live_xxx",
    api_url="https://your-custom-domain.com",
    timeout=60  # Custom timeout in seconds
)
```

### Usage Statistics

```python
warden = AIWarden()

usage = warden.get_usage()

print(f"Tier: {usage['tier']}")
print(f"Usage: {usage['usage']} / {usage['limit']}")
print(f"Remaining: {usage['limit'] - usage['usage']}")
```

---

## 🧪 Testing

Run tests with pytest:

```bash
pip install -e ".[dev]"
pytest
```

With coverage:

```bash
pytest --cov=ai_warden --cov-report=html
```

---

## 📦 Installation Options

### Basic Installation

```bash
pip install ai-warden
```

### With Async Support

```bash
pip install ai-warden[async]
```

### With Framework Support

```bash
pip install ai-warden[fastapi]
pip install ai-warden[django]
pip install ai-warden[flask]
```

### All Features

```bash
pip install ai-warden[async,fastapi,django,flask,secure]
```

### Development

```bash
pip install -e ".[dev]"
```

---

## 🤝 Contributing

Contributions are welcome! Please follow these steps:

1. Fork the repository
2. Create a feature branch: `git checkout -b feature-name`
3. Make your changes
4. Run tests: `pytest`
5. Format code: `black ai_warden/`
6. Submit a pull request

---

## 📄 License

MIT License - see [LICENSE](LICENSE) for details.

---

## 🔗 Links

- **Documentation**: https://github.com/ai-warden/ai-warden-python#readme
- **Bug Reports**: https://github.com/ai-warden/ai-warden-python/issues
- **Source Code**: https://github.com/ai-warden/ai-warden-python

---

## 🛡️ Security

If you discover a security vulnerability, please email security@ai-warden.com instead of using the issue tracker.

---

## 🙏 Acknowledgments

Built with ❤️ using:
- [Click](https://click.palletsprojects.com/) - Beautiful command-line interfaces
- [Rich](https://rich.readthedocs.io/) - Rich terminal formatting
- [Pydantic](https://pydantic-docs.helpmanual.io/) - Data validation
- [Requests](https://requests.readthedocs.io/) - HTTP client
- [httpx](https://www.python-httpx.org/) - Async HTTP client

---

**Made with 🛡️ by the AI-Warden team**
