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 = [