# Instrument with amplitude-ai

Auto-instrument a Python AI app with `amplitude-ai` in 4 phases: **Detect → Discover → Instrument → Verify**. The result is a fully instrumented app with provider wrappers, session lifecycle, multi-agent delegation (when detected), and a verification test proving correctness — all before deploying anything.

---

## Phase 1: Detect Environment

1. Read `pyproject.toml`, `requirements.txt`, or `setup.py` for dependencies
2. Detect framework: `fastapi` → FastAPI, `flask` → Flask, `django` → Django
3. Detect LLM providers: `openai`, `anthropic`, `google-generativeai`, `google-genai`, `boto3` (Bedrock), `mistralai`, `cohere`
4. Detect agent frameworks: `langchain`, `langchain-core`, `llama-index`, `openai-agents`, `crewai`, `langgraph`
5. Detect existing instrumentation: `amplitude-ai` in deps, `patch(amplitude=` or `AmplitudeAI` in source
6. Check for multi-agent signals: multiple files with LLM calls, tool definitions that call other LLM-calling functions, delegation patterns
7. Check for streaming: `stream=True` in provider calls, `async for` iteration patterns
8. Check for frontend deps: `react`, `vue`, `svelte` in a companion JS project
9. Check for message queues: `celery`, `redis`, `pika`, `boto3` SQS patterns

**Output to the developer:**

```
Detected environment:
  Runtime:     Python 3.x
  Framework:   [framework or "none"]
  Providers:   [list]
  Agent frameworks: [list or "none"]
  Existing instrumentation: [yes/no]
  Multi-agent signals: [yes/no]
  Streaming: [yes/no]
  Recommended: full instrumentation
```

**Next step:** Confirm the detection with the developer, then proceed to full instrumentation. Always instrument with agents, sessions, provider wrappers, tool tracking, and scoring. If multi-agent signals are detected, also add child agents and `run_as` delegation.

---

## Phase 2: Discover Agents and Call Sites

1. Identify files with LLM call sites (search for `chat.completions.create`, `messages.create`, `generate_content`, `invoke_model`, `converse`, `chat.complete`)
2. For each file with call sites, read the actual source and review:
   - Is it a route handler / API endpoint?
   - What provider(s) does it use?
   - Does it call other files with LLM call sites? (delegation → multi-agent)
3. Identify:
   - Agent boundaries (each distinct orchestration unit = one agent)
   - Delegation patterns (parent calls child → `run_as`)
   - Feedback handlers (thumbs up/down endpoints)
   - Tool functions (functions called by the LLM via function calling)
4. For each event emission you plan to add, trace **all code paths** that should emit the same event type. Look for error handlers, retry-exhaustion paths, timeout handlers, and fallback branches that represent the same logical operation failing — these should also emit the event (typically with `success=False` or an `error_message`).

### Multi-Agent Detection

Look for these patterns by reading the source files:
1. Do any tool functions call other functions that make LLM calls? → **delegation-as-tools** (A2A) pattern
2. Does a parent function call a function in another file that has LLM call sites? → **sequential delegation**
3. Are there multiple distinct agent roles or personas with separate system prompts?
4. Identify **orchestration wrappers** — functions that invoke sub-agents and should measure their execution. These are candidates for `@observe` / `track_span` to capture delegation latency, input/output summaries, and error status. Look for: try/except blocks around sub-agent calls, functions that dispatch to multiple agents in sequence or parallel, and any function that measures duration of a delegated operation.

Only mark the architecture as multi-agent once you've confirmed one of these patterns by reading the source.

**Output to the developer:**

```
Found N agents across M files:

Agent 1: "chat-handler"
  Description: "Handles user chat requests via streaming OpenAI GPT-4o"
  File: app/routes/chat.py
  Provider: OpenAI (chat.completions.create)
  Entry point: POST /chat

Agent 2: "recipe-agent"  (child of chat-handler, called as a tool)
  Description: "Specialized recipe planning agent called by the orchestrator"
  File: app/agents/recipe.py
  Provider: OpenAI (chat.completions.create)
  Delegation: ask_recipe_agent() tool in Agent 1 delegates to this file

Multi-agent architecture: delegation-as-tools (A2A)
  → will instrument with ai.agent().child() + session.run_as()

Proceed with instrumentation? [Review changes first / Apply / Skip]
```

**PAUSE HERE.** Let the developer review the agent names, descriptions, and structure before proceeding. They can edit names and descriptions.

---

## Phase 3: Instrument

