From 38861da5490b06d6b87799dd96af53facf1051c9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 3 Nov 2025 16:40:11 +0200 Subject: [PATCH] Go back and separate common patters in graph_common --- examples/ag-ui-research/backend/agent.py | 2 +- haiku_rag_slim/haiku/rag/a2a/__init__.py | 2 +- .../haiku/rag/graph_common/__init__.py | 5 ++ .../haiku/rag/graph_common/models.py | 42 ++++++++++++ .../haiku/rag/graph_common/prompts.py | 46 +++++++++++++ .../haiku/rag/graph_common/utils.py | 64 +++++++++++++++++++ haiku_rag_slim/haiku/rag/qa/deep/common.py | 41 ------------ .../haiku/rag/qa/deep/dependencies.py | 2 +- haiku_rag_slim/haiku/rag/qa/deep/graph.py | 49 +++++++------- haiku_rag_slim/haiku/rag/qa/deep/models.py | 23 ------- haiku_rag_slim/haiku/rag/qa/deep/prompts.py | 46 +------------ haiku_rag_slim/haiku/rag/qa/deep/state.py | 2 +- haiku_rag_slim/haiku/rag/research/__init__.py | 2 +- .../haiku/rag/research/dependencies.py | 2 +- haiku_rag_slim/haiku/rag/research/graph.py | 42 +++++++----- tests/test_deep_qa.py | 6 +- tests/test_research_graph_integration.py | 2 +- 17 files changed, 219 insertions(+), 159 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/graph_common/__init__.py create mode 100644 haiku_rag_slim/haiku/rag/graph_common/models.py create mode 100644 haiku_rag_slim/haiku/rag/graph_common/prompts.py create mode 100644 haiku_rag_slim/haiku/rag/graph_common/utils.py delete mode 100644 haiku_rag_slim/haiku/rag/qa/deep/common.py diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index 049e0dbd..5362696a 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -8,7 +8,7 @@ from pydantic_ai.ag_ui import StateDeps from haiku.rag.client import HaikuRAG from haiku.rag.config import Config -from haiku.rag.graph.common import get_model +from haiku.rag.graph_common import get_model class ResearchState(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/a2a/__init__.py b/haiku_rag_slim/haiku/rag/a2a/__init__.py index 7c3a00db..76e1ee36 100644 --- a/haiku_rag_slim/haiku/rag/a2a/__init__.py +++ b/haiku_rag_slim/haiku/rag/a2a/__init__.py @@ -6,7 +6,7 @@ import logfire from pydantic_ai import Agent, RunContext from haiku.rag.config import Config -from haiku.rag.qa.deep.common import get_model +from haiku.rag.graph_common import get_model from .context import load_message_history, save_message_history from .models import AgentDependencies, SearchResult diff --git a/haiku_rag_slim/haiku/rag/graph_common/__init__.py b/haiku_rag_slim/haiku/rag/graph_common/__init__.py new file mode 100644 index 00000000..dc47bee0 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/graph_common/__init__.py @@ -0,0 +1,5 @@ +"""Common utilities for graph implementations.""" + +from haiku.rag.graph_common.utils import get_model, log + +__all__ = ["get_model", "log"] diff --git a/haiku_rag_slim/haiku/rag/graph_common/models.py b/haiku_rag_slim/haiku/rag/graph_common/models.py new file mode 100644 index 00000000..407cb6bb --- /dev/null +++ b/haiku_rag_slim/haiku/rag/graph_common/models.py @@ -0,0 +1,42 @@ +"""Common models used across different graph implementations.""" + +from pydantic import BaseModel, Field, field_validator + + +class ResearchPlan(BaseModel): + """A structured research plan with sub-questions to explore.""" + + sub_questions: list[str] = Field( + ..., + description="Specific questions to research, phrased as complete questions", + ) + + @field_validator("sub_questions") + @classmethod + def validate_sub_questions(cls, v: list[str]) -> list[str]: + if len(v) < 1: + raise ValueError("Must have at least 1 sub-question") + if len(v) > 12: + raise ValueError("Cannot have more than 12 sub-questions") + return v + + +class SearchAnswer(BaseModel): + """Answer from a search operation with sources.""" + + query: str = Field(..., description="The question that was answered") + answer: str = Field(..., description="The comprehensive answer to the question") + context: list[str] = Field( + default_factory=list, + description="Relevant snippets that directly support the answer", + ) + sources: list[str] = Field( + default_factory=list, + description="Source URIs or titles that contributed to this answer", + ) + confidence: float = Field( + default=1.0, + description="Confidence score for this answer (0-1)", + ge=0.0, + le=1.0, + ) diff --git a/haiku_rag_slim/haiku/rag/graph_common/prompts.py b/haiku_rag_slim/haiku/rag/graph_common/prompts.py new file mode 100644 index 00000000..ed10ec98 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/graph_common/prompts.py @@ -0,0 +1,46 @@ +"""Common prompts used across different graph implementations.""" + +PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow. + +Responsibilities: +1. Understand and decompose the main question +2. Propose a minimal, high-leverage plan +3. Coordinate specialized agents to gather evidence +4. Iterate based on gaps and new findings + +Plan requirements: +- Produce at most 3 sub_questions that together cover the main question. +- Each sub_question must be a standalone, self-contained query that can run + without extra context. Include concrete entities, scope, timeframe, and any + qualifiers. Avoid ambiguous pronouns (it/they/this/that). +- Prioritize the highest-value aspects first; avoid redundancy and overlap. +- Prefer questions that are likely answerable from the current knowledge base; + if coverage is uncertain, make scopes narrower and specific. +- Order sub_questions by execution priority (most valuable first).""" + +SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist. + +Tasks: +1. Search the knowledge base for relevant evidence. +2. Analyze retrieved snippets. +3. Provide an answer strictly grounded in that evidence. + +Tool usage: +- Always call search_and_answer before drafting any answer. +- The tool returns snippets with verbatim `text`, a relevance `score`, and the + originating document identifier (document title if available, otherwise URI). +- You may call the tool multiple times to refine or broaden context, but do not + exceed 3 total calls. Favor precision over volume. +- Use scores to prioritize evidence, but include only the minimal subset of + snippet texts (verbatim) in SearchAnswer.context (typically 1-4). +- Set SearchAnswer.sources to the corresponding document identifiers for the + snippets you used (title if available, otherwise URI; one per snippet; same + order as context). Context must be text-only. +- If no relevant information is found, clearly say so and return an empty + context list and sources list. + +Answering rules: +- Be direct and specific; avoid meta commentary about the process. +- Do not include any claims not supported by the provided snippets. +- Prefer concise phrasing; avoid copying long passages. +- When evidence is partial, state the limits explicitly in the answer.""" diff --git a/haiku_rag_slim/haiku/rag/graph_common/utils.py b/haiku_rag_slim/haiku/rag/graph_common/utils.py new file mode 100644 index 00000000..f24fb06c --- /dev/null +++ b/haiku_rag_slim/haiku/rag/graph_common/utils.py @@ -0,0 +1,64 @@ +"""Common utilities for all graph implementations.""" + +from typing import Any, Protocol + +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.ollama import OllamaProvider +from pydantic_ai.providers.openai import OpenAIProvider + +from haiku.rag.config import Config + + +class HasEmitLog(Protocol): + """Protocol for objects that can emit log messages.""" + + def emit_log(self, message: str, state: Any = None) -> None: ... + + +def get_model(provider: str, model: str) -> OpenAIChatModel | str: + """ + Get a model instance for the specified provider and model name. + + Args: + provider: The model provider ("ollama", "vllm", or other) + model: The model name + + Returns: + A configured model instance + + Raises: + ValueError: If the provider is unknown + """ + if provider == "ollama": + return OpenAIChatModel( + model_name=model, + provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"), + ) + elif provider == "vllm": + return OpenAIChatModel( + model_name=model, + provider=OpenAIProvider( + base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1", + api_key="none", + ), + ) + elif provider in ("openai", "anthropic", "gemini", "groq", "bedrock"): + # These providers use string format + return f"{provider}:{model}" + else: + raise ValueError( + f"Unknown model provider: {provider}. " + f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock" + ) + + +def log(deps: HasEmitLog, state: Any, message: str) -> None: + """ + Emit a log message through the dependencies. + + Args: + deps: Dependencies object with emit_log method + state: Current state (passed to emit_log) + message: The message to log + """ + deps.emit_log(message, state) diff --git a/haiku_rag_slim/haiku/rag/qa/deep/common.py b/haiku_rag_slim/haiku/rag/qa/deep/common.py deleted file mode 100644 index 71815ba5..00000000 --- a/haiku_rag_slim/haiku/rag/qa/deep/common.py +++ /dev/null @@ -1,41 +0,0 @@ -from typing import Any, Protocol - -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider - -from haiku.rag.config import Config -from haiku.rag.qa.deep.models import SearchAnswer - - -class HasEmitLog(Protocol): - def emit_log(self, message: str, state: Any = None) -> None: ... - - -def get_model(provider: str, model: str) -> Any: - if provider == "ollama": - return OpenAIChatModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"), - ) - elif provider == "vllm": - return OpenAIChatModel( - model_name=model, - provider=OpenAIProvider( - base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1", - api_key="none", - ), - ) - else: - return f"{provider}:{model}" - - -def log(deps: HasEmitLog, state: Any, message: str) -> None: - deps.emit_log(message, state) - - -def collect_answers_reducer( - acc: list[SearchAnswer], item: SearchAnswer | None -) -> list[SearchAnswer]: - """Reducer function to collect search answers, filtering out None values.""" - return acc + [item] if item else acc diff --git a/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py b/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py index eaf7957d..f8bce190 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py @@ -2,7 +2,7 @@ from pydantic import BaseModel, Field from rich.console import Console from haiku.rag.client import HaikuRAG -from haiku.rag.qa.deep.models import SearchAnswer +from haiku.rag.graph_common.models import SearchAnswer class DeepQAContext(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/qa/deep/graph.py b/haiku_rag_slim/haiku/rag/qa/deep/graph.py index fb0a8760..a356a480 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/graph.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/graph.py @@ -3,27 +3,25 @@ from typing import Any from pydantic_ai import Agent, RunContext from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput -from pydantic_graph.beta import GraphBuilder, StepContext +from pydantic_graph.beta import Graph, GraphBuilder, StepContext +from pydantic_graph.beta.join import reduce_list_append -from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log +from haiku.rag.graph_common import get_model, log +from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer +from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT from haiku.rag.qa.deep.dependencies import DeepQADependencies -from haiku.rag.qa.deep.models import ( - DeepQAAnswer, - DeepQAEvaluation, - ResearchPlan, - SearchAnswer, -) +from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation from haiku.rag.qa.deep.prompts import ( DECISION_PROMPT, - PLAN_PROMPT, - SEARCH_AGENT_PROMPT, SYNTHESIS_PROMPT, SYNTHESIS_PROMPT_WITH_CITATIONS, ) from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState -def build_deep_qa_graph(provider: str, model: str): +def build_deep_qa_graph( + provider: str, model: str +) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]: g = GraphBuilder( state_type=DeepQAState, deps_type=DeepQADeps, @@ -82,7 +80,7 @@ def build_deep_qa_graph(provider: str, model: str): @g.step async def search_one( ctx: StepContext[DeepQAState, DeepQADeps, str], - ) -> SearchAnswer | None: + ) -> SearchAnswer: state = ctx.state deps = ctx.deps sub_q = ctx.inputs @@ -130,17 +128,22 @@ def build_deep_qa_graph(provider: str, model: str): ) try: result = await agent.run(sub_q, deps=agent_deps) + answer = result.output + if answer: + state.context.add_qa_response(answer) + preview = answer.answer[:150] + ( + "…" if len(answer.answer) > 150 else "" + ) + log(deps, state, f" [green]✓[/green] {preview}") + return answer except Exception as e: log(deps, state, f"[red]Search failed:[/red] {e}") - return None - - answer = result.output - if answer: - state.context.add_qa_response(answer) - preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "") - log(deps, state, f" [green]✓[/green] {preview}") - - return answer + failure_answer = SearchAnswer( + query=sub_q, + answer=f"Search failed after retries: {str(e)}", + confidence=0.0, + ) + return failure_answer @g.step async def get_batch( @@ -291,8 +294,8 @@ def build_deep_qa_graph(provider: str, model: str): # Build the graph structure collect_answers = g.join( - collect_answers_reducer, - initial_factory=lambda: [], + reduce_list_append, + initial_factory=list[SearchAnswer], ) g.add( diff --git a/haiku_rag_slim/haiku/rag/qa/deep/models.py b/haiku_rag_slim/haiku/rag/qa/deep/models.py index 2a7bad79..e33dfc95 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/models.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/models.py @@ -1,29 +1,6 @@ from pydantic import BaseModel, Field -class ResearchPlan(BaseModel): - main_question: str - sub_questions: list[str] - - -class SearchAnswer(BaseModel): - query: str = Field(description="The search query that was performed") - answer: str = Field(description="The answer generated based on the context") - context: list[str] = Field( - description=( - "Only the minimal set of relevant snippets (verbatim) that directly " - "support the answer" - ) - ) - sources: list[str] = Field( - description=( - "Document titles (if available) or URIs corresponding to the" - " snippets actually used in the answer (one per snippet; omit if none)" - ), - default_factory=list, - ) - - class DeepQAEvaluation(BaseModel): is_sufficient: bool = Field( description="Whether we have sufficient information to answer the question" diff --git a/haiku_rag_slim/haiku/rag/qa/deep/prompts.py b/haiku_rag_slim/haiku/rag/qa/deep/prompts.py index 0bd399b1..75f27fd7 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/prompts.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/prompts.py @@ -1,48 +1,4 @@ -PLAN_PROMPT = """You are the research orchestrator for a focused, iterative -workflow. - -Responsibilities: -1. Understand and decompose the main question -2. Propose a minimal, high‑leverage plan -3. Coordinate specialized agents to gather evidence -4. Iterate based on gaps and new findings - -Plan requirements: -- Produce at most 3 sub_questions that together cover the main question. -- Each sub_question must be a standalone, self‑contained query that can run - without extra context. Include concrete entities, scope, timeframe, and any - qualifiers. Avoid ambiguous pronouns (it/they/this/that). -- Prioritize the highest‑value aspects first; avoid redundancy and overlap. -- Prefer questions that are likely answerable from the current knowledge base; - if coverage is uncertain, make scopes narrower and specific. -- Order sub_questions by execution priority (most valuable first).""" - -SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist. - -Tasks: -1. Search the knowledge base for relevant evidence. -2. Analyze retrieved snippets. -3. Provide an answer strictly grounded in that evidence. - -Tool usage: -- Always call search_and_answer before drafting any answer. -- The tool returns snippets with verbatim `text`, a relevance `score`, and the - originating document identifier (document title if available, otherwise URI). -- You may call the tool multiple times to refine or broaden context, but do not - exceed 3 total calls. Favor precision over volume. -- Use scores to prioritize evidence, but include only the minimal subset of - snippet texts (verbatim) in SearchAnswer.context (typically 1‑4). -- Set SearchAnswer.sources to the corresponding document identifiers for the - snippets you used (title if available, otherwise URI; one per snippet; same - order as context). Context must be text‑only. -- If no relevant information is found, clearly say so and return an empty - context list and sources list. - -Answering rules: -- Be direct and specific; avoid meta commentary about the process. -- Do not include any claims not supported by the provided snippets. -- Prefer concise phrasing; avoid copying long passages. -- When evidence is partial, state the limits explicitly in the answer.""" +"""Deep QA specific prompts.""" SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers. diff --git a/haiku_rag_slim/haiku/rag/qa/deep/state.py b/haiku_rag_slim/haiku/rag/qa/deep/state.py index f76a99a3..8880da9d 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/state.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/state.py @@ -21,5 +21,5 @@ class DeepQAState: context: DeepQAContext max_sub_questions: int = 3 max_iterations: int = 2 - max_concurrency: int = 3 + max_concurrency: int = 1 iterations: int = 0 diff --git a/haiku_rag_slim/haiku/rag/research/__init__.py b/haiku_rag_slim/haiku/rag/research/__init__.py index a289a077..9406a89c 100644 --- a/haiku_rag_slim/haiku/rag/research/__init__.py +++ b/haiku_rag_slim/haiku/rag/research/__init__.py @@ -1,3 +1,3 @@ -from haiku.rag.qa.deep.models import SearchAnswer +from haiku.rag.graph_common.models import SearchAnswer from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.research.models import EvaluationResult, ResearchReport diff --git a/haiku_rag_slim/haiku/rag/research/dependencies.py b/haiku_rag_slim/haiku/rag/research/dependencies.py index 47e2b406..09376b53 100644 --- a/haiku_rag_slim/haiku/rag/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/research/dependencies.py @@ -4,7 +4,7 @@ from pydantic import BaseModel, Field from rich.console import Console from haiku.rag.client import HaikuRAG -from haiku.rag.qa.deep.models import SearchAnswer +from haiku.rag.graph_common.models import SearchAnswer from haiku.rag.research.models import ( GapRecord, InsightAnalysis, diff --git a/haiku_rag_slim/haiku/rag/research/graph.py b/haiku_rag_slim/haiku/rag/research/graph.py index bc2d46f4..a80dbc0c 100644 --- a/haiku_rag_slim/haiku/rag/research/graph.py +++ b/haiku_rag_slim/haiku/rag/research/graph.py @@ -3,11 +3,12 @@ from typing import Any from pydantic_ai import Agent, RunContext from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput -from pydantic_graph.beta import GraphBuilder, StepContext +from pydantic_graph.beta import Graph, GraphBuilder, StepContext +from pydantic_graph.beta.join import reduce_list_append -from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log -from haiku.rag.qa.deep.models import ResearchPlan, SearchAnswer -from haiku.rag.qa.deep.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT +from haiku.rag.graph_common import get_model, log +from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer +from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT from haiku.rag.research.common import ( format_analysis_for_prompt, format_context_for_prompt, @@ -26,7 +27,9 @@ from haiku.rag.research.prompts import ( from haiku.rag.research.state import ResearchDeps, ResearchState -def build_research_graph(provider: str, model: str): +def build_research_graph( + provider: str, model: str +) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: g = GraphBuilder( state_type=ResearchState, deps_type=ResearchDeps, @@ -86,7 +89,7 @@ def build_research_graph(provider: str, model: str): @g.step async def search_one( ctx: StepContext[ResearchState, ResearchDeps, str], - ) -> SearchAnswer | None: + ) -> SearchAnswer: state = ctx.state deps = ctx.deps sub_q = ctx.inputs @@ -135,17 +138,22 @@ def build_research_graph(provider: str, model: str): ) try: result = await agent.run(sub_q, deps=agent_deps) + answer = result.output + if answer: + state.context.add_qa_response(answer) + preview = answer.answer[:150] + ( + "…" if len(answer.answer) > 150 else "" + ) + log(deps, state, f" [green]✓[/green] {preview}") + return answer except Exception as e: log(deps, state, f"[red]Search failed:[/red] {e}") - return None - - answer = result.output - if answer: - state.context.add_qa_response(answer) - preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "") - log(deps, state, f" [green]✓[/green] {preview}") - - return answer + failure_answer = SearchAnswer( + query=sub_q, + answer=f"Search failed after retries: {str(e)}", + confidence=0.0, + ) + return failure_answer @g.step async def get_batch( @@ -349,8 +357,8 @@ def build_research_graph(provider: str, model: str): # Build the graph structure collect_answers = g.join( - collect_answers_reducer, - initial_factory=lambda: [], + reduce_list_append, + initial_factory=list[SearchAnswer], ) g.add( diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py index 6a6bcf87..5dab9475 100644 --- a/tests/test_deep_qa.py +++ b/tests/test_deep_qa.py @@ -2,9 +2,9 @@ import pytest from pydantic_ai.models.test import TestModel from haiku.rag.client import HaikuRAG +from haiku.rag.graph_common.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 SearchAnswer from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState @@ -16,7 +16,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): def test_model_factory(provider, model): return TestModel() - monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory) graph = build_deep_qa_graph(provider="test", model="test") @@ -50,7 +50,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path): def test_model_factory(provider, model): return TestModel() - monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory) graph = build_deep_qa_graph(provider="test", model="test") diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 69a684cd..8c4239f7 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -17,7 +17,7 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): def test_model_factory(provider, model): return TestModel() - monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory) graph = build_research_graph(provider="test", model="test")