haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/graph/research/graph.py

@@ -1,172 +1,305 @@
-from pydantic_ai import Agent
+import asyncio
+from typing import Literal
+from uuid import uuid4
+
+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_analysis_for_prompt,
-    format_context_for_prompt,
+from haiku.rag.graph.agui.emitter import (
+    emit_text_message_end,
+    emit_text_message_start,
+    emit_tool_call_args,
+    emit_tool_call_end,
+    emit_tool_call_start,
 )
-from haiku.rag.graph.research.dependencies import ResearchDependencies
+from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
 from haiku.rag.graph.research.models import (
     EvaluationResult,
-    InsightAnalysis,
+    RawSearchAnswer,
+    ResearchPlan,
     ResearchReport,
+    SearchAnswer,
 )
 from haiku.rag.graph.research.prompts import (
-    DECISION_AGENT_PROMPT,
-    INSIGHT_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 build_prompt, 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(
     config: AppConfig = Config,
+    include_plan: bool = True,
+    interactive: bool = False,
 ) -> 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)
+        interactive: Whether to include human decision nodes for HIL
 
     Returns:
         Configured Research graph
     """
-    provider = config.research.provider
-    model = config.research.model
+    model_config = config.research.model
+
+    # Build prompts with system_context if configured
+    plan_prompt = build_prompt(
+        PLAN_PROMPT
+        + "\n\nUse the gather_context tool once on the main question before planning.",
+        config,
+    )
+    search_prompt = build_prompt(SEARCH_PROMPT, config)
+    decision_prompt = build_prompt(DECISION_PROMPT, config)
+    synthesis_prompt = build_prompt(
+        config.prompts.synthesis or SYNTHESIS_PROMPT, config
+    )
     g = GraphBuilder(
         state_type=ResearchState,
         deps_type=ResearchDeps,
         output_type=ResearchReport,
     )
 
-    # Create and register the plan node using the factory
-    plan = g.step(
-        create_plan_node(
-            provider=provider,
-            model=model,
-            deps_type=ResearchDependencies,  # type: ignore[arg-type]
-            activity_message="Creating research plan",
-            output_retries=3,
-        )
-    )  # type: ignore[arg-type]
-
-    # Create and register the search_one node using the factory
-    search_one = g.step(
-        create_search_node(
-            provider=provider,
-            model=model,
-            deps_type=ResearchDependencies,  # type: ignore[arg-type]
-            with_step_wrapper=True,
-            success_message_format="Found answer with {confidence:.0%} confidence",
-            handle_exceptions=True,
-        )
-    )  # type: ignore[arg-type]
-
     @g.step
-    async def get_batch(
-        ctx: StepContext[ResearchState, ResearchDeps, 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 and process them in parallel
-        batch = list(state.context.sub_questions)
-        state.context.sub_questions.clear()
-        return batch
-
-    @g.step
-    async def analyze_insights(
-        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
-    ) -> None:
+    async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
+        """Create research plan with sub-questions."""
         state = ctx.state
         deps = ctx.deps
 
         if deps.agui_emitter:
-            deps.agui_emitter.start_step("analyze_insights")
+            deps.agui_emitter.start_step("plan")
             deps.agui_emitter.update_activity(
-                "analyzing", "Synthesizing insights and gaps"
+                "planning", {"stepName": "plan", "message": "Creating research plan"}
             )
 
         try:
-            agent = Agent(
-                model=get_model(provider, model),
-                output_type=InsightAnalysis,
-                instructions=INSIGHT_AGENT_PROMPT,
+            plan_agent = Agent(
+                model=get_model(model_config, config),
+                output_type=ResearchPlan,
+                instructions=plan_prompt,
                 retries=3,
                 output_retries=3,
                 deps_type=ResearchDependencies,
             )
 
-            context_xml = format_context_for_prompt(state.context)
+            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 = (
-                "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,
+                "Plan a focused approach for the main question.\n\n"
+                f"Main question: {state.context.original_question}"
             )
-            result = await agent.run(prompt, deps=agent_deps)
-            analysis: InsightAnalysis = result.output
 
-            state.context.integrate_analysis(analysis)
-            state.last_analysis = analysis
+            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)
 
-            # State updated with insights/gaps - emit state update and narrate
             if deps.agui_emitter:
                 deps.agui_emitter.update_state(state)
-                highlights = len(analysis.highlights) if analysis.highlights else 0
-                gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
-                resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
-                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", f"Analysis: {summary}")
+                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()
+                deps.agui_emitter.finish_step("plan")
 
     @g.step
-    async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
+    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(step_name)
+
+    @g.step
+    async def get_batch(
+        ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
+    ) -> list[str] | None:
+        """Get all remaining questions for this iteration."""
+        state = ctx.state
+
+        if not state.context.sub_questions:
+            return None
+
+        batch = list(state.context.sub_questions)
+        state.context.sub_questions.clear()
+        return batch
+
+    @g.step
+    async def decide(
+        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
+    ) -> bool:
+        """Evaluate research sufficiency and decide whether to continue."""
         state = ctx.state
         deps = ctx.deps
 
         if deps.agui_emitter:
             deps.agui_emitter.start_step("decide")
             deps.agui_emitter.update_activity(
-                "evaluating", "Evaluating research sufficiency"
+                "evaluating", {"message": "Evaluating research sufficiency"}
             )
 
         try:
             agent = Agent(
-                model=get_model(provider, model),
+                model=get_model(model_config, config),
                 output_type=EvaluationResult,
-                instructions=DECISION_AGENT_PROMPT,
+                instructions=decision_prompt,
                 retries=3,
                 output_retries=3,
                 deps_type=ResearchDependencies,
             )
 
             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
@@ -189,17 +322,28 @@ def build_research_graph(
             state.last_eval = output
             state.iterations += 1
 
+            # Get already-answered questions to avoid duplicates
+            answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
+
             for new_q in output.new_questions:
-                if new_q not in state.context.sub_questions:
-                    state.context.sub_questions.append(new_q)
+                # Skip if already in pending or already answered
+                if new_q in state.context.sub_questions:
+                    continue
+                if new_q.lower() in answered_queries:
+                    continue
+                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"
                 deps.agui_emitter.update_activity(
                     "evaluating",
-                    f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
+                    {
+                        "stepName": "decide",
+                        "message": f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
+                        "confidence": output.confidence_score,
+                        "is_sufficient": output.is_sufficient,
+                    },
                 )
 
             should_continue = (
@@ -210,26 +354,100 @@ def build_research_graph(
             return should_continue
         finally:
             if deps.agui_emitter:
-                deps.agui_emitter.finish_step()
+                deps.agui_emitter.finish_step("decide")
+
+    @g.step
+    async def human_decide(
+        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None | bool],
+    ) -> Literal["search", "synthesize"]:
+        """Wait for human decision on whether to continue searching or synthesize."""
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("human_decide")
+            deps.agui_emitter.update_state(state)
+
+        try:
+            # Emit tool call for human input wrapped in a message context
+            # This makes the tool call appear as if emitted by the LLM
+            message_id = str(uuid4())
+            tool_call_id = str(uuid4())
+
+            if deps.agui_emitter:
+                # Start an assistant message to contain the tool call
+                deps.agui_emitter.emit(emit_text_message_start(message_id))
+                # Emit tool call with parent message reference
+                deps.agui_emitter.emit(
+                    emit_tool_call_start(tool_call_id, "human_decision", message_id)
+                )
+                # Include full state for display
+                qa_responses = [
+                    {
+                        "query": qa.query,
+                        "answer": qa.answer,
+                        "confidence": qa.confidence,
+                        "citations_count": len(qa.citations),
+                    }
+                    for qa in state.context.qa_responses
+                ]
+                deps.agui_emitter.emit(
+                    emit_tool_call_args(
+                        tool_call_id,
+                        {
+                            "original_question": state.context.original_question,
+                            "sub_questions": list(state.context.sub_questions),
+                            "qa_responses": qa_responses,
+                            "iterations": state.iterations,
+                        },
+                    )
+                )
+                deps.agui_emitter.emit(emit_tool_call_end(tool_call_id))
+                # End the message after tool call
+                deps.agui_emitter.emit(emit_text_message_end(message_id))
+
+            # Wait for human input
+            if deps.human_input_queue is None:
+                raise RuntimeError("human_input_queue is required for interactive mode")
+
+            decision = await deps.human_input_queue.get()
+
+            # Process decision
+            if decision.action == "modify_questions" and decision.questions:
+                state.context.sub_questions = list(decision.questions)
+            elif decision.action == "add_questions" and decision.questions:
+                state.context.sub_questions.extend(decision.questions)
+
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+
+            if decision.action in ("search", "modify_questions", "add_questions"):
+                return "search"
+            else:
+                return "synthesize"
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step("human_decide")
 
     @g.step
     async def synthesize(
-        ctx: StepContext[ResearchState, ResearchDeps, None | bool],
+        ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
     ) -> ResearchReport:
+        """Generate final research report."""
         state = ctx.state
         deps = ctx.deps
 
         if deps.agui_emitter:
             deps.agui_emitter.start_step("synthesize")
             deps.agui_emitter.update_activity(
-                "synthesizing", "Generating final research report"
+                "synthesizing", {"message": "Generating final research report"}
             )
 
         try:
             agent = Agent(
-                model=get_model(provider, model),
+                model=get_model(model_config, config),
                 output_type=ResearchReport,
-                instructions=SYNTHESIS_AGENT_PROMPT,
+                instructions=synthesis_prompt,
                 retries=3,
                 output_retries=3,
                 deps_type=ResearchDependencies,
@@ -249,7 +467,7 @@ def build_research_graph(
             return result.output
         finally:
             if deps.agui_emitter:
-                deps.agui_emitter.finish_step()
+                deps.agui_emitter.finish_step("synthesize")
 
     # Build the graph structure
     collect_answers = g.join(
@@ -257,39 +475,76 @@ 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),
-    )
-
-    # 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(analyze_insights),
-        g.edge_from(analyze_insights).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 research")
-                .to(get_batch)
+    if interactive:
+        # Interactive mode: human decides after plan and after evaluation
+        if include_plan:
+            g.add(
+                g.edge_from(g.start_node).to(plan),
+                g.edge_from(plan).to(human_decide),
             )
-            .branch(
-                g.match(bool, matches=lambda x: not x)
-                .label("Done researching")
-                .to(synthesize)
+        else:
+            g.add(g.edge_from(g.start_node).to(human_decide))
+
+        g.add(
+            g.edge_from(human_decide).to(
+                g.decision()
+                .branch(
+                    g.match(str, matches=lambda x: x == "search")
+                    .label("Search")
+                    .to(get_batch)
+                )
+                .branch(
+                    g.match(str, matches=lambda x: x == "synthesize")
+                    .label("Synthesize")
+                    .to(synthesize)
+                )
+            ),
+            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(human_decide))
+            ),
+            g.edge_from(search_one).to(collect_answers),
+            # After search, evaluate to suggest new questions, then human decides
+            g.edge_from(collect_answers).to(decide),
+            g.edge_from(decide).to(human_decide),
+            g.edge_from(synthesize).to(g.end_node),
+        )
+    else:
+        # Non-interactive mode: automatic decision based on confidence/iterations
+        if include_plan:
+            g.add(
+                g.edge_from(g.start_node).to(plan),
+                g.edge_from(plan).to(get_batch),
             )
-        ),
-        g.edge_from(synthesize).to(g.end_node),
-    )
+        else:
+            g.add(g.edge_from(g.start_node).to(get_batch))
+
+        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),
+        )
+
+        g.add(
+            g.edge_from(decide).to(
+                g.decision()
+                .branch(
+                    g.match(bool, matches=lambda x: x)
+                    .label("Continue research")
+                    .to(get_batch)
+                )
+                .branch(
+                    g.match(bool, matches=lambda x: not x)
+                    .label("Done researching")
+                    .to(synthesize)
+                )
+            ),
+            g.edge_from(synthesize).to(g.end_node),
+        )
 
     return g.build()