danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Flatten graph structure, simplify

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-15 14:05:07 +02:00
parent 6ee0d74f92
commit 45de9bf0d5
No known key found for this signature in database
16 changed files with 402 additions and 589 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
haiku_rag_slim/haiku/rag/client.py
View file

@ -27,7 +27,7 @@ from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
logger = logging.getLogger(__name__)

5
haiku_rag_slim/haiku/rag/graph/common/__init__.py
View file

@ -1,5 +0,0 @@
"""Common utilities for graph implementations."""
from haiku.rag.utils import get_model
__all__ = ["get_model"]

107
haiku_rag_slim/haiku/rag/graph/common/models.py
View file

@ -1,107 +0,0 @@
"""Common models used across different graph implementations."""
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class ResearchPlan(BaseModel):
"""A structured research plan with sub-questions to explore."""
sub_questions: list[str] = Field(
...,
description="Specific questions to research, phrased as complete questions",
)
@field_validator("sub_questions")
@classmethod
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) < 1:
raise ValueError("Must have at least 1 sub-question")
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding."""
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
class RawSearchAnswer(BaseModel):
"""Answer to a search query with chunk references."""
query: str = Field(..., description="The question that was answered")
answer: str = Field(..., description="The answer to the question")
cited_chunks: list[str] = Field(
default_factory=list,
description="IDs of chunks used to form the answer",
)
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
class SearchAnswer(RawSearchAnswer):
"""Answer to a search query with resolved citations."""
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)
@classmethod
def from_raw(
cls,
raw: RawSearchAnswer,
search_results: "list[SearchResult]",
) -> "SearchAnswer":
"""Create SearchAnswer from RawSearchAnswer with resolved citations."""
citations = resolve_citations(raw.cited_chunks, search_results)
return cls(
query=raw.query,
answer=raw.answer,
cited_chunks=raw.cited_chunks,
confidence=raw.confidence,
citations=citations,
)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
# Build lookup by chunk_id
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for chunk_id in cited_chunk_ids:
r = by_id.get(chunk_id)
if not r:
continue
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
)
)
return citations

315
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

@ -1,315 +0,0 @@
"""Common node implementations for graph workflows."""
import asyncio
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
from pydantic_ai import Agent, RunContext
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import StepContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import RawSearchAnswer, ResearchPlan, SearchAnswer
from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.store.models import SearchResult
class GraphContext(Protocol):
"""Protocol for graph context objects."""
original_question: str
sub_questions: list[str]
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a QA response to context."""
...
class GraphState(Protocol):
"""Protocol for graph state objects."""
context: GraphContext
max_concurrency: int
search_filter: str | None
class GraphDeps(Protocol):
"""Protocol for graph dependencies."""
client: HaikuRAG
agui_emitter: AGUIEmitter[Any, Any] | None
semaphore: asyncio.Semaphore | None
class GraphAgentDeps(Protocol):
"""Protocol for agent dependencies."""
client: HaikuRAG
context: GraphContext
search_results: list[SearchResult]
def create_plan_node[AgentDepsT: GraphAgentDeps](
model_config: ModelConfig,
deps_type: type[AgentDepsT],
activity_message: str = "Creating plan",
output_retries: int | None = None,
config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
"""Create a plan node for any graph.
Args:
model_config: ModelConfig with provider, model, and settings
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies)
activity_message: Message to show during planning activity
output_retries: Number of output retries for the agent (optional)
config: AppConfig object (defaults to global Config)
Returns:
Async function that can be used as a graph step
"""
async def plan(ctx: StepContext[Any, Any, None], /) -> None:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": activity_message}
)
try:
# Build agent configuration
agent_config = {
"model": get_model(model_config, config),
"output_type": ResearchPlan,
"instructions": (
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
"retries": 3,
"deps_type": deps_type,
}
if output_retries is not None:
agent_config["output_retries"] = output_retries
plan_agent = Agent(**agent_config)
# Capture search filter for use in tool
search_filter = state.search_filter
@plan_agent.tool
async def gather_context(
ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
) -> str:
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
return "\n\n".join(r.content for r in results)
# Tool is registered via decorator above
_ = gather_context
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
# Create agent dependencies
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
# State now contains the plan - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
return plan
def create_search_node[AgentDepsT: GraphAgentDeps](
model_config: ModelConfig,
deps_type: type[AgentDepsT],
with_step_wrapper: bool = True,
success_message_format: str = "Answered: {sub_q}",
handle_exceptions: bool = False,
config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
"""Create a search_one node for any graph.
Args:
model_config: ModelConfig with provider, model, and settings
deps_type: Type of dependencies for the agent
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
success_message_format: Format string for success activity message
handle_exceptions: Whether to handle exceptions with fallback answer
config: AppConfig object (defaults to global Config)
Returns:
Async function that can be used as a graph step
"""
async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
sub_q = ctx.inputs
# Create unique step name from question text
step_name = f"search: {sub_q}"
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.start_step(step_name)
try:
# Create semaphore if not already provided
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
# Use semaphore to control concurrency
async with deps.semaphore:
return await _do_search(
state,
deps,
sub_q,
model_config,
deps_type,
success_message_format,
handle_exceptions,
config,
)
finally:
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.finish_step()
return search_one
async def _do_search[AgentDepsT: GraphAgentDeps](
state: GraphState,
deps: GraphDeps,
sub_q: str,
model_config: ModelConfig,
deps_type: type[AgentDepsT],
success_message_format: str,
handle_exceptions: bool,
config: AppConfig,
) -> SearchAnswer:
"""Internal search implementation."""
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = Agent(
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=deps_type,
)
# Capture search filter for use in tool
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[AgentDepsT], 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 ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
# Store results for citation resolution
ctx2.deps.search_results = results
# Format with metadata for agent context
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found in the knowledge base for: {query}"
return "\n\n".join(parts)
# Tool is registered via decorator above
_ = search_and_answer
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
try:
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
if raw_answer:
# Convert RawSearchAnswer to SearchAnswer with resolved citations
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
state.context.add_qa_response(answer)
# State updated with new answer - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
# Format the success message
if "{confidence" in success_message_format:
message = success_message_format.format(
sub_q=sub_q, confidence=answer.confidence
)
else:
message = success_message_format.format(sub_q=sub_q)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": message,
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
# Return empty SearchAnswer if no result
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
except Exception as e:
if handle_exceptions:
# Narrate the error
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
confidence=0.0,
)
return failure_answer
else:
raise

67
haiku_rag_slim/haiku/rag/graph/common/prompts.py
View file

@ -1,67 +0,0 @@
"""Common prompts used across different graph implementations."""
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
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
Plan requirements:
- Produce at most 3 sub_questions that together cover the main question.
- sub_questions must be a list of plain strings, where each string is a complete
question. Do NOT use objects with nested fields like {question, details}.
- 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).
Use the gather_context tool once on the main question before planning."""
SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
Process:
1. Call search_and_answer with relevant keywords from the question.
2. Review the results and their relevance scores.
3. If needed, perform follow-up searches with different keywords (max 3 total).
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Output format:
- query: Echo the question you are answering
- answer: Your concise answer based on the retrieved content
- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge.
- Use the Source and Type metadata to understand context.
- If multiple results are relevant, synthesize them coherently.
- If information is insufficient, say so clearly.
- Be concise and direct; avoid meta commentary about the process.
- Higher scores indicate more relevant results."""

