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

Do not duplicate ResearchContext in ResearchState

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-24 13:43:07 +03:00
parent 086eb56d34
commit cf9d6a116f
No known key found for this signature in database
8 changed files with 23 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -222,9 +222,9 @@ class HaikuRAGApp:
self.console.print()
graph = build_research_graph()
context = ResearchContext(original_question=question)
state = ResearchState(
question=question,
context=ResearchContext(original_question=question),
context=context,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,

4
src/haiku/rag/research/nodes/evaluate.py
View file

@ -55,8 +55,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
for insight in output.key_insights:
state.context.add_insight(insight)
for new_q in output.new_questions:
if new_q not in state.sub_questions:
state.sub_questions.append(new_q)
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
state.last_eval = output
state.iterations += 1

12
src/haiku/rag/research/nodes/plan.py
View file

@ -45,7 +45,7 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
prompt = (
"Plan a focused research approach for the main question.\n\n"
f"Main question: {state.question}"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(
@ -55,12 +55,16 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
stream=deps.stream,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.sub_questions = list(plan_result.output.sub_questions)
state.context.sub_questions = list(plan_result.output.sub_questions)
log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
log(deps, state, f" [bold]Main Question:[/bold] {state.question}")
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.sub_questions, 1):
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
return SearchDispatchNode(self.provider, self.model)

6
src/haiku/rag/research/nodes/search.py
View file

@ -24,7 +24,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
state = ctx.state
deps = ctx.deps
if not state.sub_questions:
if not state.context.sub_questions:
from haiku.rag.research.nodes.evaluate import EvaluateNode
return EvaluateNode(self.provider, self.model)
@ -32,8 +32,8 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
# Take up to max_concurrency questions and answer them concurrently
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.sub_questions and len(batch) < take:
batch.append(state.sub_questions.pop(0))
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
async def answer_one(sub_q: str) -> SearchAnswer | None:
log(

4
src/haiku/rag/research/state.py
View file

@ -1,4 +1,4 @@
from dataclasses import dataclass, field
from dataclasses import dataclass
from rich.console import Console
@ -23,9 +23,7 @@ class ResearchDeps:
@dataclass
class ResearchState:
question: str
context: ResearchContext
sub_questions: list[str] = field(default_factory=list)
iterations: int = 0
max_iterations: int = 3
max_concurrency: int = 1

6
src/haiku/rag/research/stream.py
View file

@ -34,13 +34,13 @@ class ResearchStateSnapshot:
last_sufficient = state.last_eval.is_sufficient
return cls(
question=state.question,
sub_questions=list(state.sub_questions),
question=context.original_question,
sub_questions=list(context.sub_questions),
iterations=state.iterations,
max_iterations=state.max_iterations,
max_concurrency=state.max_concurrency,
confidence_threshold=state.confidence_threshold,
pending_sub_questions=len(state.sub_questions),
pending_sub_questions=len(context.sub_questions),
answered_questions=len(context.qa_responses),
insights=list(context.insights),
gaps=list(context.gaps),

3
tests/test_research_graph.py
View file

@ -9,7 +9,6 @@ def test_build_graph_and_state():
assert graph is not None
state = ResearchState(
question="What are the key features of haiku.rag?",
context=ResearchContext(
original_question="What are the key features of haiku.rag?"
),
@ -17,7 +16,7 @@ def test_build_graph_and_state():
confidence_threshold=0.8,
)
assert state.iterations == 0
assert state.sub_questions == []
assert state.context.sub_questions == []
def test_async_loop_available():

7
tests/test_research_graph_integration.py
View file

@ -21,7 +21,6 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
graph = build_research_graph()
state = ResearchState(
question="What is haiku.rag?",
context=ResearchContext(original_question="What is haiku.rag?"),
max_iterations=1,
confidence_threshold=0.5,
@ -32,7 +31,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
) # client unused in patched nodes
async def fake_plan_run(self, ctx) -> Any:
ctx.state.sub_questions = [
ctx.state.context.sub_questions = [
"Describe haiku.rag in one sentence",
"List core components of haiku.rag",
]
@ -41,8 +40,8 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
async def fake_search_dispatch_run(self, ctx) -> Any:
# Answer all pending questions deterministically, then move to evaluation
while ctx.state.sub_questions:
q = ctx.state.sub_questions.pop(0)
while ctx.state.context.sub_questions:
q = ctx.state.context.sub_questions.pop(0)
# pydantic BaseModel kwargs not fully typed for pyright
ctx.state.context.add_qa_response(
SearchAnswer(query=q, answer="A", context=["x"], sources=["s"]) # pyright: ignore[reportCallIssue]