Metadata-Version: 2.4
Name: ai-agentree-sdk
Version: 0.1.0
Summary: Python SDK for AIAgentree — trace, audit, and govern AI agent decisions
License-Expression: MIT
Project-URL: Homepage, https://ai-agentree.com
Project-URL: Documentation, https://ai-agentree.com/docs
Project-URL: Repository, https://github.com/argumentree/argumentree
Project-URL: Changelog, https://github.com/argumentree/argumentree/blob/master/sdk/python/CHANGELOG.md
Keywords: ai,agent,decision,tracing,governance,audit
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Topic :: Software Development :: Libraries
Classifier: Programming Language :: Python :: 3
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: Programming Language :: Python :: 3.13
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests>=2.28.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: responses>=0.23; extra == "dev"
Dynamic: license-file

# ai-agentree-sdk

Python SDK for [AI Agentree](https://ai-agentree.com) — trace, audit, and govern AI agent decisions.

**Zero-configuration decision tracing for any LLM.** Connect your AI agents to AI Agentree and capture every decision trace automatically. Works with OpenAI, Anthropic, Llama, Mistral, Nemotron, and any other LLM.

## Installation

```bash
pip install ai-agentree-sdk
```

For development from source:

```bash
cd sdk/python
pip install -e ".[dev]"
```

## Quick Start — Local Mode (No API Key Needed)

Try AIAgentree locally with zero config — no backend, no API key:

```python
from aiagentree import LocalTracer, TracedAgent

# Zero-config local tracing (console + file)
tracer = LocalTracer()
client = tracer.get_client()

with TracedAgent(client, workflow_id="claim_review", entity_id="CLM-4821") as agent:
    response = agent.chat(
        llm=your_llm_client,  # OpenAI, Anthropic, Ollama, LangChain, etc.
        message="Review this insurance claim for $15,000",
    )
    # Trace printed to console + saved to agentree-traces.jsonl
    agent.export_markdown("trace.md")
    print(agent.stats())
```

## Quick Start — Cloud Mode

Connect to AI Agentree for validation, precedent search, and dashboards:

```python
from aiagentree import AgentreeClient, TracedAgent

client = AgentreeClient(
    api_key="ask_your_key_here",
    base_url="http://localhost:5000",
    tenant_id="your-tenant-id",
)

with TracedAgent(client, workflow_id="claim_review", entity_id="CLM-4821") as agent:
    response = agent.chat(
        llm=your_llm_client,
        message="Review this insurance claim for $15,000",
    )
    print(f"Trace ID: {agent.trace_id}")
```

**What happens:**
1. Your LLM receives structured prompts to output reasoning in a traceable format
2. The SDK automatically extracts inputs, reasoning steps, and decisions
3. A complete trace is submitted to AI Agentree
4. The trace is transformed into an argument tree you can inspect in the UI

## Supported LLMs

| Provider | How to Use | Example |
|----------|------------|---------|
| **OpenAI** | `TracedAgent` or `TracedLLM` wrapper | [openai_example.py](examples/openai_example.py) |
| **Anthropic** | `TracedAgent` or tool use | [anthropic_example.py](examples/anthropic_example.py) |
| **Ollama** (Llama, Mistral, Nemotron) | `TracedAgent` | [ollama_example.py](examples/ollama_example.py) |
| **LangChain** | `TracedAgent` or callback | [langchain_example.py](examples/langchain_example.py) |
| **Any LLM** | Pass a callable to `TracedAgent` | See Generic Callable below |

## Integration Methods

### Method 1: TracedAgent (Recommended)

Context manager that handles everything:

```python
from aiagentree import AgentreeClient, TracedAgent

client = AgentreeClient(api_key="...", base_url="...", tenant_id="...")

with TracedAgent(
    client=client,
    workflow_id="order_review",
    entity_type="order",
    entity_id="ORD-123",
    title="Order Review Decision",
) as agent:
    response = agent.chat(
        llm=openai_client,  # or anthropic, ollama, langchain, etc.
        message="Should we approve this $500 order from a new customer?",
        context={"order_amount": 500, "customer_type": "new"},
    )

    # Access extracted reasoning
    print(agent.extracted_reasoning)
    # {'inputs': [...], 'steps': [...], 'decision': {...}}
```

### Method 2: instrument_llm (One-Liner Wrapper)

```python
from aiagentree import AgentreeClient, instrument_llm
from openai import OpenAI

client = AgentreeClient(api_key="...", base_url="...", tenant_id="...")
traced_llm = instrument_llm(OpenAI(), client, "review")

# Use exactly like normal OpenAI — tracing is automatic
response = traced_llm.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Review order ORD-123"}],
    _entity_id="ORD-123",  # Trace metadata
)
```

### Method 3: TracedLLM Wrapper (Full Options)

```python
from openai import OpenAI
from aiagentree import AgentreeClient, TracedLLM

openai_client = OpenAI()
agentree_client = AgentreeClient(api_key="...", base_url="...", tenant_id="...")

traced_llm = TracedLLM(
    llm=openai_client,
    client=agentree_client,
    workflow_id="review",
    agent_id="order-reviewer",
    auto_transform=True,
)

response = traced_llm.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Review order ORD-123"}],
    _entity_id="ORD-123",
)
```

### Method 4: Decorator

```python
from aiagentree import AgentreeClient, traced_decision

client = AgentreeClient(api_key="...", base_url="...", tenant_id="...")

@traced_decision(client, workflow_id="approval")
def approve_request(request_id: str):
    response = my_llm.chat(f"Should we approve {request_id}?")
    return {"response": response, "entity_id": request_id}

# Automatically traced!
approve_request("REQ-456")
```

### Method 5: Manual Control

Full control over trace construction:

```python
from aiagentree import AgentreeClient

client = AgentreeClient(api_key="...", base_url="...", tenant_id="...")

trace = client.start_trace(
    agent_id="my-agent",
    workflow_id="review",
    entity_type="order",
    entity_id="ORD-123",
)

# Add events manually
trace.add_input("order_amount", 500, source="database")
trace.add_step(temp_id="s1", title="Check amount", category="financial_threshold", confidence=0.9)
trace.add_step(temp_id="s2", title="Check history", category="customer_context", confidence=0.85)
trace.add_relation("s1", "s2", is_pro=True)

trace.seal(decision_id="d1", action="approve", confidence=0.92)
trace.transform()
```

## Using with MCP (Model Context Protocol)

For Claude and other MCP-enabled agents, provide AIAgentree tools:

```python
from aiagentree import get_mcp_tools

# Get tool definitions
tools = get_mcp_tools()
# Returns: [agentree_start_trace, agentree_add_input, agentree_add_reasoning_step, agentree_seal_decision]

# Provide to your MCP-enabled agent
# Claude will call these tools as it reasons!
```

## Using with Function Calling

For OpenAI/Anthropic function calling:

```python
from aiagentree import get_openai_function_schema, get_anthropic_tool_schema

# OpenAI
functions = [get_openai_function_schema()]
response = openai.chat.completions.create(
    model="gpt-4",
    messages=[...],
    functions=functions,
    function_call={"name": "submit_decision"},
)

# Anthropic
tools = [get_anthropic_tool_schema()]
response = anthropic.messages.create(
    model="claude-3-opus-20240229",
    messages=[...],
    tools=tools,
)
```

## Prompt Templates

Use built-in prompts for consistent structured output:

```python
from aiagentree import DECISION_SYSTEM_PROMPT, format_decision_prompt

# System prompt that instructs LLM to output structured reasoning
system = DECISION_SYSTEM_PROMPT

# Format user message with context
user_message = format_decision_prompt(
    task="Review this loan application",
    context={"amount": 50000, "credit_score": 720},
)

response = llm.chat(system=system, messages=[{"role": "user", "content": user_message}])
```

## Reasoning Extraction

Extract reasoning from any LLM output (JSON or freeform):

```python
from aiagentree import ReasoningExtractor

extractor = ReasoningExtractor()

# Works with JSON output
extracted = extractor.extract('{"reasoning_steps": [...], "decision": {...}}')

# Also works with freeform text!
extracted = extractor.extract("""
1. First, I checked the order amount ($500) against our threshold
2. Then I verified the customer's history - new customer with no prior orders
3. Finally, I assessed the risk level

Decision: Approve with standard verification
""")

print(extracted)
# {'inputs': [...], 'steps': [...], 'relations': [...], 'decision': {...}}
```

## Export Methods

After tracing, export the results in multiple formats — works in both local and cloud mode:

```python
with TracedAgent(client, workflow_id="review", entity_id="CLM-4821") as agent:
    response = agent.chat(llm, "Review this claim")

    # Print raw LLM response
    agent.print_raw()

    # Export to various formats
    agent.export_json("trace.json")       # Formatted JSON
    agent.export_text("trace.txt")        # Human-readable plain text
    agent.export_markdown("trace.md")     # Markdown with headers
    agent.export_mermaid("trace.mmd")     # Mermaid diagram (steps + relations)
    agent.export_jsonl("traces.jsonl")    # Append as JSONL line (for batch analysis)

    # Get statistics
    stats = agent.stats()
    # {'word_count': 234, 'step_count': 5, 'relation_count': 3,
    #  'input_count': 2, 'categories': ['cost_benefit', 'risk_assessment'],
    #  'has_decision': True}
```

## Local-First Mode (LocalTracer)

Try AIAgentree without any backend — traces go to console + JSONL file:

```python
from aiagentree import LocalTracer

tracer = LocalTracer()                          # console + file (default)
tracer = LocalTracer(console=False)              # file only
tracer = LocalTracer(file_path="my-traces.jsonl") # custom path

client = tracer.get_client()
# Use client exactly like a cloud AgentreeClient
```

After the first trace is sealed, a one-time teaser message shows how to connect to the cloud for validation, audit trails, and dashboards.

## API Reference

### AgentreeClient

```python
client = AgentreeClient(
    api_key="ask_...",       # Required — API key
    base_url="http://...",   # Required — Backend URL
    tenant_id="...",         # Required — Tenant ID
    timeout=30,              # Optional — Request timeout
    max_retries=3,           # Optional — Retry count
)
```

### TracedAgent

```python
with TracedAgent(
    client=client,           # AgentreeClient instance
    workflow_id="...",       # Workflow identifier
    entity_type="...",       # Entity type (e.g., "order")
    entity_id="...",         # Entity ID
    agent_id="...",          # Optional agent ID
    title="...",             # Optional trace title
    auto_seal=True,          # Auto-seal on exit
    auto_transform=True,     # Auto-transform after seal
    use_structured_prompt=True,  # Prepend DECISION_SYSTEM_PROMPT (default: True)
) as agent:
    response = agent.chat(llm, message, context={...})
```

### Trace Builder Methods

| Method | Description |
|--------|-------------|
| `add_input(key, value, source=None)` | Record an input |
| `add_step(temp_id, title, category, confidence=None, text=None, ...)` | Add reasoning step |
| `add_relation(parent_temp_id, child_temp_id, is_pro=True)` | Link steps |
| `add_policy(policy_id, outcome, ...)` | Record policy evaluation |
| `set_discussion_data(title, tags=None)` | Set metadata |
| `seal(decision_id, action, confidence=None, ...)` | Complete trace (auto-validates; returns `validation_status` + `quality_score`) |
| `validate()` | Manually revalidate trace (usually not needed — seal auto-validates) |
| `transform()` | Transform to argument tree |

### Error Handling

```python
from aiagentree import AgentreeError, RateLimitError, ExtractionError

try:
    with TracedAgent(...) as agent:
        agent.chat(llm, message)
except RateLimitError as e:
    print(f"Rate limited. Retry after {e.retry_after}s")
except ExtractionError as e:
    print(f"Could not extract reasoning: {e}")
except AgentreeError as e:
    print(f"API error {e.status_code}: {e}")
```

## Examples

See the [examples/](examples/) directory:

- [basic_trace.py](examples/basic_trace.py) — Manual trace construction
- [openai_example.py](examples/openai_example.py) — OpenAI GPT integration
- [anthropic_example.py](examples/anthropic_example.py) — Anthropic Claude integration
- [ollama_example.py](examples/ollama_example.py) — Local LLMs (Llama, Mistral, Nemotron)
- [langchain_example.py](examples/langchain_example.py) — LangChain agents and chains

## How It Works — 3 Integration Stages

The SDK supports three integration stages with increasing data quality. Pick the one that fits your needs:

### Stage 1: Passive Tracing (`TracedLLM` / `instrument_llm`)

**No prompt changes.** Wraps your LLM transparently — your existing prompts are untouched. The SDK intercepts the response and extracts what it can using heuristics.

```python
traced_llm = instrument_llm(OpenAI(), client, "review")
# Your prompts are unchanged — tracing is invisible
response = traced_llm.chat.completions.create(...)
```

**Data quality:** Basic. The extractor parses freeform text with regex — bullet/numbered lists become steps, keyword matching guesses categories, no relations or confidence scores. Good for getting started without changing anything.

### Stage 2: Structured Prompts (`TracedAgent`, default)

**Adds a system prompt** that instructs the LLM to output structured JSON with reasoning steps, categories, confidence, and a decision. This is the default when using `TracedAgent`.

```python
with TracedAgent(client, workflow_id="review", entity_id="ORD-123") as agent:
    # SDK prepends DECISION_SYSTEM_PROMPT as system message
    response = agent.chat(llm, "Review this order")
```

**Data quality:** Rich. The LLM outputs JSON that the extractor parses directly — categories, confidence scores, supports/opposes relations, and structured decisions are all captured.

### Stage 3: Argument Tree Prompts (`TracedAgent` + `use_argument_prompt`)

**Uses `ARGUMENT_SYSTEM_PROMPT`** which instructs the LLM to output a full pro/con argument tree with hierarchy levels and explicit parent-child relations.

```python
with TracedAgent(client, workflow_id="review", entity_id="ORD-123",
                 use_argument_prompt=True) as agent:
    response = agent.chat(llm, "Review this order")
```

**Data quality:** Maximum. Full argument hierarchy with typed relations, hierarchy levels, and structured evidence. Best for compliance-grade audit trails.

> **Prompt overrides:** When connected to the AIAgentree backend, the SDK checks the `start_trace` response for custom prompts. If the backend provides a `prompt_overrides` field, it is used instead of the built-in defaults. This enables prompt iteration server-side without SDK updates.

### Data Quality Summary

| | Stage 1 (Passive) | Stage 2 (Structured) | Stage 3 (Argument Tree) |
|---|---|---|---|
| **Prompt changes** | None | System prompt added | System prompt added |
| **Steps** | Bullet/list heuristics | LLM-structured JSON | LLM-structured JSON |
| **Categories** | Keyword-guessed | LLM-assigned | LLM-assigned |
| **Confidence** | Not available | Per-step scores | Per-step scores |
| **Relations** | Not available | Supports/opposes | Full hierarchy |
| **Decision** | Keyword scan | Structured | Structured |

### What Happens Under the Hood

```
Your LLM Call
    │
    ▼
TracedAgent / SDK
  1. Prepend structured output prompt (Stage 2/3) or pass through (Stage 1)
  2. Call your LLM
  3. Extract reasoning (JSON parse or freeform heuristics)
  4. Submit trace events to AI Agentree
  5. Seal and transform
    │
    ▼
AI Agentree Backend
  • Validates trace completeness
  • Calculates quality score
  • Transforms to argument tree
  • Stores for audit and analysis
```

## Development

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

# Run tests
pytest

# Run with coverage
pytest --cov=aiagentree
```

## License

MIT
