From 5f4c54cd1fa5d7b49fa16e43093165899b13df65 Mon Sep 17 00:00:00 2001
From: Siddharth Wagh <siddw143@gmail.com>
Date: Mon, 25 Aug 2025 11:49:13 +0530
Subject: [PATCH] align create_deep_agent model parameters with
 create_react_agent

---
 src/deepagents/graph.py | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/deepagents/graph.py b/src/deepagents/graph.py
index fc2d4274..be52900e 100644
--- a/src/deepagents/graph.py
+++ b/src/deepagents/graph.py
@@ -2,11 +2,15 @@
 from deepagents.model import get_default_model
 from deepagents.tools import write_todos, write_file, read_file, ls, edit_file
 from deepagents.state import DeepAgentState
-from typing import Sequence, Union, Callable, Any, TypeVar, Type, Optional, Dict
+from typing import Sequence, Union, Callable, Any, TypeVar, Type, Optional, Dict, Awaitable
 from langchain_core.tools import BaseTool
-from langchain_core.language_models import LanguageModelLike
+from langchain_core.language_models import LanguageModelLike,LanguageModelInput, BaseMessage
+from langchain_core.chat_models import BaseChatModel
+from langchain_core.runnables import Runnable
 from deepagents.interrupt import create_interrupt_hook, ToolInterruptConfig
 from langgraph.types import Checkpointer
+from langgraph.runtime import Runtime
+from langgraph.typing import ContextT
 from langgraph.prebuilt import create_react_agent
 
 StateSchema = TypeVar("StateSchema", bound=DeepAgentState)
@@ -28,7 +32,19 @@
 def create_deep_agent(
     tools: Sequence[Union[BaseTool, Callable, dict[str, Any]]],
     instructions: str,
-    model: Optional[Union[str, LanguageModelLike]] = None,
+    model: Optional[Union[
+        str,
+        LanguageModelLike,
+        Callable[[StateSchema, Runtime[ContextT]], BaseChatModel],
+        Callable[[StateSchema, Runtime[ContextT]], Awaitable[BaseChatModel]],
+        Callable[
+            [StateSchema, Runtime[ContextT]], Runnable[LanguageModelInput, BaseMessage]
+        ],
+        Callable[
+            [StateSchema, Runtime[ContextT]],
+            Awaitable[Runnable[LanguageModelInput, BaseMessage]],
+        ],
+    ]] = None,
     subagents: list[SubAgent] = None,
     state_schema: Optional[StateSchemaType] = None,
     interrupt_config: Optional[ToolInterruptConfig] = None,
@@ -45,7 +61,17 @@ def create_deep_agent(
         tools: The additional tools the agent should have access to.
         instructions: The additional instructions the agent should have. Will go in
             the system prompt.
-        model: The model to use.
+        model: The language model for the agent. Supports static and dynamic
+            model selection.
+
+            - **Static model**: A chat model instance (e.g., `ChatOpenAI()`) or
+              string identifier (e.g., `"openai:gpt-4"`)
+            - **Dynamic model**: A callable with signature
+              `(state, runtime) -> BaseChatModel` that returns different models
+              based on runtime context
+              If the model has tools bound via `.bind_tools()` or other configurations,
+              the return type should be a Runnable[LanguageModelInput, BaseMessage]
+              Coroutines are also supported, allowing for asynchronous model selection.     
         subagents: The subagents to use. Each subagent should be a dictionary with the
             following keys:
                 - `name`