From bdcf81774d681e48d31b5c61595cac6013ba0ef4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Dec 2025 12:40:27 +0200 Subject: [PATCH 1/7] Remove analyze insights node, simplify research context, state, prompts & models --- .../haiku/rag/graph/research/common.py | 66 -------- .../haiku/rag/graph/research/dependencies.py | 126 +-------------- .../haiku/rag/graph/research/graph.py | 96 +----------- .../haiku/rag/graph/research/models.py | 148 ++---------------- .../haiku/rag/graph/research/prompts.py | 96 +++--------- .../haiku/rag/graph/research/state.py | 33 +--- 6 files changed, 48 insertions(+), 517 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/graph/research/common.py b/haiku_rag_slim/haiku/rag/graph/research/common.py index fcdb0820..c0b3bd96 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/common.py +++ b/haiku_rag_slim/haiku/rag/graph/research/common.py @@ -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 "" - - 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") diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index 542710dd..f5aa205a 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -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] diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 32000dce..7b1129f9 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/graph/research/models.py b/haiku_rag_slim/haiku/rag/graph/research/models.py index a132bf4b..322ca99b 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/models.py +++ b/haiku_rag_slim/haiku/rag/graph/research/models.py @@ -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): diff --git a/haiku_rag_slim/haiku/rag/graph/research/prompts.py b/haiku_rag_slim/haiku/rag/graph/research/prompts.py index 70dfbb22..d3b8aec4 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/research/prompts.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/graph/research/state.py b/haiku_rag_slim/haiku/rag/graph/research/state.py index 17e876ee..f69301dd 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/state.py +++ b/haiku_rag_slim/haiku/rag/graph/research/state.py @@ -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, From f0a7abd953d0bacf0a7c69a8f755100fba8447cc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Dec 2025 13:16:05 +0200 Subject: [PATCH 2/7] Remove deep qa graph, everything now uses the simplified research graph --- haiku_rag_slim/haiku/rag/app.py | 39 ++- haiku_rag_slim/haiku/rag/graph/__init__.py | 11 - haiku_rag_slim/haiku/rag/graph/agui/server.py | 22 +- .../haiku/rag/graph/deep_qa/__init__.py | 1 - .../haiku/rag/graph/deep_qa/dependencies.py | 29 -- .../haiku/rag/graph/deep_qa/graph.py | 250 ------------------ .../haiku/rag/graph/deep_qa/models.py | 23 -- .../haiku/rag/graph/deep_qa/prompts.py | 42 --- .../haiku/rag/graph/deep_qa/state.py | 59 ----- .../haiku/rag/graph/research/graph.py | 13 +- .../haiku/rag/graph/research/state.py | 23 +- haiku_rag_slim/haiku/rag/mcp.py | 26 +- tests/graph/test_deep_qa.py | 41 --- tests/graph/test_search_filter.py | 49 ---- tests/test_app.py | 71 +++-- tests/test_mcp.py | 9 +- 16 files changed, 120 insertions(+), 588 deletions(-) delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/models.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/deep_qa/state.py delete mode 100644 tests/graph/test_deep_qa.py diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index b7aa9428..c04cdbdf 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -329,39 +329,30 @@ class HaikuRAGApp: try: citations = [] if deep: - from haiku.rag.graph.deep_qa.dependencies import DeepQAContext - from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph - from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState - - graph = build_deep_qa_graph(config=self.config) - context = DeepQAContext(original_question=question) - state = DeepQAState.from_config(context=context, config=self.config) + graph = build_research_graph(config=self.config) + context = ResearchContext(original_question=question) + state = ResearchState.from_config( + context=context, + config=self.config, + max_iterations=2, + confidence_threshold=0.0, + ) state.search_filter = filter - deps = DeepQADeps(client=self.client) + deps = ResearchDeps(client=self.client) if verbose: - # Use AG-UI renderer to process and display events - from haiku.rag.graph.common.models import Citation - renderer = AGUIConsoleRenderer(self.console) result_dict = await renderer.render( stream_graph(graph, state, deps) ) - # Result should be a dict with 'answer' and 'citations' keys - answer = result_dict.get("answer", "") if result_dict else "" - if cite and result_dict: - # Convert dicts to Citation objects - raw_citations = result_dict.get("citations", []) - citations = [ - Citation(**c) if isinstance(c, dict) else c - for c in raw_citations - ] + answer = ( + result_dict.get("executive_summary", "") + if result_dict + else "" + ) else: - # Run without rendering events, just get the result result = await graph.run(state=state, deps=deps) - answer = result.answer - if cite: - citations = result.citations + answer = result.executive_summary else: answer, citations = await self.client.ask(question, filter=filter) diff --git a/haiku_rag_slim/haiku/rag/graph/__init__.py b/haiku_rag_slim/haiku/rag/graph/__init__.py index fff8e10a..56470996 100644 --- a/haiku_rag_slim/haiku/rag/graph/__init__.py +++ b/haiku_rag_slim/haiku/rag/graph/__init__.py @@ -1,25 +1,14 @@ -"""Graph module for haiku.rag. - -This module contains all graph-related functionality including: -- AG-UI protocol for graph streaming -- Common graph utilities and models -- Research graph implementation -- Deep QA graph implementation -""" - from haiku.rag.graph.agui import ( AGUIConsoleRenderer, AGUIEmitter, create_agui_server, stream_graph, ) -from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph from haiku.rag.graph.research.graph import build_research_graph __all__ = [ "AGUIConsoleRenderer", "AGUIEmitter", - "build_deep_qa_graph", "build_research_graph", "create_agui_server", "stream_graph", diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index 4e9418f6..1c91ff5d 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -164,9 +164,6 @@ def create_agui_server( # pragma: no cover Starlette app with research and deep ask endpoints """ from haiku.rag.client import HaikuRAG - from haiku.rag.graph.deep_qa.dependencies import DeepQAContext - from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph - from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.state import ResearchDeps, ResearchState @@ -202,26 +199,31 @@ def create_agui_server( # pragma: no cover ) return ResearchDeps(client=get_client(effective_db_path)) - # Deep ask graph factories + # Deep ask graph factories (uses research graph with quick settings) def deep_ask_graph_factory() -> Graph: - return build_deep_qa_graph(config) + return build_research_graph(config) - def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState: + def deep_ask_state_factory(input_state: dict[str, Any]) -> ResearchState: question = input_state.get("question", "") if not question: messages = input_state.get("messages", []) if messages: question = messages[0].get("content", "") - context = DeepQAContext(original_question=question) - return DeepQAState.from_config(context=context, config=config) + context = ResearchContext(original_question=question) + return ResearchState.from_config( + context=context, + config=config, + max_iterations=2, + confidence_threshold=0.0, + ) - def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps: + def deep_ask_deps_factory(input_config: dict[str, Any]) -> ResearchDeps: effective_db_path = ( db_path or input_config.get("db_path") or config.storage.data_dir / "haiku.rag.lancedb" ) - return DeepQADeps(client=get_client(effective_db_path)) + return ResearchDeps(client=get_client(effective_db_path)) # Create event stream functions for each graph type async def research_event_stream( diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py deleted file mode 100644 index aaeb1ae9..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from haiku.rag.graph.deep_qa.models import DeepQAAnswer diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py deleted file mode 100644 index efd6d92f..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py +++ /dev/null @@ -1,29 +0,0 @@ -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 - - -class DeepQAContext(BaseModel): - original_question: str = Field(description="The original question") - sub_questions: list[str] = Field( - default_factory=list, description="Decomposed sub-questions" - ) - qa_responses: list[SearchAnswer] = Field( - default_factory=list, description="QA pairs collected during answering" - ) - - def add_qa_response(self, qa: SearchAnswer) -> None: - """Add a QA response (citations already resolved).""" - self.qa_responses.append(qa) - - -class DeepQADependencies(BaseModel): - model_config = {"arbitrary_types_allowed": True} - - client: HaikuRAG = Field(description="RAG client for document operations") - context: DeepQAContext = Field(description="Shared QA context") - search_results: list[SearchResult] = Field( - default_factory=list, description="Search results for citation resolution" - ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py deleted file mode 100644 index 1df7b34f..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py +++ /dev/null @@ -1,250 +0,0 @@ -from pydantic_ai import Agent -from pydantic_ai.format_prompt import format_as_xml -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, resolve_citations -from haiku.rag.graph.common.nodes import create_plan_node, create_search_node -from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies -from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation -from haiku.rag.graph.deep_qa.prompts import DECISION_PROMPT, SYNTHESIS_PROMPT -from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState -from haiku.rag.store.models import SearchResult - - -def build_deep_qa_graph( - config: AppConfig = Config, -) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]: - """Build the Deep QA graph. - - Args: - config: AppConfig object (uses config.qa for provider, model, and graph parameters) - - Returns: - Configured Deep QA graph - """ - model_config = config.qa.model - g = GraphBuilder( - state_type=DeepQAState, - deps_type=DeepQADeps, - output_type=DeepQAAnswer, - ) - - # Create and register the plan node using the factory - plan = g.step( - create_plan_node( - model_config=model_config, - deps_type=DeepQADependencies, # type: ignore[arg-type] - activity_message="Planning approach", - output_retries=None, # Deep QA doesn't use output_retries - config=config, - ) - ) # 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, - deps_type=DeepQADependencies, # type: ignore[arg-type] - with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step - success_message_format="Answered: {sub_q}", - handle_exceptions=True, - config=config, - ) - ) # type: ignore[arg-type] - - @g.step - async def get_batch( - ctx: StepContext[DeepQAState, DeepQADeps, None | bool], - ) -> list[str] | None: - """Get all remaining questions for this iteration.""" - state = ctx.state - - if not state.context.sub_questions: - return None - - # Take ALL remaining questions - max_concurrency controls parallel execution within .map() - batch = list(state.context.sub_questions) - state.context.sub_questions.clear() - return batch - - @g.step - async def decide( - ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]], - ) -> bool: - state = ctx.state - deps = ctx.deps - - if deps.agui_emitter: - deps.agui_emitter.start_step("decide") - deps.agui_emitter.update_activity( - "evaluating", {"message": "Evaluating information sufficiency"} - ) - - try: - agent = Agent( - model=get_model(model_config, config), - output_type=DeepQAEvaluation, - instructions=DECISION_PROMPT, - retries=3, - deps_type=DeepQADependencies, - ) - - context_data = { - "original_question": state.context.original_question, - "gathered_answers": [ - { - "question": qa.query, - "answer": qa.answer, - "confidence": qa.confidence, - } - for qa in state.context.qa_responses - ], - } - context_xml = format_as_xml(context_data, root_tag="gathered_information") - - prompt = ( - "Evaluate whether we have sufficient information to answer the question.\n\n" - f"{context_xml}" - ) - - agent_deps = DeepQADependencies( - client=deps.client, - context=state.context, - ) - result = await agent.run(prompt, deps=agent_deps) - evaluation = result.output - - state.iterations += 1 - - for new_q in evaluation.new_questions: - if new_q not in state.context.sub_questions: - state.context.sub_questions.append(new_q) - - if deps.agui_emitter: - deps.agui_emitter.update_state(state) - status = "sufficient" if evaluation.is_sufficient else "insufficient" - deps.agui_emitter.update_activity( - "evaluating", - { - "stepName": "decide", - "message": f"Information {status} after {state.iterations} iteration(s)", - "is_sufficient": evaluation.is_sufficient, - "iterations": state.iterations, - }, - ) - - should_continue = ( - not evaluation.is_sufficient and state.iterations < state.max_iterations - ) - - return should_continue - finally: - if deps.agui_emitter: - deps.agui_emitter.finish_step() - - @g.step - async def synthesize( - ctx: StepContext[DeepQAState, DeepQADeps, None | bool], - ) -> DeepQAAnswer: - state = ctx.state - deps = ctx.deps - - if deps.agui_emitter: - deps.agui_emitter.start_step("synthesize") - deps.agui_emitter.update_activity( - "synthesizing", {"message": "Synthesizing final answer"} - ) - - try: - agent = Agent( - model=get_model(model_config, config), - output_type=SearchAnswer, - instructions=SYNTHESIS_PROMPT, - retries=3, - deps_type=DeepQADependencies, - ) - - context_data = { - "original_question": state.context.original_question, - "sub_answers": [ - { - "question": qa.query, - "answer": qa.answer, - "confidence": qa.confidence, - "cited_chunks": qa.cited_chunks, - } - for qa in state.context.qa_responses - ], - } - context_xml = format_as_xml(context_data, root_tag="gathered_information") - - prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}" - - agent_deps = DeepQADependencies( - client=deps.client, - context=state.context, - ) - result = await agent.run(prompt, deps=agent_deps) - llm_answer = result.output - - # Resolve citations by fetching chunks by ID - search_results = [] - for chunk_id in llm_answer.cited_chunks: - chunk = await deps.client.chunk_repository.get_by_id(chunk_id) - if chunk: - search_results.append(SearchResult.from_chunk(chunk, score=1.0)) - citations = resolve_citations(llm_answer.cited_chunks, search_results) - - if deps.agui_emitter: - deps.agui_emitter.update_activity( - "synthesizing", {"message": "Answer complete"} - ) - - return DeepQAAnswer(answer=llm_answer.answer, citations=citations) - finally: - if deps.agui_emitter: - deps.agui_emitter.finish_step() - - # Build the graph structure - collect_answers = g.join( - reduce_list_append, - initial_factory=list[SearchAnswer], - ) - - g.add( - g.edge_from(g.start_node).to(plan), - g.edge_from(plan).to(get_batch), - ) - - # Branch based on whether we have questions - g.add( - g.edge_from(get_batch).to( - g.decision() - .branch(g.match(list).label("Has questions").map().to(search_one)) - .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), - ) - - # Branch based on decision - g.add( - g.edge_from(decide).to( - g.decision() - .branch( - g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch) - ) - .branch( - g.match(bool, matches=lambda x: not x) - .label("Done with QA") - .to(synthesize) - ) - ), - g.edge_from(synthesize).to(g.end_node), - ) - - return g.build() diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py deleted file mode 100644 index ba4d2b89..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py +++ /dev/null @@ -1,23 +0,0 @@ -from pydantic import BaseModel, Field - -from haiku.rag.graph.common.models import Citation - - -class DeepQAEvaluation(BaseModel): - is_sufficient: bool = Field( - description="Whether we have sufficient information to answer the question" - ) - reasoning: str = Field(description="Explanation of the sufficiency assessment") - new_questions: list[str] = Field( - description="Additional sub-questions needed if insufficient", - default_factory=list, - ) - - -class DeepQAAnswer(BaseModel): - """Final deep QA answer with resolved citations.""" - - answer: str = Field(description="The comprehensive answer to the question") - citations: list[Citation] = Field( - default_factory=list, description="Resolved citations for the answer" - ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py deleted file mode 100644 index 9dc46615..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Deep QA specific prompts.""" - -SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers. - -Task: -- Combine the gathered information from sub-questions into a single comprehensive answer -- Answer the original question directly and completely -- Base your answer strictly on the provided evidence -- Be clear, accurate, and well-structured - -Output format: -- query: Echo the original question being answered -- answer: The complete answer to the original question (2-4 paragraphs) -- cited_chunks: List of plain strings containing chunk IDs (UUIDs only, not objects) -- confidence: A score from 0.0 to 1.0 indicating answer confidence - -Guidelines: -- Start directly with the answer - no preamble like "Based on the research..." -- Use a clear, professional tone -- Organize information logically -- If evidence is incomplete, state limitations clearly -- Do not include any claims not supported by the gathered information -- Each sub-answer includes cited_chunks IDs - include the relevant ones in your response""" - -DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question. - -Task: -- Review the original question and all gathered sub-question answers -- Determine if we have enough information to provide a comprehensive answer -- If insufficient, suggest specific new sub-questions to fill the gaps - -Output format: -- is_sufficient: Boolean indicating if we can answer the question comprehensively -- reasoning: Clear explanation of your assessment -- new_questions: List of plain strings, each a specific follow-up question (not objects) - -Guidelines: -- Be strict but reasonable in your assessment -- Focus on whether core aspects of the question are addressed -- New questions should be specific and distinct from what's been asked -- Limit new questions to 2-3 maximum -- Consider whether additional searches would meaningfully improve the answer""" diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py deleted file mode 100644 index 24dfc1a2..00000000 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py +++ /dev/null @@ -1,59 +0,0 @@ -import asyncio -from dataclasses import dataclass -from typing import TYPE_CHECKING - -from pydantic import BaseModel, Field - -from haiku.rag.client import HaikuRAG -from haiku.rag.graph.deep_qa.dependencies import DeepQAContext -from haiku.rag.graph.deep_qa.models import DeepQAAnswer - -if TYPE_CHECKING: - from haiku.rag.config.models import AppConfig - from haiku.rag.graph.agui.emitter import AGUIEmitter - - -@dataclass -class DeepQADeps: - client: HaikuRAG - agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None - semaphore: asyncio.Semaphore | None = None - - -class DeepQAState(BaseModel): - """Deep QA state for multi-agent question answering.""" - - model_config = {"arbitrary_types_allowed": True} - - context: DeepQAContext = Field(description="Shared QA context") - max_sub_questions: int = Field( - default=3, description="Maximum number of sub-questions" - ) - max_iterations: int = Field( - default=2, description="Maximum number of QA iterations" - ) - max_concurrency: int = Field( - default=1, description="Maximum parallel sub-question searches" - ) - iterations: int = Field(default=0, description="Current iteration number") - search_filter: str | None = Field( - default=None, description="SQL WHERE clause to filter search results" - ) - - @classmethod - def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState": - """Create a DeepQAState from an AppConfig. - - Args: - context: The DeepQAContext containing the question and settings - config: The AppConfig object (uses config.qa for state parameters) - - Returns: - A configured DeepQAState instance - """ - return cls( - context=context, - max_sub_questions=config.qa.max_sub_questions, - max_iterations=config.qa.max_iterations, - max_concurrency=config.qa.max_concurrency, - ) diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 7b1129f9..f263ea6c 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -19,11 +19,13 @@ from haiku.rag.graph.research.state import ResearchDeps, ResearchState def build_research_graph( config: AppConfig = Config, + include_plan: bool = True, ) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: """Build the Research graph. Args: config: AppConfig object (uses config.research for provider, model, and graph parameters) + include_plan: Whether to include the planning step (False for execute-only mode) Returns: Configured Research graph @@ -191,10 +193,13 @@ def build_research_graph( initial_factory=list[SearchAnswer], ) - g.add( - g.edge_from(g.start_node).to(plan), - g.edge_from(plan).to(get_batch), - ) + if include_plan: + g.add( + g.edge_from(g.start_node).to(plan), + g.edge_from(plan).to(get_batch), + ) + else: + g.add(g.edge_from(g.start_node).to(get_batch)) # Branch based on whether we have questions g.add( diff --git a/haiku_rag_slim/haiku/rag/graph/research/state.py b/haiku_rag_slim/haiku/rag/graph/research/state.py index f69301dd..fbf4fc39 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/state.py +++ b/haiku_rag_slim/haiku/rag/graph/research/state.py @@ -54,12 +54,27 @@ class ResearchState(BaseModel): @classmethod def from_config( - cls, context: ResearchContext, config: "AppConfig" + cls, + context: ResearchContext, + config: "AppConfig", + max_iterations: int | None = None, + confidence_threshold: float | None = None, ) -> "ResearchState": - """Create a ResearchState from an AppConfig.""" + """Create a ResearchState from an AppConfig. + + Args: + context: The ResearchContext containing the question + config: The AppConfig object + max_iterations: Override max iterations (None uses config default) + confidence_threshold: Override threshold (None uses config, 0.0 disables check) + """ return cls( context=context, - max_iterations=config.research.max_iterations, - confidence_threshold=config.research.confidence_threshold, + max_iterations=max_iterations + if max_iterations is not None + else config.research.max_iterations, + confidence_threshold=confidence_threshold + if confidence_threshold is not None + else config.research.confidence_threshold, max_concurrency=config.research.max_concurrency, ) diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 48dbaaa3..949d44b1 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -174,18 +174,26 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP: try: async with HaikuRAG(db_path, config=config) as rag: if deep: - from haiku.rag.graph.deep_qa.dependencies import DeepQAContext - from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph - from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState + from haiku.rag.graph.research.dependencies import ResearchContext + from haiku.rag.graph.research.graph import build_research_graph + from haiku.rag.graph.research.state import ( + ResearchDeps, + ResearchState, + ) - graph = build_deep_qa_graph(config=config) - context = DeepQAContext(original_question=question) - state = DeepQAState.from_config(context=context, config=config) - deps = DeepQADeps(client=rag) + graph = build_research_graph(config=config) + context = ResearchContext(original_question=question) + state = ResearchState.from_config( + context=context, + config=config, + max_iterations=2, + confidence_threshold=0.0, + ) + deps = ResearchDeps(client=rag) result = await graph.run(state=state, deps=deps) - answer = result.answer - citations = result.citations + answer = result.executive_summary + citations = [] else: answer, citations = await rag.ask(question) if cite and citations: diff --git a/tests/graph/test_deep_qa.py b/tests/graph/test_deep_qa.py deleted file mode 100644 index a87629dc..00000000 --- a/tests/graph/test_deep_qa.py +++ /dev/null @@ -1,41 +0,0 @@ -import pytest -from pydantic_ai.models.test import TestModel - -from haiku.rag.client import HaikuRAG -from haiku.rag.graph.deep_qa.dependencies import DeepQAContext -from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph -from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState - - -@pytest.mark.asyncio -async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): - """Test deep Q&A graph with mocked LLM using TestModel.""" - - # Mock get_model to return TestModel which generates valid schema-compliant data - def test_model_factory(provider, model, config=None): - return TestModel() - - # 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.deep_qa.graph.get_model", test_model_factory) - - graph = build_deep_qa_graph() - - state = DeepQAState( - context=DeepQAContext(original_question="What is haiku.rag?"), - max_sub_questions=3, - ) - - # Use real client but with TestModel for LLM calls - client = HaikuRAG(temp_db_path, create=True) - deps = DeepQADeps(client=client) - - result = await graph.run(state=state, deps=deps) - - # TestModel will generate valid structured output based on schemas - assert result.answer is not None - assert isinstance(result.answer, str) - - client.close() diff --git a/tests/graph/test_search_filter.py b/tests/graph/test_search_filter.py index 0c9e99bd..17afd7c1 100644 --- a/tests/graph/test_search_filter.py +++ b/tests/graph/test_search_filter.py @@ -2,9 +2,6 @@ import pytest from pydantic_ai.models.test import TestModel from haiku.rag.client import HaikuRAG -from haiku.rag.graph.deep_qa.dependencies import DeepQAContext -from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph -from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.state import ResearchDeps, ResearchState @@ -97,52 +94,6 @@ async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs): ) -@pytest.mark.asyncio -async def test_deep_qa_graph_uses_search_filter(monkeypatch, client_with_docs): - """Test that deep QA graph passes search_filter to search operations.""" - client, doc1_id, doc2_id = client_with_docs - - # Track search calls to verify filter is passed - search_calls = [] - original_search = client.search - - async def tracking_search(query, limit=None, search_type="hybrid", filter=None): - search_calls.append({"query": query, "filter": filter}) - return await original_search(query, limit, search_type, filter) - - client.search = tracking_search - - # Mock get_model to return TestModel - def test_model_factory(_provider, _model, _config=None): - 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.deep_qa.graph.get_model", test_model_factory) - - graph = build_deep_qa_graph() - - # Create state with search_filter - filter_clause = f"id = '{doc2_id}'" - state = DeepQAState( - context=DeepQAContext(original_question="Tell me about animals"), - max_sub_questions=2, - search_filter=filter_clause, - ) - - deps = DeepQADeps(client=client) - - await graph.run(state=state, deps=deps) - - # Verify search was called with the filter - assert len(search_calls) > 0, "Expected search to be called" - for call in search_calls: - assert call["filter"] == filter_clause, ( - f"Expected filter '{filter_clause}', got '{call['filter']}'" - ) - - @pytest.mark.asyncio async def test_search_filter_none_searches_all(monkeypatch, client_with_docs): """Test that search_filter=None searches all documents.""" diff --git a/tests/test_app.py b/tests/test_app.py index 0d848761..58188c7e 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -358,25 +358,36 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch): @pytest.mark.asyncio async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep QA.""" - from haiku.rag.graph.deep_qa.models import DeepQAAnswer + """Test asking a question with deep mode uses research graph.""" + import haiku.rag.app as app_module + from haiku.rag.graph.research.models import ResearchReport - mock_output = DeepQAAnswer(answer="Deep QA answer") + mock_output = ResearchReport( + title="Test", + executive_summary="Deep research answer", + main_findings=["Finding 1"], + conclusions=["Conclusion 1"], + sources_summary="Sources", + ) mock_graph = AsyncMock() mock_graph.run.return_value = mock_output mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() monkeypatch.setattr(app.console, "print", mock_print) + monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph) - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - with patch( - "haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph - ): - await app.ask("test question", deep=True) + with patch("haiku.rag.app.HaikuRAG") as mock_rag_class: + mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) + await app.ask("test question", deep=True) + + # Check if there was an error printed + print_calls = [str(c) for c in mock_print.call_args_list] + error_calls = [c for c in print_calls if "Error" in c] + assert not error_calls, f"Error was printed: {error_calls}" mock_graph.run.assert_called_once() call_kwargs = mock_graph.run.call_args[1] @@ -385,25 +396,31 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): @pytest.mark.asyncio async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep QA and citations (cite ignored for deep).""" - from haiku.rag.graph.deep_qa.models import DeepQAAnswer + """Test asking a question with deep mode (cite is ignored for research graph).""" + import haiku.rag.app as app_module + from haiku.rag.graph.research.models import ResearchReport - mock_output = DeepQAAnswer(answer="Deep QA answer") + mock_output = ResearchReport( + title="Test", + executive_summary="Deep research answer", + main_findings=["Finding 1"], + conclusions=["Conclusion 1"], + sources_summary="Sources", + ) mock_graph = AsyncMock() mock_graph.run.return_value = mock_output mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() monkeypatch.setattr(app.console, "print", mock_print) + monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph) - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - with patch( - "haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph - ): - await app.ask("test question", deep=True, cite=True) + with patch("haiku.rag.app.HaikuRAG") as mock_rag_class: + mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) + await app.ask("test question", deep=True, cite=True) mock_graph.run.assert_called_once() call_kwargs = mock_graph.run.call_args[1] @@ -412,9 +429,10 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): @pytest.mark.asyncio async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep QA and verbose output.""" + """Test asking a question with deep mode and verbose output.""" + import haiku.rag.app as app_module - mock_output = {"answer": "Deep QA answer", "citations": []} + mock_output = {"executive_summary": "Deep research answer"} mock_renderer = AsyncMock() mock_renderer.render.return_value = mock_output @@ -422,17 +440,16 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): mock_graph = AsyncMock() mock_client = AsyncMock() - mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() monkeypatch.setattr(app.console, "print", mock_print) + monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph) - with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): - with patch( - "haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph - ): - with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer): - await app.ask("test question", deep=True, verbose=True) + with patch("haiku.rag.app.HaikuRAG") as mock_rag_class: + mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client) + mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) + with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer): + await app.ask("test question", deep=True, verbose=True) # With verbose, it should use AGUIConsoleRenderer.render, not graph.run mock_renderer.render.assert_called_once() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 418ea339..01694757 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -243,7 +243,7 @@ async def test_mcp_delete_document(): @pytest.mark.asyncio async def test_mcp_ask_question_deep(): - """Test ask_question tool with deep=True is properly wired.""" + """Test ask_question tool with deep=True uses research graph.""" with tempfile.TemporaryDirectory() as temp_dir: db_path = Path(temp_dir) / "test.lancedb" mcp = create_mcp_server(db_path) @@ -251,7 +251,7 @@ async def test_mcp_ask_question_deep(): with ( patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class, patch( - "haiku.rag.graph.deep_qa.graph.build_deep_qa_graph" + "haiku.rag.graph.research.graph.build_research_graph" ) as mock_graph_builder, ): mock_rag = AsyncMock() @@ -260,8 +260,7 @@ async def test_mcp_ask_question_deep(): mock_graph = AsyncMock() mock_result = AsyncMock() - mock_result.answer = "Deep answer" - mock_result.citations = [] + mock_result.executive_summary = "Deep answer from research" mock_graph.run = AsyncMock(return_value=mock_result) mock_graph_builder.return_value = mock_graph @@ -273,7 +272,7 @@ async def test_mcp_ask_question_deep(): question="Deep question?", cite=False, deep=True ) - assert result == "Deep answer" + assert result == "Deep answer from research" mock_graph.run.assert_called_once() From 6e621ea34daccfe9f51532cd5bc5869380122b5b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Dec 2025 13:22:51 +0200 Subject: [PATCH 3/7] Provide full report on deep ask, modify prompt to not just summarize --- haiku_rag_slim/haiku/rag/app.py | 42 +++++++++++++------ .../haiku/rag/graph/research/prompts.py | 12 ++++-- 2 files changed, 38 insertions(+), 16 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index c04cdbdf..7f216fd0 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -329,6 +329,8 @@ class HaikuRAGApp: try: citations = [] if deep: + from haiku.rag.graph.research.models import ResearchReport + graph = build_research_graph(config=self.config) context = ResearchContext(original_question=question) state = ResearchState.from_config( @@ -345,24 +347,40 @@ class HaikuRAGApp: result_dict = await renderer.render( stream_graph(graph, state, deps) ) - answer = ( - result_dict.get("executive_summary", "") + report = ( + ResearchReport.model_validate(result_dict) if result_dict - else "" + else None ) else: - result = await graph.run(state=state, deps=deps) - answer = result.executive_summary + report = await graph.run(state=state, deps=deps) + + self.console.print(f"[bold blue]Question:[/bold blue] {question}") + self.console.print() + if report: + self.console.print("[bold green]Answer:[/bold green]") + self.console.print(Markdown(report.executive_summary)) + if report.main_findings: + self.console.print() + self.console.print("[bold cyan]Key Findings:[/bold cyan]") + for finding in report.main_findings: + self.console.print(f"• {finding}") + if report.sources_summary: + self.console.print() + self.console.print("[bold cyan]Sources:[/bold cyan]") + self.console.print(report.sources_summary) + else: + self.console.print("[yellow]No answer generated.[/yellow]") else: answer, citations = await self.client.ask(question, filter=filter) - self.console.print(f"[bold blue]Question:[/bold blue] {question}") - self.console.print() - self.console.print("[bold green]Answer:[/bold green]") - self.console.print(Markdown(answer)) - if cite and citations: - for renderable in format_citations_rich(citations): - self.console.print(renderable) + self.console.print(f"[bold blue]Question:[/bold blue] {question}") + self.console.print() + self.console.print("[bold green]Answer:[/bold green]") + self.console.print(Markdown(answer)) + if cite and citations: + for renderable in format_citations_rich(citations): + self.console.print(renderable) except Exception as e: self.console.print(f"[red]Error: {e}[/red]") diff --git a/haiku_rag_slim/haiku/rag/graph/research/prompts.py b/haiku_rag_slim/haiku/rag/graph/research/prompts.py index d3b8aec4..f91a3bbf 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/research/prompts.py @@ -20,17 +20,20 @@ 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 -research report. +research report that directly answers the original question. Goals: -1. Synthesize all gathered information into a coherent narrative. +1. Directly answer the research question using gathered evidence. 2. Present findings clearly and concisely. 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. +- executive_summary: 3-5 sentences that DIRECTLY ANSWER the original question. + Write the actual answer, not a description of what the report contains. + BAD: "This report examines the topic and presents findings..." + GOOD: "The system requires configuration X and supports features Y and Z..." - 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. @@ -42,7 +45,8 @@ All list fields must contain plain strings only, not objects. Style: - Base all content solely on the collected evidence. - Be professional, objective, and specific. -- Avoid meta commentary and refrain from speculation beyond the evidence.""" +- 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. From 557a2963410e0397a8fc05965d7e44b23ccc60b4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Dec 2025 13:43:13 +0200 Subject: [PATCH 4/7] Update ag-ui example --- CHANGELOG.md | 21 ++ examples/ag-ui-research/README.md | 10 +- examples/ag-ui-research/backend/README.md | 4 +- examples/ag-ui-research/backend/agent.py | 3 +- .../frontend/components/Agent.tsx | 31 -- .../frontend/components/StateDisplay.tsx | 329 ------------------ .../haiku/rag/graph/common/nodes.py | 2 +- 7 files changed, 28 insertions(+), 372 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2defdf9c..e986e0b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,17 @@ ### Changed - **Chunker Sets Order**: Chunkers now set `chunk.order` directly +- **Unified Research Graph**: Simplified and unified research and deep QA into a single configurable graph + - Removed `analyze_insights` node - graph now flows directly from `collect_answers` to `decide` + - Simplified `EvaluationResult` to: `is_sufficient`, `confidence_score`, `reasoning`, `new_questions` + - Simplified `ResearchContext` - removed insight/gap tracking methods + - `ask --deep` now uses research graph with `max_iterations=2`, `confidence_threshold=0.0` + - `ask --deep` output now shows executive summary, key findings, and sources + - Added `include_plan` parameter to `build_research_graph()` for plan-less execution + - Added `max_iterations` and `confidence_threshold` overrides to `ResearchState.from_config()` +- **Improved Synthesis Prompt**: Updated synthesis agent prompt to produce direct answers + - Executive summary now directly answers the question instead of describing the report + - Added explicit examples of good vs bad output style - **Evaluations Vacuum Strategy**: `populate_db` now uses periodic vacuum to prevent disk exhaustion with large datasets - Disables auto_vacuum during population, vacuums every N documents with retention=0 - New `--vacuum-interval` CLI option (default: 100) to control vacuum frequency @@ -23,6 +34,16 @@ - Added dedicated Methodology section explaining MRR, MAP, and QA Accuracy metrics - Organized results by dataset with retrieval and QA subsections +### Removed + +- **Deep QA Graph**: Removed `haiku.rag.graph.deep_qa` module entirely + - Use `build_research_graph()` with appropriate parameters instead + - `ask --deep` CLI command now uses research graph internally +- **Insight/Gap Tracking**: Removed over-engineered insight and gap tracking from research graph + - Removed `InsightRecord`, `GapRecord`, `InsightAnalysis`, `InsightStatus`, `GapSeverity` models + - Removed `format_analysis_for_prompt()` helper + - Removed `INSIGHT_AGENT_PROMPT` from prompts + ## [0.20.2] - 2025-12-12 ### Fixed diff --git a/examples/ag-ui-research/README.md b/examples/ag-ui-research/README.md index 085222a5..df6850dc 100644 --- a/examples/ag-ui-research/README.md +++ b/examples/ag-ui-research/README.md @@ -6,10 +6,9 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), ## Features -- **Multi-iteration research graph**: Automated question decomposition, search, insight extraction, and gap analysis +- **Multi-iteration research graph**: Automated question decomposition and search - **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered - **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol -- **Insight & gap tracking**: Structured insights with provenance and automatic gap identification - **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources ## Quick Start @@ -81,9 +80,8 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), - Gathers initial context about the topic 3. **Research iterations**: The graph autonomously: - Searches the knowledge base for each sub-question in parallel - - Extracts structured insights with source provenance - - Identifies information gaps and assesses confidence - - Generates new follow-up questions for gaps + - Assesses confidence in gathered information + - Generates new follow-up questions if needed - Iterates until confidence threshold is met or max iterations reached 4. **Synthesis**: Generates a comprehensive research report with: - Executive summary @@ -126,7 +124,7 @@ This example demonstrates the **agent+graph** architecture pattern: - CopilotKit for AG-UI protocol integration - Split-pane UI: chat on left, live research state on right - Real-time state synchronization via Server-Sent Events (SSE) - - `StateDisplay` component with collapsible sections for questions, insights, and gaps + - `StateDisplay` component with collapsible sections for questions and report ## Configuration diff --git a/examples/ag-ui-research/backend/README.md b/examples/ag-ui-research/backend/README.md index 9b0751e9..8b6da3df 100644 --- a/examples/ag-ui-research/backend/README.md +++ b/examples/ag-ui-research/backend/README.md @@ -15,13 +15,11 @@ The server starts on `http://localhost:8000` and uses [haiku.rag configuration]( The backend uses `create_agui_server()` from `haiku.rag.graph.agui.server` which provides: -- **Research graph execution**: Multi-iteration research workflow with insight/gap tracking +- **Research graph execution**: Multi-iteration research workflow - **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates - **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations -- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa` ## Endpoints - `GET /health` - Health check with configuration info - `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol) -- `POST /agent/deep_qa/stream` - Deep QA graph streaming endpoint (AG-UI protocol) diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index 3dce81ac..e697145f 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -56,7 +56,7 @@ How to decide: - "Tell me about Y" → Use run_research tool When you use run_research, the graph will decompose questions, search the knowledge base, -extract insights, and generate a comprehensive report. +and generate a comprehensive report. Be friendly and conversational in all responses.""", ) @@ -100,7 +100,6 @@ Main Findings: Conclusions: {chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])} -Total insights gathered: {len(state.context.insights)} Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"} Iterations completed: {state.iterations} diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx index 8a98a3dd..75429a91 100644 --- a/examples/ag-ui-research/frontend/components/Agent.tsx +++ b/examples/ag-ui-research/frontend/components/Agent.tsx @@ -6,26 +6,6 @@ import "@copilotkit/react-ui/styles.css"; import DocumentSelector from "./DocumentSelector"; import StateDisplay from "./StateDisplay"; -interface InsightRecord { - id: string; - summary: string; - status: string; - notes?: string; - supporting_sources: string[]; - originating_questions: string[]; -} - -interface GapRecord { - id: string; - description: string; - severity: string; - blocking: boolean; - resolved: boolean; - notes?: string; - supporting_sources: string[]; - resolved_by: string[]; -} - interface Citation { document_id: string; chunk_id: string; @@ -48,14 +28,10 @@ interface ResearchContext { original_question: string; sub_questions: string[]; qa_responses: SearchAnswer[]; - insights: InsightRecord[]; - gaps: GapRecord[]; } interface EvaluationResult { - key_insights: string[]; new_questions: string[]; - gaps: string[]; confidence_score: number; is_sufficient: boolean; reasoning: string; @@ -78,10 +54,6 @@ interface ResearchState { confidence_threshold: number; max_concurrency: number; last_eval: EvaluationResult | null; - last_analysis: { - insights_extracted: InsightRecord[]; - gaps_identified: GapRecord[]; - } | null; result?: ResearchReport; current_activity?: string; current_activity_message?: string; @@ -96,15 +68,12 @@ function AgentContent() { original_question: "", sub_questions: [], qa_responses: [], - insights: [], - gaps: [], }, iterations: 0, max_iterations: 3, confidence_threshold: 0.8, max_concurrency: 1, last_eval: null, - last_analysis: null, documentFilter: [], }, }); diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx index c0ad3291..387b4a57 100644 --- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx +++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx @@ -11,26 +11,6 @@ interface VisualGroundingState { error: string | null; } -interface InsightRecord { - id: string; - summary: string; - status: string; - notes?: string; - supporting_sources: string[]; - originating_questions: string[]; -} - -interface GapRecord { - id: string; - description: string; - severity: string; - blocking: boolean; - resolved: boolean; - notes?: string; - supporting_sources: string[]; - resolved_by: string[]; -} - interface Citation { document_id: string; chunk_id: string; @@ -53,14 +33,10 @@ interface ResearchContext { original_question: string; sub_questions: string[]; qa_responses: SearchAnswer[]; - insights: InsightRecord[]; - gaps: GapRecord[]; } interface EvaluationResult { - key_insights: string[]; new_questions: string[]; - gaps: string[]; confidence_score: number; is_sufficient: boolean; reasoning: string; @@ -83,10 +59,6 @@ interface ResearchState { confidence_threshold: number; max_concurrency: number; last_eval: EvaluationResult | null; - last_analysis: { - insights_extracted: InsightRecord[]; - gaps_identified: GapRecord[]; - } | null; result?: ResearchReport; current_activity?: string; current_activity_message?: string; @@ -101,8 +73,6 @@ export default function StateDisplay({ state }: StateDisplayProps) { Record >({ questions: true, - insights: true, - gaps: true, report: true, }); @@ -694,305 +664,6 @@ export default function StateDisplay({ state }: StateDisplayProps) { )} - {/* Insights */} - {state.context.insights.length > 0 && ( -
- - {expandedSections.insights && ( -
- {state.context.insights.map((insight) => ( -
-
- - {insight.status} - - - {insight.supporting_sources.length} sources - -
-
- -
- {insight.notes && ( -
- -
- )} - {insight.supporting_sources.length > 0 && ( -
- Sources: - {insight.supporting_sources.map((source, srcIdx) => ( - - {srcIdx > 0 && ", "} - {source} - - ))} -
- )} -
- ))} -
- )} -
- )} - - {/* Knowledge Gaps */} - {state.context.gaps.length > 0 && ( -
- - {expandedSections.gaps && ( -
- {state.context.gaps.map((gap) => ( -
-
-
- - {gap.severity} - - {gap.blocking && ( - - Blocking - - )} - {gap.resolved && ( - - Resolved - - )} -
-
-
- -
- {gap.notes && ( -
- -
- )} - {gap.resolved && gap.resolved_by.length > 0 && ( -
- Resolved by: - {gap.resolved_by.map((source, srcIdx) => ( - - {srcIdx > 0 && ", "} - {source} - - ))} -
- )} -
- ))} -
- )} -
- )} - {/* Final Report */} {state.result && (
Date: Mon, 15 Dec 2025 13:43:24 +0200 Subject: [PATCH 5/7] Update docs --- docs/agents.md | 267 +++++++++++++++---------------------------------- docs/server.md | 3 +- 2 files changed, 80 insertions(+), 190 deletions(-) diff --git a/docs/agents.md b/docs/agents.md index 4d218383..59c48482 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -1,156 +1,57 @@ # Agents -Three agentic flows are provided by haiku.rag: +Two agentic flows are provided by haiku.rag: -- Simple QA Agent — a focused question answering agent -- Deep QA Agent — multi-agent question decomposition for complex questions -- Research Multi‑Agent — a multi‑step, analyzable research workflow +- **Simple QA Agent** — a focused question answering agent +- **Research Graph** — a multi-step research workflow with question decomposition -For an interactive example using Pydantic AI and AG-UI, see the [Interactive Research Assistant](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research) example ([demo video](https://vimeo.com/1128874386)). The demo uses a knowledge base containing haiku.rag's code and documentation. +For an interactive example using Pydantic AI and AG-UI, see the [Interactive Research Assistant](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research) example ([demo video](https://vimeo.com/1128874386)). See [QA and Research Configuration](configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings. - -### Simple QA Agent +## Simple QA Agent The simple QA agent answers a single question using the knowledge base. It retrieves relevant chunks, optionally expands context around them, and asks the model to answer strictly based on that context. Key points: - Uses a single `search_documents` tool to fetch relevant chunks -- Can be run with or without inline citations in the prompt (citations prefer - document titles when present, otherwise URIs) +- Can be run with or without inline citations in the prompt - Returns a plain string answer -Python usage: +**CLI usage:** + +```bash +haiku-rag ask "What is climate change?" + +# With citations +haiku-rag ask "What is climate change?" --cite + +# Deep mode (uses research graph with optimized settings) +haiku-rag ask "What are the main features of haiku.rag?" --deep +``` + +**Python usage:** ```python from haiku.rag.client import HaikuRAG from haiku.rag.qa.agent import QuestionAnswerAgent async with HaikuRAG(path_to_db) as client: - # Choose a provider and model (see Configuration for env defaults) agent = QuestionAnswerAgent( client=client, - provider="openai", # or "ollama", "vllm", etc. + provider="openai", model="gpt-4o-mini", - use_citations=False, # set True to bias prompt towards citing sources + use_citations=False, ) answer = await agent.answer("What is climate change?") print(answer) ``` -### Deep QA Agent +## Research Graph -Deep QA is a multi-agent system that decomposes complex questions into sub-questions, answers them in batches, evaluates sufficiency, and iterates if needed before synthesizing a final answer. It's lighter than the full research workflow but more powerful than the simple QA agent. - -```mermaid ---- -title: Deep QA graph ---- -stateDiagram-v2 - [*] --> plan - plan --> get_batch - get_batch --> search_one: Has questions (map) - get_batch --> synthesize: No questions - search_one --> collect_answers - collect_answers --> decide - decide --> get_batch: Continue QA - decide --> synthesize: Done with QA - synthesize --> [*] -``` - -Key nodes: - -- **plan**: Decomposes the question into focused sub-questions using a presearch tool -- **get_batch**: Retrieves remaining sub-questions for the current iteration -- **search_one**: Answers a single sub-question using the knowledge base (mapped in parallel) -- **collect_answers**: Aggregates search results from parallel executions -- **decide**: Evaluates if sufficient information has been gathered or if more iterations are needed -- **synthesize**: Generates the final comprehensive answer from all gathered information - -Key differences from Research: - -- **Simpler evaluation**: Uses sufficiency check (not confidence + insight analysis) -- **Direct answers**: Returns just the answer (not a full research report) -- **Question-focused**: Optimized for answering specific questions, not open-ended research -- **Supports citations**: Can include inline source citations like `[document.md]` -- **Configurable iterations**: Control max_iterations (default: 2) and max_concurrency (default: 1) - -Note on parallel execution: -- The `search_one` node is mapped over all questions in a batch -- Parallelism is controlled via `max_concurrency` -- All questions in an iteration are processed before evaluation - -CLI usage: - -```bash -# Deep QA without citations -haiku-rag ask "What are the main features of haiku.rag?" --deep - -# Deep QA with citations -haiku-rag ask "What are the main features of haiku.rag?" --deep --cite -``` - -Python usage: - -```python -from haiku.rag.client import HaikuRAG -from haiku.rag.config import Config -from haiku.rag.graph.deep_qa.dependencies import DeepQAContext -from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph -from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState - -async with HaikuRAG(path_to_db) as client: - # Use global config (recommended) - graph = build_deep_qa_graph(config=Config) - context = DeepQAContext( - original_question="What are the main features of haiku.rag?", - use_citations=True - ) - state = DeepQAState.from_config(context=context, config=Config) - deps = DeepQADeps(client=client) - - result = await graph.run( - state=state, - deps=deps - ) - - print(result.answer) - print(result.sources) -``` - -Alternative usage with custom config: - -```python -# Create a custom config with different settings -from haiku.rag.config.models import AppConfig, QAConfig - -custom_config = AppConfig( - qa=QAConfig( - provider="openai", - model="gpt-4o-mini", - max_sub_questions=5, - max_iterations=3, - max_concurrency=2, - ) -) - -graph = build_deep_qa_graph(config=custom_config) -context = DeepQAContext( - original_question="What are the main features of haiku.rag?", - use_citations=True -) -state = DeepQAState.from_config(context=context, config=custom_config) -deps = DeepQADeps(client=client) - -result = await graph.run(state=state, deps=deps) -``` - -### Research Graph - -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. +The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report. ```mermaid --- @@ -162,47 +63,49 @@ stateDiagram-v2 get_batch --> search_one: Has questions (map) get_batch --> synthesize: No questions search_one --> collect_answers - collect_answers --> analyze_insights - analyze_insights --> decide + collect_answers --> decide decide --> get_batch: Continue research decide --> synthesize: Done researching synthesize --> [*] ``` -Key nodes: +**Key nodes:** -- **plan**: Builds up to 3 standalone sub‑questions (uses an internal presearch tool) -- **get_batch**: Retrieves remaining sub‑questions for the current iteration -- **search_one**: Answers a single sub‑question using the KB with minimal, verbatim context (mapped in parallel) +- **plan**: Builds up to 3 standalone sub-questions (uses an internal presearch tool) +- **get_batch**: Retrieves remaining sub-questions for the current iteration +- **search_one**: Answers a single sub-question using the KB (mapped in parallel) - **collect_answers**: Aggregates search results from parallel executions -- **analyze_insights**: Synthesizes fresh insights, updates gaps, and suggests new sub-questions -- **decide**: Checks sufficiency/confidence thresholds and determines whether to continue research +- **decide**: Evaluates confidence and determines whether to continue or synthesize - **synthesize**: Generates a final structured research report -Primary models: +**Primary models:** -- `SearchAnswer` — one per sub‑question (query, answer, context, sources) -- `InsightRecord` / `GapRecord` — structured tracking of findings and open issues -- `InsightAnalysis` — output of the analysis stage (insights, gaps, commentary) -- `EvaluationResult` — insights, new questions, sufficiency, confidence +- `SearchAnswer` — one per sub-question (query, answer, confidence, citations) +- `EvaluationResult` — confidence score, new questions, sufficiency assessment - `ResearchReport` — final report (title, executive summary, findings, conclusions, …) -Note on parallel execution: +**Parallel execution:** + - The `search_one` node is mapped over all questions in a batch - Parallelism is controlled via `max_concurrency` -- Analysis and decision nodes process results after each batch completes +- Decision nodes process results after each batch completes -CLI usage: +### CLI Usage ```bash -# Basic usage (uses config from file or defaults) +# Basic usage +haiku-rag research "How does haiku.rag organize and query documents?" + +# With verbose output (shows progress) haiku-rag research "How does haiku.rag organize and query documents?" --verbose -# With custom config file -haiku-rag --config my-research-config.yaml research "How does haiku.rag organize and query documents?" --verbose +# With document filter +haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'" ``` -Python usage (blocking result): +### Python Usage + +**Basic example:** ```python from haiku.rag.client import HaikuRAG @@ -212,41 +115,25 @@ from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.state import ResearchDeps, ResearchState async with HaikuRAG(path_to_db) as client: - # Use global config (recommended) graph = build_research_graph(config=Config) - question = "What are the main drivers and trends of global temperature anomalies since 1990?" - context = ResearchContext(original_question=question) + context = ResearchContext(original_question="What are the main features?") state = ResearchState.from_config(context=context, config=Config) deps = ResearchDeps(client=client) - result = await graph.run( - state=state, - deps=deps, - ) + report = await graph.run(state=state, deps=deps) - report = result print(report.title) print(report.executive_summary) ``` -### Filtering Documents - -Both Research and Deep QA graphs support restricting searches to specific documents via the `search_filter` parameter. Set it to a SQL WHERE clause before running: - -```python -state = ResearchState.from_config(context=context, config=Config) -# Only search documents with these IDs -state.search_filter = "id IN ('doc-123', 'doc-456')" - -result = await graph.run(state=state, deps=deps) -``` - -The filter applies to all search operations in the graph (context gathering and sub-question searches). See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax. - -Alternative usage with custom config: +**With custom config:** ```python +from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig, ResearchConfig +from haiku.rag.graph.research.dependencies import ResearchContext +from haiku.rag.graph.research.graph import build_research_graph +from haiku.rag.graph.research.state import ResearchDeps, ResearchState custom_config = AppConfig( research=ResearchConfig( @@ -258,28 +145,28 @@ custom_config = AppConfig( ) ) -graph = build_research_graph(config=custom_config) -context = ResearchContext(original_question=question) -state = ResearchState.from_config(context=context, config=custom_config) -deps = ResearchDeps(client=client) +async with HaikuRAG(path_to_db) as client: + graph = build_research_graph(config=custom_config) + context = ResearchContext(original_question="What are the main features?") + state = ResearchState.from_config(context=context, config=custom_config) + deps = ResearchDeps(client=client) -result = await graph.run(state=state, deps=deps) + report = await graph.run(state=state, deps=deps) ``` -Python usage (streamed AG-UI events): +**Streaming AG-UI events:** ```python -from haiku.rag.graph.agui import stream_graph from haiku.rag.client import HaikuRAG from haiku.rag.config import Config +from haiku.rag.graph.agui import stream_graph from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.state import ResearchDeps, ResearchState async with HaikuRAG(path_to_db) as client: graph = build_research_graph(config=Config) - question = "What are the main drivers and trends of global temperature anomalies since 1990?" - context = ResearchContext(original_question=question) + context = ResearchContext(original_question="What are the main features?") state = ResearchState.from_config(context=context, config=Config) deps = ResearchDeps(client=client) @@ -287,21 +174,25 @@ async with HaikuRAG(path_to_db) as client: if event["type"] == "STEP_STARTED": print(f"Starting step: {event['stepName']}") elif event["type"] == "ACTIVITY_SNAPSHOT": - # Activity events include structured data alongside messages - content = event['content'] + content = event["content"] print(f" {content['message']}") - - # Different activity types have different structured fields - if 'confidence' in content: + if "confidence" in content: print(f" Confidence: {content['confidence']:.0%}") - if 'sub_questions' in content: - for q in content['sub_questions']: - print(f" - {q}") - if 'insights' in content: - print(f" New insights: {len(content['insights'])}") elif event["type"] == "RUN_FINISHED": - print("\nResearch complete!\n") - result = event["result"] - print(result["title"]) - print(result["executive_summary"]) + report = event["result"] + print(report["executive_summary"]) ``` + +### Filtering Documents + +Restrict searches to specific documents via the `search_filter` parameter: + +```python +# Set filter before running the graph +state = ResearchState.from_config(context=context, config=Config) +state.search_filter = "id IN ('doc-123', 'doc-456')" + +report = await graph.run(state=state, deps=deps) +``` + +The filter applies to all search operations in the graph. See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax. diff --git a/docs/server.md b/docs/server.md index f312620d..68e837be 100644 --- a/docs/server.md +++ b/docs/server.md @@ -239,8 +239,7 @@ data: {"type":"RUN_FINISHED","threadId":"abc123","runId":"xyz789","result":{"tit - Additional structured fields depending on the activity type: - **Planning**: `sub_questions` (list of strings) - **Searching**: `query` (string), `confidence` (float, on completion), `error` (string, on failure) - - **Analyzing**: `insights` (list of insight objects), `gaps` (list of gap objects), `resolved_gaps` (list of strings) - - **Evaluating**: `confidence` (float), `is_sufficient` (boolean) for research; `is_sufficient` (boolean), `iterations` (int) for deep QA + - **Evaluating**: `confidence` (float), `is_sufficient` (boolean), `new_questions` (list of strings) The `message` field is always present for simple rendering, while structured fields enable richer UI features like displaying lists, charts, and detailed status information. From 45de9bf0d537f50bb874a08df8d79403a7bb1df4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 15 Dec 2025 14:05:07 +0200 Subject: [PATCH 6/7] Flatten graph structure, simplify --- haiku_rag_slim/haiku/rag/client.py | 2 +- .../haiku/rag/graph/common/__init__.py | 5 - .../haiku/rag/graph/common/models.py | 107 ------ .../haiku/rag/graph/common/nodes.py | 315 ------------------ .../haiku/rag/graph/common/prompts.py | 67 ---- .../haiku/rag/graph/research/__init__.py | 7 +- .../haiku/rag/graph/research/common.py | 29 -- .../haiku/rag/graph/research/dependencies.py | 10 +- .../haiku/rag/graph/research/graph.py | 248 ++++++++++++-- .../haiku/rag/graph/research/models.py | 105 +++++- .../haiku/rag/graph/research/prompts.py | 82 ++++- haiku_rag_slim/haiku/rag/qa/agent.py | 4 +- haiku_rag_slim/haiku/rag/utils.py | 2 +- tests/graph/test_research_graph.py | 2 - tests/graph/test_search_filter.py | 4 - tests/test_app.py | 2 +- 16 files changed, 402 insertions(+), 589 deletions(-) delete mode 100644 haiku_rag_slim/haiku/rag/graph/common/__init__.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/common/models.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/common/nodes.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/common/prompts.py delete mode 100644 haiku_rag_slim/haiku/rag/graph/research/common.py diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 1e57195a..edd3e9f9 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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__) diff --git a/haiku_rag_slim/haiku/rag/graph/common/__init__.py b/haiku_rag_slim/haiku/rag/graph/common/__init__.py deleted file mode 100644 index 199aabe6..00000000 --- a/haiku_rag_slim/haiku/rag/graph/common/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Common utilities for graph implementations.""" - -from haiku.rag.utils import get_model - -__all__ = ["get_model"] diff --git a/haiku_rag_slim/haiku/rag/graph/common/models.py b/haiku_rag_slim/haiku/rag/graph/common/models.py deleted file mode 100644 index e559096e..00000000 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ /dev/null @@ -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 diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py deleted file mode 100644 index f751a0a1..00000000 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ /dev/null @@ -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 diff --git a/haiku_rag_slim/haiku/rag/graph/common/prompts.py b/haiku_rag_slim/haiku/rag/graph/common/prompts.py deleted file mode 100644 index ed8d49c2..00000000 --- a/haiku_rag_slim/haiku/rag/graph/common/prompts.py +++ /dev/null @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/graph/research/__init__.py b/haiku_rag_slim/haiku/rag/graph/research/__init__.py index 60f53a90..36f46a92 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/__init__.py +++ b/haiku_rag_slim/haiku/rag/graph/research/__init__.py @@ -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, +) diff --git a/haiku_rag_slim/haiku/rag/graph/research/common.py b/haiku_rag_slim/haiku/rag/graph/research/common.py deleted file mode 100644 index c0b3bd96..00000000 --- a/haiku_rag_slim/haiku/rag/graph/research/common.py +++ /dev/null @@ -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") diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index f5aa205a..a7ecbe22 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index f263ea6c..4e190665 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/graph/research/models.py b/haiku_rag_slim/haiku/rag/graph/research/models.py index 322ca99b..922132f4 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/models.py +++ b/haiku_rag_slim/haiku/rag/graph/research/models.py @@ -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): diff --git a/haiku_rag_slim/haiku/rag/graph/research/prompts.py b/haiku_rag_slim/haiku/rag/graph/research/prompts.py index f91a3bbf..6bfa8ed4 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/research/prompts.py @@ -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).""" diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index 03b6c0d5..d0ab6600 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -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): diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 9be4011f..9889b680 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -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( diff --git a/tests/graph/test_research_graph.py b/tests/graph/test_research_graph.py index badb79e6..66f7482e 100644 --- a/tests/graph/test_research_graph.py +++ b/tests/graph/test_research_graph.py @@ -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() diff --git a/tests/graph/test_search_filter.py b/tests/graph/test_search_filter.py index 17afd7c1..d106033a 100644 --- a/tests/graph/test_search_filter.py +++ b/tests/graph/test_search_filter.py @@ -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() diff --git a/tests/test_app.py b/tests/test_app.py index 58188c7e..65da4e78 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -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 = [ From 921893cff56e6b63039781e404bbdd79f158275f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Dec 2025 10:10:03 +0200 Subject: [PATCH 7/7] Remove deep ask from ag-ui server --- docs/server.md | 25 ++----- haiku_rag_slim/haiku/rag/graph/agui/server.py | 71 ++++--------------- 2 files changed, 18 insertions(+), 78 deletions(-) diff --git a/docs/server.md b/docs/server.md index 68e837be..33109877 100644 --- a/docs/server.md +++ b/docs/server.md @@ -119,7 +119,7 @@ URLs are also supported - the content is fetched and converted to markdown. ## AG-UI Server -The AG-UI server provides HTTP streaming of both research and deep ask graph execution using Server-Sent Events (SSE). +The AG-UI server provides HTTP streaming of research graph execution using Server-Sent Events (SSE). ### Starting the AG-UI Server @@ -131,7 +131,6 @@ This starts an HTTP server (default: http://0.0.0.0:8000) that exposes: - `GET /health` - Health check endpoint - `POST /v1/research/stream` - Research graph streaming endpoint -- `POST /v1/deep-ask/stream` - Deep ask graph streaming endpoint ### Configuration @@ -154,9 +153,9 @@ agui: - **cors_methods**: Allowed HTTP methods (default: `["GET", "POST", "OPTIONS"]`) - **cors_headers**: Allowed headers (default: `["*"]`) -### Using the Streaming Endpoints +### Using the Streaming Endpoint -Both endpoints accept POST requests with the same AG-UI RunAgentInput format and stream AG-UI events. +The endpoint accepts POST requests with AG-UI RunAgentInput format and streams AG-UI events. **Request format:** ```json @@ -171,7 +170,7 @@ Both endpoints accept POST requests with the same AG-UI RunAgentInput format and } ``` -**Research endpoint example:** +**Example:** ```bash curl -X POST http://localhost:8000/v1/research/stream \ -H "Content-Type: application/json" \ @@ -183,24 +182,12 @@ curl -X POST http://localhost:8000/v1/research/stream \ --no-buffer ``` -**Deep ask endpoint example:** -```bash -curl -X POST http://localhost:8000/v1/deep-ask/stream \ - -H "Content-Type: application/json" \ - -d '{ - "state": { - "question": "How does haiku.rag handle document chunking?", - "use_citations": true - } - }' \ - --no-buffer -``` - The `--no-buffer` flag ensures curl displays events as they arrive instead of buffering them. **Note:** The `state` object can include: - `question`: The question to answer (required) -- `use_citations`: Enable citations in deep ask responses (optional, deep ask only) +- `max_iterations`: Maximum research iterations (optional, defaults to config) +- `confidence_threshold`: Confidence threshold for early termination (optional, defaults to config) **Response:** Server-Sent Events stream with AG-UI protocol events: - `RUN_STARTED` - Graph execution started diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index 1c91ff5d..1ccba823 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -154,14 +154,14 @@ def format_sse_event(event: AGUIEvent) -> str: def create_agui_server( # pragma: no cover config: "AppConfig", db_path: Path | None = None ) -> Starlette: - """Create AG-UI server with both research and deep ask endpoints. + """Create AG-UI server with research endpoint. Args: - config: Application config with research and qa settings + config: Application config with research settings db_path: Optional database path override Returns: - Starlette app with research and deep ask endpoints + Starlette app with research endpoint """ from haiku.rag.client import HaikuRAG from haiku.rag.graph.research.dependencies import ResearchContext @@ -189,7 +189,14 @@ def create_agui_server( # pragma: no cover if messages: question = messages[0].get("content", "") context = ResearchContext(original_question=question) - return ResearchState.from_config(context=context, config=config) + max_iterations = input_state.get("max_iterations") + confidence_threshold = input_state.get("confidence_threshold") + return ResearchState.from_config( + context=context, + config=config, + max_iterations=max_iterations, + confidence_threshold=confidence_threshold, + ) def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps: effective_db_path = ( @@ -199,33 +206,7 @@ def create_agui_server( # pragma: no cover ) return ResearchDeps(client=get_client(effective_db_path)) - # Deep ask graph factories (uses research graph with quick settings) - def deep_ask_graph_factory() -> Graph: - return build_research_graph(config) - - def deep_ask_state_factory(input_state: dict[str, Any]) -> ResearchState: - question = input_state.get("question", "") - if not question: - messages = input_state.get("messages", []) - if messages: - question = messages[0].get("content", "") - context = ResearchContext(original_question=question) - return ResearchState.from_config( - context=context, - config=config, - max_iterations=2, - confidence_threshold=0.0, - ) - - def deep_ask_deps_factory(input_config: dict[str, Any]) -> ResearchDeps: - effective_db_path = ( - db_path - or input_config.get("db_path") - or config.storage.data_dir / "haiku.rag.lancedb" - ) - return ResearchDeps(client=get_client(effective_db_path)) - - # Create event stream functions for each graph type + # Create event stream function async def research_event_stream( input_data: RunAgentInput, ) -> AsyncIterator[str]: @@ -238,18 +219,6 @@ def create_agui_server( # pragma: no cover event_data = format_sse_event(event) yield event_data - async def deep_ask_event_stream( - input_data: RunAgentInput, - ) -> AsyncIterator[str]: - """Generate SSE event stream from deep ask graph execution.""" - graph = deep_ask_graph_factory() - initial_state = deep_ask_state_factory(input_data.state) - deps = deep_ask_deps_factory(input_data.config) - - async for event in stream_graph(graph, initial_state, deps): - event_data = format_sse_event(event) - yield event_data - # Endpoint handlers async def stream_research(request: Request) -> StreamingResponse: """Research graph streaming endpoint.""" @@ -266,21 +235,6 @@ def create_agui_server( # pragma: no cover }, ) - async def stream_deep_ask(request: Request) -> StreamingResponse: - """Deep ask graph streaming endpoint.""" - body = await request.json() - input_data = RunAgentInput(**body) - - return StreamingResponse( - deep_ask_event_stream(input_data), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "X-Accel-Buffering": "no", - }, - ) - async def health_check(_: Request) -> JSONResponse: """Health check endpoint.""" return JSONResponse({"status": "healthy"}) @@ -288,7 +242,6 @@ def create_agui_server( # pragma: no cover # Define routes routes = [ Route("/v1/research/stream", stream_research, methods=["POST"]), - Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]), Route("/health", health_check, methods=["GET"]), ]