diff --git a/libs/langgraph/langgraph/agent/middleware/dynamic_model.py b/libs/langgraph/langgraph/agent/middleware/dynamic_model.py new file mode 100644 index 000000000..a3352f748 --- /dev/null +++ b/libs/langgraph/langgraph/agent/middleware/dynamic_model.py @@ -0,0 +1,28 @@ +from langgraph.agent.types import AgentState, AgentUpdate, AgentJump, AgentMiddleware +from langgraph.agent.types import ModelRequest +from typing import Literal +from dataclasses import dataclass +from langchain_core.language_models.chat_models import BaseChatModel + + +class DynamicModelMiddleware(AgentMiddleware): + """Selects different models based on task complexity""" + + def __init__( + self, + basic_model: BaseChatModel, + complex_model: BaseChatModel, + message_threshold: int = 5, + ): + self.basic_model = basic_model + self.complex_model = complex_model + self.message_threshold = message_threshold + + def modify_model_request( + self, request: ModelRequest, state: AgentState + ) -> ModelRequest: + if len(state.messages) > self.message_threshold: + request.model = self.complex_model + else: + request.model = self.basic_model + return request diff --git a/libs/langgraph/langgraph/agent/middleware/dynamic_prep.py b/libs/langgraph/langgraph/agent/middleware/dynamic_prep.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/langgraph/langgraph/agent/middleware/limiting_calls.py b/libs/langgraph/langgraph/agent/middleware/model_calls.py similarity index 100% rename from libs/langgraph/langgraph/agent/middleware/limiting_calls.py rename to libs/langgraph/langgraph/agent/middleware/model_calls.py