69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
from pydantic_ai import RunContext
|
|
from pydantic_ai.format_prompt import format_as_xml
|
|
from pydantic_ai.run import AgentRunResult
|
|
|
|
from haiku.rag.research.base import BaseResearchAgent, SearchAnswer
|
|
from haiku.rag.research.dependencies import ResearchDependencies
|
|
from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
|
|
|
|
|
|
class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
|
|
"""Agent specialized in answering questions using RAG search."""
|
|
|
|
def __init__(self, provider: str, model: str) -> None:
|
|
super().__init__(provider, model, output_type=SearchAnswer)
|
|
|
|
async def run(
|
|
self, prompt: str, deps: ResearchDependencies, **kwargs
|
|
) -> AgentRunResult[SearchAnswer]:
|
|
"""Execute the agent and persist the QA pair in shared context.
|
|
|
|
Pydantic AI enforces `SearchAnswer` as the output model; we just store
|
|
the QA response with the last search results as sources.
|
|
"""
|
|
console = deps.console
|
|
if console:
|
|
console.print(f"\t{prompt}")
|
|
|
|
result = await super().run(prompt, deps, **kwargs)
|
|
deps.context.add_qa_response(result.output)
|
|
deps.context.sub_questions.remove(prompt)
|
|
if console:
|
|
answer = result.output.answer
|
|
answer_preview = answer[:150] + "…" if len(answer) > 150 else answer
|
|
console.log(f"\n [green]✓[/green] {answer_preview}")
|
|
|
|
return result
|
|
|
|
def get_system_prompt(self) -> str:
|
|
return SEARCH_AGENT_PROMPT
|
|
|
|
def register_tools(self) -> None:
|
|
"""Register search-specific tools."""
|
|
|
|
@self.agent.tool
|
|
async def search_and_answer(
|
|
ctx: RunContext[ResearchDependencies],
|
|
query: str,
|
|
limit: int = 5,
|
|
) -> str:
|
|
"""Search the KB and return a concise context pack."""
|
|
search_results = await ctx.deps.client.search(query, limit=limit)
|
|
expanded = await ctx.deps.client.expand_context(search_results)
|
|
|
|
snippet_entries = [
|
|
{
|
|
"text": chunk.content,
|
|
"score": score,
|
|
"document_uri": (chunk.document_uri or ""),
|
|
}
|
|
for chunk, score in expanded
|
|
]
|
|
|
|
# Return an XML-formatted payload with the question and snippets.
|
|
if snippet_entries:
|
|
return format_as_xml(snippet_entries, root_tag="snippets")
|
|
else:
|
|
return (
|
|
f"No relevant information found in the knowledge base for: {query}"
|
|
)
|