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()