From c45d507213a240f5a4f64867358cd9f78995ad0d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 21 Nov 2025 10:57:14 +0200 Subject: [PATCH] Allow custom config when building research/deep ask graphs. Closes #149 --- CHANGELOG.md | 4 ++++ haiku_rag_slim/haiku/rag/graph/common/nodes.py | 12 ++++++++++-- haiku_rag_slim/haiku/rag/graph/common/utils.py | 10 +++++++--- haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py | 6 ++++-- haiku_rag_slim/haiku/rag/graph/research/graph.py | 8 +++++--- tests/graph/test_deep_qa.py | 4 ++-- tests/graph/test_research_graph.py | 2 +- 7 files changed, 33 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 08d65490..44bc6f51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ ### Fixed - **AG-UI Activity Events**: Activity events now correctly use structured dict content instead of strings +- **Graph Configuration**: Graph builder functions now properly accept and use non-global config (#149) + - `build_research_graph()` and `build_deep_qa_graph()` now pass config to all agents and model creation + - `get_model()` utility function accepts `config` parameter (defaults to global Config) + - Allows creating multiple graphs with different configurations in the same application ## [0.17.2] - 2025-11-19 diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index bc2da554..858170d8 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -10,6 +10,8 @@ from pydantic_ai.output import ToolOutput from pydantic_graph.beta import StepContext from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.config.models import AppConfig from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.common import get_model from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer @@ -55,6 +57,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( deps_type: type[AgentDepsT], activity_message: str = "Creating plan", output_retries: int | None = None, + config: AppConfig = Config, ) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]: """Create a plan node for any graph. @@ -64,6 +67,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies) activity_message: Message to show during planning activity output_retries: Number of output retries for the agent (optional) + config: AppConfig object (defaults to global Config) Returns: Async function that can be used as a graph step @@ -80,7 +84,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( try: # Build agent configuration agent_config = { - "model": get_model(provider, model), + "model": get_model(provider, model, config), "output_type": ResearchPlan, "instructions": ( PLAN_PROMPT @@ -136,6 +140,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps]( with_step_wrapper: bool = True, success_message_format: str = "Answered: {sub_q}", handle_exceptions: bool = False, + config: AppConfig = Config, ) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]: """Create a search_one node for any graph. @@ -146,6 +151,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps]( with_step_wrapper: Whether to wrap with agui_emitter start/finish step success_message_format: Format string for success activity message handle_exceptions: Whether to handle exceptions with fallback answer + config: AppConfig object (defaults to global Config) Returns: Async function that can be used as a graph step @@ -178,6 +184,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps]( deps_type, success_message_format, handle_exceptions, + config, ) finally: if deps.agui_emitter and with_step_wrapper: @@ -195,6 +202,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( deps_type: type[AgentDepsT], success_message_format: str, handle_exceptions: bool, + config: AppConfig, ) -> SearchAnswer: """Internal search implementation.""" if deps.agui_emitter: @@ -203,7 +211,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( ) agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=ToolOutput(SearchAnswer, max_retries=3), instructions=SEARCH_AGENT_PROMPT, retries=3, diff --git a/haiku_rag_slim/haiku/rag/graph/common/utils.py b/haiku_rag_slim/haiku/rag/graph/common/utils.py index 19c424c8..c2711915 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/utils.py +++ b/haiku_rag_slim/haiku/rag/graph/common/utils.py @@ -5,15 +5,19 @@ from pydantic_ai.providers.ollama import OllamaProvider from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.config import Config +from haiku.rag.config.models import AppConfig -def get_model(provider: str, model: str) -> OpenAIChatModel | str: +def get_model( + provider: str, model: str, config: AppConfig = Config +) -> 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 + config: AppConfig object (defaults to global Config) Returns: A configured model instance @@ -24,13 +28,13 @@ def get_model(provider: str, model: str) -> OpenAIChatModel | str: if provider == "ollama": return OpenAIChatModel( model_name=model, - provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"), + 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", + base_url=f"{config.providers.vllm.research_base_url or config.providers.vllm.qa_base_url}/v1", api_key="none", ), ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py index 07b6d649..0fa4bdee 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py @@ -45,6 +45,7 @@ def build_deep_qa_graph( deps_type=DeepQADependencies, # type: ignore[arg-type] activity_message="Planning approach", output_retries=None, # Deep QA doesn't use output_retries + config=config, ) ) # type: ignore[arg-type] @@ -57,6 +58,7 @@ def build_deep_qa_graph( with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step success_message_format="Answered: {sub_q}", handle_exceptions=True, + config=config, ) ) # type: ignore[arg-type] @@ -90,7 +92,7 @@ def build_deep_qa_graph( try: agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=DeepQAEvaluation, instructions=DECISION_PROMPT, retries=3, @@ -168,7 +170,7 @@ def build_deep_qa_graph( ) agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=DeepQAAnswer, instructions=prompt_template, retries=3, diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 048c258f..6187f8cf 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -52,6 +52,7 @@ def build_research_graph( deps_type=ResearchDependencies, # type: ignore[arg-type] activity_message="Creating research plan", output_retries=3, + config=config, ) ) # type: ignore[arg-type] @@ -64,6 +65,7 @@ def build_research_graph( with_step_wrapper=True, success_message_format="Found answer with {confidence:.0%} confidence", handle_exceptions=True, + config=config, ) ) # type: ignore[arg-type] @@ -97,7 +99,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=InsightAnalysis, instructions=INSIGHT_AGENT_PROMPT, retries=3, @@ -155,7 +157,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=EvaluationResult, instructions=DECISION_AGENT_PROMPT, retries=3, @@ -231,7 +233,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model), + model=get_model(provider, model, config), output_type=ResearchReport, instructions=SYNTHESIS_AGENT_PROMPT, retries=3, diff --git a/tests/graph/test_deep_qa.py b/tests/graph/test_deep_qa.py index 582b3242..3c7e4e13 100644 --- a/tests/graph/test_deep_qa.py +++ b/tests/graph/test_deep_qa.py @@ -13,7 +13,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): """Test deep Q&A graph with mocked LLM using TestModel.""" # Mock get_model to return TestModel which generates valid schema-compliant data - def test_model_factory(provider, model): + def test_model_factory(provider, model, config=None): return TestModel() monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory) @@ -47,7 +47,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path): """Test deep Q&A with citations enabled using TestModel.""" # Mock get_model to return TestModel - def test_model_factory(provider, model): + def test_model_factory(provider, model, config=None): return TestModel() monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory) diff --git a/tests/graph/test_research_graph.py b/tests/graph/test_research_graph.py index f452b022..889ae92b 100644 --- a/tests/graph/test_research_graph.py +++ b/tests/graph/test_research_graph.py @@ -36,7 +36,7 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): """Test research graph with mocked LLM using AG-UI events.""" # Mock get_model to return TestModel which generates valid schema-compliant data - def test_model_factory(_provider, _model): + def test_model_factory(_provider, _model, _config=None): return TestModel() monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)