ai-agentree-sdk

ai-agentree-sdk

Python SDK for AI Agentree — 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

pip install ai-agentree-sdk

For development from source:

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:

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:

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. Sealing the trace automatically builds the argument tree — the Decision Packet is retrievable immediately after seal (no separate transform call; an explicit transform() is an idempotent no-op on an already-transformed trace)

Supported LLMs

Integration Methods

Method 1: TracedAgent (Recommended)

Context manager that handles everything:

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)

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)

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

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:

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:

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:

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:

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):

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:

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:

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

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

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

Error Handling

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/ directory:

• basic_trace.py — Manual trace construction
• openai_example.py — OpenAI GPT integration
• anthropic_example.py — Anthropic Claude integration
• ollama_example.py — Local LLMs (Llama, Mistral, Nemotron)
• 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.

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.

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.

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

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

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

# Run tests
pytest

# Run with coverage
pytest --cov=aiagentree

License

MIT