Agentic example with and without agui

docs/python.md

@@ -455,4 +455,4 @@ async with HaikuRAG("path/to/db.lancedb") as client:
     result = await agent.run("What are the main findings?", deps=deps)
 ```
 
-See [Toolsets](tools.md) for the full API reference and composition guide.
+See [Toolsets](tools.md) for the full API reference and composition guide, and the [`examples/`](https://github.com/ggozad/haiku.rag/tree/main/examples) directory for runnable scripts.

docs/tools.md

@@ -245,6 +245,8 @@ deps = AgentDeps(client=client, tool_context=context)
 
 Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests.
 
+For complete runnable examples, see [`examples/custom_agent.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent.py) (standalone) and [`examples/custom_agent_agui.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent_agui.py) (AG-UI streaming server).
+
 All toolsets respect session-level document filters when a `SessionState` is registered in the context. This means setting `SessionState.document_filter` restricts all tools simultaneously.
 
 ## AG-UI State Management

examples/README.md

@@ -11,3 +11,23 @@ Complete Docker setup for running haiku.rag with all services:
 - MCP server for AI assistant integration
 
 See `docker/README.md` for setup instructions.
+
+## Custom Agent
+
+**Script:** `custom_agent.py`
+
+Composes `search`, `qa`, and `document` toolsets into a pydantic-ai `Agent` using `AgentDeps` and `prepare_context`. Shows how to run queries and inspect accumulated state (citations, QA history).
+
+```bash
+uv run python examples/custom_agent.py /path/to/db.lancedb
+```
+
+## Custom Agent with AG-UI Streaming
+
+**Script:** `custom_agent_agui.py`
+
+A Starlette app that serves an AG-UI streaming endpoint using composed toolsets, `AgentDeps`, and `ToolContextCache` for multi-session support.
+
+```bash
+DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000
+```

examples/custom_agent.py

@@ -0,0 +1,74 @@
+"""Custom agent using haiku.rag composable toolsets.
+
+Demonstrates how to compose search, QA, and document toolsets into a
+pydantic-ai Agent using AgentDeps and prepare_context.
+
+Requirements:
+    - An Ollama instance running locally (default embedder)
+    - An Anthropic API key (for the QA model) or adjust the model below
+    - A haiku.rag database with documents already ingested
+
+Usage:
+    uv run python examples/custom_agent.py /path/to/db.lancedb
+"""
+
+import asyncio
+import sys
+
+from pydantic_ai import Agent
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.tools import (
+    AgentDeps,
+    ToolContext,
+    create_document_toolset,
+    create_qa_toolset,
+    create_search_toolset,
+    prepare_context,
+)
+
+
+async def main(db_path: str) -> None:
+    async with HaikuRAG(db_path) as client:
+        # Compose toolsets into an agent
+        config = client.config
+        search_toolset = create_search_toolset(config)
+        qa_toolset = create_qa_toolset(config)
+        document_toolset = create_document_toolset(config)
+
+        agent = Agent(
+            "anthropic:claude-haiku-4-5-20251001",
+            deps_type=AgentDeps,
+            output_type=str,
+            instructions=(
+                "You are a helpful assistant with access to a knowledge base. "
+                "Use the available tools to answer questions."
+            ),
+            toolsets=[search_toolset, qa_toolset, document_toolset],
+        )
+
+        # Prepare a shared ToolContext
+        context = ToolContext()
+        prepare_context(context, features=["search", "documents", "qa"])
+
+        deps = AgentDeps(client=client, tool_context=context)
+
+        print("Custom agent ready. Ctrl+C to exit.\n")
+        while True:
+            try:
+                user_input = input("You: ").strip()
+            except (EOFError, KeyboardInterrupt):
+                break
+
+            if not user_input:
+                continue
+
+            result = await agent.run(user_input, deps=deps)
+            print(f"\nAgent: {result.output}\n")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print(f"Usage: {sys.argv[0]} <db_path>")
+        sys.exit(1)
+    asyncio.run(main(sys.argv[1]))

examples/custom_agent_agui.py

@@ -0,0 +1,116 @@
+"""Custom agent with AG-UI streaming.
+
+A Starlette app that composes haiku.rag toolsets into an AG-UI compatible
+agent. Multi-session support via ToolContextCache.
+
+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 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.client import HaikuRAG
+from haiku.rag.config.models import AppConfig
+from haiku.rag.tools import (
+    AgentDeps,
+    ToolContextCache,
+    create_qa_toolset,
+    create_search_toolset,
+    prepare_context,
+)
+
+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)
+
+AGUI_STATE_KEY = "my_app"
+
+config = AppConfig()
+
+# ToolContextCache maintains per-thread state across requests
+context_cache = ToolContextCache()
+
+# Singleton client
+_client: HaikuRAG | None = None
+
+
+def get_client() -> HaikuRAG:
+    global _client
+    if _client is None:
+        _client = HaikuRAG(db_path=db_path)
+    return _client
+
+
+# Create the agent once at module level
+agent = Agent(
+    "anthropic:claude-haiku-4-5-20251001",
+    deps_type=AgentDeps,
+    output_type=str,
+    instructions=(
+        "You are a helpful assistant with access to a knowledge base. "
+        "Use the search and ask tools to answer questions."
+    ),
+    toolsets=[
+        create_search_toolset(config),
+        create_qa_toolset(config),
+    ],
+)
+
+
+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)
+
+    thread_id = getattr(run_input, "thread_id", None) or "default"
+    context, is_new = context_cache.get_or_create(thread_id)
+    if is_new:
+        prepare_context(
+            context,
+            features=["search", "qa"],
+            state_key=AGUI_STATE_KEY,
+        )
+
+    deps = AgentDeps(client=get_client(), tool_context=context)
+
+    adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
+    event_stream = adapter.run_stream(deps=deps)
+    sse_event_stream = adapter.encode_stream(event_stream)
+
+    return StreamingResponse(
+        sse_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"]),
+    ],
+)