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

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

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-30 14:51:18 +03:00
parent 49eac86259
commit a6275fcf25
No known key found for this signature in database
8 changed files with 69 additions and 70 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
docs/agents.md
View file

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

4
src/haiku/rag/app.py
View file

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

2
src/haiku/rag/qa/deep/__init__.py
View file

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

20
src/haiku/rag/qa/deep/graph.py
View file

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

4
src/haiku/rag/qa/deep/models.py
View file

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

32
src/haiku/rag/qa/deep/nodes.py
View file

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

12
tests/test_app.py
View file

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

50
tests/test_deep_qa.py
View file

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