a04a16c717
HAIKU_RAG_DB and DB_PATH are gone: a capability covers what the configuration places or the db_path it is given, and the app backend and the AG-UI example load their configuration as the CLI does. The compose files point HAIKU_RAG_CONFIG_PATH at the mounted haiku.rag.yaml, which places the database at /data where DB_VOLUME is mounted; the backend refuses a configured set since it serves one database. The chat scopes a selection by source only over a set and names databases on filter rows only across several. Docstrings, docs and test fixtures stop describing an unnamed database; every database a search, listing or citation reports carries a name.
94 lines
3 KiB
Python
94 lines
3 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:
|
|
uv run uvicorn examples.custom_agent_agui:app --reload --port 8000
|
|
|
|
The configuration places the database (HAIKU_RAG_CONFIG_PATH, or
|
|
./haiku.rag.yaml).
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
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
|
|
|
|
capability = create_capability(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"]),
|
|
],
|
|
)
|