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

Remove deep qa graph, everything now uses the simplified research graph

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-15 13:16:05 +02:00
parent bdcf81774d
commit f0a7abd953
No known key found for this signature in database
16 changed files with 120 additions and 588 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
haiku_rag_slim/haiku/rag/app.py
View file

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

11
haiku_rag_slim/haiku/rag/graph/__init__.py
View file

@ -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",

22
haiku_rag_slim/haiku/rag/graph/agui/server.py
View file

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

1
haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py
View file

@ -1 +0,0 @@
from haiku.rag.graph.deep_qa.models import DeepQAAnswer

29
haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py
View file

@ -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"
)

250
haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py
View file

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

23
haiku_rag_slim/haiku/rag/graph/deep_qa/models.py
View file

@ -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"
)

42
haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py
View file

@ -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"""

59
haiku_rag_slim/haiku/rag/graph/deep_qa/state.py
View file

@ -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,
)

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

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

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

@ -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,
)

26
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -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:

41
tests/graph/test_deep_qa.py
View file

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

49
tests/graph/test_search_filter.py
View file

@ -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."""

71
tests/test_app.py
View file

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

9
tests/test_mcp.py
View file

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