Rich Agent Stream Events

from ai_infra import Agent
from ai_infra.llm.streaming import StreamConfig

config = StreamConfig(
    visibility="detailed",
    cost_per_input_token=0.000003,   # $3 per 1M input tokens
    cost_per_output_token=0.000015,  # $15 per 1M output tokens
)

total_cost = 0.0
async for event in agent.astream("Summarize the auth module", stream_config=config):
    if event.type == "usage":
        total_cost += event.cost or 0.0
        print(f"  Tokens: {event.input_tokens} in / {event.output_tokens} out")
        if event.cost:
            print(f"  Cost: ${event.cost:.6f}")
    elif event.type == "token":
        print(event.content, end="", flush=True)

print(f"\nTotal cost: ${total_cost:.6f}")

@app.post("/chat")
async def chat(message: str):
    total_input = 0
    total_output = 0

    async def generate():
        nonlocal total_input, total_output
        async for event in agent.astream(message, stream_config=config):
            if event.type == "usage":
                total_input += event.input_tokens or 0
                total_output += event.output_tokens or 0
            yield f"data: {json.dumps(event.to_dict())}\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"X-Token-Count": f"{total_input}+{total_output}"},
    )

turn_start (turn_id=1)
  reasoning → tool_start → tool_end
turn_end (turn_id=1, tools_called=1)

turn_start (turn_id=2)
  reasoning → tool_start → tool_end → tool_start → tool_end
turn_end (turn_id=2, tools_called=2)

turn_start (turn_id=3)
  token token token ...
turn_end (turn_id=3, tools_called=0)
done

async for event in agent.astream("Refactor the auth module"):
    match event.type:
        case "turn_start":
            ui.show_spinner(f"Turn {event.turn_id}...")
        case "turn_end":
            if event.tools_called:
                ui.update_status(f"Used {event.tools_called} tools")
            ui.hide_spinner()
        case "token":
            ui.append_text(event.content)
        case "done":
            ui.finalize()

turns = []
async for event in agent.astream(prompt, stream_config=config):
    if event.type == "turn_start":
        turns.append({"id": event.turn_id, "tools": 0})
    elif event.type == "turn_end":
        turns[-1]["tools"] = event.tools_called or 0
        progress = f"Turn {event.turn_id} complete ({event.tools_called} tools)"
        yield {"type": "progress", "message": progress, "turns": len(turns)}

async for event in agent.astream("Fix the login bug"):
    match event.type:
        case "intent":
            ui.set_status(event.content)  # "Reading files"
        case "tool_start":
            ui.show_tool_badge(event.tool)
        case "token":
            ui.clear_status()
            ui.append_text(event.content)

config = StreamConfig(
    tool_intent_map={
        "search_docs": "Searching documentation",
        "query_db": "Querying database",
        "run_tests": "Running test suite",
        "deploy": "Deploying to staging",
    }
)
async for event in agent.astream(prompt, stream_config=config):
    if event.type == "intent":
        print(f"Agent: {event.content}")

# Direct list
[{"id": 1, "title": "Fetch data", "status": "completed"}, ...]

# Dict with items/todos/todoList/todo_items key
{"items": [{"id": 1, "title": "Fetch data", "status": "completed"}, ...]}
{"todos": [...]}
{"todoList": [...]}

async for event in agent.astream("Plan the migration"):
    if event.type == "todo":
        for item in event.todo_items:
            icon = "[x]" if item["status"] == "completed" else "[ ]"
            print(f"  {icon} {item['title']}")
    elif event.type == "token":
        print(event.content, end="", flush=True)

function TaskList({ todoItems }: { todoItems: TodoItem[] }) {
  return (
    <ul>
      {todoItems.map((item) => (
        <li key={item.id} className={item.status === "completed" ? "done" : ""}>
          <input type="checkbox" checked={item.status === "completed"} readOnly />
          {item.title}
        </li>
      ))}
    </ul>
  );
}

// In your SSE handler:
if (event.type === "todo") {
  setTodoItems(event.todo_items);
}

config = StreamConfig(todo_tool_name="update_tasks")
async for event in agent.astream(prompt, stream_config=config):
    ...