68 lines
2.4 KiB
Python
68 lines
2.4 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 import Config
|
|
from haiku.rag.config.models import ModelConfig
|
|
from haiku.rag.graph.common import get_model
|
|
from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations
|
|
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
|
|
from haiku.rag.store.models import SearchResult
|
|
|
|
|
|
class Dependencies(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
client: HaikuRAG
|
|
search_results: list[SearchResult] = []
|
|
|
|
|
|
class QuestionAnswerAgent:
|
|
def __init__(
|
|
self,
|
|
client: HaikuRAG,
|
|
model_config: ModelConfig,
|
|
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 = 5,
|
|
) -> 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)
|
|
results = await ctx.deps.client.expand_context(results)
|
|
# Store results for citation resolution
|
|
ctx.deps.search_results = results
|
|
# Format with chunk IDs
|
|
parts = []
|
|
for r in results:
|
|
parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")
|
|
return "\n\n".join(parts) if parts else "No results found."
|
|
|
|
async def answer(self, question: str) -> tuple[str, list[Citation]]:
|
|
"""Answer a question using the RAG system.
|
|
|
|
Returns:
|
|
Tuple of (answer text, list of resolved citations)
|
|
"""
|
|
deps = Dependencies(client=self._client)
|
|
result = await self._agent.run(question, deps=deps)
|
|
citations = resolve_citations(result.output.cited_chunks, deps.search_results)
|
|
return result.output.answer, citations
|