771ac9c96c
The README feature list and the overview stopped at the analysis capability. Both examples and the app backend composed agents without the capabilities the documentation recommends alongside an evidence capability. custom_agent.py ran each input as an independent agent run, so it needed a state dict and a carried history before compaction could mean anything there: without state the evidence record is empty, and earlier evidence would reduce to receipts retaining nothing.
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
"""Custom agent with AG-UI streaming.
|
|
|
|
A Starlette app that serves an AG-UI streaming endpoint using the
|
|
haiku.rag's native Pydantic AI RAG capability.
|
|
|
|
Requirements:
|
|
- An Ollama instance running locally (default embedder)
|
|
- An Anthropic API key (for the QA model) or adjust the model below
|
|
|
|
Usage:
|
|
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from ag_ui.core import EventType, StateSnapshotEvent
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.ui import SSE_CONTENT_TYPE
|
|
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
|
from starlette.applications import Starlette
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse, Response, StreamingResponse
|
|
from starlette.routing import Route
|
|
|
|
from haiku.rag.capabilities.compaction import create_capability as compaction
|
|
from haiku.rag.capabilities.policy import create_capability as citation_policy
|
|
from haiku.rag.capabilities.rag import RAGState, create_capability
|
|
|
|
db_path = os.environ.get("DB_PATH")
|
|
if not db_path:
|
|
print(
|
|
"Set DB_PATH environment variable to your haiku.rag database", file=sys.stderr
|
|
)
|
|
sys.exit(1)
|
|
|
|
capability = create_capability(db_path=Path(db_path), defer_loading=False)
|
|
|
|
|
|
@dataclass
|
|
class AppDeps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
agent = Agent(
|
|
"anthropic:claude-haiku-4-5-20251001",
|
|
# The client returns the state snapshot with every run, so earlier questions are
|
|
# reduced to the evidence they cited and every answer declares its grounding.
|
|
capabilities=[capability, compaction(), citation_policy()],
|
|
deps_type=AppDeps,
|
|
)
|
|
|
|
|
|
async def stream_chat(request: Request) -> Response:
|
|
body = await request.body()
|
|
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
|
run_input = AGUIAdapter.build_run_input(body)
|
|
|
|
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
|
|
|
|
incoming_state = run_input.state if isinstance(run_input.state, dict) else {}
|
|
incoming_state.setdefault("rag", RAGState().model_dump(mode="json"))
|
|
deps = AppDeps(state=incoming_state)
|
|
|
|
async def event_stream():
|
|
async def with_final_state():
|
|
async for event in adapter.run_stream(deps=deps):
|
|
if getattr(event, "type", None) == EventType.RUN_FINISHED:
|
|
yield StateSnapshotEvent(
|
|
type=EventType.STATE_SNAPSHOT,
|
|
snapshot=deps.state,
|
|
)
|
|
yield event
|
|
|
|
async for chunk in adapter.encode_stream(with_final_state()):
|
|
yield chunk
|
|
|
|
return StreamingResponse(
|
|
event_stream(),
|
|
media_type=accept,
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
async def health_check(_: Request) -> JSONResponse:
|
|
return JSONResponse({"status": "healthy"})
|
|
|
|
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/v1/chat/stream", stream_chat, methods=["POST"]),
|
|
Route("/health", health_check, methods=["GET"]),
|
|
],
|
|
)
|