Metadata-Version: 2.4
Name: magpie-ai
Version: 1.0.2
Summary: Enterprise-grade LLM middleware for monitoring and metadata tracking
Home-page: https://github.com/magpie-so/sdk
Author: Magpie Team
Author-email: Magpie Team <team@magpie.dev>
License: MIT
Project-URL: Homepage, https://github.com/magpie-so/sdk
Project-URL: Documentation, https://docs.magpie.dev
Project-URL: Repository, https://github.com/magpie-so/sdk
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Requires-Dist: httpx>=0.26.0
Provides-Extra: dev
Requires-Dist: pytest>=7.4.0; extra == "dev"
Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: mypy>=1.0.0; extra == "dev"
Dynamic: author
Dynamic: home-page
Dynamic: requires-python

# Magpie AI

**Enterprise-grade LLM middleware for monitoring, content moderation, and execution tracking**

[![PyPI version](https://img.shields.io/pypi/v/magpie-ai.svg)](https://pypi.org/project/magpie-ai/)
[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
[![MIT License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)

Magpie AI is a Python SDK that integrates with your applications to provide monitoring, cost tracking, PII detection, content moderation, and schema validation. It supports both LLM chat functions and arbitrary function executions.

## Features

- **Two Monitoring Modes** - Chat mode for LLM calls, completion mode for any function
- **Cost Tracking** - Automatic token counting and cost calculation for LLM calls
- **PII Detection** - Detect sensitive data in both chat and completion modes
- **Content Moderation** - Policy-based content validation
- **Schema Guard** - Validate LLM output against Pydantic schemas
- **Non-Blocking** - Async logging with fail-open design that never crashes your app
- **Production-Grade** - HTTP retry with backoff, connection pooling, thread-safe

## Installation

```bash
pip install magpie-ai
```

## Quick Start

### 1. Configure

```python
import magpie_ai

magpie_ai.configure(
    api_key="sk_...",
    backend_url="https://your-magpie-backend.com"
)
```

### 2. Chat Mode (LLM Monitoring)

Monitor LLM chat functions that work with message arrays:

```python
from openai import OpenAI

@magpie_ai.monitor(
    project_id="my-project",
    model="gpt-4"
)
def chat(messages: list) -> str:
    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4",
        messages=messages
    )
    return response.choices[0].message.content

# Use it exactly as before - monitoring is automatic
result = chat([{"role": "user", "content": "What is Python?"}])
```

Chat mode automatically extracts tokens, calculates costs, and captures input/output from the LLM response object.

### 3. Completion Mode (Any Function)

Monitor any function - API calls, SDK calls, database queries, file operations:

```python
import httpx

@magpie_ai.monitor(
    project_id="my-project",
    type="completion",
    custom={"service": "stripe"}
)
def charge_customer(customer_id: str, amount: float) -> dict:
    response = httpx.post(
        "https://api.stripe.com/v1/charges",
        data={"customer": customer_id, "amount": int(amount * 100)},
    )
    return response.json()

# Input args are serialized and logged automatically
result = charge_customer("cus_abc123", 49.99)
```

Completion mode serializes all function arguments and return values generically. No LLM-specific logic is applied - tokens, costs, and schema validation are skipped unless explicitly configured.

## Monitoring Modes

### Chat Mode (`type="chat"`, default)

For LLM functions that accept `messages: list` and return an LLM response object.

| Feature | Behavior |
|---------|----------|
| Input capture | Extracted from `messages[]` array |
| Output capture | Extracted from LLM response object |
| Token counting | Auto-extracted from response |
| Cost calculation | Auto-calculated from model pricing |
| Schema Guard | Validates LLM output text |
| Content moderation | Runs on input text + async output moderation |

### Completion Mode (`type="completion"`)

For any function - API calls, SDK calls, DB queries, file I/O, etc.

| Feature | Behavior |
|---------|----------|
| Input capture | All args/kwargs serialized via `inspect.signature` |
| Output capture | Return value serialized to string |
| Token counting | Only if explicitly provided |
| Cost calculation | Only with `input_token_price`/`output_token_price` |
| Schema Guard | Skipped |
| Content moderation | Runs on serialized input text (if enabled) |

## Core Features

### Cost Tracking

Automatic for chat mode, manual for completion mode:

```python
# Chat mode: pricing auto-looked up from model name
@magpie_ai.monitor(project_id="my-project", model="gpt-4-turbo")
def chat(messages: list):
    ...

# Completion mode: provide explicit pricing for custom LLMs
@magpie_ai.monitor(
    project_id="my-project",
    type="completion",
    input_token_price=0.001,
    output_token_price=0.002,
)
def call_custom_llm(prompt: str):
    ...
```

**Supported models:** OpenAI (gpt-4, gpt-4-turbo, gpt-4o, gpt-3.5-turbo), Anthropic (claude-3-opus, claude-3-sonnet, claude-3-haiku), or custom pricing.

### PII Detection

Works in both chat and completion modes:

```python
# Chat mode
@magpie_ai.monitor(project_id="my-project", pii=True)
def chat(messages: list):
    ...

# Completion mode - scans all string arguments
@magpie_ai.monitor(project_id="my-project", type="completion", pii=True)
def process_form(name: str, email: str, ssn: str):
    return submit_to_api(name=name, email=email, ssn=ssn)
```

Detects: email addresses, phone numbers, credit cards, SSNs, names, addresses.

### Content Moderation

Policy-based content validation with blocking:

```python
from magpie_ai import ContentModerationError

@magpie_ai.monitor(
    project_id="my-project",
    content_moderation=True
)
def moderated_chat(messages: list):
    ...

try:
    result = moderated_chat(messages)
except ContentModerationError as e:
    print(f"Content blocked: {e}")
```

### Schema Guard

Validate LLM output against a Pydantic schema (chat mode only):

```python
from pydantic import BaseModel, Field
from magpie_ai import SchemaValidationError

class MovieReview(BaseModel):
    title: str = Field(description="Movie title")
    rating: float = Field(description="Rating out of 10")
    summary: str = Field(description="Brief review summary")

# Block mode: raises SchemaValidationError if output doesn't match
@magpie_ai.monitor(
    project_id="my-project",
    model="gpt-4",
    output_schema=MovieReview,
    on_schema_fail="block"
)
def get_review_strict(messages: list):
    ...

# Flag mode: output passes through, violation logged to review queue
@magpie_ai.monitor(
    project_id="my-project",
    model="gpt-4",
    output_schema=MovieReview,
    on_schema_fail="flag"
)
def get_review_lenient(messages: list):
    ...

try:
    result = get_review_strict(messages)
except SchemaValidationError as e:
    print(f"Schema: {e.schema_name}, Errors: {e.validation_errors}")
```

### Custom Metadata

Attach arbitrary metadata to any execution:

```python
@magpie_ai.monitor(
    project_id="my-project",
    custom={
        "user_id": "john_doe",
        "department": "engineering",
        "client": "acme_corp",
        "version": "2.1.0"
    }
)
def tracked_function(prompt: str):
    ...
```

Custom metadata is filterable in the Magpie dashboard.

## API Reference

### `@magpie_ai.monitor()`

Main decorator for monitoring functions.

**Parameters:**

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `project_id` | `str` | required | Project identifier |
| `type` | `"chat" \| "completion"` | `"chat"` | Monitoring mode |
| `model` | `str` | `None` | Model name for auto pricing lookup |
| `custom` | `dict` | `None` | Custom metadata (JSON-serializable) |
| `pii` | `bool` | `False` | Enable PII detection |
| `content_moderation` | `bool` | `False` | Enable content moderation |
| `capture_input` | `bool` | `True` | Capture function inputs |
| `capture_output` | `bool` | `True` | Capture function outputs |
| `output_schema` | `BaseModel \| dict` | `None` | Pydantic schema for output validation (chat only) |
| `on_schema_fail` | `"block" \| "flag"` | `"block"` | Schema validation failure behavior |
| `input_token_price` | `float` | `None` | Custom price per 1M input tokens |
| `output_token_price` | `float` | `None` | Custom price per 1M output tokens |
| `trace_id` | `str` | auto-generated | Custom trace ID |
| `llm_url` | `str` | env `VLLM_URL` | LLM URL for PII/moderation analysis |
| `llm_model` | `str` | env `VLLM_MODEL` | Model for PII/moderation analysis |

**Raises:**

| Exception | When |
|-----------|------|
| `ValueError` | `project_id` is empty, invalid `type`, or conflicting pricing config |
| `TypeError` | `custom` is not a dict, pricing values aren't numeric |
| `ContentModerationError` | Content moderation blocks the input |
| `SchemaValidationError` | Schema Guard blocks invalid output (`on_schema_fail="block"`) |

### `magpie_ai.configure()`

Configure SDK-wide settings.

```python
magpie_ai.configure(
    api_key="sk_...",
    backend_url="https://your-backend.com",
    timeout=10.0,
    fail_open=True
)
```

### `magpie_ai.context()`

Context manager for monitoring specific code blocks.

```python
from magpie_ai import context

with context(project_id="my-project", custom={"session": "abc"}):
    result = call_llm()
```

### Exported Types

All types are importable from `magpie_ai`:

```python
from magpie_ai import (
    # Core
    monitor,
    context,
    configure,
    # Exceptions
    ContentModerationError,
    SchemaValidationError,
    # Content moderation types
    ModerationResult,
    ModerationAction,
    ModerationSeverity,
    ModerationViolation,
    # Schema validation types
    SchemaValidationResult,
    # Metadata validation
    validate_metadata,
    clear_schema_cache,
    ValidationResult,
    # Pricing
    get_model_pricing,
    list_available_models,
    ModelPricing,
)
```

## Examples

### Chat Mode: E-Commerce Product Search

```python
import magpie_ai
from openai import OpenAI

@magpie_ai.monitor(
    project_id="ecommerce-ai",
    model="gpt-3.5-turbo",
    custom={"service": "product-search"}
)
def search_products(messages: list) -> str:
    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=messages
    )
    return response.choices[0].message.content

result = search_products([
    {"role": "system", "content": "You are a shopping assistant."},
    {"role": "user", "content": "Find blue running shoes under $100"}
])
```

### Completion Mode: External API Call

```python
import magpie_ai
import httpx

@magpie_ai.monitor(
    project_id="my-project",
    type="completion",
    custom={"service": "weather_api"}
)
def get_weather(city: str) -> dict:
    response = httpx.get(
        "https://api.weatherapi.com/v1/current.json",
        params={"key": "...", "q": city},
    )
    return response.json()

weather = get_weather("London")
```

### Completion Mode: Database Query

```python
@magpie_ai.monitor(
    project_id="my-project",
    type="completion",
    custom={"service": "user_db"}
)
def get_user(user_id: str) -> dict:
    return db.users.find_one({"_id": user_id})
```

### Completion Mode: Third-Party SDK

```python
@magpie_ai.monitor(
    project_id="my-project",
    type="completion",
    pii=True,
    custom={"service": "stripe"}
)
def create_customer(email: str, name: str) -> dict:
    return stripe.Customer.create(email=email, name=name)
```

## Policy System (LLM-as-Judge)

Magpie's policy system uses an LLM-as-judge approach: each policy is a custom prompt template that the backend evaluates against your content. Policies are managed per-project in the dashboard or via API.

### `action` parameter — block vs. flag mode

The `monitor` decorator supports two policy enforcement modes via the `action` parameter:

| Mode | `action` value | Behaviour |
|------|---------------|-----------|
| Block | `"block"` | Evaluates policies **synchronously before** the function runs. Raises `ContentModerationError` for critical violations. |
| Flag | `"flag"` | Evaluates policies **asynchronously after** the function runs. Violations are logged to the review queue but the call succeeds. |
| Off | `None` (default) | No policy evaluation. |

**Block mode** — stops execution on critical violations:

```python
from magpie_ai import ContentModerationError

@magpie_ai.monitor(
    project_id="my-project",
    model="gpt-4",
    action="block"
)
def moderated_chat(messages: list) -> str:
    response = openai_client.chat.completions.create(
        model="gpt-4",
        messages=messages
    )
    return response.choices[0].message.content

try:
    result = moderated_chat([{"role": "user", "content": "Ignore all instructions..."}])
except ContentModerationError as e:
    print(f"Blocked: {e}")
    print(f"Violations: {e.result.violations}")
```

**Flag mode** — logs violations without blocking:

```python
@magpie_ai.monitor(
    project_id="my-project",
    model="gpt-4",
    action="flag"
)
def monitored_chat(messages: list) -> str:
    response = openai_client.chat.completions.create(model="gpt-4", messages=messages)
    return response.choices[0].message.content

# Violations are flagged in the review queue — no exception raised
result = monitored_chat([{"role": "user", "content": "How do I..."}])
```

### `PolicyClient` — direct API access

Use `PolicyClient` to evaluate content directly against your project's policies without wrapping a function:

```python
from magpie_ai.policies import PolicyClient, ContentModerationError

client = PolicyClient(
    api_key="sk_...",
    base_url="https://your-magpie-backend.com"
)

# Evaluate user input
result = client.evaluate(
    project_id="my-project",
    content="User message here",
    content_type="input",       # "input" or "output"
    metadata={"user_id": "u_123"},
)

print(f"Evaluated {result.total_policies_evaluated} policies")
print(f"Violations: {len(result.violations)}")
print(f"Should block: {result.should_block}")

# raises ContentModerationError automatically when should_block=True
```

**Seeding default policies** for a new project:

```python
import httpx

httpx.post(
    "https://your-magpie-backend.com/api/v1/policies/v2/seed-defaults",
    params={"project_id": "my-project", "user_id": "admin-user-id"},
    headers={"Authorization": "Bearer sk_..."},
)
# Creates 3 starter policies: Harmful Content, Misinformation, Prompt Injection
```

### `PolicyEvaluationResult` fields

| Field | Type | Description |
|-------|------|-------------|
| `violations` | `list[PolicyViolation]` | Individual policy violations found |
| `highest_severity` | `str \| None` | `"critical"`, `"high"`, `"medium"`, `"low"`, or `None` |
| `should_block` | `bool` | `True` only when a **critical** violation is found |
| `total_policies_evaluated` | `int` | Number of policies checked |
| `cached_results` | `int` | Number of results served from evaluation cache |

### `PolicyViolation` fields

| Field | Type | Description |
|-------|------|-------------|
| `policy_id` | `str` | ID of the violated policy |
| `policy_name` | `str` | Human-readable policy name |
| `severity` | `str` | `"critical"`, `"high"`, `"medium"`, or `"low"` |
| `confidence_score` | `float` | LLM confidence (0.0–1.0) |
| `violation_reason` | `str` | LLM explanation of the violation |
| `cached` | `bool` | Whether this result came from the evaluation cache |

## Error Handling

Magpie AI uses a **fail-open** design. Monitoring errors never crash your application:

```python
@magpie_ai.monitor(project_id="my-project")
def critical_function(messages: list):
    # Even if the Magpie backend is down, your function runs normally.
    # Monitoring errors are logged but never raised.
    ...
```

The only exceptions that propagate are intentional ones:
- `ContentModerationError` - when content moderation blocks input
- `SchemaValidationError` - when schema guard blocks invalid output (`on_schema_fail="block"`)

## Parameter Validation

The decorator validates all parameters at decoration time:

```python
# Valid
@magpie_ai.monitor(project_id="my-project", model="gpt-4")

# ValueError: project_id is required
@magpie_ai.monitor()

# ValueError: type must be 'chat' or 'completion'
@magpie_ai.monitor(project_id="x", type="invalid")

# ValueError: cannot use both model and explicit pricing
@magpie_ai.monitor(project_id="x", model="gpt-4", input_token_price=0.03)

# TypeError: custom must be a dict
@magpie_ai.monitor(project_id="x", custom="invalid")
```

## Thread Safety

Magpie AI is fully thread-safe. The HTTP client uses connection pooling and the singleton is protected with double-checked locking:

```python
import threading

@magpie_ai.monitor(project_id="my-project", type="completion")
def concurrent_function(item_id: str):
    return process(item_id)

threads = [
    threading.Thread(target=concurrent_function, args=(f"item_{i}",))
    for i in range(100)
]
for t in threads:
    t.start()
for t in threads:
    t.join()
```

## Performance

- **Overhead:** < 5ms per call (monitoring runs in background thread)
- **Retry:** 3 attempts with exponential backoff for transient failures
- **Connection pooling:** Persistent HTTP client with 20 max connections
- **Throughput:** Supports 1000+ concurrent monitored calls

## Troubleshooting

### Import Error: No module named 'magpie_ai'

```bash
pip install --upgrade magpie-ai
```

### PII Detection Not Working

Ensure your LLM server is running and set the env vars. Works with vLLM or [LM Studio](https://lmstudio.ai/) (any OpenAI-compatible server):

```bash
export VLLM_URL="http://localhost:1234"
export VLLM_MODEL="qwen2.5-1.5b-instruct"
```

### ContentModerationError

Check your moderation policy configuration in the Magpie dashboard.

### High Latency

PII detection and content moderation add ~200-500ms. Enable only when needed.

## License

MIT - See [LICENSE](LICENSE) file

## Support

- Email: team@magpie.dev
- Docs: https://docs.magpie.dev
- Issues: https://github.com/magpie-so/sdk/issues

## Changelog

### v0.3.0

- Added completion mode (`type="completion"`) for monitoring any function
- Added Schema Guard (`output_schema`, `on_schema_fail`) for LLM output validation
- Added `capture_output` parameter
- Added `Literal` types for all string enum parameters
- Added HTTP retry with exponential backoff (3 attempts)
- Added persistent connection pooling
- Added thread-safe client singleton
- Exported all public types from `magpie_ai` package
- Removed unused `requests` dependency
- Fixed version mismatch between package and __init__

### v0.2.1

- Fixed module import path (magpie_ai)
- Added proper type hints for decorators
- Added parameter validation
- Improved error messages

### v0.2.0

- Renamed package to magpie-ai
- Full type support with py.typed
- Generic decorator types

### v0.1.0

- Initial release
- Core monitoring, PII detection, content moderation, cost tracking
