From cf9d6a116f880c13c9bd968a36eb99e79a891d9d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Sep 2025 13:43:07 +0300 Subject: [PATCH] Do not duplicate ResearchContext in ResearchState --- src/haiku/rag/app.py | 4 ++-- src/haiku/rag/research/nodes/evaluate.py | 4 ++-- src/haiku/rag/research/nodes/plan.py | 12 ++++++++---- src/haiku/rag/research/nodes/search.py | 6 +++--- src/haiku/rag/research/state.py | 4 +--- src/haiku/rag/research/stream.py | 6 +++--- tests/test_research_graph.py | 3 +-- tests/test_research_graph_integration.py | 7 +++---- 8 files changed, 23 insertions(+), 23 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 3edce247..0a0302f9 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -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, diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py index 23181136..46634b52 100644 --- a/src/haiku/rag/research/nodes/evaluate.py +++ b/src/haiku/rag/research/nodes/evaluate.py @@ -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 diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py index 0f3726cb..63612a55 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/research/nodes/plan.py @@ -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) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py index fddc1b9e..c7e471a9 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/research/nodes/search.py @@ -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( diff --git a/src/haiku/rag/research/state.py b/src/haiku/rag/research/state.py index ad0920b2..238accf6 100644 --- a/src/haiku/rag/research/state.py +++ b/src/haiku/rag/research/state.py @@ -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 diff --git a/src/haiku/rag/research/stream.py b/src/haiku/rag/research/stream.py index 0df6b8fb..9ecf2d80 100644 --- a/src/haiku/rag/research/stream.py +++ b/src/haiku/rag/research/stream.py @@ -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), diff --git a/tests/test_research_graph.py b/tests/test_research_graph.py index 8b50a6f3..b986eeba 100644 --- a/tests/test_research_graph.py +++ b/tests/test_research_graph.py @@ -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(): diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 66c411f5..64863f9f 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -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]