diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 1914cc541..f142c7865 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -248,6 +248,591 @@ def _validate_chat_history( raise ValueError(error_message) +class _AgentBuilder: + """Internal scaffolding utility for constructing configurable ReAct-style agent graphs. + + This class encapsulates configuration logic, validation, and dynamic construction + of nodes and edges based on user input, allowing users to build structured agent + behaviors without manually wiring up the entire workflow. + """ + + 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, + **deprecated_kwargs: Any, + ) -> None: + # Handle deprecated config_schema parameter + if ( + config_schema := deprecated_kwargs.pop("config_schema", MISSING) + ) is not MISSING: + warn( + "`config_schema` is no longer supported. Use `context_schema` instead.", + category=LangGraphDeprecatedSinceV10, + ) + + if context_schema is not None: + context_schema = config_schema + + # Validate version parameter + if version not in ("v1", "v2"): + raise ValueError( + f"Invalid version {version}. Supported versions are 'v1' and 'v2'." + ) + + # Validate state schema requirements + if state_schema is not None: + required_keys = {"messages", "remaining_steps"} + if response_format is not None: + required_keys.add("structured_response") + + schema_keys = set(get_type_hints(state_schema)) + if missing_keys := required_keys - set(schema_keys): + raise ValueError(f"Missing required key(s) {missing_keys} in state_schema") + + # Set default state schema based on response_format + if state_schema is None: + state_schema = ( + AgentStateWithStructuredResponse + if response_format is not None + else AgentState + ) + + # Store all parameters as instance variables + 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 + + # Process tools (ToolNode vs sequence) + if isinstance(tools, ToolNode): + self.tool_classes = list(tools.tools_by_name.values()) + self.tool_node = tools + self.llm_builtin_tools = [] + else: + self.llm_builtin_tools = [t for t in tools if isinstance(t, dict)] + self.tool_node = ToolNode([t for t in tools if not isinstance(t, dict)]) + self.tool_classes = list(self.tool_node.tools_by_name.values()) + + # Determine model characteristics + self.is_dynamic_model = not isinstance(model, (str, Runnable)) and callable(model) + self.is_async_dynamic_model = self.is_dynamic_model and inspect.iscoroutinefunction(model) + + # Identify tools with return_direct behavior + self.should_return_direct = {t.name for t in self.tool_classes if t.return_direct} + + # Initialize other state + self.static_model: Optional[Runnable] = None + self.entrypoint: str = "agent" + self._resolve_model: Optional[Callable] = None + self._aresolve_model: Optional[Callable] = None + self.post_model_hook_router: Optional[Callable] = None + + def _validate_state_schema(self) -> None: + """Validate custom state schema requirements.""" + pass # TODO: Implement + + def _setup_model_and_tools(self) -> None: + """Handle model resolution and tool binding.""" + if not self.is_dynamic_model: + # Static model initialization + if isinstance(self.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." + ) + + self.model = cast(BaseChatModel, init_chat_model(self.model)) + + # Tool binding + if ( + _should_bind_tools(self.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 + ): + self.model = cast(BaseChatModel, self.model).bind_tools( + self.tool_classes + self.llm_builtin_tools # type: ignore[operator] + ) + + # Create prompt runnable + self.static_model = _get_prompt_runnable(self.prompt) | self.model # type: ignore[operator] + else: + # For dynamic models, we'll create the runnable at runtime + self.static_model = None + + # Create model resolution functions + def _resolve_model( + state: StateSchema, runtime: Runtime[ContextT] + ) -> LanguageModelLike: + """Resolve the model to use, handling both static and dynamic models.""" + if self.is_dynamic_model: + return _get_prompt_runnable(self.prompt) | self.model(state, runtime) # type: ignore[operator] + else: + return self.static_model + + async def _aresolve_model( + state: StateSchema, runtime: Runtime[ContextT] + ) -> LanguageModelLike: + """Async resolve the model to use, handling both static and dynamic models.""" + if self.is_async_dynamic_model: + resolved_model = await self.model(state, runtime) # type: ignore[misc,operator] + return _get_prompt_runnable(self.prompt) | resolved_model + elif self.is_dynamic_model: + return _get_prompt_runnable(self.prompt) | self.model(state, runtime) # type: ignore[operator] + else: + return self.static_model + + # Store as instance methods + self._resolve_model = _resolve_model + self._aresolve_model = _aresolve_model + + def _create_model_node(self) -> RunnableCallable: + """Create the core LLM interaction node.""" + def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool: + has_tool_calls = isinstance(response, AIMessage) and response.tool_calls + all_tools_return_direct = ( + all(call["name"] in self.should_return_direct for call in response.tool_calls) + if isinstance(response, AIMessage) + else False + ) + remaining_steps = _get_state_value(state, "remaining_steps", None) + is_last_step = _get_state_value(state, "is_last_step", False) + return ( + (remaining_steps is None and is_last_step and has_tool_calls) + or ( + remaining_steps is not None + and remaining_steps < 1 + and all_tools_return_direct + ) + or (remaining_steps is not None and remaining_steps < 2 and has_tool_calls) + ) + + def _get_model_input_state(state: StateSchema) -> StateSchema: + if self.pre_model_hook is not None: + messages = ( + _get_state_value(state, "llm_input_messages") + ) or _get_state_value(state, "messages") + error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}" + else: + messages = _get_state_value(state, "messages") + error_msg = ( + f"Expected input to call_model to have 'messages' key, but got {state}" + ) + + if messages is None: + raise ValueError(error_msg) + + _validate_chat_history(messages) + # we're passing messages under `messages` key, as this is expected by the prompt + if isinstance(self.state_schema, type) and issubclass(self.state_schema, BaseModel): + state.messages = messages # type: ignore + else: + state["messages"] = messages # type: ignore + + return state + + # Define the function that calls the model + def call_model( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + if self.is_async_dynamic_model: + msg = ( + "Async model callable provided but agent invoked synchronously. " + "Use agent.ainvoke() or agent.astream(), or " + "provide a sync model callable." + ) + raise RuntimeError(msg) + + model_input = _get_model_input_state(state) + + if self.is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + dynamic_model = self._resolve_model(state, runtime) + response = cast(AIMessage, dynamic_model.invoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, self.static_model.invoke(model_input, config)) # type: ignore[union-attr] + + # add agent name to the AIMessage + response.name = self.name + + if _are_more_steps_needed(state, response): + return { + "messages": [ + AIMessage( + id=response.id, + content="Sorry, need more steps to process this request.", + ) + ] + } + # We return a list, because this will get added to the existing list + return {"messages": [response]} + + async def acall_model( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + model_input = _get_model_input_state(state) + + if self.is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + # (supports both sync and async) + dynamic_model = await self._aresolve_model(state, runtime) + response = cast(AIMessage, await dynamic_model.ainvoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, await self.static_model.ainvoke(model_input, config)) # type: ignore[union-attr] + + # add agent name to the AIMessage + response.name = self.name + if _are_more_steps_needed(state, response): + return { + "messages": [ + AIMessage( + id=response.id, + content="Sorry, need more steps to process this request.", + ) + ] + } + # We return a list, because this will get added to the existing list + return {"messages": [response]} + + # Create input schema + input_schema: StateSchemaType + if self.pre_model_hook is not None: + # Dynamically create a schema that inherits from state_schema and adds 'llm_input_messages' + if isinstance(self.state_schema, type) and issubclass(self.state_schema, BaseModel): + # For Pydantic schemas + from pydantic import create_model + + input_schema = create_model( + "CallModelInputSchema", + llm_input_messages=(list[AnyMessage], ...), + __base__=self.state_schema, + ) + else: + # For TypedDict schemas + class CallModelInputSchema(self.state_schema): # type: ignore + llm_input_messages: list[AnyMessage] + + input_schema = CallModelInputSchema + else: + input_schema = self.state_schema + + return RunnableCallable(call_model, acall_model, input_schema=input_schema) + + def _create_structured_response_node(self) -> Optional[RunnableCallable]: + """Create structured output formatting node if needed.""" + if self.response_format is None: + return None + + def generate_structured_response( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + if self.is_async_dynamic_model: + msg = ( + "Async model callable provided but agent invoked synchronously. " + "Use agent.ainvoke() or agent.astream(), or provide a sync model callable." + ) + raise RuntimeError(msg) + + messages = _get_state_value(state, "messages") + structured_response_schema = self.response_format + if isinstance(self.response_format, tuple): + system_prompt, structured_response_schema = self.response_format + messages = [SystemMessage(content=system_prompt)] + list(messages) + + resolved_model = self._resolve_model(state, runtime) + model_with_structured_output = _get_model( + resolved_model + ).with_structured_output( + cast(StructuredResponseSchema, structured_response_schema) + ) + response = model_with_structured_output.invoke(messages, config) + return {"structured_response": response} + + async def agenerate_structured_response( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + messages = _get_state_value(state, "messages") + structured_response_schema = self.response_format + if isinstance(self.response_format, tuple): + system_prompt, structured_response_schema = self.response_format + messages = [SystemMessage(content=system_prompt)] + list(messages) + + resolved_model = await self._aresolve_model(state, runtime) + model_with_structured_output = _get_model( + resolved_model + ).with_structured_output( + cast(StructuredResponseSchema, structured_response_schema) + ) + response = await model_with_structured_output.ainvoke(messages, config) + return {"structured_response": response} + + return RunnableCallable( + generate_structured_response, + agenerate_structured_response, + ) + + def _create_model_router(self) -> Callable: + """Create execution flow routing after model call.""" + # Define the function that determines whether to continue or not + def should_continue(state: StateSchema) -> Union[str, list[Send]]: + messages = _get_state_value(state, "messages") + last_message = messages[-1] + # If there is no function call, then we finish + if not isinstance(last_message, AIMessage) or not last_message.tool_calls: + if self.post_model_hook is not None: + return "post_model_hook" + elif self.response_format is not None: + return "generate_structured_response" + else: + return END + # Otherwise if there is, we continue + else: + if self.version == "v1": + return "tools" + elif self.version == "v2": + if self.post_model_hook is not None: + return "post_model_hook" + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in last_message.tool_calls + ] + + return should_continue + + def _create_tools_router(self) -> Optional[Callable]: + """Create post-tool-call routing based on return_direct.""" + if not self.should_return_direct: + return None + + def route_tool_responses(state: StateSchema) -> str: + for m in reversed(_get_state_value(state, "messages")): + if not isinstance(m, ToolMessage): + break + if m.name in self.should_return_direct: + return END + + # handle a case of parallel tool calls where + # the tool w/ `return_direct` was executed in a different `Send` + if isinstance(m, AIMessage) and m.tool_calls: + if any(call["name"] in self.should_return_direct for call in m.tool_calls): + return END + + return self.entrypoint + + return route_tool_responses + + def _setup_hooks(self, workflow: StateGraph) -> None: + """Add pre/post model hooks to the workflow.""" + # Optionally add a pre-model hook node that will be called + # every time before the "agent" (LLM-calling node) + if self.pre_model_hook is not None: + workflow.add_node("pre_model_hook", self.pre_model_hook) # type: ignore[arg-type] + workflow.add_edge("pre_model_hook", "agent") + self.entrypoint = "pre_model_hook" + else: + self.entrypoint = "agent" + + # Add a post model hook node if post_model_hook is provided + if self.post_model_hook is not None: + workflow.add_node("post_model_hook", self.post_model_hook) # type: ignore[arg-type] + + def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]: + """Route to the next node after post_model_hook. + + Routes to one of: + * "tools": if there are pending tool calls without a corresponding message. + * "generate_structured_response": if no pending tool calls exist and response_format is specified. + * END: if no pending tool calls exist and no response_format is specified. + """ + + messages = _get_state_value(state, "messages") + tool_messages = [ + m.tool_call_id for m in messages if isinstance(m, ToolMessage) + ] + last_ai_message = next( + m for m in reversed(messages) if isinstance(m, AIMessage) + ) + pending_tool_calls = [ + c for c in last_ai_message.tool_calls if c["id"] not in tool_messages + ] + + if pending_tool_calls: + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in pending_tool_calls + ] + elif isinstance(messages[-1], ToolMessage): + return self.entrypoint + elif self.response_format is not None: + return "generate_structured_response" + else: + return END + + # Store for use in build method + self.post_model_hook_router = post_model_hook_router + + def build(self) -> CompiledStateGraph: + """Assemble the full graph based on all options.""" + # Setup models and tools first + self._setup_model_and_tools() + + tool_calling_enabled = len(self.tool_classes) > 0 + + if not tool_calling_enabled: + # Non-tool-calling workflow + workflow = StateGraph(state_schema=self.state_schema, context_schema=self.context_schema) + model_node = self._create_model_node() + workflow.add_node("agent", model_node) + + # Set up hooks + self._setup_hooks(workflow) + workflow.set_entry_point(self.entrypoint) + + if self.post_model_hook is not None: + workflow.add_edge("agent", "post_model_hook") + + # Add structured response node if needed + structured_response_node = self._create_structured_response_node() + if structured_response_node is not None: + workflow.add_node("generate_structured_response", structured_response_node) + if self.post_model_hook is not None: + workflow.add_edge("post_model_hook", "generate_structured_response") + else: + workflow.add_edge("agent", "generate_structured_response") + + return workflow.compile( + checkpointer=self.checkpointer, + store=self.store, + interrupt_before=self.interrupt_before, + interrupt_after=self.interrupt_after, + debug=self.debug, + name=self.name, + ) + + # Tool-calling workflow + workflow = StateGraph( + state_schema=self.state_schema or AgentState, context_schema=self.context_schema + ) + + # Add nodes + model_node = self._create_model_node() + workflow.add_node("agent", model_node) + workflow.add_node("tools", self.tool_node) + + # Set up hooks + self._setup_hooks(workflow) + workflow.set_entry_point(self.entrypoint) + + # Prepare path mappings + agent_paths = [] + post_model_hook_paths = [self.entrypoint, "tools"] + + if self.post_model_hook is not None: + agent_paths.append("post_model_hook") + workflow.add_edge("agent", "post_model_hook") + else: + agent_paths.append("tools") + + # Add a structured output node if response_format is provided + structured_response_node = self._create_structured_response_node() + if structured_response_node is not None: + workflow.add_node("generate_structured_response", structured_response_node) + if self.post_model_hook is not None: + post_model_hook_paths.append("generate_structured_response") + else: + agent_paths.append("generate_structured_response") + else: + if self.post_model_hook is not None: + post_model_hook_paths.append(END) + else: + agent_paths.append(END) + + # Add conditional edges for post model hook + if self.post_model_hook is not None: + workflow.add_conditional_edges( + "post_model_hook", + self.post_model_hook_router, # type: ignore[arg-type] + path_map=post_model_hook_paths, + ) + + # Add conditional edges for agent + model_router = self._create_model_router() + workflow.add_conditional_edges( + "agent", + model_router, # type: ignore[arg-type] + path_map=agent_paths, + ) + + # Add conditional edges for tools if needed + tools_router = self._create_tools_router() + if tools_router is not None: + workflow.add_conditional_edges( + "tools", tools_router, path_map=[self.entrypoint, END] + ) + else: + workflow.add_edge("tools", self.entrypoint) + + # Finally, compile it! + return workflow.compile( + checkpointer=self.checkpointer, + store=self.store, + interrupt_before=self.interrupt_before, + interrupt_after=self.interrupt_after, + debug=self.debug, + name=self.name, + ) + + def create_react_agent( model: Union[ str, @@ -455,491 +1040,33 @@ def create_react_agent( print(chunk) ``` """ - if ( - config_schema := deprecated_kwargs.pop("config_schema", MISSING) - ) is not MISSING: - warn( - "`config_schema` is no longer supported. Use `context_schema` instead.", - category=LangGraphDeprecatedSinceV10, - ) + # Validate deprecated kwargs + if deprecated_kwargs: + # Only config_schema is handled in _AgentBuilder.__init__ + pass - if context_schema is not None: - context_schema = config_schema - - if version not in ("v1", "v2"): - raise ValueError( - f"Invalid version {version}. Supported versions are 'v1' and 'v2'." - ) - - if state_schema is not None: - required_keys = {"messages", "remaining_steps"} - if response_format is not None: - required_keys.add("structured_response") - - schema_keys = set(get_type_hints(state_schema)) - if missing_keys := required_keys - set(schema_keys): - raise ValueError(f"Missing required key(s) {missing_keys} in state_schema") - - if state_schema is None: - state_schema = ( - AgentStateWithStructuredResponse - if response_format is not None - else AgentState - ) - - llm_builtin_tools: list[dict] = [] - if isinstance(tools, ToolNode): - tool_classes = list(tools.tools_by_name.values()) - tool_node = tools - else: - llm_builtin_tools = [t for t in tools if isinstance(t, dict)] - tool_node = ToolNode([t for t in tools if not isinstance(t, dict)]) - tool_classes = list(tool_node.tools_by_name.values()) - - is_dynamic_model = not isinstance(model, (str, Runnable)) and callable(model) - is_async_dynamic_model = is_dynamic_model and inspect.iscoroutinefunction(model) - - tool_calling_enabled = len(tool_classes) > 0 - - if not is_dynamic_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, tool_classes, num_builtin=len(llm_builtin_tools)) # type: ignore[arg-type] - and len(tool_classes + llm_builtin_tools) > 0 - ): - model = cast(BaseChatModel, model).bind_tools( - tool_classes + llm_builtin_tools # type: ignore[operator] - ) - - static_model: Optional[Runnable] = _get_prompt_runnable(prompt) | model # type: ignore[operator] - else: - # For dynamic models, we'll create the runnable at runtime - static_model = None - - # If any of the tools are configured to return_directly after running, - # our graph needs to check if these were called - should_return_direct = {t.name for t in tool_classes if t.return_direct} - - def _resolve_model( - state: StateSchema, runtime: Runtime[ContextT] - ) -> LanguageModelLike: - """Resolve the model to use, handling both static and dynamic models.""" - if is_dynamic_model: - return _get_prompt_runnable(prompt) | model(state, runtime) # type: ignore[operator] - else: - return static_model - - async def _aresolve_model( - state: StateSchema, runtime: Runtime[ContextT] - ) -> LanguageModelLike: - """Async resolve the model to use, handling both static and dynamic models.""" - if is_async_dynamic_model: - resolved_model = await model(state, runtime) # type: ignore[misc,operator] - return _get_prompt_runnable(prompt) | resolved_model - elif is_dynamic_model: - return _get_prompt_runnable(prompt) | model(state, runtime) # type: ignore[operator] - else: - return static_model - - def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool: - has_tool_calls = isinstance(response, AIMessage) and response.tool_calls - all_tools_return_direct = ( - all(call["name"] in should_return_direct for call in response.tool_calls) - if isinstance(response, AIMessage) - else False - ) - remaining_steps = _get_state_value(state, "remaining_steps", None) - is_last_step = _get_state_value(state, "is_last_step", False) - return ( - (remaining_steps is None and is_last_step and has_tool_calls) - or ( - remaining_steps is not None - and remaining_steps < 1 - and all_tools_return_direct - ) - or (remaining_steps is not None and remaining_steps < 2 and has_tool_calls) - ) - - def _get_model_input_state(state: StateSchema) -> StateSchema: - if pre_model_hook is not None: - messages = ( - _get_state_value(state, "llm_input_messages") - ) or _get_state_value(state, "messages") - error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}" - else: - messages = _get_state_value(state, "messages") - error_msg = ( - f"Expected input to call_model to have 'messages' key, but got {state}" - ) - - if messages is None: - raise ValueError(error_msg) - - _validate_chat_history(messages) - # we're passing messages under `messages` key, as this is expected by the prompt - if isinstance(state_schema, type) and issubclass(state_schema, BaseModel): - state.messages = messages # type: ignore - else: - state["messages"] = messages # type: ignore - - return state - - # Define the function that calls the model - def call_model( - state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig - ) -> StateSchema: - if is_async_dynamic_model: - msg = ( - "Async model callable provided but agent invoked synchronously. " - "Use agent.ainvoke() or agent.astream(), or " - "provide a sync model callable." - ) - raise RuntimeError(msg) - - model_input = _get_model_input_state(state) - - if is_dynamic_model: - # Resolve dynamic model at runtime and apply prompt - dynamic_model = _resolve_model(state, runtime) - response = cast(AIMessage, dynamic_model.invoke(model_input, config)) # type: ignore[arg-type] - else: - response = cast(AIMessage, static_model.invoke(model_input, config)) # type: ignore[union-attr] - - # add agent name to the AIMessage - response.name = name - - if _are_more_steps_needed(state, response): - return { - "messages": [ - AIMessage( - id=response.id, - content="Sorry, need more steps to process this request.", - ) - ] - } - # We return a list, because this will get added to the existing list - return {"messages": [response]} - - async def acall_model( - state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig - ) -> StateSchema: - model_input = _get_model_input_state(state) - - if is_dynamic_model: - # Resolve dynamic model at runtime and apply prompt - # (supports both sync and async) - dynamic_model = await _aresolve_model(state, runtime) - response = cast(AIMessage, await dynamic_model.ainvoke(model_input, config)) # type: ignore[arg-type] - else: - response = cast(AIMessage, await static_model.ainvoke(model_input, config)) # type: ignore[union-attr] - - # add agent name to the AIMessage - response.name = name - if _are_more_steps_needed(state, response): - return { - "messages": [ - AIMessage( - id=response.id, - content="Sorry, need more steps to process this request.", - ) - ] - } - # We return a list, because this will get added to the existing list - return {"messages": [response]} - - input_schema: StateSchemaType - if pre_model_hook is not None: - # Dynamically create a schema that inherits from state_schema and adds 'llm_input_messages' - if isinstance(state_schema, type) and issubclass(state_schema, BaseModel): - # For Pydantic schemas - from pydantic import create_model - - input_schema = create_model( - "CallModelInputSchema", - llm_input_messages=(list[AnyMessage], ...), - __base__=state_schema, - ) - else: - # For TypedDict schemas - class CallModelInputSchema(state_schema): # type: ignore - llm_input_messages: list[AnyMessage] - - input_schema = CallModelInputSchema - else: - input_schema = state_schema - - def generate_structured_response( - state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig - ) -> StateSchema: - if is_async_dynamic_model: - msg = ( - "Async model callable provided but agent invoked synchronously. " - "Use agent.ainvoke() or agent.astream(), or provide a sync model callable." - ) - raise RuntimeError(msg) - - messages = _get_state_value(state, "messages") - structured_response_schema = response_format - if isinstance(response_format, tuple): - system_prompt, structured_response_schema = response_format - messages = [SystemMessage(content=system_prompt)] + list(messages) - - resolved_model = _resolve_model(state, runtime) - model_with_structured_output = _get_model( - resolved_model - ).with_structured_output( - cast(StructuredResponseSchema, structured_response_schema) - ) - response = model_with_structured_output.invoke(messages, config) - return {"structured_response": response} - - async def agenerate_structured_response( - state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig - ) -> StateSchema: - messages = _get_state_value(state, "messages") - structured_response_schema = response_format - if isinstance(response_format, tuple): - system_prompt, structured_response_schema = response_format - messages = [SystemMessage(content=system_prompt)] + list(messages) - - resolved_model = await _aresolve_model(state, runtime) - model_with_structured_output = _get_model( - resolved_model - ).with_structured_output( - cast(StructuredResponseSchema, structured_response_schema) - ) - response = await model_with_structured_output.ainvoke(messages, config) - return {"structured_response": response} - - if not tool_calling_enabled: - # Define a new graph - workflow = StateGraph(state_schema=state_schema, context_schema=context_schema) - workflow.add_node( - "agent", - RunnableCallable(call_model, acall_model), - input_schema=input_schema, - ) - if pre_model_hook is not None: - workflow.add_node("pre_model_hook", pre_model_hook) # type: ignore[arg-type] - workflow.add_edge("pre_model_hook", "agent") - entrypoint = "pre_model_hook" - else: - entrypoint = "agent" - - workflow.set_entry_point(entrypoint) - - if post_model_hook is not None: - workflow.add_node("post_model_hook", post_model_hook) # type: ignore[arg-type] - workflow.add_edge("agent", "post_model_hook") - - if response_format is not None: - workflow.add_node( - "generate_structured_response", - RunnableCallable( - generate_structured_response, - agenerate_structured_response, - ), - ) - if post_model_hook is not None: - workflow.add_edge("post_model_hook", "generate_structured_response") - else: - workflow.add_edge("agent", "generate_structured_response") - - return workflow.compile( - checkpointer=checkpointer, - store=store, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - debug=debug, - name=name, - ) - - # Define the function that determines whether to continue or not - def should_continue(state: StateSchema) -> Union[str, list[Send]]: - messages = _get_state_value(state, "messages") - last_message = messages[-1] - # If there is no function call, then we finish - if not isinstance(last_message, AIMessage) or not last_message.tool_calls: - if post_model_hook is not None: - return "post_model_hook" - elif response_format is not None: - return "generate_structured_response" - else: - return END - # Otherwise if there is, we continue - else: - if version == "v1": - return "tools" - elif version == "v2": - if post_model_hook is not None: - return "post_model_hook" - return [ - Send( - "tools", - ToolCallWithContext( - __type="tool_call_with_context", - tool_call=tool_call, - state=state, - ), - ) - for tool_call in last_message.tool_calls - ] - - # Define a new graph - workflow = StateGraph( - state_schema=state_schema or AgentState, context_schema=context_schema - ) - - # Define the two nodes we will cycle between - workflow.add_node( - "agent", - RunnableCallable(call_model, acall_model), - input_schema=input_schema, - ) - workflow.add_node("tools", tool_node) - - # Optionally add a pre-model hook node that will be called - # every time before the "agent" (LLM-calling node) - if pre_model_hook is not None: - workflow.add_node("pre_model_hook", pre_model_hook) # type: ignore[arg-type] - workflow.add_edge("pre_model_hook", "agent") - entrypoint = "pre_model_hook" - else: - entrypoint = "agent" - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point(entrypoint) - - agent_paths = [] - post_model_hook_paths = [entrypoint, "tools"] - - # Add a post model hook node if post_model_hook is provided - if post_model_hook is not None: - workflow.add_node("post_model_hook", post_model_hook) # type: ignore[arg-type] - agent_paths.append("post_model_hook") - workflow.add_edge("agent", "post_model_hook") - else: - agent_paths.append("tools") - - # Add a structured output node if response_format is provided - if response_format is not None: - workflow.add_node( - "generate_structured_response", - RunnableCallable( - generate_structured_response, - agenerate_structured_response, - ), - ) - if post_model_hook is not None: - post_model_hook_paths.append("generate_structured_response") - else: - agent_paths.append("generate_structured_response") - else: - if post_model_hook is not None: - post_model_hook_paths.append(END) - else: - agent_paths.append(END) - - if post_model_hook is not None: - - def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]: - """Route to the next node after post_model_hook. - - Routes to one of: - * "tools": if there are pending tool calls without a corresponding message. - * "generate_structured_response": if no pending tool calls exist and response_format is specified. - * END: if no pending tool calls exist and no response_format is specified. - """ - - messages = _get_state_value(state, "messages") - tool_messages = [ - m.tool_call_id for m in messages if isinstance(m, ToolMessage) - ] - last_ai_message = next( - m for m in reversed(messages) if isinstance(m, AIMessage) - ) - pending_tool_calls = [ - c for c in last_ai_message.tool_calls if c["id"] not in tool_messages - ] - - if pending_tool_calls: - return [ - Send( - "tools", - ToolCallWithContext( - __type="tool_call_with_context", - tool_call=tool_call, - state=state, - ), - ) - for tool_call in pending_tool_calls - ] - elif isinstance(messages[-1], ToolMessage): - return entrypoint - elif response_format is not None: - return "generate_structured_response" - else: - return END - - workflow.add_conditional_edges( - "post_model_hook", - post_model_hook_router, # type: ignore[arg-type] - path_map=post_model_hook_paths, - ) - - workflow.add_conditional_edges( - "agent", - should_continue, # type: ignore[arg-type] - path_map=agent_paths, - ) - - def route_tool_responses(state: StateSchema) -> str: - for m in reversed(_get_state_value(state, "messages")): - if not isinstance(m, ToolMessage): - break - if m.name in should_return_direct: - return END - - # handle a case of parallel tool calls where - # the tool w/ `return_direct` was executed in a different `Send` - if isinstance(m, AIMessage) and m.tool_calls: - if any(call["name"] in should_return_direct for call in m.tool_calls): - return END - - return entrypoint - - if should_return_direct: - workflow.add_conditional_edges( - "tools", route_tool_responses, path_map=[entrypoint, END] - ) - else: - workflow.add_edge("tools", entrypoint) - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - return workflow.compile( + # Create builder and delegate all logic to it + builder = _AgentBuilder( + model=model, + tools=tools, + prompt=prompt, + response_format=response_format, + pre_model_hook=pre_model_hook, + post_model_hook=post_model_hook, + state_schema=state_schema, + context_schema=context_schema, checkpointer=checkpointer, store=store, interrupt_before=interrupt_before, interrupt_after=interrupt_after, debug=debug, + version=version, name=name, + **deprecated_kwargs, ) + return builder.build() + # Keep for backwards compatibility create_tool_calling_executor = create_react_agent