### Step 3a: Install dependencies

```bash
pip install amplitude-ai
```

### Step 3b: Create bootstrap file

Create `app/amplitude_bootstrap.py` (or the project's conventional config/lib path):

**Choose `content_mode` based on privacy needs:**

- **`full`** — captures full prompt/response text. Best for debugging and enrichment. Always pair with `redact_pii=True` unless the customer has explicit consent.
- **`metadata_only`** — captures token counts, latency, model, cost, but no text. Use for sensitive PII or regulated data.
- **`customer_enriched`** — no text captured by default, but the customer can send enriched summaries via `track_session_enrichment()`.

```python
from amplitude import Amplitude
from amplitude_ai import AmplitudeAI, AIConfig, OpenAI

amplitude = Amplitude("YOUR_AMPLITUDE_API_KEY")
ai = AmplitudeAI(
    amplitude=amplitude,
    config=AIConfig(content_mode="full", redact_pii=True),
)

# One wrapped client per provider detected in Phase 1
openai_client = OpenAI(amplitude=ai, api_key="sk-...")

# Add more providers as detected:
# from amplitude_ai import Anthropic
# anthropic_client = Anthropic(amplitude=ai, api_key="sk-ant-...")
```

Add `AMPLITUDE_AI_API_KEY` to `.env.example`. Check `.gitignore` includes `.env`.

> **Note:** If you cannot modify provider instantiation sites, use `wrap(existing_client, amplitude=ai)` to instrument an existing client, or `patch(amplitude=amplitude)` for zero-code verification. These capture fewer event types — always prefer provider wrappers when possible.

### Step 3c: Swap provider imports

Replace direct provider instantiation with imports from the bootstrap file:

**Before:**
```python
import openai
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
```

**After:**
```python
from app.amplitude_bootstrap import openai_client as client
```

### Step 3d: Add session context

Wrap route handlers with agent + session:

```python
from app.amplitude_bootstrap import ai

agent = ai.agent("chat-handler", description="Handles user chat requests via OpenAI GPT-4o")

@app.post("/chat")
async def chat(request: Request):
    data = await request.json()
    async with agent.session(user_id=data["user_id"], session_id=data["session_id"]) as s:
        s.track_user_message(data["message"])
        response = client.chat.completions.create(model="gpt-4o", messages=data["messages"])
        return {"content": response.choices[0].message.content}
    # Session auto-flushes in serverless (Lambda, Vercel, Netlify, etc.)
    # For non-serverless, or tracking outside session context, call: ai.flush()
```

If multi-agent signals were detected, add delegation with `run_as`:

```python
orchestrator = ai.agent("shopping-agent", description="Orchestrates shopping requests")
recipe_agent = orchestrator.child("recipe-agent", description="Finds recipes")

# Inside parent's session context:
async with s.arun_as(recipe_agent) as cs:
    cs.track_user_message(delegated_query)
    result = client.chat.completions.create(model="gpt-4o", messages=[...])
```

### Step 3e: Track tools and explicit AI responses

**Tool tracking** with the `@tool` decorator:

```python
from amplitude_ai import tool

@tool
def search_products(query: str) -> dict:
    return db.search(query)

# Inside session context, just call the decorated function:
result = search_products(query)
# [Agent] Tool Call event automatically emitted with duration, success, input/output
```

**Explicit AI response capture** (when provider wrappers can't auto-capture):

```python
s.track_ai_message(
    content=completed_message.content,
    model="gpt-4o",
    provider="openai",
    latency_ms=1200.0,
    input_tokens=usage.prompt_tokens,
    output_tokens=usage.completion_tokens,
    total_tokens=usage.total_tokens,
)
```

### Step 3f: Add span tracking for orchestration

When a parent agent delegates work to a child agent, wrap the delegation call with span tracking to measure latency and capture errors. Look for existing try/except blocks around sub-agent execution — these are natural places to add span tracking with both success and error paths:

```python
from amplitude_ai import observe

# Option A: decorator on orchestration functions
@observe
async def run_sub_agent(prompt: str) -> str:
    return await sub_agent.execute(prompt)

# Option B: explicit tracking when you need more control
run_start = time.monotonic()
try:
    result = await sub_agent.execute(prompt)
    duration_ms = (time.monotonic() - run_start) * 1000
    s.track_span(
        span_name=child_agent_name,
        latency_ms=duration_ms,
        input_state={"prompt": prompt[:1000]},
        output_state={"response": result[:1000]},
    )
except Exception as e:
    duration_ms = (time.monotonic() - run_start) * 1000
    s.track_span(
        span_name=child_agent_name,
        latency_ms=duration_ms,
        is_error=True,
        error_message=f"{type(e).__name__}: {e}",
    )
    raise
```

### Step 3g: Add scoring

If feedback handlers were detected (thumbs up/down endpoints), add scoring. Check whether the handler receives a `message_id` — if so, target the specific message for finer-grained scoring. Otherwise fall back to session-level scoring:

```python
target_id = message_id if message_id else session_id
target_type = "message" if message_id else "session"
ai.score(
    user_id=user_id, name="user-feedback", value=1 if thumbs_up else 0,
    target_id=target_id, target_type=target_type, source="user",
    session_id=session_id, comment=feedback_text,
)
```

### Step 3h: Streaming session lifecycle

If the app uses streaming, the session must stay open until the stream is fully consumed:

```python
# WRONG: session ends when generator is created, not consumed
async with agent.session(user_id=uid, session_id=sid) as s:
    stream = client.chat.completions.create(model="gpt-4o", messages=msgs, stream=True)
    return StreamingResponse(stream)  # session closes here, before stream consumed!

# CORRECT: session stays open until stream is fully consumed
async with agent.session(user_id=uid, session_id=sid) as s:
    stream = client.chat.completions.create(model="gpt-4o", messages=msgs, stream=True)
    async def generate():
        for chunk in stream:
            yield chunk.choices[0].delta.content or ""
    return StreamingResponse(generate())
    # session stays open because we're still inside the 'async with' block
```

### Step 3i: Cross-service context propagation

If the app uses message queues or has a multi-service architecture:

```python
from amplitude_ai.propagation import inject_context, extract_context

# Sender side: include Amplitude context in message headers
headers = inject_context()  # reads from active session context via contextvars
await queue.send(payload, headers=headers)

# Receiver side: restore context from headers
ctx = extract_context(message.headers)
# Use the extracted session_id and user_id to continue the session
with agent.session(session_id=ctx.get("session_id"), user_id=ctx.get("user_id")) as s:
    ...
```

### Step 3j: Browser-server session linking

If frontend deps were detected, extract browser IDs from request headers:

```python
browser_session_id = request.headers.get("x-amplitude-session-id")
device_id = request.headers.get("x-amplitude-device-id")
with agent.session(user_id=uid, browser_session_id=browser_session_id, device_id=device_id) as s:
    ...
```

### Step 3k: Framework-specific notes

**FastAPI/Starlette**: Use the built-in middleware for automatic session context:
```python
from amplitude_ai import AmplitudeAIMiddleware

app.add_middleware(
    AmplitudeAIMiddleware,
    amplitude_ai=ai,
    user_id_resolver=lambda request: request.headers.get("x-user-id"),
)
```

**Flask / Django**: No built-in middleware. Wrap route handlers manually with `agent.session()`.

---

## Phase 4: Verify

### Step 4a: Create verification test

Create `test_amplitude_verify.py` that verifies:
- Each agent emits `[Agent] User Message` with correct `[Agent] Agent ID`
- Sessions are properly closed (`assert_session_closed`)
- Multi-agent delegation preserves session ID across `run_as`

```python
from amplitude_ai.testing import MockAmplitudeAI

def test_agent_emits_events():
    mock = MockAmplitudeAI()
    agent = mock.agent("test-agent")
    with agent.session(user_id="u1", session_id="s1") as s:
        s.track_user_message("hello")
    mock.assert_event_tracked("[Agent] User Message", agent_id="test-agent")
    mock.assert_session_closed("s1")

def test_multi_agent_delegation():
    mock = MockAmplitudeAI()
    parent = mock.agent("parent-agent")
    child = parent.child("child-agent")
    with parent.session(user_id="u1", session_id="s1") as s:
        with s.run_as(child) as cs:
            cs.track_user_message("delegated task")
    child_events = mock.events_for_agent("child-agent")
    assert len(child_events) > 0
    assert child_events[0]["event_properties"]["[Agent] Session ID"] == "s1"
```

### Step 4b: Run verification

```bash
pytest test_amplitude_verify.py -v
```

### Step 4c: Run doctor

```bash
amplitude-ai-doctor
```

### Step 4d: Run project checks

```bash
pytest              # Existing tests still pass
mypy .              # Type checking passes (if used)
```

### Step 4e: Show confidence report

```
Verification complete:
  Doctor checks:        5/5 passed
  Event sequence test:  PASSED (N events captured)
  Existing tests:       PASSED

  Content mode: full (PII redacted)

Next steps:
  1. Set AMPLITUDE_AI_API_KEY in your environment
  2. Keep test_amplitude_verify.py for CI regression testing
  3. Deploy and verify live events in Amplitude
```

---

## API Quick Reference

### Core API

| API | What it does |
|-----|-------------|
| `AmplitudeAI(amplitude=amp, config=cfg)` | Initialize SDK |
| `AIConfig(content_mode=, redact_pii=, debug=)` | Privacy/debug config |
| `ai.agent(agent_id, description=, user_id=)` | Create bound agent |
| `agent.child(agent_id, description=)` | Create child agent |
| `agent.session(user_id=, session_id=, auto_flush=)` | Create session (`auto_flush` auto-detects serverless) |
| `with agent.session(...) as s:` | Sync session context (auto-flushes in serverless) |
| `async with agent.session(...) as s:` | Async session context |
| `with s.run_as(child_agent) as cs:` | Delegate to child agent (sync) |
| `async with s.arun_as(child_agent) as cs:` | Delegate to child agent (async) |
| `ai.flush()` | Flush events (serverless) |
| `ai.shutdown()` | Flush and release resources |

### Tracking Methods (on session `s`)

| Method | Event Emitted |
|--------|--------------|
| `s.track_user_message(content)` | `[Agent] User Message` |
| `s.track_ai_message(content, model, provider, latency_ms)` | `[Agent] AI Response` |
| `s.track_tool_call(tool_name, latency_ms, success)` | `[Agent] Tool Call` |
| `s.score(name, value, target_id)` | `[Agent] Score` |
| `s.track_session_enrichment(enrichments)` | `[Agent] Session Enrichment` |

### Decorators

| Decorator | Event Emitted | Usage |
|-----------|--------------|-------|
| `@tool` | `[Agent] Tool Call` | Decorate tool functions |
| `@observe` | `[Agent] Span` | Decorate any function for observability |

### Provider Wrappers

All imported from `amplitude_ai`:

| Provider | Constructor |
|----------|------------|
| OpenAI | `OpenAI(amplitude=ai, api_key=...)` |
| AsyncOpenAI | `AsyncOpenAI(amplitude=ai, api_key=...)` |
| Anthropic | `Anthropic(amplitude=ai, api_key=...)` |
| Gemini | `Gemini(amplitude=ai, api_key=...)` |
| AzureOpenAI | `AzureOpenAI(amplitude=ai, azure_endpoint=..., api_key=...)` |
| Bedrock | `Bedrock(amplitude=ai)` |
| Mistral | `Mistral(amplitude=ai, api_key=...)` |

### Other APIs

| API | Usage |
|-----|-------|
| `patch(amplitude=amp)` / `unpatch()` | Zero-code instrumentation |
| `wrap(client, amplitude=ai)` | Wrap existing provider client |
| `inject_context()` / `extract_context(headers)` | Cross-service propagation (from `amplitude_ai.propagation`) |
| `AmplitudeAIMiddleware` | FastAPI/Starlette middleware |
| `MockAmplitudeAI` (from `amplitude_ai.testing`) | Deterministic test double |

---

## Ecosystem-Specific Guidance

### LangChain / LangGraph
- Use `AmplitudeCallbackHandler` for LangChain — LLM calls via LangChain are captured through the callback
- LangGraph graph orchestration events (node transitions, checkpoints) are not yet instrumented

### OpenAI Assistants API
- Provider wrappers do NOT auto-instrument the Assistants API (async/polling-based)
- Use manual tracking: `track_user_message()` when creating a message, `track_ai_message()` when polling

### AWS Bedrock
- Use the `Bedrock` wrapper which instruments both `invoke_model` and `converse` APIs
- Streaming via `invoke_model_with_response_stream` is supported

---

## Safety Rules

- **Never modify unrelated files.** Only touch files with LLM call sites and the bootstrap file.
- **Never duplicate instrumentation.** Check for existing `patch()` or wrapper calls before adding new ones.
- **Pause before Phase 3.** Always show the discovery report and get developer confirmation.
- **Prefer additive changes.** Add imports and wrappers rather than rewriting entire files.
- **Keep content mode explicit.** Default is `full` + `redact_pii=True`. Never silently downgrade.
- **Preserve existing tests.** Instrumentation must not break the test suite.
- **Idempotent.** Running this twice should not double-instrument. Detect existing instrumentation and skip.
