Update deep ask graph

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

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

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

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,

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(

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

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

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)