Streaming Guide

from ai_infra import Agent

agent = Agent(tools=[my_tool])

# Basic streaming - print tokens as they arrive
async for event in agent.astream("What is the weather in NYC?"):
    if event.type == "token":
        print(event.content, end="", flush=True)

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from ai_infra import Agent
import json

app = FastAPI()
agent = Agent(tools=[search_docs])

@app.post("/chat")
async def chat(message: str):
    async def generate():
        async for event in agent.astream(message, visibility="detailed"):
            yield f"data: {json.dumps(event.to_dict())}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

from fastapi import WebSocket

@app.websocket("/chat")
async def chat_ws(websocket: WebSocket):
    await websocket.accept()
    message = await websocket.receive_text()

    async for event in agent.astream(message):
        await websocket.send_json(event.to_dict())

    await websocket.close()

async for event in agent.astream("Search for auth"):
    if event.type == "tool_end" and event.result_structured:
        # event.result is a dict, not a string
        results = event.result["results"]
        query = event.result["query"]
        for r in results:
            print(f"{r['source']}: {r['score']}")
    elif event.type == "tool_end":
        # event.result is a string
        print(event.result)

# Text result (result_structured=False)
{"type": "tool_end", "result": "text output...", ...}

# Structured result (result_structured=True)
{"type": "tool_end", "structured_result": {"results": [...], "query": "..."}, ...}

from ai_infra import Agent, atemporary_api_key

async with atemporary_api_key("openai", user_key):
    async for event in agent.astream(prompt):
        yield event.to_dict()

from ai_infra import Agent, load_mcp_tools_cached

tools = await load_mcp_tools_cached("http://localhost:8000/mcp")
agent = Agent(tools=tools)

async for event in agent.astream("What is the refund policy?"):
    if event.type == "token":
        print(event.content, end="", flush=True)

async for event in agent.astream("Search docs", visibility="detailed"):
    if event.type == "tool_start":
        print(event.arguments)

config = {
    "configurable": {"thread_id": "user-123"},
    "tags": ["production"],
}
async for event in agent.astream("Continue our conversation", config=config):
    ...

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()
agent = Agent(tools=[search_docs])

@app.post("/chat")
async def chat(message: str, provider: str, api_key: str):
    async def generate():
        async with atemporary_api_key(provider, api_key):
            async for event in agent.astream(message):
                yield f"data: {event.to_dict()}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

from fastapi import WebSocket

@router.websocket("/chat")
async def chat_ws(ws: WebSocket):
    await ws.accept()
    async for event in agent.astream("Hello"):
        await ws.send_json(event.to_dict())

MCP_URL = "http://localhost:8000/mcp"

async def stream_with_mcp(message: str, api_key: str):
    tools = await load_mcp_tools_cached(MCP_URL)
    agent = Agent(tools=tools)
    async with atemporary_api_key("openai", api_key):
        async for event in agent.astream(message, visibility="debug"):
            yield event

### Result 1 (svc-infra: auth.md)
[snippet about authentication]
---
### Result 2 (ai-infra: core/llm.md)
[snippet about LLM usage]

import re
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from ai_infra import Agent, atemporary_api_key, load_mcp_tools_cached

app = FastAPI()

def parse_doc_results(result: str) -> list[dict]:
    """Parse search_docs output into structured data."""
    pattern = r"### Result \d+ \((.+?): (.+?)\)"
    matches = re.finditer(pattern, result)
    return [
        {"package": m.group(1), "path": m.group(2)}
        for m in matches
    ]

@app.post("/chat")
async def chat(message: str, provider: str, api_key: str):
    tools = await load_mcp_tools_cached("http://localhost:8000/mcp")
    agent = Agent(tools=tools)

    async def generate():
        async with atemporary_api_key(provider, api_key):
            async for event in agent.astream(message, visibility="detailed"):
                # Regular events pass through
                if event.type != "tool_end":
                    yield f"data: {json.dumps(event.to_dict())}\n\n"
                    continue

                # Parse tool results for search tools
                if "search" in event.tool and event.result:
                    docs = parse_doc_results(event.result)

                    # Emit custom event with structured data
                    yield f"data: {json.dumps({
                        'type': 'tool_results',
                        'tool': event.tool,
                        'docs': docs,  # [{"package": "svc-infra", "path": "auth.md"}, ...]
                        'latency_ms': event.latency_ms,
                    })}\n\n"
                else:
                    # Regular tool_end event
                    yield f"data: {json.dumps(event.to_dict())}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

// Handle tool results in chat UI
if (event.type === "tool_results") {
    event.docs.forEach((doc: {package: string, path: string}) => {
        const url = `/${doc.package}/${doc.path.replace('.md', '')}`;
        // Create clickable chip/link in UI
        showClickableChip(doc.package, doc.path, url);
    });
}