From a6275fcf25534c6973742cbfbfffa7491bef1d84 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 30 Sep 2025 14:51:18 +0300 Subject: [PATCH] Rename nodes to make it clear they are part of Q/A --- docs/agents.md | 15 +++++----- src/haiku/rag/app.py | 4 +-- src/haiku/rag/qa/deep/__init__.py | 2 +- src/haiku/rag/qa/deep/graph.py | 20 ++++++------- src/haiku/rag/qa/deep/models.py | 4 +-- src/haiku/rag/qa/deep/nodes.py | 32 ++++++++++---------- tests/test_app.py | 12 ++++---- tests/test_deep_qa.py | 50 +++++++++++++++---------------- 8 files changed, 69 insertions(+), 70 deletions(-) diff --git a/docs/agents.md b/docs/agents.md index d926258f..9e9c5372 100644 --- a/docs/agents.md +++ b/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 - DeepSearchDispatchNode --> DeepSearchDispatchNode - DeepSearchDispatchNode --> DeepDecisionNode - DeepDecisionNode --> DeepSearchDispatchNode - DeepDecisionNode --> DeepSynthesizeNode - DeepSynthesizeNode --> [*] + DeepQAPlanNode --> DeepQASearchDispatchNode + DeepQASearchDispatchNode --> DeepQADecisionNode + DeepQADecisionNode --> DeepQASearchDispatchNode + DeepQADecisionNode --> DeepQASynthesizeNode + DeepQASynthesizeNode --> [*] ``` 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 ) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index c5786a96..ebf18515 100644 --- a/src/haiku/rag/app.py +++ b/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, ) diff --git a/src/haiku/rag/qa/deep/__init__.py b/src/haiku/rag/qa/deep/__init__.py index 38782a45..0dc102f6 100644 --- a/src/haiku/rag/qa/deep/__init__.py +++ b/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 diff --git a/src/haiku/rag/qa/deep/graph.py b/src/haiku/rag/qa/deep/graph.py index a5a90afd..f5701ebd 100644 --- a/src/haiku/rag/qa/deep/graph.py +++ b/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, - DeepPlanNode, - DeepSearchDispatchNode, - DeepSynthesizeNode, + DeepQADecisionNode, + DeepQAPlanNode, + DeepQASearchDispatchNode, + 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, - DeepSearchDispatchNode, - DeepDecisionNode, - DeepSynthesizeNode, + DeepQAPlanNode, + DeepQASearchDispatchNode, + DeepQADecisionNode, + DeepQASynthesizeNode, ] ) diff --git a/src/haiku/rag/qa/deep/models.py b/src/haiku/rag/qa/deep/models.py index 933c3616..e33dfc95 100644 --- a/src/haiku/rag/qa/deep/models.py +++ b/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", diff --git a/src/haiku/rag/qa/deep/nodes.py b/src/haiku/rag/qa/deep/nodes.py index d97284c8..9461efe2 100644 --- a/src/haiku/rag/qa/deep/nodes.py +++ b/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, diff --git a/tests/test_app.py b/tests/test_app.py index 289f5b95..6cd549c9 100644 --- a/tests/test_app.py +++ b/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 diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py index 1f6f8587..ee9b1a1d 100644 --- a/tests/test_deep_qa.py +++ b/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, - DeepPlanNode, - DeepSearchDispatchNode, - DeepSynthesizeNode, + DeepQADecisionNode, + DeepQAPlanNode, + DeepQASearchDispatchNode, + 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(DeepSearchDispatchNode, "run", fake_search_dispatch_run) - monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run) - monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run) + monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_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(DeepSearchDispatchNode, "run", fake_search_dispatch_run) - monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run) - monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run) + monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_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