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

Remove analyze insights node, simplify research context, state, prompts & models

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-15 12:40:27 +02:00
parent 50b7fb4461
commit bdcf81774d
No known key found for this signature in database
6 changed files with 48 additions and 517 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,12 +1,10 @@
from pydantic_ai import format_as_xml
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import InsightAnalysis
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,
@ -27,69 +25,5 @@ def format_context_for_prompt(context: ResearchContext) -> str:
}
for qa in context.qa_responses
],
"insights": [
{
"id": insight.id,
"summary": insight.summary,
"status": insight.status.value,
"supporting_sources": insight.supporting_sources,
"originating_questions": insight.originating_questions,
"notes": insight.notes,
}
for insight in context.insights
],
"gaps": [
{
"id": gap.id,
"description": gap.description,
"severity": gap.severity.value,
"blocking": gap.blocking,
"resolved": gap.resolved,
"resolved_by": gap.resolved_by,
"supporting_sources": gap.supporting_sources,
"notes": gap.notes,
}
for gap in context.gaps
],
}
return format_as_xml(context_data, root_tag="research_context")
def format_analysis_for_prompt(
analysis: InsightAnalysis | None,
) -> str:
"""Format the latest insight analysis as XML for prompts."""
if analysis is None:
return "<latest_analysis />"
data = {
"commentary": analysis.commentary,
"highlights": [
{
"id": insight.id,
"summary": insight.summary,
"status": insight.status.value,
"supporting_sources": insight.supporting_sources,
"originating_questions": insight.originating_questions,
"notes": insight.notes,
}
for insight in analysis.highlights
],
"gap_assessments": [
{
"id": gap.id,
"description": gap.description,
"severity": gap.severity.value,
"blocking": gap.blocking,
"resolved": gap.resolved,
"resolved_by": gap.resolved_by,
"supporting_sources": gap.supporting_sources,
"notes": gap.notes,
}
for gap in analysis.gap_assessments
],
"resolved_gaps": analysis.resolved_gaps,
"new_questions": analysis.new_questions,
}
return format_as_xml(data, root_tag="latest_analysis")

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

@ -1,14 +1,7 @@
from collections.abc import Iterable
from pydantic import BaseModel, Field, PrivateAttr
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.models import (
GapRecord,
InsightAnalysis,
InsightRecord,
)
from haiku.rag.store.models import SearchResult
@ -22,121 +15,11 @@ class ResearchContext(BaseModel):
qa_responses: list[SearchAnswer] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
insights: list[InsightRecord] = Field(
default_factory=list, description="Key insights discovered"
)
gaps: list[GapRecord] = Field(
default_factory=list, description="Identified information gaps"
)
# Private dict indexes for O(1) lookups
_insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict)
_gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict)
def model_post_init(self, __context: object) -> None:
"""Build indexes after initialization."""
self._insights_by_id = {ins.id: ins for ins in self.insights}
self._gaps_by_id = {gap.id: gap for gap in self.gaps}
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a structured QA response (citations already resolved)."""
"""Add a structured QA response."""
self.qa_responses.append(qa)
def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
"""Merge one or more insights into the shared context with deduplication."""
merged: list[InsightRecord] = []
for record in records:
candidate = InsightRecord.model_validate(record)
existing = self._insights_by_id.get(candidate.id)
if existing:
# Update existing insight
existing.summary = candidate.summary
existing.status = candidate.status
if candidate.notes:
existing.notes = candidate.notes
existing.supporting_sources = _merge_unique(
existing.supporting_sources, candidate.supporting_sources
)
existing.originating_questions = _merge_unique(
existing.originating_questions, candidate.originating_questions
)
merged.append(existing)
else:
# Add new insight
new_insight = candidate.model_copy(deep=True)
self.insights.append(new_insight)
self._insights_by_id[new_insight.id] = new_insight
merged.append(new_insight)
return merged
def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]:
"""Merge one or more gap records into the shared context with deduplication."""
merged: list[GapRecord] = []
for record in records:
candidate = GapRecord.model_validate(record)
existing = self._gaps_by_id.get(candidate.id)
if existing:
# Update existing gap
existing.description = candidate.description
existing.severity = candidate.severity
existing.blocking = candidate.blocking
existing.resolved = candidate.resolved
if candidate.notes:
existing.notes = candidate.notes
existing.supporting_sources = _merge_unique(
existing.supporting_sources, candidate.supporting_sources
)
existing.resolved_by = _merge_unique(
existing.resolved_by, candidate.resolved_by
)
merged.append(existing)
else:
# Add new gap
new_gap = candidate.model_copy(deep=True)
self.gaps.append(new_gap)
self._gaps_by_id[new_gap.id] = new_gap
merged.append(new_gap)
return merged
def mark_gap_resolved(
self, identifier: str, resolved_by: Iterable[str] | None = None
) -> GapRecord | None:
"""Mark a gap as resolved by identifier."""
gap = self._gaps_by_id.get(identifier)
if gap is None:
return None
gap.resolved = True
gap.blocking = False
if resolved_by:
gap.resolved_by = _merge_unique(gap.resolved_by, list(resolved_by))
return gap
def integrate_analysis(self, analysis: InsightAnalysis) -> None:
"""Apply an analysis result to the shared context."""
merged_insights: list[InsightRecord] = []
if analysis.highlights:
merged_insights = self.upsert_insights(analysis.highlights)
analysis.highlights = merged_insights
if analysis.gap_assessments:
merged_gaps = self.upsert_gaps(analysis.gap_assessments)
analysis.gap_assessments = merged_gaps
if analysis.resolved_gaps:
resolved_by_list = (
[ins.id for ins in merged_insights] if merged_insights else None
)
for resolved in analysis.resolved_gaps:
self.mark_gap_resolved(resolved, resolved_by=resolved_by_list)
for question in analysis.new_questions:
if question not in self.sub_questions:
self.sub_questions.append(question)
class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context."""
@ -148,8 +31,3 @@ class ResearchDependencies(BaseModel):
search_results: list[SearchResult] = Field(
default_factory=list, description="Search results for citation resolution"
)
def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
"""Merge two iterables preserving order while removing duplicates."""
return [k for k in dict.fromkeys([*existing, *incoming]) if k]

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

