Metadata-Version: 2.4
Name: snowglobe-sdk
Version: 0.5.0a0
Summary: SDK for the Snowglobe API
Author-email: Guardrails AI <contact@guardrailsai.com>
License: MIT License
        
        Copyright (c) 2024 Guardrails AI
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Requires-Python: <4,>=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: httpx>=0.27.0
Requires-Dist: pydantic<3.0.0,>=2.0.0
Requires-Dist: tenacity>=9.0.0
Provides-Extra: dev
Requires-Dist: ruff; extra == "dev"
Requires-Dist: pre-commit>=4.1.0; extra == "dev"
Requires-Dist: coverage>=7.6.12; extra == "dev"
Requires-Dist: pyright[nodejs]>=1.1.396; extra == "dev"
Dynamic: license-file

# Snowglobe Python SDK

Python SDK for the [Snowglobe](https://guardrailsai.com/snowglobe) API.

## Requirements

- Python 3.11+

## Installation

```bash
pip install snowglobe-sdk
```

## Initialization

```python
import asyncio
from snowglobe.sdk import SnowglobeClient

client = SnowglobeClient(
    api_key="your-api-key",
)
```

**With an organization ID** (multi-tenant):

```python
client = SnowglobeClient(
    api_key="your-api-key",
    organization_id="org-123",
)
```

**With a custom base URL:**

```python
client = SnowglobeClient(
    api_key="your-api-key",
    base_url="https://your-instance.snowglobe.guardrailsai.com",
)
```

**With a custom `httpx.AsyncClient`** (e.g. to configure timeouts or proxies):

```python
import httpx
from snowglobe.sdk import SnowglobeClient

http = httpx.AsyncClient(timeout=30.0)
client = SnowglobeClient(api_key="your-api-key", http_client=http)
```

All methods are `async` and must be awaited inside an async context.

```python
async def main():
    client = SnowglobeClient(api_key="your-api-key")
    agents = await client.agents.list_agents()
    print(agents)

asyncio.run(main())
```

---

## Namespaces

The client exposes four namespaced API objects:

| Namespace | Attribute | Description |
|---|---|---|
| Agents | `client.agents` | Create and manage AI agents (applications) |
| Simulations | `client.simulations` | Run and inspect simulations |
| Metrics | `client.metrics` | Browse risk/metric definitions |
| Documents | `client.documents` | Upload context documents |

---

## `client.agents`

Manage AI agents (applications under test).

### `create_agent(body)`

Create a new agent.

```python
agent = await client.agents.create_agent({
    "name": "My Chatbot",
    "icon": "🤖",
    "description": "Customer support bot",
    "connection_info": {
        "api_key_ref": "OPENAI_API_KEY",
        "model_name": "gpt-4o",
        "system_prompt": "You are a helpful assistant.",
    },
})
print(agent.id)
```

### `list_agents()`

List all agents in the organization.

```python
agents = await client.agents.list_agents()
for agent in agents:
    print(agent.id, agent.name)
```

### `get_agent(id)`

Fetch a single agent by ID.

```python
agent = await client.agents.get_agent("agent-123")
```

### `update_agent(id, body)`

Update an existing agent.

```python
agent = await client.agents.update_agent("agent-123", {
    "name": "Updated Name",
    "description": "New description",
})
```

### `delete_agent(id)`

Delete an agent.

```python
await client.agents.delete_agent("agent-123")
```

---

## `client.simulations`

Run simulations and retrieve results.

### `create_simulation(body, *, as_draft=None)`

Create a new simulation.

```python
simulation = await client.simulations.create_simulation({
    "name": "Q4 Red Team",
    "role": "assistant",
    "is_template": False,
    "risks": [{"id": "risk-123", "name": "Bias", "type": "LLM", "version": 1}],
})
print(simulation.id, simulation.state)
```

Pass `as_draft="true"` to save without launching.

### `list_simulations(*, limit=None)`

List simulations, optionally capped at `limit` results.

```python
simulations = await client.simulations.list_simulations(limit=20)
```

### `get_simulation(id, *, exclude_calculated_status=None)`

Fetch a simulation by ID.

```python
sim = await client.simulations.get_simulation("sim-123")
print(sim.state, sim.statistics)
```

### `update_simulation(id, body, *, as_draft=None)`

Update a simulation's configuration.

```python
sim = await client.simulations.update_simulation("sim-123", {
    "name": "Updated Sim Name",
})
```

### `delete_simulation(id)`

Delete a simulation.

```python
await client.simulations.delete_simulation("sim-123")
```

### `update_simulation_settings(id, body, *, access=None)`

Update settings (e.g. auto-approve personas) for a simulation.

```python
settings = await client.simulations.update_simulation_settings("sim-123", {
    "autoApprovePersonas": True,
})
```

### `download_simulation_data(id)`

Download the full conversation dataset for a completed simulation.

```python
data = await client.simulations.download_simulation_data("sim-123")
for conversation in data:
    print(conversation.persona, len(conversation.messages), "turns")
```

### `get_simulation_diversity_metrics(experiment_id)`

Retrieve diversity metrics for a simulation.

```python
metrics = await client.simulations.get_simulation_diversity_metrics("sim-123")
```

### `create_simulation_diversity_metrics(experiment_id, body)`

Create a diversity metrics record for a simulation.

```python
metrics = await client.simulations.create_simulation_diversity_metrics(
    "sim-123", body
)
```

### `update_simulation_diversity_metrics(experiment_id, body)`

Update an existing diversity metrics record.

```python
metrics = await client.simulations.update_simulation_diversity_metrics(
    "sim-123", body
)
```

### `create_conversation_group(body)`

Group conversations for analysis.

```python
group = await client.simulations.create_conversation_group({
    "name": "Jailbreak attempts",
    "experiment_id": "sim-123",
})
```

### `update_conversation_group(id, body)`

Update findings or metadata on a conversation group.

```python
group = await client.simulations.update_conversation_group("group-123", {
    "findings": {"severity": "high", "note": "Consistent prompt injection pattern"},
})
```

### `get_test(id, test_id, *, include_embedding=None)`

Fetch a single test (conversation turn) with its risk evaluations.

```python
test = await client.simulations.get_test("sim-123", "test-456")
print(test.prompt, test.response)
for evaluation in test.risk_evaluations:
    print(evaluation.risk_type, evaluation.risk_triggered)
```

### `batch_create_risk_evaluations(id, test_id, body, *, tests=None)`

Submit risk evaluation results for a test in bulk.

```python
evaluations = await client.simulations.batch_create_risk_evaluations(
    "sim-123",
    "test-456",
    [
        {
            "risk_type": "bias",
            "risk_triggered": False,
            "confidence": 95,
            "judge_response": "No bias detected.",
        }
    ],
)
```

### `list_conversations_for_test(id, test_id, *, start_from=None, app_id=None, include_adaptability_messages=None)`

List the full conversation history associated with a test.

```python
conversations = await client.simulations.list_conversations_for_test(
    "sim-123", "test-456"
)
for conv in conversations:
    for msg in conv.messages:
        print(f"[{msg.role}] {msg.content}")
```

---

## `client.metrics`

Browse risk/metric definitions available in the platform.

### `list_metrics(*, lineage_id=None, version=None)`

List all available risk metrics, optionally filtered.

```python
risks = await client.metrics.list_metrics()
for risk in risks:
    print(risk.name, risk.type, f"v{risk.version}")
```

Filter by lineage:

```python
risks = await client.metrics.list_metrics(lineage_id="lineage-abc", version="2")
```

### `get_metric(id)`

Fetch a single risk metric by ID.

```python
risk = await client.metrics.get_metric("risk-123")
print(risk.name, risk.promptSource)
```

---

## `client.documents`

Upload reference documents to associate with agents or simulations.

### `upload_documents(body)`

Request pre-signed upload URLs for one or more files.

```python
response = await client.documents.upload_documents({
    "fileNames": ["policy.pdf", "faq.txt"],
})
for url_info in response.urls:
    print(url_info.id, url_info.url)
    # PUT your file bytes to url_info.url directly
```

---

## Return Types

All methods return instances of the method's respective return type as a pydantic model. The full model definitions live in `snowglobe.sdk.models` and can be imported directly for type annotations.

```python
from snowglobe.sdk.models import Agent, Simulation, Risk
```

## Error Handling

Network or API errors raise `httpx.HTTPStatusError`. All methods automatically retry up to 5 times with exponential backoff (powered by [tenacity](https://tenacity.readthedocs.io)).

```python
import httpx

try:
    agent = await client.agents.get_agent("nonexistent-id")
except httpx.HTTPStatusError as e:
    print(e.response.status_code, e.response.text)
```

---

## Core Components

### 1. Agent Creation

Agents represent the AI systems you want to test. Each agent requires a name, icon, and connection info describing how Snowglobe should reach your LLM.

```python
agent = await client.agents.create_agent({
    "name": "customer_support_agent",
    "description": "Customer support agent for Amatto's pizza restaurant",
    "icon": "pizza",
    "connection_info": {
        "provider": "OpenAI",
        "endpoint": "",
        "model_name": "openai/gpt-4o",
        "api_key_ref": "OPENAI_API_KEY",
        "system_prompt": "You are a helpful expert customer support agent for Amatto's pizza.",
        "extra_body": [],
        "extra_headers": [],
    },
})
agent_id = agent.id
```

> **Note:** Agents using a custom code integration via `snowglobe-connect` require a two-step setup:
> 1. Create the agent with the API as shown above.
> 2. Map the agent's ID in your `snowglobe-connect` deployment's `agents.json` file.
> 3. Run `snowglobe-connect start`.

### 2. Simulation Configuration

Simulations define how conversations are generated and evaluated against your agent.

```python
simulation_config = {
    "name": "continuous integration simulation",
    "role": "Customer support agent for Amatto's pizza restaurant",
    "user_description": "",
    "use_cases": "",
    "generation_status": "pending",
    "evaluation_status": "pending",
    "validation_status": "pending",
    "source_data": {
        "docs": {
            "misc": [],
            "knowledge_base": [],
            "historical_data": []
        },
        "evaluation_configuration": {
            "No Financial Advice": {
                "id": "e5af8dee-6d8d-4144-b754-204d24879ec9",
                "name": "No Financial Advice",
                "version": 1,
                "metadata": {},
            },
        },
        "generation_configuration": {
            "max_topics": 1,
            "max_personas": 4,
            "branching_factor": 25,
            "max_conversations": 500,
            "max_conversation_length": 4,
            "continue_conversations_from_adapted_messages": False,
            "data_gen_mode": "coverage_focused_v3",
            "intent": "",
            "persona_topic_generators": [
                {
                    "name": "app_description_system_prompt",
                    "settings": {
                        "max_personas": 4
                    }
                }
            ],
            "min_conversation_length": 1
        }
    },
    "is_template": False
}
```

---

## CI Integration Workflow

### Step 1: Create and Configure Agent

```python
async def create_agent(client, agent_config):
    """Create a new agent for testing."""
    agent = await client.agents.create_agent(agent_config)
    print(f"Agent created successfully: {agent.id}")
    return agent.id
```

### Step 2: Launch Simulation

```python
async def launch_simulation(client, agent_id, simulation_config):
    """Create and launch a new simulation."""
    config = {
        **simulation_config,
        "application_id": agent_id,
        "app_id": agent_id,
    }
    simulation = await client.simulations.create_simulation(config)
    simulation_id = simulation.id

    # Auto-approve personas for CI automation
    await client.simulations.update_simulation_settings(
        simulation_id, {"autoApprovePersonas": True}
    )
    return simulation_id
```

### Step 3: Monitor Simulation Progress

```python
import asyncio

async def wait_for_completion(client, simulation_id, timeout_minutes=20):
    """Poll simulation until completion or timeout."""
    max_attempts = timeout_minutes * 6  # poll every 10 seconds

    for _ in range(max_attempts):
        sim = await client.simulations.get_simulation(simulation_id)
        state_num = sim.state_num
        print(f"Simulation state: {state_num}")

        # state_num >= 17 indicates completion (see Simulation States below)
        if state_num is not None and state_num >= 17:
            print("Simulation completed successfully")
            return True

        await asyncio.sleep(10)

    raise TimeoutError("Simulation timed out before completion")
```

### Step 4: Retrieve Results

```python
import json

async def download_results(client, simulation_id):
    """Download simulation results for analysis."""
    data = await client.simulations.download_simulation_data(simulation_id)

    filename = f"{simulation_id}_results.json"
    with open(filename, "w") as f:
        json.dump([d.model_dump() for d in data], f, indent=2)

    print(f"Results saved to {filename}")
    return filename
```

---

## Complete CI Integration Example

```python
import asyncio
import json
import os

from snowglobe.sdk import SnowglobeClient

async def run_agent_simulation(agent_config, simulation_config):
    """Complete workflow for running agent simulations in CI."""
    client = SnowglobeClient(
        api_key=os.environ["SNOWGLOBE_API_KEY"],
        organization_id=os.environ.get("SNOWGLOBE_ORG_ID"),
    )

    try:
        agent_id = await create_agent(client, agent_config)
        simulation_id = await launch_simulation(client, agent_id, simulation_config)
        await wait_for_completion(client, simulation_id)
        results_file = await download_results(client, simulation_id)

        return {
            "success": True,
            "agent_id": agent_id,
            "simulation_id": simulation_id,
            "results_file": results_file,
        }

    except Exception as e:
        print(f"Simulation failed: {e}")
        return {"success": False, "error": str(e)}


asyncio.run(
    run_agent_simulation(agent_config, simulation_config)
)
```

### Error Handling

```python
async def robust_simulation_run(agent_config, simulation_config):
    """Simulation run with comprehensive error handling."""
    import httpx

    try:
        return await run_agent_simulation(agent_config, simulation_config)
    except ValidationError as e:
        print(f"Configuration validation failed: {e}")
        return {"success": False, "error": "validation", "details": str(e)}
    except httpx.HTTPStatusError as e:
        print(f"API error {e.response.status_code}: {e.response.text}")
        return {"success": False, "error": "api_error", "details": str(e)}
    except TimeoutError as e:
        print(f"Simulation timed out: {e}")
        return {"success": False, "error": "timeout", "details": str(e)}
    except Exception as e:
        print(f"Unexpected error: {e}")
        return {"success": False, "error": "unexpected", "details": str(e)}
```

### CI Pipeline Integration

```yaml
# Example GitHub Actions workflow
name: Agent Testing
on: [push, pull_request]

jobs:
  test-agent:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
      - name: Install dependencies
        run: pip install snowglobe-sdk
      - name: Run agent simulation
        env:
          SNOWGLOBE_API_KEY: ${{ secrets.SNOWGLOBE_API_KEY }}
          SNOWGLOBE_ORG_ID: ${{ secrets.SNOWGLOBE_ORG_ID }}
          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
        run: python test_agent_simulation.py
```

---

## Simulation States

The `state_num` field on a simulation indicates its current phase:

| `state_num` | State name | Description |
|---|---|---|
| 0–2 | Draft / Queued | Simulation created, waiting to start |
| 3–5 | Experiment started | Initialization and setup |
| 6–8 | Generation in progress | Persona, topic, and conversation generation |
| 9–11 | Evaluation in progress | Agent responses being judged against risks |
| 12–14 | Validation in progress | Results validated |
| 15–16 | Adaptation in progress | Adapted (adversarial) conversations being generated |
| 17+ | Experiment completed | Results available; `download_simulation_data` can be called |

Poll `sim.state_num` and wait for `>= 17` before downloading results.

---

## Troubleshooting

### Common Issues

**Authentication errors**
- Verify your API key and organization ID are correct.
- Confirm the `x-api-key` header is being sent (the SDK sets this automatically from `api_key`).
- Check network connectivity to your control plane URL.

**Simulation failures**
- Review the agent's `connection_info` — all required fields (`api_key_ref`, `model_name`, `system_prompt`) must be present.
- Verify the LLM provider API key referenced by `api_key_ref` is valid and has sufficient quota.
- Check that `source_data.generation_configuration` values are within acceptable ranges.

**Timeout issues**
- Increase `timeout_minutes` in `wait_for_completion` for large simulations.
- Reduce `max_personas`, `max_topics`, or `branching_factor` for faster runs.

### Debugging Tips

```python
async def debug_simulation(client, simulation_id):
    """Print detailed simulation status for debugging."""
    sim = await client.simulations.get_simulation(simulation_id)

    print(f"State:             {sim.state} (num={sim.state_num})")
    print(f"Generation status: {sim.generation_status}")
    print(f"Evaluation status: {sim.evaluation_status}")
    print(f"Validation status: {sim.validation_status}")
    print(f"Status reason:     {sim.status_reason}")
    print(f"Statistics:        {sim.statistics}")
```
