91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
import asyncio
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from pydantic_ai import Agent, RunContext
|
|
from pydantic_ai.format_prompt import format_as_xml
|
|
from pydantic_ai.output import ToolOutput
|
|
from pydantic_graph import BaseNode, GraphRunContext
|
|
|
|
from haiku.rag.research.common import get_model, log
|
|
from haiku.rag.research.dependencies import ResearchDependencies
|
|
from haiku.rag.research.models import ResearchReport, SearchAnswer
|
|
from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
|
|
from haiku.rag.research.state import ResearchDeps, ResearchState
|
|
|
|
|
|
@dataclass
|
|
class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
|
provider: str
|
|
model: str
|
|
|
|
async def run(
|
|
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
|
|
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
if not state.sub_questions:
|
|
from haiku.rag.research.nodes.evaluate import EvaluateNode
|
|
|
|
return EvaluateNode(self.provider, self.model)
|
|
|
|
# Take up to max_concurrency questions and answer them concurrently
|
|
take = max(1, state.max_concurrency)
|
|
batch: list[str] = []
|
|
while state.sub_questions and len(batch) < take:
|
|
batch.append(state.sub_questions.pop(0))
|
|
|
|
async def answer_one(sub_q: str) -> SearchAnswer | None:
|
|
log(
|
|
deps.console,
|
|
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
|
|
)
|
|
agent = Agent(
|
|
model=get_model(self.provider, self.model),
|
|
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
|
instructions=SEARCH_AGENT_PROMPT,
|
|
retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
@agent.tool
|
|
async def search_and_answer(
|
|
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 5
|
|
) -> str:
|
|
search_results = await ctx2.deps.client.search(query, limit=limit)
|
|
expanded = await ctx2.deps.client.expand_context(search_results)
|
|
|
|
entries: list[dict[str, Any]] = [
|
|
{
|
|
"text": chunk.content,
|
|
"score": score,
|
|
"document_uri": (chunk.document_uri or ""),
|
|
}
|
|
for chunk, score in expanded
|
|
]
|
|
if not entries:
|
|
return f"No relevant information found in the knowledge base for: {query}"
|
|
|
|
return format_as_xml(entries, root_tag="snippets")
|
|
|
|
agent_deps = ResearchDependencies(
|
|
client=deps.client, context=state.context, console=deps.console
|
|
)
|
|
try:
|
|
result = await agent.run(sub_q, deps=agent_deps)
|
|
except Exception as e:
|
|
log(deps.console, f"[red]Search failed:[/red] {e}")
|
|
return None
|
|
|
|
return result.output
|
|
|
|
answers = await asyncio.gather(*(answer_one(q) for q in batch))
|
|
for ans in answers:
|
|
if ans is None:
|
|
continue
|
|
state.context.add_qa_response(ans)
|
|
if deps.console:
|
|
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
|
|
log(deps.console, f" [green]✓[/green] {preview}")
|
|
|
|
return SearchDispatchNode(self.provider, self.model)
|