Merge pull request #151 from ggozad/fix/graph-config

Allow custom config when building research/deep ask graphs.

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
 

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,

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

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,

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,

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)

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)