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.
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
"""Custom agent using the native haiku.rag RAG capability.
|
|
|
|
Demonstrates composing native Pydantic AI capabilities into an agent, and what a
|
|
multi-turn conversation needs to carry between runs.
|
|
|
|
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 dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.messages import ModelMessage
|
|
|
|
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 create_capability as rag
|
|
|
|
|
|
@dataclass
|
|
class Deps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
async def main(db_path: str) -> None:
|
|
agent = Agent(
|
|
"anthropic:claude-haiku-4-5-20251001",
|
|
capabilities=[
|
|
rag(db_path=Path(db_path), defer_loading=False),
|
|
compaction(),
|
|
citation_policy(),
|
|
],
|
|
deps_type=Deps,
|
|
)
|
|
|
|
# One state dict and one history for the whole session. The capabilities read
|
|
# both: the state holds what was retrieved and cited, and the message counts
|
|
# are how they tell one question from the next.
|
|
deps = Deps()
|
|
messages: list[ModelMessage] = []
|
|
|
|
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, message_history=messages)
|
|
messages = list(result.all_messages())
|
|
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]))
|