7
haiku_rag_slim/haiku/rag/graph/research/__init__.py
View file

@ -1,3 +1,6 @@
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
from haiku.rag.graph.research.models import (
EvaluationResult,
ResearchReport,
SearchAnswer,
)

29
haiku_rag_slim/haiku/rag/graph/research/common.py
View file

@ -1,29 +0,0 @@
from pydantic_ai import format_as_xml
from haiku.rag.graph.research.dependencies import ResearchContext
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,
"confidence": qa.confidence,
"sources": [
{
"document_uri": c.document_uri,
"document_title": c.document_title,
"page_numbers": c.page_numbers,
"headings": c.headings,
}
for c in qa.citations
],
}
for qa in context.qa_responses
],
}
return format_as_xml(context_data, root_tag="research_context")

10
haiku_rag_slim/haiku/rag/graph/research/dependencies.py
View file

@ -1,9 +1,13 @@
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.graph.research.models import SearchAnswer
class ResearchContext(BaseModel):
"""Context shared across research agents."""
@ -12,11 +16,11 @@ class ResearchContext(BaseModel):
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list[SearchAnswer] = Field(
qa_responses: list[Any] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
def add_qa_response(self, qa: SearchAnswer) -> None:
def add_qa_response(self, qa: "SearchAnswer") -> None:
"""Add a structured QA response."""
self.qa_responses.append(qa)

248
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

@ -1,20 +1,54 @@
from pydantic_ai import Agent
import asyncio
from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from pydantic_graph.beta.join import reduce_list_append
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.research.common import format_context_for_prompt
from haiku.rag.graph.research.dependencies import ResearchDependencies
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import (
EvaluationResult,
RawSearchAnswer,
ResearchPlan,
ResearchReport,
SearchAnswer,
)
from haiku.rag.graph.research.prompts import (
DECISION_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
DECISION_PROMPT,
PLAN_PROMPT,
SEARCH_PROMPT,
SYNTHESIS_PROMPT,
)
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
from haiku.rag.utils import get_model
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,
"confidence": qa.confidence,
"sources": [
{
"document_uri": c.document_uri,
"document_title": c.document_title,
"page_numbers": c.page_numbers,
"headings": c.headings,
}
for c in qa.citations
],
}
for qa in context.qa_responses
],
}
return format_as_xml(context_data, root_tag="research_context")
def build_research_graph(
@ -37,26 +71,172 @@ def build_research_graph(
output_type=ResearchReport,
)
plan = g.step(
create_plan_node(
model_config=model_config,
deps_type=ResearchDependencies, # type: ignore[arg-type]
activity_message="Creating research plan",
output_retries=3,
config=config,
)
) # type: ignore[arg-type]
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
"""Create research plan with sub-questions."""
state = ctx.state
deps = ctx.deps
search_one = g.step(
create_search_node(
model_config=model_config,
deps_type=ResearchDependencies, # type: ignore[arg-type]
with_step_wrapper=True,
success_message_format="Found answer with {confidence:.0%} confidence",
handle_exceptions=True,
config=config,
)
) # type: ignore[arg-type]
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": "Creating research plan"}
)
try:
plan_agent = Agent(
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
return "\n\n".join(r.content for r in results)
_ = gather_context
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single sub-question using the knowledge base."""
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
step_name = f"search: {sub_q}"
if deps.agui_emitter:
deps.agui_emitter.start_step(step_name)
try:
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = Agent(
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=SEARCH_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
"""Search the knowledge base for relevant documents."""
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
ctx2.deps.search_results = results
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found for: {query}"
return "\n\n".join(parts)
_ = search_and_answer
agent_deps = ResearchDependencies(
client=deps.client, context=state.context
)
try:
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
if raw_answer:
answer = SearchAnswer.from_raw(
raw_answer, agent_deps.search_results
)
state.context.add_qa_response(answer)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Found answer with {answer.confidence:.0%} confidence",
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
except Exception as e:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
return SearchAnswer(
query=sub_q,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def get_batch(
@ -76,6 +256,7 @@ def build_research_graph(
async def decide(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> bool:
"""Evaluate research sufficiency and decide whether to continue."""
state = ctx.state
deps = ctx.deps
@ -89,7 +270,7 @@ def build_research_graph(
agent = Agent(
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
instructions=DECISION_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
@ -152,6 +333,7 @@ def build_research_graph(
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
@ -165,7 +347,7 @@ def build_research_graph(
agent = Agent(
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT,
instructions=SYNTHESIS_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
@ -201,7 +383,6 @@ def build_research_graph(
else:
g.add(g.edge_from(g.start_node).to(get_batch))
# Branch based on whether we have questions
g.add(
g.edge_from(get_batch).to(
g.decision()
@ -209,12 +390,9 @@ def build_research_graph(
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(
decide
), # Direct: collect → decide (no analyze_insights)
g.edge_from(collect_answers).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()

105
haiku_rag_slim/haiku/rag/graph/research/models.py
View file

@ -1,4 +1,107 @@
from pydantic import BaseModel, Field
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class ResearchPlan(BaseModel):
"""A structured research plan with sub-questions to explore."""
sub_questions: list[str] = Field(
...,
description="Specific questions to research, phrased as complete questions",
)
@field_validator("sub_questions")
@classmethod
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) < 1:
raise ValueError("Must have at least 1 sub-question")
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding."""
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
class RawSearchAnswer(BaseModel):
"""Answer to a search query with chunk references."""
query: str = Field(..., description="The question that was answered")
answer: str = Field(..., description="The answer to the question")
cited_chunks: list[str] = Field(
default_factory=list,
description="IDs of chunks used to form the answer",
)
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
class SearchAnswer(RawSearchAnswer):
"""Answer to a search query with resolved citations."""
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)
@classmethod
def from_raw(
cls,
raw: RawSearchAnswer,
search_results: "list[SearchResult]",
) -> "SearchAnswer":
"""Create SearchAnswer from RawSearchAnswer with resolved citations."""
citations = resolve_citations(raw.cited_chunks, search_results)
return cls(
query=raw.query,
answer=raw.answer,
cited_chunks=raw.cited_chunks,
confidence=raw.confidence,
citations=citations,
)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for chunk_id in cited_chunk_ids:
r = by_id.get(chunk_id)
if not r:
continue
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
)
)
return citations
class EvaluationResult(BaseModel):

