75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
from pydantic import BaseModel
|
|
from pydantic_ai import Agent, RunContext
|
|
from pydantic_ai.output import ToolOutput
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig, ModelConfig
|
|
from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
|
|
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.utils import get_model
|
|
|
|
|
|
class Dependencies(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
client: HaikuRAG
|
|
search_results: list[SearchResult] = []
|
|
search_filter: str | None = None
|
|
|
|
|
|
class QuestionAnswerAgent:
|
|
def __init__(
|
|
self,
|
|
client: HaikuRAG,
|
|
model_config: ModelConfig,
|
|
config: AppConfig | None = None,
|
|
system_prompt: str | None = None,
|
|
):
|
|
self._client = client
|
|
model_obj = get_model(model_config, config)
|
|
|
|
self._agent = Agent(
|
|
model=model_obj,
|
|
deps_type=Dependencies,
|
|
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
|
|
instructions=system_prompt or QA_SYSTEM_PROMPT,
|
|
retries=3,
|
|
)
|
|
|
|
@self._agent.tool
|
|
async def search_documents(
|
|
ctx: RunContext[Dependencies],
|
|
query: str,
|
|
limit: int | None = None,
|
|
) -> str:
|
|
"""Search the knowledge base for relevant documents.
|
|
|
|
Returns results with chunk IDs and relevance scores.
|
|
Reference results by their chunk_id in cited_chunks.
|
|
"""
|
|
results = await ctx.deps.client.search(
|
|
query, limit=limit, filter=ctx.deps.search_filter
|
|
)
|
|
results = await ctx.deps.client.expand_context(results)
|
|
# Store results for citation resolution
|
|
ctx.deps.search_results = results
|
|
# Format with metadata for agent context
|
|
parts = [r.format_for_agent() for r in results]
|
|
return "\n\n".join(parts) if parts else "No results found."
|
|
|
|
async def answer(
|
|
self, question: str, filter: str | None = None
|
|
) -> tuple[str, list[Citation]]:
|
|
"""Answer a question using the RAG system.
|
|
|
|
Args:
|
|
question: The question to answer
|
|
filter: SQL WHERE clause to filter documents
|
|
|
|
Returns:
|
|
Tuple of (answer text, list of resolved citations)
|
|
"""
|
|
deps = Dependencies(client=self._client, search_filter=filter)
|
|
result = await self._agent.run(question, deps=deps)
|
|
citations = resolve_citations(result.output.cited_chunks, deps.search_results)
|
|
return result.output.answer, citations
|