Rename nodes to make it clear they are part of Q/A

docs/agents.md

@@ -47,12 +47,11 @@ Deep QA is a multi-agent system that decomposes complex questions into sub-quest
 title: Deep QA graph
 ---
 stateDiagram-v2
-  DeepPlanNode --> DeepSearchDispatchNode
+  DeepQAPlanNode --> DeepQASearchDispatchNode
-  DeepSearchDispatchNode --> DeepSearchDispatchNode
+  DeepQASearchDispatchNode --> DeepQADecisionNode
-  DeepSearchDispatchNode --> DeepDecisionNode
+  DeepQADecisionNode --> DeepQASearchDispatchNode
-  DeepDecisionNode --> DeepSearchDispatchNode
+  DeepQADecisionNode --> DeepQASynthesizeNode
-  DeepDecisionNode --> DeepSynthesizeNode
+  DeepQASynthesizeNode --> [*]
-  DeepSynthesizeNode --> [*]
 ```
 
 Key nodes:
@@ -86,7 +85,7 @@ Python usage:
 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.nodes import DeepPlanNode
+from haiku.rag.qa.deep.nodes import DeepQAPlanNode
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 async with HaikuRAG(path_to_db) as client:
@@ -104,7 +103,7 @@ async with HaikuRAG(path_to_db) as client:
     deps = DeepQADeps(client=client)
 
     result = await graph.run(
-        start_node=DeepPlanNode(provider="openai", model="gpt-4o-mini"),
+        start_node=DeepQAPlanNode(provider="openai", model="gpt-4o-mini"),
         state=state,
         deps=deps
     )

src/haiku/rag/app.py

@@ -208,7 +208,7 @@ class HaikuRAGApp:
 
                     from haiku.rag.qa.deep.dependencies import DeepQAContext
                     from haiku.rag.qa.deep.graph import build_deep_qa_graph
-                    from haiku.rag.qa.deep.nodes import DeepPlanNode
+                    from haiku.rag.qa.deep.nodes import DeepQAPlanNode
                     from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
                     graph = build_deep_qa_graph()
@@ -220,7 +220,7 @@ class HaikuRAGApp:
                         client=self.client, console=Console() if verbose else None
                     )
 
-                    start_node = DeepPlanNode(
+                    start_node = DeepQAPlanNode(
                         provider=Config.QA_PROVIDER,
                         model=Config.QA_MODEL,
                     )

src/haiku/rag/qa/deep/__init__.py

@@ -1 +1 @@
-from haiku.rag.qa.deep.models import DeepAnswer
+from haiku.rag.qa.deep.models import DeepQAAnswer

src/haiku/rag/qa/deep/graph.py

@@ -1,21 +1,21 @@
 from pydantic_graph import Graph
 
-from haiku.rag.qa.deep.models import DeepAnswer
+from haiku.rag.qa.deep.models import DeepQAAnswer
 from haiku.rag.qa.deep.nodes import (
-    DeepDecisionNode,
+    DeepQADecisionNode,
-    DeepPlanNode,
+    DeepQAPlanNode,
-    DeepSearchDispatchNode,
+    DeepQASearchDispatchNode,
-    DeepSynthesizeNode,
+    DeepQASynthesizeNode,
 )
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 
-def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepAnswer]:
+def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepQAAnswer]:
     return Graph(
         nodes=[
-            DeepPlanNode,
+            DeepQAPlanNode,
-            DeepSearchDispatchNode,
+            DeepQASearchDispatchNode,
-            DeepDecisionNode,
+            DeepQADecisionNode,
-            DeepSynthesizeNode,
+            DeepQASynthesizeNode,
         ]
     )

src/haiku/rag/qa/deep/models.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 
 
-class DeepEvaluation(BaseModel):
+class DeepQAEvaluation(BaseModel):
     is_sufficient: bool = Field(
         description="Whether we have sufficient information to answer the question"
     )
@@ -12,7 +12,7 @@ class DeepEvaluation(BaseModel):
     )
 
 
-class DeepAnswer(BaseModel):
+class DeepQAAnswer(BaseModel):
     answer: str = Field(description="The comprehensive answer to the question")
     sources: list[str] = Field(
         description="Document titles or URIs used to generate the answer",

src/haiku/rag/qa/deep/nodes.py

@@ -11,7 +11,7 @@ from haiku.rag.graph.common import get_model, log
 from haiku.rag.graph.models import ResearchPlan, SearchAnswer
 from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.qa.deep.dependencies import DeepQADependencies
-from haiku.rag.qa.deep.models import DeepAnswer, DeepEvaluation
+from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
 from haiku.rag.qa.deep.prompts import (
     DECISION_PROMPT,
     SYNTHESIS_PROMPT,
@@ -21,13 +21,13 @@ from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 
 @dataclass
-class DeepPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
+class DeepQAPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
     provider: str
     model: str
 
     async def run(
         self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
-    ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]:
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
         state = ctx.state
         deps = ctx.deps
 
@@ -77,22 +77,22 @@ class DeepPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
         for i, sq in enumerate(state.context.sub_questions, 1):
             log(deps, state, f"      {i}. {sq}")
 
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
 
 @dataclass
-class DeepSearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
+class DeepQASearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
     provider: str
     model: str
 
     async def run(
         self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
-    ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]:
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
         state = ctx.state
         deps = ctx.deps
 
         if not state.context.sub_questions:
-            return DeepDecisionNode(self.provider, self.model)
+            return DeepQADecisionNode(self.provider, self.model)
 
         # Take up to max_concurrency questions and answer them concurrently
         take = max(1, state.max_concurrency)
@@ -157,17 +157,17 @@ class DeepSearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
             preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
             log(deps, state, f"   [green]✓[/green] {preview}")
 
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
 
 @dataclass
-class DeepDecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
+class DeepQADecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
     provider: str
     model: str
 
     async def run(
         self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
-    ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]:
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
         state = ctx.state
         deps = ctx.deps
 
@@ -179,7 +179,7 @@ class DeepDecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
 
         agent = Agent(
             model=get_model(self.provider, self.model),
-            output_type=DeepEvaluation,
+            output_type=DeepQAEvaluation,
             instructions=DECISION_PROMPT,
             retries=3,
             deps_type=DeepQADependencies,
@@ -236,24 +236,24 @@ class DeepDecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
                     f"\n[bold yellow]⚠️  Reached max iterations ({state.max_iterations})[/bold yellow]",
                 )
             log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
-            return DeepSynthesizeNode(self.provider, self.model)
+            return DeepQASynthesizeNode(self.provider, self.model)
 
         log(
             deps,
             state,
             f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
         )
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
 
 @dataclass
-class DeepSynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
+class DeepQASynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
     provider: str
     model: str
 
     async def run(
         self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
-    ) -> End[DeepAnswer]:
+    ) -> End[DeepQAAnswer]:
         state = ctx.state
         deps = ctx.deps
 
@@ -271,7 +271,7 @@ class DeepSynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]):
 
         agent = Agent(
             model=get_model(self.provider, self.model),
-            output_type=DeepAnswer,
+            output_type=DeepQAAnswer,
             instructions=prompt_template,
             retries=3,
             deps_type=DeepQADependencies,

tests/test_app.py

@@ -268,9 +268,9 @@ 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.qa.deep.models import DeepAnswer
+    from haiku.rag.qa.deep.models import DeepQAAnswer
 
-    mock_output = DeepAnswer(answer="Deep QA answer", sources=["test.md"])
+    mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
     mock_result = MagicMock()
     mock_result.output = mock_output
 
@@ -298,9 +298,9 @@ 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."""
-    from haiku.rag.qa.deep.models import DeepAnswer
+    from haiku.rag.qa.deep.models import DeepQAAnswer
 
-    mock_output = DeepAnswer(
+    mock_output = DeepQAAnswer(
         answer="Deep QA answer with citations [test.md]", sources=["test.md"]
     )
     mock_result = MagicMock()
@@ -330,9 +330,9 @@ 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 DeepAnswer
+    from haiku.rag.qa.deep.models import DeepQAAnswer
 
-    mock_output = DeepAnswer(answer="Deep QA answer", sources=["test.md"])
+    mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
     mock_result = MagicMock()
     mock_result.output = mock_output
 

tests/test_deep_qa.py

@@ -5,12 +5,12 @@ import pytest
 from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.qa.deep.dependencies import DeepQAContext
 from haiku.rag.qa.deep.graph import build_deep_qa_graph
-from haiku.rag.qa.deep.models import DeepAnswer
+from haiku.rag.qa.deep.models import DeepQAAnswer
 from haiku.rag.qa.deep.nodes import (
-    DeepDecisionNode,
+    DeepQADecisionNode,
-    DeepPlanNode,
+    DeepQAPlanNode,
-    DeepSearchDispatchNode,
+    DeepQASearchDispatchNode,
-    DeepSynthesizeNode,
+    DeepQASynthesizeNode,
 )
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
@@ -32,11 +32,11 @@ async def test_deep_qa_graph_end_to_end(monkeypatch):
             "Describe haiku.rag in one sentence",
             "List core components of haiku.rag",
         ]
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
     async def fake_search_dispatch_run(self, ctx) -> Any:
         if not ctx.state.context.sub_questions:
-            return DeepDecisionNode(self.provider, self.model)
+            return DeepQADecisionNode(self.provider, self.model)
 
         batch = ctx.state.context.sub_questions[:]
         ctx.state.context.sub_questions.clear()
@@ -50,28 +50,28 @@ async def test_deep_qa_graph_end_to_end(monkeypatch):
                     sources=["test.md"],
                 )
             )
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
     async def fake_decision_run(self, ctx) -> Any:
         ctx.state.iterations += 1
-        return DeepSynthesizeNode(self.provider, self.model)
+        return DeepQASynthesizeNode(self.provider, self.model)
 
     async def fake_synthesize_run(self, ctx) -> Any:
         from pydantic_graph import End
 
         return End(
-            DeepAnswer(
+            DeepQAAnswer(
                 answer="haiku.rag is a RAG system with components A, B, C.",
                 sources=["test.md"],
             )
         )
 
-    monkeypatch.setattr(DeepPlanNode, "run", fake_plan_run)
+    monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run)
-    monkeypatch.setattr(DeepSearchDispatchNode, "run", fake_search_dispatch_run)
+    monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run)
-    monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run)
+    monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run)
-    monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run)
+    monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run)
 
-    start = DeepPlanNode(provider="ollama", model="test")
+    start = DeepQAPlanNode(provider="ollama", model="test")
     result = await graph.run(start_node=start, state=state, deps=deps)
 
     assert result.output.answer == "haiku.rag is a RAG system with components A, B, C."
@@ -91,11 +91,11 @@ async def test_deep_qa_with_citations(monkeypatch):
 
     async def fake_plan_run(self, ctx) -> Any:
         ctx.state.context.sub_questions = ["What is Python used for?"]
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
     async def fake_search_dispatch_run(self, ctx) -> Any:
         if not ctx.state.context.sub_questions:
-            return DeepDecisionNode(self.provider, self.model)
+            return DeepQADecisionNode(self.provider, self.model)
 
         batch = ctx.state.context.sub_questions[:]
         ctx.state.context.sub_questions.clear()
@@ -109,28 +109,28 @@ async def test_deep_qa_with_citations(monkeypatch):
                     sources=["python.md"],
                 )
             )
-        return DeepSearchDispatchNode(self.provider, self.model)
+        return DeepQASearchDispatchNode(self.provider, self.model)
 
     async def fake_decision_run(self, ctx) -> Any:
         ctx.state.iterations += 1
-        return DeepSynthesizeNode(self.provider, self.model)
+        return DeepQASynthesizeNode(self.provider, self.model)
 
     async def fake_synthesize_run(self, ctx) -> Any:
         from pydantic_graph import End
 
         return End(
-            DeepAnswer(
+            DeepQAAnswer(
                 answer="Python is a programming language [python.md].",
                 sources=["python.md"],
             )
         )
 
-    monkeypatch.setattr(DeepPlanNode, "run", fake_plan_run)
+    monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run)
-    monkeypatch.setattr(DeepSearchDispatchNode, "run", fake_search_dispatch_run)
+    monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run)
-    monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run)
+    monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run)
-    monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run)
+    monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run)
 
-    start = DeepPlanNode(provider="ollama", model="test")
+    start = DeepQAPlanNode(provider="ollama", model="test")
     result = await graph.run(start_node=start, state=state, deps=deps)
 
     assert "[python.md]" in result.output.answer