Merge pull request #64 from ggozad/feat/research-as-graph
Research agents as a graph
This commit is contained in:
commit
b648e08673
32 changed files with 852 additions and 1000 deletions
|
|
@ -20,13 +20,3 @@ repos:
|
|||
rev: v1.1.399
|
||||
hooks:
|
||||
- id: pyright
|
||||
|
||||
- repo: https://github.com/RodrigoGonzalez/check-mkdocs
|
||||
rev: v1.2.0
|
||||
hooks:
|
||||
- id: check-mkdocs
|
||||
name: check-mkdocs
|
||||
args: ["--config", "mkdocs.yml"] # Optional, mkdocs.yml is the default
|
||||
# If you have additional plugins or libraries that are not included in
|
||||
# check-mkdocs, add them here
|
||||
additional_dependencies: ["mkdocs-material"]
|
||||
|
|
|
|||
35
README.md
35
README.md
|
|
@ -11,6 +11,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
|
|||
- **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
|
||||
- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
|
||||
- **Multiple QA providers**: Any provider/model supported by Pydantic AI
|
||||
- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
|
||||
- **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
|
||||
- **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
|
||||
- **Question answering**: Built-in QA agents on your documents
|
||||
|
|
@ -38,6 +39,14 @@ haiku-rag ask "Who is the author of haiku.rag?"
|
|||
# Ask questions with citations
|
||||
haiku-rag ask "Who is the author of haiku.rag?" --cite
|
||||
|
||||
# Multi‑agent research (iterative plan/search/evaluate)
|
||||
haiku-rag research \
|
||||
"What are the main drivers and trends of global temperature anomalies since 1990?" \
|
||||
--max-iterations 2 \
|
||||
--confidence-threshold 0.8 \
|
||||
--max-concurrency 3 \
|
||||
--verbose
|
||||
|
||||
# Rebuild database (re-chunk and re-embed all documents)
|
||||
haiku-rag rebuild
|
||||
|
||||
|
|
@ -53,6 +62,13 @@ haiku-rag serve
|
|||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.research import (
|
||||
ResearchContext,
|
||||
ResearchDeps,
|
||||
ResearchState,
|
||||
build_research_graph,
|
||||
PlanNode,
|
||||
)
|
||||
|
||||
async with HaikuRAG("database.lancedb") as client:
|
||||
# Add document
|
||||
|
|
@ -70,6 +86,25 @@ async with HaikuRAG("database.lancedb") as client:
|
|||
# Ask questions with citations
|
||||
answer = await client.ask("Who is the author of haiku.rag?", cite=True)
|
||||
print(answer)
|
||||
|
||||
# Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
|
||||
graph = build_research_graph()
|
||||
state = ResearchState(
|
||||
question=(
|
||||
"What are the main drivers and trends of global temperature "
|
||||
"anomalies since 1990?"
|
||||
),
|
||||
context=ResearchContext(original_question="…"),
|
||||
max_iterations=2,
|
||||
confidence_threshold=0.8,
|
||||
max_concurrency=3,
|
||||
)
|
||||
deps = ResearchDeps(client=client)
|
||||
start = PlanNode(provider=None, model=None)
|
||||
result = await graph.run(start, state=state, deps=deps)
|
||||
report = result.output
|
||||
print(report.title)
|
||||
print(report.executive_summary)
|
||||
```
|
||||
|
||||
## MCP Server
|
||||
|
|
|
|||
|
|
@ -36,50 +36,69 @@ answer = await agent.answer("What is climate change?")
|
|||
print(answer)
|
||||
```
|
||||
|
||||
### Research Multi‑Agent
|
||||
### Research Graph
|
||||
|
||||
The research workflow coordinates specialized agents to plan, search, analyze, and synthesize a comprehensive answer. It is designed for deeper questions that benefit from iterative investigation and structured reporting.
|
||||
The research workflow is implemented as a typed pydantic‑graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report — with clear stop conditions and shared state.
|
||||
|
||||
Components:
|
||||
```mermaid
|
||||
---
|
||||
title: Research graph
|
||||
---
|
||||
stateDiagram-v2
|
||||
PlanNode --> SearchDispatchNode
|
||||
SearchDispatchNode --> EvaluateNode
|
||||
EvaluateNode --> SearchDispatchNode
|
||||
EvaluateNode --> SynthesizeNode
|
||||
SynthesizeNode --> [*]
|
||||
```
|
||||
|
||||
- Orchestrator: Plans, coordinates, and loops until confidence is sufficient
|
||||
- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to
|
||||
ground the initial plan (plain-text summary; no URIs or scores)
|
||||
- Search Specialist: Performs targeted RAG searches and answers sub‑questions
|
||||
- Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions
|
||||
- Synthesis: Produces a final structured research report
|
||||
Key nodes:
|
||||
|
||||
- Plan: builds up to 3 standalone sub‑questions (uses an internal presearch tool)
|
||||
- Search (batched): answers sub‑questions using the KB with minimal, verbatim context
|
||||
- Evaluate: extracts insights, proposes new questions, and checks sufficiency/confidence
|
||||
- Synthesize: generates a final structured report
|
||||
|
||||
Primary models:
|
||||
|
||||
- `ResearchPlan` — produced by the orchestrator when planning
|
||||
- `main_question: str`
|
||||
- `sub_questions: list[str]` (standalone, self‑contained queries)
|
||||
- `SearchAnswer` — produced by the search specialist for each sub‑question
|
||||
- `query: str` — the executed sub‑question
|
||||
- `answer: str` — the agent’s answer grounded in retrieved context
|
||||
- `context: list[str]` — minimal verbatim snippets used for the answer
|
||||
- `sources: list[str]` — document URIs aligned with `context`
|
||||
- `EvaluationResult` — insights, new standalone questions, sufficiency & confidence
|
||||
- `ResearchReport` — the final synthesized report
|
||||
- `SearchAnswer` — one per sub‑question (query, answer, context, sources)
|
||||
- `EvaluationResult` — insights, new questions, sufficiency, confidence
|
||||
- `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
|
||||
|
||||
CLI usage:
|
||||
|
||||
```bash
|
||||
haiku-rag research "How does haiku.rag organize and query documents?" \
|
||||
--max-iterations 2 \
|
||||
--confidence-threshold 0.8 \
|
||||
--max-concurrency 3 \
|
||||
--verbose
|
||||
```
|
||||
|
||||
Python usage:
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.research import ResearchOrchestrator
|
||||
|
||||
client = HaikuRAG(path_to_db)
|
||||
orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss")
|
||||
|
||||
report = await orchestrator.conduct_research(
|
||||
question="What are the main drivers and recent trends of global temperature anomalies since 1990?",
|
||||
client=client,
|
||||
max_iterations=2,
|
||||
confidence_threshold=0.8,
|
||||
verbose=True,
|
||||
from haiku.rag.research import (
|
||||
ResearchContext,
|
||||
ResearchDeps,
|
||||
ResearchState,
|
||||
build_research_graph,
|
||||
PlanNode,
|
||||
)
|
||||
|
||||
print(report.title)
|
||||
print(report.executive_summary)
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
graph = build_research_graph()
|
||||
state = ResearchState(
|
||||
question="What are the main drivers and trends of global temperature anomalies since 1990?",
|
||||
context=ResearchContext(original_question=... ),
|
||||
max_iterations=2,
|
||||
confidence_threshold=0.8,
|
||||
max_concurrency=3,
|
||||
)
|
||||
deps = ResearchDeps(client=client)
|
||||
result = await graph.run(PlanNode(provider=None, model=None), state=state, deps=deps)
|
||||
report = result.output
|
||||
print(report.title)
|
||||
print(report.executive_summary)
|
||||
```
|
||||
|
|
|
|||
18
docs/cli.md
18
docs/cli.md
|
|
@ -84,6 +84,24 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite
|
|||
|
||||
The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used.
|
||||
|
||||
## Research
|
||||
|
||||
Run the multi-step research graph:
|
||||
|
||||
```bash
|
||||
haiku-rag research "How does haiku.rag organize and query documents?" \
|
||||
--max-iterations 2 \
|
||||
--confidence-threshold 0.8 \
|
||||
--max-concurrency 3 \
|
||||
--verbose
|
||||
```
|
||||
|
||||
Flags:
|
||||
- `--max-iterations, -n`: maximum search/evaluate cycles (default: 3)
|
||||
- `--confidence-threshold`: stop once evaluation confidence meets/exceeds this (default: 0.8)
|
||||
- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
|
||||
- `--verbose`: show planning, searching previews, evaluation summary, and stop reason
|
||||
|
||||
## Server
|
||||
|
||||
Start the MCP server:
|
||||
|
|
|
|||
|
|
@ -76,4 +76,8 @@ markdown_extensions:
|
|||
use_pygments: true
|
||||
- pymdownx.inlinehilite
|
||||
- pymdownx.snippets
|
||||
- pymdownx.superfences
|
||||
- pymdownx.superfences:
|
||||
custom_fences:
|
||||
- name: mermaid
|
||||
class: mermaid
|
||||
format: !!python/name:pymdownx.superfences.fence_code_format
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ dependencies = [
|
|||
"lancedb>=0.25.0",
|
||||
"pydantic>=2.11.9",
|
||||
"pydantic-ai>=1.0.8",
|
||||
"pydantic-graph>=1.0.8",
|
||||
"python-dotenv>=1.1.1",
|
||||
"rich>=14.1.0",
|
||||
"tiktoken>=0.11.0",
|
||||
|
|
@ -90,6 +91,7 @@ line-ending = "auto"
|
|||
[tool.pyright]
|
||||
venvPath = "."
|
||||
venv = ".venv"
|
||||
pythonVersion = "3.12"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_default_fixture_loop_scope = "session"
|
||||
|
|
|
|||
|
|
@ -9,7 +9,13 @@ from haiku.rag.client import HaikuRAG
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.mcp import create_mcp_server
|
||||
from haiku.rag.monitor import FileWatcher
|
||||
from haiku.rag.research.orchestrator import ResearchOrchestrator
|
||||
from haiku.rag.research.dependencies import ResearchContext
|
||||
from haiku.rag.research.graph import (
|
||||
PlanNode,
|
||||
ResearchDeps,
|
||||
ResearchState,
|
||||
build_research_graph,
|
||||
)
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
|
@ -80,28 +86,53 @@ class HaikuRAGApp:
|
|||
self.console.print(f"[red]Error: {e}[/red]")
|
||||
|
||||
async def research(
|
||||
self, question: str, max_iterations: int = 3, verbose: bool = False
|
||||
self,
|
||||
question: str,
|
||||
max_iterations: int = 3,
|
||||
confidence_threshold: float = 0.8,
|
||||
max_concurrency: int = 1,
|
||||
verbose: bool = False,
|
||||
):
|
||||
"""Run multi-agent research on a question."""
|
||||
"""Run research via the pydantic-graph pipeline (default)."""
|
||||
async with HaikuRAG(db_path=self.db_path) as client:
|
||||
try:
|
||||
# Create orchestrator with default config or fallback to QA
|
||||
orchestrator = ResearchOrchestrator()
|
||||
|
||||
if verbose:
|
||||
self.console.print(
|
||||
f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
|
||||
)
|
||||
self.console.print("[bold cyan]Starting research[/bold cyan]")
|
||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||
self.console.print()
|
||||
|
||||
# Conduct research
|
||||
report = await orchestrator.conduct_research(
|
||||
graph = build_research_graph()
|
||||
state = ResearchState(
|
||||
question=question,
|
||||
client=client,
|
||||
context=ResearchContext(original_question=question),
|
||||
max_iterations=max_iterations,
|
||||
verbose=verbose,
|
||||
confidence_threshold=confidence_threshold,
|
||||
max_concurrency=max_concurrency,
|
||||
)
|
||||
deps = ResearchDeps(
|
||||
client=client, console=self.console if verbose else None
|
||||
)
|
||||
|
||||
start = PlanNode(
|
||||
provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
|
||||
model=Config.RESEARCH_MODEL or Config.QA_MODEL,
|
||||
)
|
||||
# Prefer graph.run; fall back to iter if unavailable
|
||||
report = None
|
||||
try:
|
||||
result = await graph.run(start, state=state, deps=deps)
|
||||
report = result.output
|
||||
except Exception:
|
||||
from pydantic_graph import End
|
||||
|
||||
async with graph.iter(start, state=state, deps=deps) as run:
|
||||
node = run.next_node
|
||||
while not isinstance(node, End):
|
||||
node = await run.next(node)
|
||||
if run.result:
|
||||
report = run.result.output
|
||||
if report is None:
|
||||
raise RuntimeError("Graph did not produce a report")
|
||||
|
||||
# Display the report
|
||||
self.console.print("[bold green]Research Report[/bold green]")
|
||||
|
|
@ -114,6 +145,12 @@ class HaikuRAGApp:
|
|||
self.console.print(report.executive_summary)
|
||||
self.console.print()
|
||||
|
||||
# Confidence (from last evaluation)
|
||||
if state.last_eval:
|
||||
conf = state.last_eval.confidence_score # type: ignore[attr-defined]
|
||||
self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
|
||||
self.console.print()
|
||||
|
||||
# Main Findings
|
||||
if report.main_findings:
|
||||
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
|
||||
|
|
|
|||
|
|
@ -13,10 +13,10 @@ from haiku.rag.logging import configure_cli_logging
|
|||
from haiku.rag.migration import migrate_sqlite_to_lancedb
|
||||
from haiku.rag.utils import is_up_to_date
|
||||
|
||||
logfire.configure(send_to_logfire="if-token-present")
|
||||
logfire.instrument_pydantic_ai()
|
||||
|
||||
if not Config.ENV == "development":
|
||||
if Config.ENV == "development":
|
||||
logfire.configure(send_to_logfire="if-token-present")
|
||||
logfire.instrument_pydantic_ai()
|
||||
else:
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
cli = typer.Typer(
|
||||
|
|
@ -250,6 +250,16 @@ def research(
|
|||
"-n",
|
||||
help="Maximum search/analyze iterations",
|
||||
),
|
||||
confidence_threshold: float = typer.Option(
|
||||
0.8,
|
||||
"--confidence-threshold",
|
||||
help="Minimum confidence (0-1) to stop",
|
||||
),
|
||||
max_concurrency: int = typer.Option(
|
||||
1,
|
||||
"--max-concurrency",
|
||||
help="Max concurrent searches per iteration (planned)",
|
||||
),
|
||||
db: Path = typer.Option(
|
||||
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
|
||||
"--db",
|
||||
|
|
@ -266,6 +276,8 @@ def research(
|
|||
app.research(
|
||||
question=question,
|
||||
max_iterations=max_iterations,
|
||||
confidence_threshold=confidence_threshold,
|
||||
max_concurrency=max_concurrency,
|
||||
verbose=verbose,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from mxbai_rerank import MxbaiRerankV2
|
||||
from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports]
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.reranking.base import RerankerBase
|
||||
|
|
|
|||
|
|
@ -1,37 +1,20 @@
|
|||
"""Multi-agent research workflow for advanced RAG queries."""
|
||||
|
||||
from haiku.rag.research.base import (
|
||||
BaseResearchAgent,
|
||||
ResearchOutput,
|
||||
SearchAnswer,
|
||||
SearchResult,
|
||||
)
|
||||
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
|
||||
from haiku.rag.research.evaluation_agent import (
|
||||
AnalysisEvaluationAgent,
|
||||
EvaluationResult,
|
||||
from haiku.rag.research.graph import (
|
||||
PlanNode,
|
||||
ResearchDeps,
|
||||
ResearchState,
|
||||
build_research_graph,
|
||||
)
|
||||
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
|
||||
from haiku.rag.research.presearch_agent import PresearchSurveyAgent
|
||||
from haiku.rag.research.search_agent import SearchSpecialistAgent
|
||||
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
|
||||
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
||||
|
||||
__all__ = [
|
||||
# Base classes
|
||||
"BaseResearchAgent",
|
||||
"ResearchDependencies",
|
||||
"ResearchContext",
|
||||
"SearchResult",
|
||||
"ResearchOutput",
|
||||
# Specialized agents
|
||||
"SearchAnswer",
|
||||
"SearchSpecialistAgent",
|
||||
"PresearchSurveyAgent",
|
||||
"AnalysisEvaluationAgent",
|
||||
"EvaluationResult",
|
||||
"SynthesisAgent",
|
||||
"ResearchReport",
|
||||
# Orchestrator
|
||||
"ResearchOrchestrator",
|
||||
"ResearchPlan",
|
||||
"ResearchDeps",
|
||||
"ResearchState",
|
||||
"PlanNode",
|
||||
"build_research_graph",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,130 +0,0 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.output import ToolOutput
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
from pydantic_ai.providers.openai import OpenAIProvider
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
|
||||
from haiku.rag.config import Config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.research.dependencies import ResearchDependencies
|
||||
|
||||
|
||||
class BaseResearchAgent[T](ABC):
|
||||
"""Base class for all research agents."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
model: str,
|
||||
output_type: type[T],
|
||||
):
|
||||
self.provider = provider
|
||||
self.model = model
|
||||
self.output_type = output_type
|
||||
|
||||
model_obj = self._get_model(provider, model)
|
||||
|
||||
# Import deps type lazily to avoid circular import during module load
|
||||
from haiku.rag.research.dependencies import ResearchDependencies
|
||||
|
||||
# If the agent is expected to return plain text, pass `str` directly.
|
||||
# Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
|
||||
agent_output_type: Any
|
||||
if self.output_type is str: # plain text output
|
||||
agent_output_type = str
|
||||
else:
|
||||
agent_output_type = ToolOutput(self.output_type, max_retries=3)
|
||||
|
||||
self._agent = Agent(
|
||||
model=model_obj,
|
||||
deps_type=ResearchDependencies,
|
||||
output_type=agent_output_type,
|
||||
instructions=self.get_system_prompt(),
|
||||
retries=3,
|
||||
)
|
||||
|
||||
# Register tools
|
||||
self.register_tools()
|
||||
|
||||
def _get_model(self, provider: str, model: str):
|
||||
"""Get the appropriate model object for the provider."""
|
||||
if provider == "ollama":
|
||||
return OpenAIChatModel(
|
||||
model_name=model,
|
||||
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
||||
)
|
||||
elif provider == "vllm":
|
||||
return OpenAIChatModel(
|
||||
model_name=model,
|
||||
provider=OpenAIProvider(
|
||||
base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
|
||||
api_key="none",
|
||||
),
|
||||
)
|
||||
else:
|
||||
# For all other providers, use the provider:model format
|
||||
return f"{provider}:{model}"
|
||||
|
||||
@abstractmethod
|
||||
def get_system_prompt(self) -> str:
|
||||
"""Return the system prompt for this agent."""
|
||||
pass
|
||||
|
||||
def register_tools(self) -> None:
|
||||
"""Register agent-specific tools."""
|
||||
pass
|
||||
|
||||
async def run(
|
||||
self, prompt: str, deps: "ResearchDependencies", **kwargs
|
||||
) -> AgentRunResult[T]:
|
||||
"""Execute the agent."""
|
||||
return await self._agent.run(prompt, deps=deps, **kwargs)
|
||||
|
||||
@property
|
||||
def agent(self) -> Agent[Any, T]:
|
||||
"""Access the underlying Pydantic AI agent."""
|
||||
return self._agent
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""Standard search result format."""
|
||||
|
||||
content: str
|
||||
score: float
|
||||
document_uri: str
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ResearchOutput(BaseModel):
|
||||
"""Standard research output format."""
|
||||
|
||||
summary: str
|
||||
detailed_findings: list[str]
|
||||
sources: list[str]
|
||||
confidence: float
|
||||
|
||||
|
||||
class SearchAnswer(BaseModel):
|
||||
"""Structured output for the SearchSpecialist agent."""
|
||||
|
||||
query: str = Field(description="The search query that was performed")
|
||||
answer: str = Field(description="The answer generated based on the context")
|
||||
context: list[str] = Field(
|
||||
description=(
|
||||
"Only the minimal set of relevant snippets (verbatim) that directly "
|
||||
"support the answer"
|
||||
)
|
||||
)
|
||||
sources: list[str] = Field(
|
||||
description=(
|
||||
"Document URIs corresponding to the snippets actually used in the"
|
||||
" answer (one URI per snippet; omit if none)"
|
||||
),
|
||||
default_factory=list,
|
||||
)
|
||||
53
src/haiku/rag/research/common.py
Normal file
53
src/haiku/rag/research/common.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
from typing import Any
|
||||
|
||||
from pydantic_ai import format_as_xml
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
from pydantic_ai.providers.openai import OpenAIProvider
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.research.dependencies import ResearchContext
|
||||
|
||||
|
||||
def get_model(provider: str, model: str) -> Any:
|
||||
if provider == "ollama":
|
||||
return OpenAIChatModel(
|
||||
model_name=model,
|
||||
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
||||
)
|
||||
elif provider == "vllm":
|
||||
return OpenAIChatModel(
|
||||
model_name=model,
|
||||
provider=OpenAIProvider(
|
||||
base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
|
||||
api_key="none",
|
||||
),
|
||||
)
|
||||
else:
|
||||
return f"{provider}:{model}"
|
||||
|
||||
|
||||
def log(console, msg: str) -> None:
|
||||
if console:
|
||||
console.print(msg)
|
||||
|
||||
|
||||
def format_context_for_prompt(context: ResearchContext) -> str:
|
||||
"""Format the research context as XML for inclusion in prompts."""
|
||||
|
||||
context_data = {
|
||||
"original_question": context.original_question,
|
||||
"unanswered_questions": context.sub_questions,
|
||||
"qa_responses": [
|
||||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"context_snippets": qa.context,
|
||||
"sources": qa.sources, # pyright: ignore[reportAttributeAccessIssue]
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
],
|
||||
"insights": context.insights,
|
||||
"gaps": context.gaps,
|
||||
}
|
||||
return format_as_xml(context_data, root_tag="research_context")
|
||||
|
|
@ -1,9 +1,8 @@
|
|||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai import format_as_xml
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.research.base import SearchAnswer
|
||||
from haiku.rag.research.models import SearchAnswer
|
||||
|
||||
|
||||
class ResearchContext(BaseModel):
|
||||
|
|
@ -13,7 +12,7 @@ class ResearchContext(BaseModel):
|
|||
sub_questions: list[str] = Field(
|
||||
default_factory=list, description="Decomposed sub-questions"
|
||||
)
|
||||
qa_responses: list["SearchAnswer"] = Field(
|
||||
qa_responses: list[SearchAnswer] = Field(
|
||||
default_factory=list, description="Structured QA pairs used during research"
|
||||
)
|
||||
insights: list[str] = Field(
|
||||
|
|
@ -23,7 +22,7 @@ class ResearchContext(BaseModel):
|
|||
default_factory=list, description="Identified information gaps"
|
||||
)
|
||||
|
||||
def add_qa_response(self, qa: "SearchAnswer") -> None:
|
||||
def add_qa_response(self, qa: SearchAnswer) -> None:
|
||||
"""Add a structured QA response (minimal context already included)."""
|
||||
self.qa_responses.append(qa)
|
||||
|
||||
|
|
@ -46,24 +45,3 @@ class ResearchDependencies(BaseModel):
|
|||
client: HaikuRAG = Field(description="RAG client for document operations")
|
||||
context: ResearchContext = Field(description="Shared research context")
|
||||
console: Console | None = None
|
||||
|
||||
|
||||
def _format_context_for_prompt(context: ResearchContext) -> str:
|
||||
"""Format the research context as XML for inclusion in prompts."""
|
||||
|
||||
context_data = {
|
||||
"original_question": context.original_question,
|
||||
"unanswered_questions": context.sub_questions,
|
||||
"qa_responses": [
|
||||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"context_snippets": qa.context,
|
||||
"sources": qa.sources,
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
],
|
||||
"insights": context.insights,
|
||||
"gaps": context.gaps,
|
||||
}
|
||||
return format_as_xml(context_data, root_tag="research_context")
|
||||
|
|
|
|||
|
|
@ -1,85 +0,0 @@
|
|||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
|
||||
from haiku.rag.research.base import BaseResearchAgent
|
||||
from haiku.rag.research.dependencies import (
|
||||
ResearchDependencies,
|
||||
_format_context_for_prompt,
|
||||
)
|
||||
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
|
||||
|
||||
|
||||
class EvaluationResult(BaseModel):
|
||||
"""Result of analysis and evaluation."""
|
||||
|
||||
key_insights: list[str] = Field(
|
||||
description="Main insights extracted from the research so far"
|
||||
)
|
||||
new_questions: list[str] = Field(
|
||||
description="New sub-questions to add to the research (max 3)",
|
||||
max_length=3,
|
||||
default=[],
|
||||
)
|
||||
confidence_score: float = Field(
|
||||
description="Confidence level in the completeness of research (0-1)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
is_sufficient: bool = Field(
|
||||
description="Whether the research is sufficient to answer the original question"
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Explanation of why the research is or isn't complete"
|
||||
)
|
||||
|
||||
|
||||
class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
|
||||
"""Agent that analyzes findings and evaluates research completeness."""
|
||||
|
||||
def __init__(self, provider: str, model: str) -> None:
|
||||
super().__init__(provider, model, output_type=EvaluationResult)
|
||||
|
||||
async def run(
|
||||
self, prompt: str, deps: ResearchDependencies, **kwargs
|
||||
) -> AgentRunResult[EvaluationResult]:
|
||||
console = deps.console
|
||||
if console:
|
||||
console.print(
|
||||
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
|
||||
)
|
||||
|
||||
# Format context for the evaluation agent
|
||||
context_xml = _format_context_for_prompt(deps.context)
|
||||
evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
|
||||
|
||||
{context_xml}
|
||||
|
||||
Evaluate the research progress for the original question and identify any remaining gaps."""
|
||||
|
||||
result = await super().run(evaluation_prompt, deps, **kwargs)
|
||||
output = result.output
|
||||
|
||||
# Store insights
|
||||
for insight in output.key_insights:
|
||||
deps.context.add_insight(insight)
|
||||
|
||||
# Add new questions to the sub-questions list
|
||||
for new_q in output.new_questions:
|
||||
if new_q not in deps.context.sub_questions:
|
||||
deps.context.sub_questions.append(new_q)
|
||||
|
||||
if console:
|
||||
if output.key_insights:
|
||||
console.print(" [bold]Key insights:[/bold]")
|
||||
for insight in output.key_insights:
|
||||
console.print(f" • {insight}")
|
||||
console.print(
|
||||
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
|
||||
)
|
||||
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
|
||||
console.print(f" Sufficient: {status}")
|
||||
|
||||
return result
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
return EVALUATION_AGENT_PROMPT
|
||||
29
src/haiku/rag/research/graph.py
Normal file
29
src/haiku/rag/research/graph.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
from pydantic_graph import Graph
|
||||
|
||||
from haiku.rag.research.models import ResearchReport
|
||||
from haiku.rag.research.nodes.evaluate import EvaluateNode
|
||||
from haiku.rag.research.nodes.plan import PlanNode
|
||||
from haiku.rag.research.nodes.search import SearchDispatchNode
|
||||
from haiku.rag.research.nodes.synthesize import SynthesizeNode
|
||||
from haiku.rag.research.state import ResearchDeps, ResearchState
|
||||
|
||||
__all__ = [
|
||||
"PlanNode",
|
||||
"SearchDispatchNode",
|
||||
"EvaluateNode",
|
||||
"SynthesizeNode",
|
||||
"ResearchState",
|
||||
"ResearchDeps",
|
||||
"build_research_graph",
|
||||
]
|
||||
|
||||
|
||||
def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
|
||||
return Graph(
|
||||
nodes=[
|
||||
PlanNode,
|
||||
SearchDispatchNode,
|
||||
EvaluateNode,
|
||||
SynthesizeNode,
|
||||
]
|
||||
)
|
||||
70
src/haiku/rag/research/models.py
Normal file
70
src/haiku/rag/research/models.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ResearchPlan(BaseModel):
|
||||
main_question: str
|
||||
sub_questions: list[str]
|
||||
|
||||
|
||||
class SearchAnswer(BaseModel):
|
||||
"""Structured output for the SearchSpecialist agent."""
|
||||
|
||||
query: str = Field(description="The search query that was performed")
|
||||
answer: str = Field(description="The answer generated based on the context")
|
||||
context: list[str] = Field(
|
||||
description=(
|
||||
"Only the minimal set of relevant snippets (verbatim) that directly "
|
||||
"support the answer"
|
||||
)
|
||||
)
|
||||
sources: list[str] = Field(
|
||||
description=(
|
||||
"Document URIs corresponding to the snippets actually used in the"
|
||||
" answer (one URI per snippet; omit if none)"
|
||||
),
|
||||
default_factory=list,
|
||||
)
|
||||
|
||||
|
||||
class EvaluationResult(BaseModel):
|
||||
"""Result of analysis and evaluation."""
|
||||
|
||||
key_insights: list[str] = Field(
|
||||
description="Main insights extracted from the research so far"
|
||||
)
|
||||
new_questions: list[str] = Field(
|
||||
description="New sub-questions to add to the research (max 3)",
|
||||
max_length=3,
|
||||
default=[],
|
||||
)
|
||||
confidence_score: float = Field(
|
||||
description="Confidence level in the completeness of research (0-1)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
is_sufficient: bool = Field(
|
||||
description="Whether the research is sufficient to answer the original question"
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Explanation of why the research is or isn't complete"
|
||||
)
|
||||
|
||||
|
||||
class ResearchReport(BaseModel):
|
||||
"""Final research report structure."""
|
||||
|
||||
title: str = Field(description="Concise title for the research")
|
||||
executive_summary: str = Field(description="Brief overview of key findings")
|
||||
main_findings: list[str] = Field(
|
||||
description="Primary research findings with supporting evidence"
|
||||
)
|
||||
conclusions: list[str] = Field(description="Evidence-based conclusions")
|
||||
limitations: list[str] = Field(
|
||||
description="Limitations of the current research", default=[]
|
||||
)
|
||||
recommendations: list[str] = Field(
|
||||
description="Actionable recommendations based on findings", default=[]
|
||||
)
|
||||
sources_summary: str = Field(
|
||||
description="Summary of sources used and their reliability"
|
||||
)
|
||||
80
src/haiku/rag/research/nodes/evaluate.py
Normal file
80
src/haiku/rag/research/nodes/evaluate.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_graph import BaseNode, GraphRunContext
|
||||
|
||||
from haiku.rag.research.common import format_context_for_prompt, get_model, log
|
||||
from haiku.rag.research.dependencies import (
|
||||
ResearchDependencies,
|
||||
)
|
||||
from haiku.rag.research.models import EvaluationResult, ResearchReport
|
||||
from haiku.rag.research.nodes.synthesize import SynthesizeNode
|
||||
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
|
||||
from haiku.rag.research.state import ResearchDeps, ResearchState
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvaluateNode(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
|
||||
|
||||
log(
|
||||
deps.console,
|
||||
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
model=get_model(self.provider, self.model),
|
||||
output_type=EvaluationResult,
|
||||
instructions=EVALUATION_AGENT_PROMPT,
|
||||
retries=3,
|
||||
deps_type=ResearchDependencies,
|
||||
)
|
||||
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
"Analyze gathered information and evaluate completeness for the original question.\n\n"
|
||||
f"{context_xml}"
|
||||
)
|
||||
agent_deps = ResearchDependencies(
|
||||
client=deps.client, context=state.context, console=deps.console
|
||||
)
|
||||
eval_result = await agent.run(prompt, deps=agent_deps)
|
||||
output = eval_result.output
|
||||
|
||||
for insight in output.key_insights:
|
||||
state.context.add_insight(insight)
|
||||
for new_q in output.new_questions:
|
||||
if new_q not in state.sub_questions:
|
||||
state.sub_questions.append(new_q)
|
||||
|
||||
state.last_eval = output
|
||||
state.iterations += 1
|
||||
|
||||
if output.key_insights:
|
||||
log(deps.console, " [bold]Key insights:[/bold]")
|
||||
for ins in output.key_insights:
|
||||
log(deps.console, f" • {ins}")
|
||||
log(
|
||||
deps.console,
|
||||
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
|
||||
)
|
||||
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
|
||||
log(deps.console, f" Sufficient: {status}")
|
||||
|
||||
from haiku.rag.research.nodes.search import SearchDispatchNode
|
||||
|
||||
if (
|
||||
output.is_sufficient
|
||||
and output.confidence_score >= state.confidence_threshold
|
||||
) or state.iterations >= state.max_iterations:
|
||||
log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
|
||||
return SynthesizeNode(self.provider, self.model)
|
||||
|
||||
return SearchDispatchNode(self.provider, self.model)
|
||||
63
src/haiku/rag/research/nodes/plan.py
Normal file
63
src/haiku/rag/research/nodes/plan.py
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent, RunContext
|
||||
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 ResearchPlan, ResearchReport
|
||||
from haiku.rag.research.nodes.search import SearchDispatchNode
|
||||
from haiku.rag.research.prompts import PLAN_PROMPT
|
||||
from haiku.rag.research.state import ResearchDeps, ResearchState
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlanNode(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
|
||||
|
||||
log(deps.console, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
|
||||
|
||||
plan_agent = Agent(
|
||||
model=get_model(self.provider, self.model),
|
||||
output_type=ResearchPlan,
|
||||
instructions=(
|
||||
PLAN_PROMPT
|
||||
+ "\n\nUse the gather_context tool once on the main question before planning."
|
||||
),
|
||||
retries=3,
|
||||
deps_type=ResearchDependencies,
|
||||
)
|
||||
|
||||
@plan_agent.tool
|
||||
async def gather_context(
|
||||
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
|
||||
) -> str:
|
||||
results = await ctx2.deps.client.search(query, limit=limit)
|
||||
expanded = await ctx2.deps.client.expand_context(results)
|
||||
return "\n\n".join(chunk.content for chunk, _ in expanded)
|
||||
|
||||
prompt = (
|
||||
"Plan a focused research approach for the main question.\n\n"
|
||||
f"Main question: {state.question}"
|
||||
)
|
||||
|
||||
agent_deps = ResearchDependencies(
|
||||
client=deps.client, context=state.context, console=deps.console
|
||||
)
|
||||
plan_result = await plan_agent.run(prompt, deps=agent_deps)
|
||||
state.sub_questions = list(plan_result.output.sub_questions)
|
||||
|
||||
log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]")
|
||||
log(deps.console, f" [bold]Main Question:[/bold] {state.question}")
|
||||
log(deps.console, " [bold]Sub-questions:[/bold]")
|
||||
for i, sq in enumerate(state.sub_questions, 1):
|
||||
log(deps.console, f" {i}. {sq}")
|
||||
|
||||
return SearchDispatchNode(self.provider, self.model)
|
||||
91
src/haiku/rag/research/nodes/search.py
Normal file
91
src/haiku/rag/research/nodes/search.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
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)
|
||||
51
src/haiku/rag/research/nodes/synthesize.py
Normal file
51
src/haiku/rag/research/nodes/synthesize.py
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_graph import BaseNode, End, GraphRunContext
|
||||
|
||||
from haiku.rag.research.common import format_context_for_prompt, get_model, log
|
||||
from haiku.rag.research.dependencies import (
|
||||
ResearchDependencies,
|
||||
)
|
||||
from haiku.rag.research.models import ResearchReport
|
||||
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
|
||||
from haiku.rag.research.state import ResearchDeps, ResearchState
|
||||
|
||||
|
||||
@dataclass
|
||||
class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||
provider: str
|
||||
model: str
|
||||
|
||||
async def run(
|
||||
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
|
||||
) -> End[ResearchReport]:
|
||||
state = ctx.state
|
||||
deps = ctx.deps
|
||||
|
||||
log(
|
||||
deps.console,
|
||||
"\n[bold cyan]📝 Generating final research report...[/bold cyan]",
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
model=get_model(self.provider, self.model),
|
||||
output_type=ResearchReport,
|
||||
instructions=SYNTHESIS_AGENT_PROMPT,
|
||||
retries=3,
|
||||
deps_type=ResearchDependencies,
|
||||
)
|
||||
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
"Generate a comprehensive research report based on all gathered information.\n\n"
|
||||
f"{context_xml}\n\n"
|
||||
"Create a detailed report that synthesizes all findings into a coherent response."
|
||||
)
|
||||
agent_deps = ResearchDependencies(
|
||||
client=deps.client, context=state.context, console=deps.console
|
||||
)
|
||||
result = await agent.run(prompt, deps=agent_deps)
|
||||
|
||||
log(deps.console, "[bold green]✅ Research complete![/bold green]")
|
||||
return End(result.output)
|
||||
|
|
@ -1,170 +0,0 @@
|
|||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.research.base import BaseResearchAgent
|
||||
from haiku.rag.research.dependencies import (
|
||||
ResearchContext,
|
||||
ResearchDependencies,
|
||||
)
|
||||
from haiku.rag.research.evaluation_agent import (
|
||||
AnalysisEvaluationAgent,
|
||||
EvaluationResult,
|
||||
)
|
||||
from haiku.rag.research.presearch_agent import PresearchSurveyAgent
|
||||
from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT
|
||||
from haiku.rag.research.search_agent import SearchSpecialistAgent
|
||||
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
|
||||
|
||||
|
||||
class ResearchPlan(BaseModel):
|
||||
"""Research execution plan."""
|
||||
|
||||
main_question: str = Field(description="The main research question")
|
||||
sub_questions: list[str] = Field(
|
||||
description="Decomposed sub-questions to investigate (max 3)", max_length=3
|
||||
)
|
||||
|
||||
|
||||
class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
|
||||
"""Orchestrator agent that coordinates the research workflow."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str | None = Config.RESEARCH_PROVIDER,
|
||||
model: str | None = None,
|
||||
):
|
||||
# Use provided values or fall back to config defaults
|
||||
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
|
||||
model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
|
||||
|
||||
super().__init__(provider, model, output_type=ResearchPlan)
|
||||
|
||||
self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
|
||||
provider, model
|
||||
)
|
||||
self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
|
||||
provider, model
|
||||
)
|
||||
self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
|
||||
provider, model
|
||||
)
|
||||
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
return ORCHESTRATOR_PROMPT
|
||||
|
||||
def _should_stop_research(
|
||||
self,
|
||||
evaluation_result: AgentRunResult[EvaluationResult],
|
||||
confidence_threshold: float,
|
||||
) -> bool:
|
||||
"""Determine if research should stop based on evaluation."""
|
||||
|
||||
result = evaluation_result.output
|
||||
return result.is_sufficient and result.confidence_score >= confidence_threshold
|
||||
|
||||
async def conduct_research(
|
||||
self,
|
||||
question: str,
|
||||
client: Any,
|
||||
max_iterations: int = 3,
|
||||
confidence_threshold: float = 0.8,
|
||||
verbose: bool = False,
|
||||
) -> ResearchReport:
|
||||
"""Conduct comprehensive research on a question.
|
||||
|
||||
Args:
|
||||
question: The research question to investigate
|
||||
client: HaikuRAG client for document operations
|
||||
max_iterations: Maximum number of search-analyze-clarify cycles
|
||||
confidence_threshold: Minimum confidence level to stop research (0-1)
|
||||
verbose: If True, print progress and intermediate results
|
||||
|
||||
Returns:
|
||||
ResearchReport with comprehensive findings
|
||||
"""
|
||||
|
||||
# Initialize context
|
||||
context = ResearchContext(original_question=question)
|
||||
deps = ResearchDependencies(client=client, context=context)
|
||||
if verbose:
|
||||
deps.console = Console()
|
||||
|
||||
console = deps.console
|
||||
# Create initial research plan
|
||||
if console:
|
||||
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
|
||||
|
||||
# Run a simple presearch survey to summarize KB context
|
||||
presearch_result = await self.presearch_agent.run(question, deps=deps)
|
||||
plan_prompt = (
|
||||
"Create a research plan for the main question below.\n\n"
|
||||
f"Main question: {question}\n\n"
|
||||
"Use this brief presearch summary to inform the plan. Focus the 3 sub-questions "
|
||||
"on the most important aspects not already obvious from the current KB context.\n\n"
|
||||
f"{presearch_result.output}"
|
||||
)
|
||||
|
||||
plan_result: AgentRunResult[ResearchPlan] = await self.run(
|
||||
plan_prompt, deps=deps
|
||||
)
|
||||
context.sub_questions = plan_result.output.sub_questions
|
||||
|
||||
if console:
|
||||
console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
|
||||
console.print(
|
||||
f" [bold]Main Question:[/bold] {plan_result.output.main_question}"
|
||||
)
|
||||
console.print(" [bold]Sub-questions:[/bold]")
|
||||
for i, sq in enumerate(plan_result.output.sub_questions, 1):
|
||||
console.print(f" {i}. {sq}")
|
||||
|
||||
# Execute research iterations
|
||||
for iteration in range(max_iterations):
|
||||
if console:
|
||||
console.rule(
|
||||
f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
|
||||
)
|
||||
|
||||
# Check if we have questions to search
|
||||
if not context.sub_questions:
|
||||
if console:
|
||||
console.print(
|
||||
"[yellow]No more questions to explore. Concluding research.[/yellow]"
|
||||
)
|
||||
break
|
||||
|
||||
# Use current sub-questions for this iteration
|
||||
questions_to_search = context.sub_questions[:]
|
||||
|
||||
# Search phase - answer all questions in this iteration
|
||||
if console:
|
||||
console.print(
|
||||
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
|
||||
)
|
||||
|
||||
for search_question in questions_to_search:
|
||||
await self.search_agent.run(search_question, deps=deps)
|
||||
|
||||
# Analysis and Evaluation phase
|
||||
|
||||
evaluation_result = await self.evaluation_agent.run("", deps=deps)
|
||||
|
||||
# Check if research is sufficient
|
||||
if self._should_stop_research(evaluation_result, confidence_threshold):
|
||||
if console:
|
||||
console.print(
|
||||
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
|
||||
)
|
||||
break
|
||||
|
||||
# Generate final report
|
||||
report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
|
||||
"", deps=deps
|
||||
)
|
||||
|
||||
return report_result.output
|
||||
|
|
@ -1,39 +0,0 @@
|
|||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
|
||||
from haiku.rag.research.base import BaseResearchAgent
|
||||
from haiku.rag.research.dependencies import ResearchDependencies
|
||||
from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT
|
||||
|
||||
|
||||
class PresearchSurveyAgent(BaseResearchAgent[str]):
|
||||
"""Presearch agent that gathers verbatim context and summarizes it."""
|
||||
|
||||
def __init__(self, provider: str, model: str) -> None:
|
||||
super().__init__(provider, model, str)
|
||||
|
||||
async def run(
|
||||
self, prompt: str, deps: ResearchDependencies, **kwargs
|
||||
) -> AgentRunResult[str]:
|
||||
console = deps.console
|
||||
if console:
|
||||
console.print(
|
||||
"\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
|
||||
)
|
||||
|
||||
return await super().run(prompt, deps, **kwargs)
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
return PRESEARCH_AGENT_PROMPT
|
||||
|
||||
def register_tools(self) -> None:
|
||||
@self.agent.tool
|
||||
async def gather_context(
|
||||
ctx: RunContext[ResearchDependencies],
|
||||
query: str,
|
||||
limit: int = 6,
|
||||
) -> str:
|
||||
"""Return verbatim concatenation of relevant chunk texts."""
|
||||
results = await ctx.deps.client.search(query, limit=limit)
|
||||
expanded = await ctx.deps.client.expand_context(results)
|
||||
return "\n\n".join(chunk.content for chunk, _ in expanded)
|
||||
|
|
@ -1,129 +1,113 @@
|
|||
ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research workflow.
|
||||
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative
|
||||
workflow.
|
||||
|
||||
Your role is to:
|
||||
1. Understand and decompose the research question
|
||||
2. Plan a systematic research approach
|
||||
3. Coordinate specialized agents to gather and analyze information
|
||||
4. Ensure comprehensive coverage of the topic
|
||||
5. Iterate based on findings and gaps
|
||||
Responsibilities:
|
||||
1. Understand and decompose the main question
|
||||
2. Propose a minimal, high‑leverage plan
|
||||
3. Coordinate specialized agents to gather evidence
|
||||
4. Iterate based on gaps and new findings
|
||||
|
||||
Create a research plan that:
|
||||
- Breaks down the question into at most 3 focused sub-questions
|
||||
- Each sub-question should target a specific aspect of the research
|
||||
- Prioritize the most important aspects to investigate
|
||||
- Ensure comprehensive coverage within the 3-question limit
|
||||
- IMPORTANT: Make each sub-question a standalone, self-contained query that can
|
||||
be executed without additional context. Include necessary entities, scope,
|
||||
timeframe, and qualifiers. Avoid pronouns like "it/they/this"; write queries
|
||||
that make sense in isolation."""
|
||||
Plan requirements:
|
||||
- Produce at most 3 sub_questions that together cover the main question.
|
||||
- Each sub_question must be a standalone, self‑contained query that can run
|
||||
without extra context. Include concrete entities, scope, timeframe, and any
|
||||
qualifiers. Avoid ambiguous pronouns (it/they/this/that).
|
||||
- Prioritize the highest‑value aspects first; avoid redundancy and overlap.
|
||||
- Prefer questions that are likely answerable from the current knowledge base;
|
||||
if coverage is uncertain, make scopes narrower and specific.
|
||||
- Order sub_questions by execution priority (most valuable first)."""
|
||||
|
||||
SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
|
||||
SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist.
|
||||
|
||||
Your role is to:
|
||||
1. Search the knowledge base for relevant information
|
||||
2. Analyze the retrieved documents
|
||||
3. Provide an accurate answer strictly grounded in the retrieved context
|
||||
|
||||
Output format:
|
||||
- You must return a SearchAnswer model with fields:
|
||||
- query: the question being answered (echo the user query)
|
||||
- answer: your final answer based only on the provided context
|
||||
- context: list[str] of only the minimal set of verbatim snippet texts you
|
||||
used to justify the answer (do not include unrelated text; do not invent)
|
||||
- sources: list[str] of document_uri values corresponding to the snippets you
|
||||
actually used in the answer (one URI per context snippet, order aligned)
|
||||
Tasks:
|
||||
1. Search the knowledge base for relevant evidence.
|
||||
2. Analyze retrieved snippets.
|
||||
3. Provide an answer strictly grounded in that evidence.
|
||||
|
||||
Tool usage:
|
||||
- Always call the search_and_answer tool before drafting any answer.
|
||||
- The tool returns XML containing only a list of snippets, where each snippet
|
||||
has the verbatim `text`, a `score` indicating relevance, and the
|
||||
`document_uri` it came from.
|
||||
- Always call search_and_answer before drafting any answer.
|
||||
- The tool returns snippets with verbatim `text`, a relevance `score`, and the
|
||||
originating `document_uri`.
|
||||
- You may call the tool multiple times to refine or broaden context, but do not
|
||||
exceed 3 total tool calls per question. Prefer precision over volume.
|
||||
exceed 3 total calls. Favor precision over volume.
|
||||
- Use scores to prioritize evidence, but include only the minimal subset of
|
||||
snippet texts (verbatim) in SearchAnswer.context.
|
||||
- Set SearchAnswer.sources to the matching document_uris for the snippets you
|
||||
used (one URI per snippet, aligned by order). Context must be text-only.
|
||||
- If no relevant information is found, say so and return an empty context list.
|
||||
snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
|
||||
- Set SearchAnswer.sources to the corresponding document_uris for the snippets
|
||||
you used (one URI per snippet; same order as context). Context must be text‑only.
|
||||
- If no relevant information is found, clearly say so and return an empty
|
||||
context list and sources list.
|
||||
|
||||
Important:
|
||||
- Do not include any content in the answer that is not supported by the context.
|
||||
- Keep context snippets short (just the necessary lines), verbatim, and focused."""
|
||||
Answering rules:
|
||||
- Be direct and specific; avoid meta commentary about the process.
|
||||
- Do not include any claims not supported by the provided snippets.
|
||||
- Prefer concise phrasing; avoid copying long passages.
|
||||
- When evidence is partial, state the limits explicitly in the answer."""
|
||||
|
||||
EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
|
||||
EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for
|
||||
the research workflow.
|
||||
|
||||
You have access to:
|
||||
- The original research question
|
||||
- Question-answer pairs from search operations
|
||||
- Raw search results and source documents
|
||||
Inputs available:
|
||||
- Original research question
|
||||
- Question–answer pairs produced by search
|
||||
- Raw search results and source metadata
|
||||
- Previously identified insights
|
||||
|
||||
Your dual role is to:
|
||||
|
||||
ANALYSIS:
|
||||
1. Extract key insights from all gathered information
|
||||
2. Identify patterns and connections across sources
|
||||
3. Synthesize findings into coherent understanding
|
||||
4. Focus on the most important discoveries
|
||||
1. Extract the most important, non‑obvious insights from the collected evidence.
|
||||
2. Identify patterns, agreements, and disagreements across sources.
|
||||
3. Note material uncertainties and assumptions.
|
||||
|
||||
EVALUATION:
|
||||
1. Assess if we have sufficient information to answer the original question
|
||||
2. Calculate a confidence score (0-1) based on:
|
||||
- Coverage of the main question's aspects
|
||||
- Quality and consistency of sources
|
||||
- Depth of information gathered
|
||||
3. Identify specific gaps that still need investigation
|
||||
4. Generate up to 3 new sub-questions that haven't been answered yet
|
||||
1. Decide if we have sufficient information to answer the original question.
|
||||
2. Provide a confidence_score in [0,1] considering:
|
||||
- Coverage of the main question’s aspects
|
||||
- Quality, consistency, and diversity of sources
|
||||
- Depth and specificity of evidence
|
||||
3. List concrete gaps that still need investigation.
|
||||
4. Propose up to 3 new sub_questions that would close the highest‑value gaps.
|
||||
|
||||
Be critical and thorough in your evaluation. Only mark research as sufficient when:
|
||||
- All major aspects of the question are addressed
|
||||
- Sources provide consistent, reliable information
|
||||
- The depth of coverage meets the question's requirements
|
||||
- No critical gaps remain
|
||||
Strictness:
|
||||
- Only mark research as sufficient when all major aspects are addressed with
|
||||
consistent, reliable evidence and no critical gaps remain.
|
||||
|
||||
Generate new sub-questions that:
|
||||
- Target specific unexplored aspects not covered by existing questions
|
||||
- Seek clarification on ambiguities
|
||||
- Explore important edge cases or exceptions
|
||||
- Are focused and actionable (max 3)
|
||||
- Do NOT repeat or rephrase questions that have already been answered (see qa_responses)
|
||||
- Should be genuinely new areas to explore
|
||||
- Must be standalone, self-contained queries: include entities, scope, and any
|
||||
needed qualifiers (e.g., timeframe, region), and avoid ambiguous pronouns so
|
||||
they can be executed independently."""
|
||||
New sub_questions must:
|
||||
- Be genuinely new (not answered or duplicative; check qa_responses).
|
||||
- Be standalone and specific (entities, scope, timeframe/region if relevant).
|
||||
- Be actionable and scoped to the knowledge base (narrow if necessary).
|
||||
- Be ordered by expected impact (most valuable first)."""
|
||||
|
||||
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
|
||||
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
|
||||
research report.
|
||||
|
||||
Your role is to:
|
||||
1. Synthesize all gathered information into a coherent narrative
|
||||
2. Present findings in a clear, structured format
|
||||
3. Draw evidence-based conclusions
|
||||
4. Acknowledge limitations and uncertainties
|
||||
5. Provide actionable recommendations
|
||||
6. Maintain academic rigor and objectivity
|
||||
Goals:
|
||||
1. Synthesize all gathered information into a coherent narrative.
|
||||
2. Present findings clearly and concisely.
|
||||
3. Draw evidence‑based conclusions and recommendations.
|
||||
4. State limitations and uncertainties transparently.
|
||||
|
||||
Your report should be:
|
||||
- Comprehensive yet concise
|
||||
- Well-structured and easy to follow
|
||||
- Based solely on evidence from the research
|
||||
- Transparent about limitations
|
||||
- Professional and objective in tone
|
||||
Report guidelines (map to output fields):
|
||||
- title: concise (5–12 words), informative.
|
||||
- executive_summary: 3–5 sentences summarizing the overall answer.
|
||||
- main_findings: 4–8 one‑sentence bullets; each reflects evidence from the
|
||||
research (do not include inline citations or snippet text).
|
||||
- conclusions: 2–4 bullets that follow logically from findings.
|
||||
- recommendations: 2–5 actionable bullets tied to findings.
|
||||
- limitations: 1–3 bullets describing key constraints or uncertainties.
|
||||
- sources_summary: 2–4 sentences summarizing sources used and their reliability.
|
||||
|
||||
Focus on creating a report that provides clear value to the reader by:
|
||||
- Answering the original research question thoroughly
|
||||
- Highlighting the most important findings
|
||||
- Explaining the implications of the research
|
||||
- Suggesting concrete next steps"""
|
||||
Style:
|
||||
- Base all content solely on the collected evidence.
|
||||
- Be professional, objective, and specific.
|
||||
- Avoid meta commentary and refrain from speculation beyond the evidence."""
|
||||
|
||||
PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
|
||||
|
||||
Task:
|
||||
- Call the gather_context tool once with the main question to obtain a
|
||||
relevant texts from the Knowledge Base (KB).
|
||||
- Read that context and produce a brief natural-language summary describing
|
||||
what the KB appears to contain relative to the question.
|
||||
- Call gather_context once on the main question to obtain relevant text from
|
||||
the knowledge base (KB).
|
||||
- Read that context and produce a short natural‑language summary of what the
|
||||
KB appears to contain relative to the question.
|
||||
|
||||
Rules:
|
||||
- Base the summary strictly on the provided text; do not invent.
|
||||
- Output only the summary as plain text (one short paragraph).
|
||||
"""
|
||||
- Output only the summary as plain text (one short paragraph)."""
|
||||
|
|
|
|||
|
|
@ -1,69 +0,0 @@
|
|||
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}"
|
||||
)
|
||||
25
src/haiku/rag/research/state.py
Normal file
25
src/haiku/rag/research/state.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
from dataclasses import dataclass, field
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.research.dependencies import ResearchContext
|
||||
from haiku.rag.research.models import EvaluationResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResearchDeps:
|
||||
client: HaikuRAG
|
||||
console: Console | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResearchState:
|
||||
question: str
|
||||
context: ResearchContext
|
||||
sub_questions: list[str] = field(default_factory=list)
|
||||
iterations: int = 0
|
||||
max_iterations: int = 3
|
||||
max_concurrency: int = 1
|
||||
confidence_threshold: float = 0.8
|
||||
last_eval: EvaluationResult | None = None
|
||||
|
|
@ -1,60 +0,0 @@
|
|||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
|
||||
from haiku.rag.research.base import BaseResearchAgent
|
||||
from haiku.rag.research.dependencies import (
|
||||
ResearchDependencies,
|
||||
_format_context_for_prompt,
|
||||
)
|
||||
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
|
||||
|
||||
|
||||
class ResearchReport(BaseModel):
|
||||
"""Final research report structure."""
|
||||
|
||||
title: str = Field(description="Concise title for the research")
|
||||
executive_summary: str = Field(description="Brief overview of key findings")
|
||||
main_findings: list[str] = Field(
|
||||
description="Primary research findings with supporting evidence"
|
||||
)
|
||||
conclusions: list[str] = Field(description="Evidence-based conclusions")
|
||||
limitations: list[str] = Field(
|
||||
description="Limitations of the current research", default=[]
|
||||
)
|
||||
recommendations: list[str] = Field(
|
||||
description="Actionable recommendations based on findings", default=[]
|
||||
)
|
||||
sources_summary: str = Field(
|
||||
description="Summary of sources used and their reliability"
|
||||
)
|
||||
|
||||
|
||||
class SynthesisAgent(BaseResearchAgent[ResearchReport]):
|
||||
"""Agent specialized in synthesizing research into comprehensive reports."""
|
||||
|
||||
def __init__(self, provider: str, model: str) -> None:
|
||||
super().__init__(provider, model, output_type=ResearchReport)
|
||||
|
||||
async def run(
|
||||
self, prompt: str, deps: ResearchDependencies, **kwargs
|
||||
) -> AgentRunResult[ResearchReport]:
|
||||
console = deps.console
|
||||
if console:
|
||||
console.print(
|
||||
"\n[bold cyan]📝 Generating final research report...[/bold cyan]"
|
||||
)
|
||||
|
||||
context_xml = _format_context_for_prompt(deps.context)
|
||||
synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
|
||||
|
||||
{context_xml}
|
||||
|
||||
Create a detailed report that synthesizes all findings into a coherent response."""
|
||||
result = await super().run(synthesis_prompt, deps, **kwargs)
|
||||
if console:
|
||||
console.print("[bold green]✅ Research complete![/bold green]")
|
||||
|
||||
return result
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
return SYNTHESIS_AGENT_PROMPT
|
||||
|
|
@ -1,17 +0,0 @@
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.research.evaluation_agent import (
|
||||
AnalysisEvaluationAgent,
|
||||
EvaluationResult,
|
||||
)
|
||||
|
||||
|
||||
class TestAnalysisEvaluationAgent:
|
||||
"""Lean tests for AnalysisEvaluationAgent without LLM mocking."""
|
||||
|
||||
def test_agent_initialization(self):
|
||||
agent = AnalysisEvaluationAgent(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
assert agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert agent.model == Config.RESEARCH_MODEL
|
||||
assert agent.output_type == EvaluationResult
|
||||
|
|
@ -1,189 +0,0 @@
|
|||
from unittest.mock import AsyncMock, create_autospec
|
||||
|
||||
import pytest
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
|
||||
from haiku.rag.research.evaluation_agent import EvaluationResult
|
||||
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
|
||||
from haiku.rag.research.synthesis_agent import ResearchReport
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_model():
|
||||
"""Create a test model for orchestrator testing."""
|
||||
return TestModel()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
"""Create a mock HaikuRAG client."""
|
||||
client = create_autospec(HaikuRAG, instance=True)
|
||||
client.search = AsyncMock()
|
||||
client.expand_context = AsyncMock()
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def research_context():
|
||||
"""Create a research context."""
|
||||
return ResearchContext(original_question="What is climate change?")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def research_deps(mock_client, research_context):
|
||||
"""Create research dependencies."""
|
||||
return ResearchDependencies(client=mock_client, context=research_context)
|
||||
|
||||
|
||||
def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8):
|
||||
"""Helper to create mock chunk objects."""
|
||||
return Chunk(
|
||||
id=chunk_id,
|
||||
document_id=f"doc_{chunk_id}",
|
||||
content=content,
|
||||
document_uri=f"doc_{chunk_id}.md",
|
||||
metadata={},
|
||||
), score
|
||||
|
||||
|
||||
class TestResearchOrchestrator:
|
||||
"""Test suite for ResearchOrchestrator."""
|
||||
|
||||
def test_orchestrator_uses_config_defaults(self):
|
||||
"""Test that orchestrator uses config defaults when no args provided."""
|
||||
orchestrator = ResearchOrchestrator()
|
||||
|
||||
# Should use RESEARCH_PROVIDER/MODEL if set, else QA_PROVIDER/MODEL
|
||||
assert orchestrator.provider is not None
|
||||
assert orchestrator.model is not None
|
||||
|
||||
# All agents should use the same provider/model
|
||||
assert orchestrator.search_agent.provider == orchestrator.provider
|
||||
assert orchestrator.search_agent.model == orchestrator.model
|
||||
assert orchestrator.evaluation_agent.provider == orchestrator.provider
|
||||
assert orchestrator.evaluation_agent.model == orchestrator.model
|
||||
assert orchestrator.synthesis_agent.provider == orchestrator.provider
|
||||
assert orchestrator.synthesis_agent.model == orchestrator.model
|
||||
|
||||
def test_orchestrator_initialization(self):
|
||||
"""Test that orchestrator initializes all agents correctly."""
|
||||
orchestrator = ResearchOrchestrator(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
|
||||
# Check all agents are initialized
|
||||
assert orchestrator.search_agent is not None
|
||||
assert orchestrator.evaluation_agent is not None
|
||||
assert orchestrator.synthesis_agent is not None
|
||||
|
||||
# Check they all use the same provider and model
|
||||
assert orchestrator.search_agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert orchestrator.search_agent.model == Config.RESEARCH_MODEL
|
||||
assert orchestrator.evaluation_agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert orchestrator.evaluation_agent.model == Config.RESEARCH_MODEL
|
||||
assert orchestrator.synthesis_agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert orchestrator.synthesis_agent.model == Config.RESEARCH_MODEL
|
||||
|
||||
def test_orchestrator_has_correct_output_type(self):
|
||||
"""Test that orchestrator's output type is ResearchPlan."""
|
||||
orchestrator = ResearchOrchestrator(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
assert orchestrator.output_type == ResearchPlan
|
||||
|
||||
def test_orchestrator_has_no_tools(self):
|
||||
"""Test that orchestrator no longer registers tools (direct agent calls now)."""
|
||||
orchestrator = ResearchOrchestrator(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
|
||||
# Get the tools from the agent
|
||||
tools = orchestrator.agent._function_toolset.tools
|
||||
tool_names = list(tools.keys())
|
||||
|
||||
# Should have no tools since we call agents directly now
|
||||
assert len(tool_names) == 0
|
||||
|
||||
def test_should_stop_research_logic(self):
|
||||
"""Test the stopping logic based on EvaluationResult."""
|
||||
orchestrator = ResearchOrchestrator(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
|
||||
# Create mock evaluation results
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Sufficient research result
|
||||
sufficient_result = MagicMock()
|
||||
sufficient_result.output = EvaluationResult(
|
||||
key_insights=["Climate is changing", "Human activity is the cause"],
|
||||
new_questions=[],
|
||||
confidence_score=0.9,
|
||||
is_sufficient=True,
|
||||
reasoning="All aspects covered comprehensively",
|
||||
)
|
||||
|
||||
# Insufficient research result
|
||||
insufficient_result = MagicMock()
|
||||
insufficient_result.output = EvaluationResult(
|
||||
key_insights=["Some data found"],
|
||||
new_questions=[
|
||||
"What about economic impacts?",
|
||||
"Regional variations?",
|
||||
],
|
||||
confidence_score=0.4,
|
||||
is_sufficient=False,
|
||||
reasoning="Major gaps remain in understanding",
|
||||
)
|
||||
|
||||
# Test with sufficient research (threshold 0.8)
|
||||
assert orchestrator._should_stop_research(sufficient_result, 0.8)
|
||||
|
||||
# Test with insufficient research
|
||||
assert not orchestrator._should_stop_research(insufficient_result, 0.8)
|
||||
|
||||
# Test with high confidence but below threshold
|
||||
sufficient_result.output.confidence_score = 0.75
|
||||
assert not orchestrator._should_stop_research(sufficient_result, 0.8)
|
||||
|
||||
# Test with is_sufficient=False even with high confidence
|
||||
insufficient_result.output.confidence_score = 0.95
|
||||
assert not orchestrator._should_stop_research(insufficient_result, 0.8)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conduct_research_workflow(self, test_model, mock_client):
|
||||
"""Test the basic research workflow using TestModel."""
|
||||
orchestrator = ResearchOrchestrator(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
|
||||
# Setup mock client returns
|
||||
mock_chunks = [
|
||||
create_mock_chunk("1", "Climate change information"),
|
||||
]
|
||||
mock_client.search.return_value = mock_chunks
|
||||
mock_client.expand_context.return_value = mock_chunks
|
||||
|
||||
# Use TestModel for all agents
|
||||
with orchestrator.agent.override(model=test_model):
|
||||
with orchestrator.search_agent.agent.override(model=test_model):
|
||||
with orchestrator.evaluation_agent.agent.override(model=test_model):
|
||||
with orchestrator.synthesis_agent.agent.override(model=test_model):
|
||||
# Run the research
|
||||
report = await orchestrator.conduct_research(
|
||||
"What is climate change?", mock_client, max_iterations=1
|
||||
)
|
||||
|
||||
# Verify we got a valid report structure
|
||||
assert isinstance(report, ResearchReport)
|
||||
assert report.title
|
||||
assert report.executive_summary
|
||||
assert isinstance(report.main_findings, list)
|
||||
assert isinstance(report.conclusions, list)
|
||||
assert isinstance(report.limitations, list)
|
||||
assert isinstance(report.recommendations, list)
|
||||
assert report.sources_summary
|
||||
|
|
@ -1,14 +0,0 @@
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
|
||||
|
||||
|
||||
class TestSearchSpecialistAgent:
|
||||
"""Lean tests for SearchSpecialistAgent without LLM mocking."""
|
||||
|
||||
def test_agent_initialization(self):
|
||||
agent = SearchSpecialistAgent(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
assert agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert agent.model == Config.RESEARCH_MODEL
|
||||
assert agent.output_type is SearchAnswer
|
||||
|
|
@ -1,14 +0,0 @@
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
|
||||
|
||||
|
||||
class TestSynthesisAgent:
|
||||
"""Lean tests for SynthesisAgent without LLM mocking."""
|
||||
|
||||
def test_agent_initialization(self):
|
||||
agent = SynthesisAgent(
|
||||
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
|
||||
)
|
||||
assert agent.provider == Config.RESEARCH_PROVIDER
|
||||
assert agent.model == Config.RESEARCH_MODEL
|
||||
assert agent.output_type == ResearchReport
|
||||
26
tests/test_research_graph.py
Normal file
26
tests/test_research_graph.py
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
import asyncio
|
||||
|
||||
from haiku.rag.research.dependencies import ResearchContext
|
||||
from haiku.rag.research.graph import ResearchState, build_research_graph
|
||||
|
||||
|
||||
def test_build_graph_and_state():
|
||||
graph = build_research_graph()
|
||||
assert graph is not None
|
||||
|
||||
state = ResearchState(
|
||||
question="What are the key features of haiku.rag?",
|
||||
context=ResearchContext(
|
||||
original_question="What are the key features of haiku.rag?"
|
||||
),
|
||||
max_iterations=1,
|
||||
confidence_threshold=0.8,
|
||||
)
|
||||
assert state.iterations == 0
|
||||
assert state.sub_questions == []
|
||||
|
||||
|
||||
def test_async_loop_available():
|
||||
# Ensure an event loop can be created in test env
|
||||
loop = asyncio.new_event_loop()
|
||||
loop.close()
|
||||
89
tests/test_research_graph_integration.py
Normal file
89
tests/test_research_graph_integration.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.research.dependencies import ResearchContext
|
||||
from haiku.rag.research.graph import (
|
||||
EvaluateNode,
|
||||
PlanNode,
|
||||
ResearchDeps,
|
||||
ResearchState,
|
||||
SearchDispatchNode,
|
||||
SynthesizeNode,
|
||||
build_research_graph,
|
||||
)
|
||||
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
|
||||
graph = build_research_graph()
|
||||
|
||||
state = ResearchState(
|
||||
question="What is haiku.rag?",
|
||||
context=ResearchContext(original_question="What is haiku.rag?"),
|
||||
max_iterations=1,
|
||||
confidence_threshold=0.5,
|
||||
max_concurrency=2,
|
||||
)
|
||||
deps = ResearchDeps(
|
||||
client=cast(Any, None), console=None
|
||||
) # client unused in patched nodes
|
||||
|
||||
async def fake_plan_run(self, ctx) -> Any:
|
||||
ctx.state.sub_questions = [
|
||||
"Describe haiku.rag in one sentence",
|
||||
"List core components of haiku.rag",
|
||||
]
|
||||
return SearchDispatchNode(self.provider, self.model)
|
||||
|
||||
async def fake_search_dispatch_run(self, ctx) -> Any:
|
||||
# Answer all pending questions deterministically, then move to evaluation
|
||||
while ctx.state.sub_questions:
|
||||
q = ctx.state.sub_questions.pop(0)
|
||||
# pydantic BaseModel kwargs not fully typed for pyright
|
||||
ctx.state.context.add_qa_response(
|
||||
SearchAnswer(query=q, answer="A", context=["x"], sources=["s"]) # pyright: ignore[reportCallIssue]
|
||||
)
|
||||
return EvaluateNode(self.provider, self.model)
|
||||
|
||||
async def fake_evaluate_run(self, ctx) -> Any:
|
||||
ctx.state.last_eval = EvaluationResult(
|
||||
key_insights=["ok"],
|
||||
new_questions=[],
|
||||
confidence_score=1.0,
|
||||
is_sufficient=True,
|
||||
reasoning="done",
|
||||
)
|
||||
ctx.state.iterations += 1
|
||||
return SynthesizeNode(self.provider, self.model)
|
||||
|
||||
async def fake_synthesize_run(self, ctx) -> Any:
|
||||
report = ResearchReport(
|
||||
title="Haiku RAG",
|
||||
executive_summary="...",
|
||||
main_findings=["f1"],
|
||||
conclusions=["c1"],
|
||||
limitations=[],
|
||||
recommendations=[],
|
||||
sources_summary="s",
|
||||
)
|
||||
from pydantic_graph import End
|
||||
|
||||
return End(report)
|
||||
|
||||
monkeypatch.setattr(PlanNode, "run", fake_plan_run, raising=False)
|
||||
monkeypatch.setattr(
|
||||
SearchDispatchNode, "run", fake_search_dispatch_run, raising=False
|
||||
)
|
||||
monkeypatch.setattr(EvaluateNode, "run", fake_evaluate_run, raising=False)
|
||||
monkeypatch.setattr(SynthesizeNode, "run", fake_synthesize_run, raising=False)
|
||||
|
||||
start = PlanNode(provider="test", model="test")
|
||||
|
||||
result = await graph.run(start, state=state, deps=deps)
|
||||
report = result.output
|
||||
|
||||
assert isinstance(report, ResearchReport)
|
||||
assert report.title == "Haiku RAG"
|
||||
assert len(state.context.qa_responses) == 2
|
||||
Loading…
Reference in a new issue