@ -7,19 +7,11 @@ 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_analysis_for_prompt,
format_context_for_prompt,
)
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,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
from haiku.rag.graph.research.prompts import (
DECISION_AGENT_PROMPT,
INSIGHT_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
)
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@ -43,7 +35,6 @@ def build_research_graph(
output_type=ResearchReport,
)
# Create and register the plan node using the factory
plan = g.step(
create_plan_node(
model_config=model_config,
@ -54,7 +45,6 @@ def build_research_graph(
)
) # type: ignore[arg-type]
# Create and register the search_one node using the factory
search_one = g.step(
create_search_node(
model_config=model_config,
@ -76,84 +66,14 @@ def build_research_graph(
if not state.context.sub_questions:
return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def analyze_insights(
async def decide(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> None:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("analyze_insights")
deps.agui_emitter.update_activity(
"analyzing", {"message": "Synthesizing insights and gaps"}
)
try:
agent = Agent(
model=get_model(model_config, config),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
# State updated with insights/gaps - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
highlights = len(analysis.highlights)
gaps = len(analysis.gap_assessments)
resolved = len(analysis.resolved_gaps)
parts = []
if highlights:
parts.append(f"{highlights} insights")
if gaps:
parts.append(f"{gaps} gaps")
if resolved:
parts.append(f"{resolved} resolved")
summary = ", ".join(parts) if parts else "No updates"
deps.agui_emitter.update_activity(
"analyzing",
{
"stepName": "analyze_insights",
"message": f"Analysis: {summary}",
"insights": [
h.model_dump(mode="json") for h in analysis.highlights
],
"gaps": [
g.model_dump(mode="json") for g in analysis.gap_assessments
],
"resolved_gaps": list(analysis.resolved_gaps),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
) -> bool:
state = ctx.state
deps = ctx.deps
@ -174,11 +94,9 @@ def build_research_graph(
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_xml,
]
if state.last_eval is not None:
prev = state.last_eval
@ -205,7 +123,6 @@ def build_research_graph(
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
# State updated with evaluation - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
sufficient = "Yes" if output.is_sufficient else "No"
@ -287,8 +204,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(analyze_insights),
g.edge_from(analyze_insights).to(decide),
g.edge_from(collect_answers).to(
decide
), # Direct: collect → decide (no analyze_insights)
)
# Branch based on decision

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

@ -1,149 +1,25 @@
import uuid
from enum import Enum
from pydantic import BaseModel, Field, field_validator
def _deduplicate_list(items: list[str]) -> list[str]:
"""Remove duplicates while preserving order."""
return list(dict.fromkeys(items))
class InsightStatus(str, Enum):
OPEN = "open"
VALIDATED = "validated"
TENTATIVE = "tentative"
class GapSeverity(str, Enum):
LOW = "low"
MEDIUM = "medium"
HIGH = "high"
class TrackedRecord(BaseModel):
"""Base model for tracked entities with sources and metadata."""
model_config = {"validate_assignment": True}
id: str = Field(
default_factory=lambda: str(uuid.uuid4())[:8],
description="Unique identifier for the record",
)
supporting_sources: list[str] = Field(
default_factory=list,
description="Source identifiers backing this record",
)
notes: str | None = Field(
default=None,
description="Optional elaboration or caveats",
)
@field_validator("supporting_sources", mode="before")
@classmethod
def deduplicate_sources(cls, v: list[str]) -> list[str]:
"""Ensure supporting_sources has no duplicates."""
return _deduplicate_list(v) if v else []
class InsightRecord(TrackedRecord):
"""Structured insight with provenance and lifecycle metadata."""
summary: str = Field(description="Concise description of the insight")
status: InsightStatus = Field(
default=InsightStatus.OPEN,
description="Lifecycle status for the insight",
)
originating_questions: list[str] = Field(
default_factory=list,
description="Research sub-questions that produced this insight",
)
@field_validator("originating_questions", mode="before")
@classmethod
def deduplicate_questions(cls, v: list[str]) -> list[str]:
"""Ensure originating_questions has no duplicates."""
return _deduplicate_list(v) if v else []
class GapRecord(TrackedRecord):
"""Structured representation of an identified research gap."""
description: str = Field(description="Concrete statement of what is missing")
severity: GapSeverity = Field(
default=GapSeverity.MEDIUM,
description="Severity of the gap for answering the main question",
)
blocking: bool = Field(
default=True,
description="Whether this gap blocks a confident answer",
)
resolved: bool = Field(
default=False,
description="Flag indicating if the gap has been resolved",
)
resolved_by: list[str] = Field(
default_factory=list,
description="Insight IDs or notes explaining how the gap was closed",
)
@field_validator("resolved_by", mode="before")
@classmethod
def deduplicate_resolved_by(cls, v: list[str]) -> list[str]:
"""Ensure resolved_by has no duplicates."""
return _deduplicate_list(v) if v else []
class InsightAnalysis(BaseModel):
"""Output of the insight aggregation agent."""
highlights: list[InsightRecord] = Field(
default_factory=list,
description="New or updated insights discovered this iteration",
)
gap_assessments: list[GapRecord] = Field(
default_factory=list,
description="New or updated gap records based on current evidence",
)
resolved_gaps: list[str] = Field(
default_factory=list,
description="Gap identifiers or descriptions considered resolved",
)
new_questions: list[str] = Field(
default_factory=list,
max_length=3,
description="Up to three follow-up sub-questions to pursue next",
)
commentary: str = Field(
description="Short narrative summary of the incremental findings",
)
from pydantic import BaseModel, Field
class EvaluationResult(BaseModel):
"""Result of analysis and evaluation."""
"""Result of research sufficiency 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=[],
)
gaps: list[str] = Field(
description="Concrete information gaps that remain", default_factory=list
)
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"
)
confidence_score: float = Field(
ge=0.0,
le=1.0,
description="Confidence level in the completeness of research (0-1)",
)
reasoning: str = Field(
description="Explanation of why the research is or isn't complete"
)
new_questions: list[str] = Field(
default_factory=list,
max_length=3,
description="New sub-questions to add to the research (max 3)",
)
class ResearchReport(BaseModel):

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

@ -1,75 +1,23 @@
INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
research workflow.
DECISION_AGENT_PROMPT = """You are the research evaluator responsible for assessing
whether gathered evidence sufficiently answers the research question.
Inputs available:
- Original research question and sub-questions
- Question–answer pairs with supporting snippets and sources
- Existing insights and gaps (with status metadata)
- Original research question
- Question-answer pairs with supporting sources
- Previous evaluation (if any)
Tasks:
1. Extract new or refined insights that advance understanding of the question.
2. Update gap status, creating new gap entries when necessary and marking
resolved ones explicitly.
3. Suggest up to 3 high-impact follow-up sub_questions that would close the
most important remaining gaps.
1. Assess whether the collected evidence answers the original question.
2. Provide a confidence_score in [0,1] reflecting coverage and evidence quality.
3. Optionally propose up to 3 new sub-questions if important gaps remain.
Output format (map directly to fields):
- highlights: list of insights with fields {summary, status, supporting_sources,
originating_questions, notes}. Use status one of {validated, open, tentative}.
supporting_sources and originating_questions must be lists of plain strings.
- gap_assessments: list of gaps with fields {description, severity, blocking,
resolved, resolved_by, supporting_sources, notes}. Severity must be one of
{low, medium, high}. resolved_by and supporting_sources must be lists of plain strings.
- resolved_gaps: list of plain strings (identifiers or descriptions for gaps now closed).
- new_questions: list of plain strings, up to 3 standalone questions (no duplicates).
- commentary: 1–3 sentences summarizing what changed this round.
Output fields:
- is_sufficient: true when the question is adequately answered
- confidence_score: numeric in [0,1]
- reasoning: brief explanation of the assessment
- new_questions: list of follow-up questions (max 3), only if needed
All list fields must contain plain strings only, not objects.
Guidance:
- Be concise and avoid repeating previously recorded information unless it
changed materially.
- For supporting_sources, use only the document_uri strings from the sources.
- Only propose new sub_questions that directly address remaining gaps.
- When marking a gap as resolved, ensure the rationale is clear via
resolved_by or notes."""
DECISION_AGENT_PROMPT = """You are the research governor responsible for making
stop/go decisions.
Inputs available:
- Original research question and current plan
- Full insight ledger with status metadata
- Up-to-date gap tracker, including resolved indicators
- Latest insight analysis summary (highlights, gap changes, new questions)
- Previous evaluation decision (if any)
Tasks:
1. Determine whether the collected evidence now answers the original question.
2. Provide a confidence_score in [0,1] that reflects coverage, evidence quality,
and agreement across sources.
3. List the highest-priority gaps that still block a confident answer. Reference
existing gap descriptions rather than inventing new ones.
4. Optionally propose up to 3 new sub_questions only if they are not already in
the current backlog.
Strictness:
- Only mark research as sufficient when every critical aspect of the main
question is addressed with reliable, corroborated evidence.
- Treat unresolved high-severity or blocking gaps as a hard stop.
Output fields must line up with EvaluationResult:
- key_insights: list of plain strings, concise bullet-ready statements.
- new_questions: list of plain strings, follow-up sub-questions (max 3).
- gaps: list of plain strings, remaining blockers (reuse wording from tracked gaps).
- confidence_score: numeric in [0,1].
- is_sufficient: true only when no blocking gaps remain.
- reasoning: short narrative tying the decision to evidence coverage.
All list fields must contain plain strings only, not objects.
Remember: prefer maintaining continuity with the structured context over
introducing new terminology."""
Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
research report.
@ -77,16 +25,16 @@ research report.
Goals:
1. Synthesize all gathered information into a coherent narrative.
2. Present findings clearly and concisely.
3. Draw evidence‑based conclusions and recommendations.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
Report guidelines (map to output fields):
- title: concise (5–12 words), informative.
- executive_summary: 3–5 sentences summarizing the overall answer.
- main_findings: list of plain strings, 4–8 one‑sentence bullets reflecting evidence.
- conclusions: list of plain strings, 2–4 bullets following logically from findings.
- recommendations: list of plain strings, 2–5 actionable bullets tied to findings.
- limitations: list of plain strings, 1–3 bullets describing constraints or uncertainties.
- title: concise (5-12 words), informative.
- executive_summary: 3-5 sentences summarizing the overall answer.
- main_findings: list of plain strings, 4-8 one-sentence bullets reflecting evidence.
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects.
@ -101,7 +49,7 @@ 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
- Read that context and produce a short natural-language summary of what the
KB appears to contain relative to the question.
Rules:

33
haiku_rag_slim/haiku/rag/graph/research/state.py
View file

@ -6,11 +6,7 @@ from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import (
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@ -26,12 +22,7 @@ class ResearchDeps:
semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
"""Emit a log message through AG-UI events.
Args:
message: The message to log
state: Optional state to include in state update
"""
"""Emit a log message through AG-UI events."""
if self.agui_emitter:
self.agui_emitter.log(message)
if state:
@ -39,15 +30,12 @@ class ResearchDeps:
class ResearchState(BaseModel):
"""Research graph state model.
Fully JSON-serializable Pydantic model suitable for AG-UI state synchronization.
"""
"""Research graph state model."""
model_config = {"arbitrary_types_allowed": True}
context: ResearchContext = Field(
description="Shared research context with questions, insights, and gaps"
description="Shared research context with questions and QA responses"
)
iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
@ -60,9 +48,6 @@ class ResearchState(BaseModel):
last_eval: EvaluationResult | None = Field(
default=None, description="Last evaluation result"
)
last_analysis: InsightAnalysis | None = Field(
default=None, description="Last insight analysis"
)
search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results"
)
@ -71,15 +56,7 @@ class ResearchState(BaseModel):
def from_config(
cls, context: ResearchContext, config: "AppConfig"
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
Args:
context: The ResearchContext containing the question and settings
config: The AppConfig object (uses config.research for state parameters)
Returns:
A configured ResearchState instance
"""
"""Create a ResearchState from an AppConfig."""
return cls(
context=context,
max_iterations=config.research.max_iterations,