diff --git a/docs/server.md b/docs/server.md index e53232df..74c039c4 100644 --- a/docs/server.md +++ b/docs/server.md @@ -96,7 +96,7 @@ URLs are also supported for web content. ## AG-UI Server -The AG-UI server provides HTTP streaming of research graph execution using Server-Sent Events (SSE). +The AG-UI server provides HTTP streaming of both research and deep ask graph execution using Server-Sent Events (SSE). ### Starting the AG-UI Server @@ -107,7 +107,8 @@ haiku-rag serve --agui This starts an HTTP server (default: http://0.0.0.0:8000) that exposes: - `GET /health` - Health check endpoint -- `POST /v1/agent/stream` - Research graph streaming endpoint +- `POST /v1/research/stream` - Research graph streaming endpoint +- `POST /v1/deep-ask/stream` - Deep ask graph streaming endpoint ### Configuration @@ -123,9 +124,9 @@ agui: See [Configuration](configuration.md#ag-ui-server-configuration) for all available options. -### Using the Streaming Endpoint +### Using the Streaming Endpoints -The `/v1/agent/stream` endpoint accepts POST requests with research parameters and streams AG-UI events: +Both endpoints accept POST requests with the same AG-UI RunAgentInput format and stream AG-UI events. **Request format:** ```json @@ -133,24 +134,33 @@ The `/v1/agent/stream` endpoint accepts POST requests with research parameters a "threadId": "optional-thread-id", "runId": "optional-run-id", "state": { - "context": { - "original_question": "What are the key features of haiku.rag?" - } + "question": "What are the key features of haiku.rag?" }, "messages": [], "config": {} } ``` -**Example with curl:** +**Research endpoint example:** ```bash -curl -X POST http://localhost:8000/v1/agent/stream \ +curl -X POST http://localhost:8000/v1/research/stream \ -H "Content-Type: application/json" \ -d '{ "state": { - "context": { - "original_question": "What are the key features of haiku.rag?" - } + "question": "What are the key features of haiku.rag?" + } + }' \ + --no-buffer +``` + +**Deep ask endpoint example:** +```bash +curl -X POST http://localhost:8000/v1/deep-ask/stream \ + -H "Content-Type: application/json" \ + -d '{ + "state": { + "question": "How does haiku.rag handle document chunking?", + "use_citations": true } }' \ --no-buffer @@ -158,6 +168,10 @@ curl -X POST http://localhost:8000/v1/agent/stream \ The `--no-buffer` flag ensures curl displays events as they arrive instead of buffering them. +**Note:** The `state` object can include: +- `question`: The question to answer (required) +- `use_citations`: Enable citations in deep ask responses (optional, deep ask only) + **Response:** Server-Sent Events stream with AG-UI protocol events: - `RUN_STARTED` - Graph execution started - `STATE_SNAPSHOT` - Current state snapshot diff --git a/haiku_rag_slim/haiku/rag/agui/__init__.py b/haiku_rag_slim/haiku/rag/agui/__init__.py index 29ecd62d..fab7a016 100644 --- a/haiku_rag_slim/haiku/rag/agui/__init__.py +++ b/haiku_rag_slim/haiku/rag/agui/__init__.py @@ -21,7 +21,7 @@ from haiku.rag.agui.events import ( from haiku.rag.agui.server import ( RunAgentInput, create_agui_app, - create_research_server, + create_agui_server, format_sse_event, ) from haiku.rag.agui.state import compute_state_delta @@ -34,7 +34,7 @@ __all__ = [ "RunAgentInput", "compute_state_delta", "create_agui_app", - "create_research_server", + "create_agui_server", "emit_activity", "emit_activity_delta", "emit_run_error", diff --git a/haiku_rag_slim/haiku/rag/agui/server.py b/haiku_rag_slim/haiku/rag/agui/server.py index e0f94765..ff1be032 100644 --- a/haiku_rag_slim/haiku/rag/agui/server.py +++ b/haiku_rag_slim/haiku/rag/agui/server.py @@ -146,19 +146,20 @@ def format_sse_event(event: AGUIEvent) -> str: return f"data: {event_json}\n\n" -def create_research_server(config: Any, db_path: Any | None = None) -> Starlette: - """Create AG-UI server for research graph. - - This is a convenience function specifically for the research graph. +def create_agui_server(config: Any, db_path: Any | None = None) -> Starlette: + """Create AG-UI server with both research and deep ask endpoints. Args: - config: Application config with research settings + config: Application config with research and qa settings db_path: Optional database path override Returns: - Starlette app configured for research graph + Starlette app with research and deep ask endpoints """ from haiku.rag.client import HaikuRAG + from haiku.rag.qa.deep.dependencies import DeepQAContext + from haiku.rag.qa.deep.graph import build_deep_qa_graph + from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.graph import build_research_graph from haiku.rag.research.state import ResearchDeps, ResearchState @@ -166,44 +167,139 @@ def create_research_server(config: Any, db_path: Any | None = None) -> Starlette # Store client reference for proper lifecycle management _client_cache: dict[str, HaikuRAG] = {} - def graph_factory() -> Graph: - """Create research graph instance.""" + def get_client(effective_db_path: Any) -> HaikuRAG: + """Get or create cached client.""" + path_key = str(effective_db_path) + if path_key not in _client_cache: + _client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=config) + return _client_cache[path_key] + + # Research graph factories + def research_graph_factory() -> Graph: return build_research_graph(config) - def state_factory(input_state: dict[str, Any]) -> ResearchState: - """Create research state from input.""" - # Extract question from input state or messages + def research_state_factory(input_state: dict[str, Any]) -> ResearchState: question = input_state.get("question", "") if not question: - # Try to get from first message if available messages = input_state.get("messages", []) if messages: question = messages[0].get("content", "") - - # Create context and state context = ResearchContext(original_question=question) return ResearchState.from_config(context=context, config=config) - def deps_factory(input_config: dict[str, Any]) -> ResearchDeps: - """Create research dependencies.""" - # Use provided db_path or fallback to config + def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps: effective_db_path = ( db_path or input_config.get("db_path") or config.storage.data_dir / "haiku.rag.lancedb" ) + return ResearchDeps(client=get_client(effective_db_path)) - # Reuse existing client if available - path_key = str(effective_db_path) - if path_key not in _client_cache: - _client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=config) + # Deep ask graph factories + def deep_ask_graph_factory() -> Graph: + return build_deep_qa_graph(config) - return ResearchDeps(client=_client_cache[path_key]) + def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState: + question = input_state.get("question", "") + if not question: + messages = input_state.get("messages", []) + if messages: + question = messages[0].get("content", "") + use_citations = input_state.get("use_citations", False) + context = DeepQAContext(original_question=question, use_citations=use_citations) + return DeepQAState.from_config(context=context, config=config) - # Use AG-UI config from app config - return create_agui_app( - graph_factory=graph_factory, - state_factory=state_factory, - deps_factory=deps_factory, - config=config.agui, + def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps: + 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)) + + # Create event stream functions for each graph type + async def research_event_stream( + input_data: RunAgentInput, + ) -> AsyncIterator[str]: + """Generate SSE event stream from research graph execution.""" + graph = research_graph_factory() + initial_state = research_state_factory(input_data.state) + deps = research_deps_factory(input_data.config) + + async for event in stream_graph(graph, initial_state, deps): + event_data = format_sse_event(event) + yield event_data + + async def deep_ask_event_stream( + input_data: RunAgentInput, + ) -> AsyncIterator[str]: + """Generate SSE event stream from deep ask graph execution.""" + graph = deep_ask_graph_factory() + initial_state = deep_ask_state_factory(input_data.state) + deps = deep_ask_deps_factory(input_data.config) + + async for event in stream_graph(graph, initial_state, deps): + event_data = format_sse_event(event) + yield event_data + + # Endpoint handlers + async def stream_research(request: Request) -> StreamingResponse: + """Research graph streaming endpoint.""" + body = await request.json() + input_data = RunAgentInput(**body) + + return StreamingResponse( + research_event_stream(input_data), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + async def stream_deep_ask(request: Request) -> StreamingResponse: + """Deep ask graph streaming endpoint.""" + body = await request.json() + input_data = RunAgentInput(**body) + + return StreamingResponse( + deep_ask_event_stream(input_data), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + async def health_check(_: Request) -> JSONResponse: + """Health check endpoint.""" + return JSONResponse({"status": "healthy"}) + + # Define routes + routes = [ + Route("/v1/research/stream", stream_research, methods=["POST"]), + Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]), + Route("/health", health_check, methods=["GET"]), + ] + + # Configure CORS middleware + middleware = [ + Middleware( + CORSMiddleware, + allow_origins=config.agui.cors_origins, + allow_credentials=config.agui.cors_credentials, + allow_methods=config.agui.cors_methods, + allow_headers=config.agui.cors_headers, + ) + ] + + # Create Starlette app + app = Starlette( + routes=routes, + middleware=middleware, + debug=False, ) + + return app diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 40a20fe6..1ce07208 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -216,8 +216,6 @@ class HaikuRAGApp: async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client: try: if deep: - from rich.console import Console - from haiku.rag.qa.deep.dependencies import DeepQAContext from haiku.rag.qa.deep.graph import build_deep_qa_graph from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState @@ -227,12 +225,22 @@ class HaikuRAGApp: original_question=question, use_citations=cite ) state = DeepQAState.from_config(context=context, config=self.config) - deps = DeepQADeps( - client=self.client, console=Console() if verbose else None - ) + deps = DeepQADeps(client=self.client) - result = await graph.run(state=state, deps=deps) - answer = result.answer + if verbose: + # Use AG-UI renderer to process and display events + from haiku.rag.agui import AGUIConsoleRenderer + + renderer = AGUIConsoleRenderer(self.console) + result_dict = await renderer.render( + stream_graph(graph, state, deps) + ) + # Result should be a dict with 'answer' key + answer = result_dict.get("answer", "") if result_dict else "" + else: + # Run without rendering events, just get the result + result = await graph.run(state=state, deps=deps) + answer = result.answer else: answer = await self.client.ask(question, cite=cite) @@ -489,12 +497,12 @@ class HaikuRAGApp: async def run_agui(): import uvicorn - from haiku.rag.agui import create_research_server + from haiku.rag.agui import create_agui_server logger.info( f"Starting AG-UI server on {self.config.agui.host}:{self.config.agui.port}" ) - app = create_research_server(self.config, db_path=self.db_path) + app = create_agui_server(self.config, db_path=self.db_path) config = uvicorn.Config( app=app, host=self.config.agui.host, diff --git a/haiku_rag_slim/haiku/rag/qa/deep/graph.py b/haiku_rag_slim/haiku/rag/qa/deep/graph.py index 9ded283c..4ae92bc3 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/graph.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/graph.py @@ -8,7 +8,7 @@ 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, log +from haiku.rag.graph_common import get_model from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT from haiku.rag.qa.deep.dependencies import DeepQADependencies @@ -45,49 +45,52 @@ def build_deep_qa_graph( state = ctx.state deps = ctx.deps - log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]") + if deps.agui_emitter: + deps.agui_emitter.start_step("plan") + deps.agui_emitter.update_activity("planning", "Planning approach") - plan_agent = Agent( - model=get_model(provider, model), - output_type=ResearchPlan, - instructions=( - PLAN_PROMPT - + "\n\nUse the gather_context tool once on the main question before planning." - ), - retries=3, - deps_type=DeepQADependencies, - ) + try: + plan_agent = Agent( + model=get_model(provider, model), + output_type=ResearchPlan, + instructions=( + PLAN_PROMPT + + "\n\nUse the gather_context tool once on the main question before planning." + ), + retries=3, + deps_type=DeepQADependencies, + ) - @plan_agent.tool - async def gather_context( - ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6 - ) -> str: - results = await ctx2.deps.client.search(query, limit=limit) - expanded = await ctx2.deps.client.expand_context(results) - return "\n\n".join(chunk.content for chunk, _ in expanded) + @plan_agent.tool + async def gather_context( + ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6 + ) -> str: + results = await ctx2.deps.client.search(query, limit=limit) + expanded = await ctx2.deps.client.expand_context(results) + return "\n\n".join(chunk.content for chunk, _ in expanded) - prompt = ( - "Plan a focused approach for the main question.\n\n" - f"Main question: {state.context.original_question}" - ) + prompt = ( + "Plan a focused approach for the main question.\n\n" + f"Main question: {state.context.original_question}" + ) - agent_deps = DeepQADependencies( - client=deps.client, - context=state.context, - console=deps.console, - ) - plan_result = await plan_agent.run(prompt, deps=agent_deps) - state.context.sub_questions = list(plan_result.output.sub_questions) + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=None, + ) + plan_result = await plan_agent.run(prompt, deps=agent_deps) + state.context.sub_questions = list(plan_result.output.sub_questions) - log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]") - log( - deps, - state, - f" [bold]Main Question:[/bold] {state.context.original_question}", - ) - log(deps, state, " [bold]Sub-questions:[/bold]") - for i, sq in enumerate(state.context.sub_questions, 1): - log(deps, state, f" {i}. {sq}") + if deps.agui_emitter: + deps.agui_emitter.update_state(state) + count = len(state.context.sub_questions) + deps.agui_emitter.update_activity( + "planning", f"Created plan with {count} sub-questions" + ) + finally: + if deps.agui_emitter: + deps.agui_emitter.finish_step() @g.step async def search_one( @@ -112,11 +115,8 @@ def build_deep_qa_graph( deps: DeepQADeps, sub_q: str, ) -> SearchAnswer: - log( - deps, - state, - f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}", - ) + if deps.agui_emitter: + deps.agui_emitter.update_activity("searching", f"Searching: {sub_q}") agent = Agent( model=get_model(provider, model), @@ -151,20 +151,18 @@ def build_deep_qa_graph( agent_deps = DeepQADependencies( client=deps.client, context=state.context, - console=deps.console, + console=None, ) try: result = await agent.run(sub_q, deps=agent_deps) answer = result.output if answer: state.context.add_qa_response(answer) - preview = answer.answer[:150] + ( - "…" if len(answer.answer) > 150 else "" - ) - log(deps, state, f" [green]✓[/green] {preview}") + if deps.agui_emitter: + deps.agui_emitter.update_state(state) + deps.agui_emitter.update_activity("searching", f"Answered: {sub_q}") return answer except Exception as e: - log(deps, state, f"[red]Search failed:[/red] {e}") failure_answer = SearchAnswer( query=sub_q, answer=f"Search failed after retries: {str(e)}", @@ -194,81 +192,69 @@ def build_deep_qa_graph( state = ctx.state deps = ctx.deps - log( - deps, - state, - "\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]", - ) - - agent = Agent( - model=get_model(provider, model), - 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, - "sources": qa.sources, - } - 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, - console=deps.console, - ) - result = await agent.run(prompt, deps=agent_deps) - evaluation = result.output - - state.iterations += 1 - - log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}") - status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]" - log(deps, state, f" Sufficient: {status}") - - for new_q in evaluation.new_questions: - if new_q not in state.context.sub_questions: - state.context.sub_questions.append(new_q) - - if evaluation.new_questions: - log(deps, state, " [cyan]New questions:[/cyan]") - for question in evaluation.new_questions: - log(deps, state, f" • {question}") - - should_continue = ( - not evaluation.is_sufficient and state.iterations < state.max_iterations - ) - - if not should_continue: - if state.iterations >= state.max_iterations: - log( - deps, - state, - f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]", - ) - log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]") - else: - log( - deps, - state, - f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]", + if deps.agui_emitter: + deps.agui_emitter.start_step("decide") + deps.agui_emitter.update_activity( + "evaluating", "Evaluating information sufficiency" ) - return should_continue + try: + agent = Agent( + model=get_model(provider, model), + 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, + "sources": qa.sources, + } + 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, + console=None, + ) + 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", + f"Information {status} after {state.iterations} iteration(s)", + ) + + 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( @@ -277,50 +263,56 @@ def build_deep_qa_graph( state = ctx.state deps = ctx.deps - log( - deps, - state, - "\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]", - ) + if deps.agui_emitter: + deps.agui_emitter.start_step("synthesize") + deps.agui_emitter.update_activity( + "synthesizing", "Synthesizing final answer" + ) - prompt_template = ( - SYNTHESIS_PROMPT_WITH_CITATIONS - if state.context.use_citations - else SYNTHESIS_PROMPT - ) + try: + prompt_template = ( + SYNTHESIS_PROMPT_WITH_CITATIONS + if state.context.use_citations + else SYNTHESIS_PROMPT + ) - agent = Agent( - model=get_model(provider, model), - output_type=DeepQAAnswer, - instructions=prompt_template, - retries=3, - deps_type=DeepQADependencies, - ) + agent = Agent( + model=get_model(provider, model), + output_type=DeepQAAnswer, + instructions=prompt_template, + retries=3, + deps_type=DeepQADependencies, + ) - context_data = { - "original_question": state.context.original_question, - "sub_answers": [ - { - "question": qa.query, - "answer": qa.answer, - "sources": qa.sources, - } - for qa in state.context.qa_responses - ], - } - context_xml = format_as_xml(context_data, root_tag="gathered_information") + context_data = { + "original_question": state.context.original_question, + "sub_answers": [ + { + "question": qa.query, + "answer": qa.answer, + "sources": qa.sources, + } + 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}" + prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}" - agent_deps = DeepQADependencies( - client=deps.client, - context=state.context, - console=deps.console, - ) - result = await agent.run(prompt, deps=agent_deps) + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=None, + ) + result = await agent.run(prompt, deps=agent_deps) - log(deps, state, "[bold green]✅ Answer complete![/bold green]") - return result.output + if deps.agui_emitter: + deps.agui_emitter.update_activity("synthesizing", "Answer complete") + + return result.output + finally: + if deps.agui_emitter: + deps.agui_emitter.finish_step() # Build the graph structure collect_answers = g.join( diff --git a/haiku_rag_slim/haiku/rag/qa/deep/state.py b/haiku_rag_slim/haiku/rag/qa/deep/state.py index 0e07098e..a340d42d 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/state.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/state.py @@ -1,8 +1,8 @@ import asyncio from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from rich.console import Console +from pydantic import BaseModel, Field from haiku.rag.client import HaikuRAG from haiku.rag.qa.deep.dependencies import DeepQAContext @@ -14,21 +14,26 @@ if TYPE_CHECKING: @dataclass class DeepQADeps: client: HaikuRAG - console: Console | None = None + agui_emitter: Any | None = None semaphore: asyncio.Semaphore | None = None - def emit_log(self, message: str, state: "DeepQAState | None" = None) -> None: - if self.console: - self.console.print(message) +class DeepQAState(BaseModel): + """Deep QA state for multi-agent question answering.""" -@dataclass -class DeepQAState: - context: DeepQAContext - max_sub_questions: int = 3 - max_iterations: int = 2 - max_concurrency: int = 1 - iterations: int = 0 + 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") @classmethod def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState": diff --git a/tests/test_app.py b/tests/test_app.py index 8d3b003a..51153aa4 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -401,12 +401,13 @@ 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.""" - from haiku.rag.qa.deep.models import DeepQAAnswer - mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"]) + mock_output = {"answer": "Deep QA answer", "sources": ["test.md"]} + + mock_renderer = AsyncMock() + mock_renderer.render.return_value = mock_output mock_graph = AsyncMock() - mock_graph.run.return_value = mock_output mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -418,8 +419,11 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): with patch( "haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph ): - await app.ask("test question", deep=True, verbose=True) + with patch( + "haiku.rag.agui.AGUIConsoleRenderer", return_value=mock_renderer + ): + await app.ask("test question", deep=True, verbose=True) - mock_graph.run.assert_called_once() - call_kwargs = mock_graph.run.call_args[1] - assert call_kwargs["deps"].console is not None + # With verbose, it should use AGUIConsoleRenderer.render, not graph.run + mock_renderer.render.assert_called_once() + mock_graph.run.assert_not_called() diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py index c220fcda..b6757498 100644 --- a/tests/test_deep_qa.py +++ b/tests/test_deep_qa.py @@ -30,7 +30,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): # Use real client but with TestModel for LLM calls client = HaikuRAG(temp_db_path) - deps = DeepQADeps(client=client, console=None) + deps = DeepQADeps(client=client) result = await graph.run(state=state, deps=deps) @@ -62,7 +62,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path): # Use real client but with TestModel for LLM calls client = HaikuRAG(temp_db_path) - deps = DeepQADeps(client=client, console=None) + deps = DeepQADeps(client=client) result = await graph.run(state=state, deps=deps)