diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 1914cc541..82a457777 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -248,6 +248,144 @@ def _validate_chat_history( raise ValueError(error_message) +class _AgentBuilder: + """Internal builder class for constructing ReAct-style agent graphs.""" + + def __init__( + self, + model: Union[ + str, + LanguageModelLike, + Callable[[StateSchema, Runtime[ContextT]], BaseChatModel], + Callable[[StateSchema, Runtime[ContextT]], Awaitable[BaseChatModel]], + ], + tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], + *, + prompt: Optional[Prompt] = None, + response_format: Optional[ + Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] + ] = None, + pre_model_hook: Optional[RunnableLike] = None, + post_model_hook: Optional[RunnableLike] = None, + state_schema: Optional[StateSchemaType] = None, + context_schema: Optional[Type[Any]] = None, + checkpointer: Optional[Checkpointer] = None, + store: Optional[BaseStore] = None, + interrupt_before: Optional[list[str]] = None, + interrupt_after: Optional[list[str]] = None, + debug: bool = False, + version: Literal["v1", "v2"] = "v2", + name: Optional[str] = None, + ): + """Initialize the agent builder with all parameters.""" + # Store all parameters + self.model = model + self.tools = tools + self.prompt = prompt + self.response_format = response_format + self.pre_model_hook = pre_model_hook + self.post_model_hook = post_model_hook + self.state_schema = state_schema + self.context_schema = context_schema + self.checkpointer = checkpointer + self.store = store + self.interrupt_before = interrupt_before + self.interrupt_after = interrupt_after + self.debug = debug + self.version = version + self.name = name + + # Initialize derived attributes + self.tool_node: ToolNode + self.tool_classes: list[Type[BaseTool]] + self.llm_builtin_tools: list[dict] = [] + self.is_dynamic_model: bool + self.is_async_dynamic_model: bool + self.tool_calling_enabled: bool + self.static_model: Optional[Runnable] = None + self.should_return_direct: set[str] = set() + + # Run validation and setup + self._validate_version() + self._validate_state_schema() + self._validate_and_setup_tools() + self._validate_and_setup_model() + + def _validate_version(self) -> None: + """Validate the version parameter.""" + if self.version not in ("v1", "v2"): + raise ValueError( + f"Invalid version {self.version}. Supported versions are 'v1' and 'v2'." + ) + + def _validate_state_schema(self) -> None: + """Validate and set up the state schema.""" + if self.state_schema is not None: + required_keys = {"messages", "remaining_steps"} + if self.response_format is not None: + required_keys.add("structured_response") + + schema_keys = set(get_type_hints(self.state_schema)) + if missing_keys := required_keys - set(schema_keys): + raise ValueError(f"Missing required key(s) {missing_keys} in state_schema") + + if self.state_schema is None: + self.state_schema = ( + AgentStateWithStructuredResponse + if self.response_format is not None + else AgentState + ) + + def _validate_and_setup_tools(self) -> None: + """Validate and set up tools and tool node.""" + if isinstance(self.tools, ToolNode): + self.tool_classes = list(self.tools.tools_by_name.values()) + self.tool_node = self.tools + else: + self.llm_builtin_tools = [t for t in self.tools if isinstance(t, dict)] + self.tool_node = ToolNode([t for t in self.tools if not isinstance(t, dict)]) + self.tool_classes = list(self.tool_node.tools_by_name.values()) + + self.tool_calling_enabled = len(self.tool_classes) > 0 + + # If any of the tools are configured to return_directly after running, + # our graph needs to check if these were called + self.should_return_direct = {t.name for t in self.tool_classes if t.return_direct} + + def _validate_and_setup_model(self) -> None: + """Validate and set up the model.""" + self.is_dynamic_model = not isinstance(self.model, (str, Runnable)) and callable(self.model) + self.is_async_dynamic_model = self.is_dynamic_model and inspect.iscoroutinefunction(self.model) + + if not self.is_dynamic_model: + model = self.model + if isinstance(model, str): + try: + from langchain.chat_models import ( # type: ignore[import-not-found] + init_chat_model, + ) + except ImportError: + raise ImportError( + "Please install langchain (`pip install langchain`) to " + "use ':' string syntax for `model` parameter." + ) + + model = cast(BaseChatModel, init_chat_model(model)) + + if ( + _should_bind_tools(model, self.tool_classes, num_builtin=len(self.llm_builtin_tools)) # type: ignore[arg-type] + and len(self.tool_classes + self.llm_builtin_tools) > 0 + ): + model = cast(BaseChatModel, model).bind_tools( + self.tool_classes + self.llm_builtin_tools # type: ignore[operator] + ) + + self.static_model = _get_prompt_runnable(self.prompt) | model # type: ignore[operator] + else: + # For dynamic models, we'll create the runnable at runtime + self.static_model = None + + def create_react_agent( model: Union[ str, @@ -952,3 +1090,4 @@ __all__ = [ "AgentStateWithStructuredResponse", "AgentStateWithStructuredResponsePydantic", ] +