82
haiku_rag_slim/haiku/rag/graph/research/prompts.py
View file

@ -1,4 +1,70 @@
DECISION_AGENT_PROMPT = """You are the research evaluator responsible for assessing
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
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
Plan requirements:
- Produce at most 3 sub_questions that together cover the main question.
- sub_questions must be a list of plain strings, where each string is a complete
question. Do NOT use objects with nested fields like {question, details}.
- 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).
Use the gather_context tool once on the main question before planning."""
SEARCH_PROMPT = """You are a search and question-answering specialist.
Process:
1. Call search_and_answer with relevant keywords from the question.
2. Review the results and their relevance scores.
3. If needed, perform follow-up searches with different keywords (max 3 total).
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Output format:
- query: Echo the question you are answering
- answer: Your concise answer based on the retrieved content
- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge.
- Use the Source and Type metadata to understand context.
- If multiple results are relevant, synthesize them coherently.
- If information is insufficient, say so clearly.
- Be concise and direct; avoid meta commentary about the process.
- Higher scores indicate more relevant results."""
DECISION_PROMPT = """You are the research evaluator responsible for assessing
whether gathered evidence sufficiently answers the research question.
Inputs available:
@ -19,7 +85,7 @@ Output fields:
Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
research report that directly answers the original question.
Goals:
@ -47,15 +113,3 @@ Style:
- Be professional, objective, and specific.
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
Instead, state the actual information directly."""
PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
Task:
- 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)."""

4
haiku_rag_slim/haiku/rag/qa/agent.py
View file

@ -5,10 +5,10 @@ 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.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):

2
haiku_rag_slim/haiku/rag/utils.py
View file

@ -8,7 +8,7 @@ from packaging.version import Version, parse
if TYPE_CHECKING:
from rich.console import RenderableType
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
def apply_common_settings(

2
tests/graph/test_research_graph.py
View file

@ -18,8 +18,6 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
# Patch all locations where get_model is imported
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()

4
tests/graph/test_search_filter.py
View file

@ -67,8 +67,6 @@ async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()
@ -114,8 +112,6 @@ async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()

2
tests/test_app.py
View file

@ -310,7 +310,7 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with citations."""
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
mock_answer = "Test answer with citations"
mock_citations = [