diff --git a/docs/python.md b/docs/python.md index ab158e82..a89b173e 100644 --- a/docs/python.md +++ b/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. diff --git a/docs/tools.md b/docs/tools.md index 6468e825..257577e9 100644 --- a/docs/tools.md +++ b/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 diff --git a/examples/README.md b/examples/README.md index afb7fd39..74613d23 100644 --- a/examples/README.md +++ b/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 +``` diff --git a/examples/custom_agent.py b/examples/custom_agent.py new file mode 100644 index 00000000..91b4c51d --- /dev/null +++ b/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]} ") + sys.exit(1) + asyncio.run(main(sys.argv[1])) diff --git a/examples/custom_agent_agui.py b/examples/custom_agent_agui.py new file mode 100644 index 00000000..4e4a928c --- /dev/null +++ b/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"]), + ], +)