diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 4a73f1a08..2804990b6 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -78,13 +78,13 @@ class StateGraph(Graph): node_inboxes = { # we take any value written to channel because all writers # write the entire state as of that step, which is equal for all - f"{key}:inbox": AnyValue(Any) - for key in self.nodes + f"{key}:inbox": AnyValue(self.schema) + for key in list(self.nodes) + [START] } node_outboxes = { # we clear outbox channels after each step key: EphemeralValue(Any) - for key in self.nodes + for key in list(self.nodes) + [START] } for key in self.nodes: @@ -114,7 +114,12 @@ class StateGraph(Graph): return CompiledGraph( graph=self, nodes=nodes, - channels={**self.channels, **node_inboxes, **node_outboxes}, + channels={ + **self.channels, + **node_inboxes, + **node_outboxes, + END: LastValue(self.schema), + }, input=f"{START}:inbox", output=END, hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys, diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 249ae2c92..ac713e9ec 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -113,9 +113,12 @@ ''' # --- # name: test_conditional_graph_state - '{"title": "LangGraphOutput"}' + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' # --- # name: test_conditional_graph_state.1 + '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' +# --- +# name: test_conditional_graph_state.2 ''' { "nodes": [ @@ -123,14 +126,228 @@ "id": "__start__", "type": "schema", "data": { - "title": "LangGraphInput" + "title": "LangGraphInput", + "$ref": "#/definitions/AgentState", + "definitions": { + "AgentAction": { + "title": "AgentAction", + "description": "A full description of an action for an ActionAgent to execute.", + "type": "object", + "properties": { + "tool": { + "title": "Tool", + "type": "string" + }, + "tool_input": { + "title": "Tool Input", + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + }, + "log": { + "title": "Log", + "type": "string" + }, + "type": { + "title": "Type", + "default": "AgentAction", + "enum": [ + "AgentAction" + ], + "type": "string" + } + }, + "required": [ + "tool", + "tool_input", + "log" + ] + }, + "AgentFinish": { + "title": "AgentFinish", + "description": "The final return value of an ActionAgent.", + "type": "object", + "properties": { + "return_values": { + "title": "Return Values", + "type": "object" + }, + "log": { + "title": "Log", + "type": "string" + }, + "type": { + "title": "Type", + "default": "AgentFinish", + "enum": [ + "AgentFinish" + ], + "type": "string" + } + }, + "required": [ + "return_values", + "log" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "agent_outcome": { + "title": "Agent Outcome", + "anyOf": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "$ref": "#/definitions/AgentFinish" + } + ] + }, + "intermediate_steps": { + "title": "Intermediate Steps", + "type": "array", + "items": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "type": "string" + } + ] + } + } + } + } + } } }, { "id": "__end__", "type": "schema", "data": { - "title": "LangGraphOutput" + "title": "LangGraphOutput", + "$ref": "#/definitions/AgentState", + "definitions": { + "AgentAction": { + "title": "AgentAction", + "description": "A full description of an action for an ActionAgent to execute.", + "type": "object", + "properties": { + "tool": { + "title": "Tool", + "type": "string" + }, + "tool_input": { + "title": "Tool Input", + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + }, + "log": { + "title": "Log", + "type": "string" + }, + "type": { + "title": "Type", + "default": "AgentAction", + "enum": [ + "AgentAction" + ], + "type": "string" + } + }, + "required": [ + "tool", + "tool_input", + "log" + ] + }, + "AgentFinish": { + "title": "AgentFinish", + "description": "The final return value of an ActionAgent.", + "type": "object", + "properties": { + "return_values": { + "title": "Return Values", + "type": "object" + }, + "log": { + "title": "Log", + "type": "string" + }, + "type": { + "title": "Type", + "default": "AgentFinish", + "enum": [ + "AgentFinish" + ], + "type": "string" + } + }, + "required": [ + "return_values", + "log" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "agent_outcome": { + "title": "Agent Outcome", + "anyOf": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "$ref": "#/definitions/AgentFinish" + } + ] + }, + "intermediate_steps": { + "title": "Intermediate Steps", + "type": "array", + "items": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "type": "string" + } + ] + } + } + } + } + } } }, { @@ -203,7 +420,7 @@ } ''' # --- -# name: test_conditional_graph_state.2 +# name: test_conditional_graph_state.3 ''' +-----------+ | __start__ | @@ -228,3 +445,1266 @@ +---------+ +-------+ ''' # --- +# name: test_message_graph + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph.1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "description": "Message for priming AI behavior, usually passed in as the first of a sequence\\nof input messages.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "description": "Message for passing the result of executing a function back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "description": "Message for passing the result of executing a tool back to a model.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput", + "type": "array", + "items": { + "anyOf": [ + { + "$ref": "#/definitions/AIMessage" + }, + { + "$ref": "#/definitions/HumanMessage" + }, + { + "$ref": "#/definitions/ChatMessage" + }, + { + "$ref": "#/definitions/SystemMessage" + }, + { + "$ref": "#/definitions/FunctionMessage" + }, + { + "$ref": "#/definitions/ToolMessage" + } + ] + }, + "definitions": { + "AIMessage": { + "title": "AIMessage", + "description": "Message from an AI.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "ai", + "enum": [ + "ai" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "example": { + "title": "Example", + "default": false, + "type": "boolean" + } + }, + "required": [ + "content" + ] + }, + "HumanMessage": { + "title": "HumanMessage", + "description": "Message from a human.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "human", + "enum": [ + "human" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "example": { + "title": "Example", + "default": false, + "type": "boolean" + } + }, + "required": [ + "content" + ] + }, + "ChatMessage": { + "title": "ChatMessage", + "description": "Message that can be assigned an arbitrary speaker (i.e. role).", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "chat", + "enum": [ + "chat" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "role": { + "title": "Role", + "type": "string" + } + }, + "required": [ + "content", + "role" + ] + }, + "SystemMessage": { + "title": "SystemMessage", + "description": "Message for priming AI behavior, usually passed in as the first of a sequence\nof input messages.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "system", + "enum": [ + "system" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content" + ] + }, + "FunctionMessage": { + "title": "FunctionMessage", + "description": "Message for passing the result of executing a function back to a model.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "function", + "enum": [ + "function" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "name" + ] + }, + "ToolMessage": { + "title": "ToolMessage", + "description": "Message for passing the result of executing a tool back to a model.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "tool", + "enum": [ + "tool" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "tool_call_id": { + "title": "Tool Call Id", + "type": "string" + } + }, + "required": [ + "content", + "tool_call_id" + ] + } + } + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput", + "type": "array", + "items": { + "anyOf": [ + { + "$ref": "#/definitions/AIMessage" + }, + { + "$ref": "#/definitions/HumanMessage" + }, + { + "$ref": "#/definitions/ChatMessage" + }, + { + "$ref": "#/definitions/SystemMessage" + }, + { + "$ref": "#/definitions/FunctionMessage" + }, + { + "$ref": "#/definitions/ToolMessage" + } + ] + }, + "definitions": { + "AIMessage": { + "title": "AIMessage", + "description": "Message from an AI.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "ai", + "enum": [ + "ai" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "example": { + "title": "Example", + "default": false, + "type": "boolean" + } + }, + "required": [ + "content" + ] + }, + "HumanMessage": { + "title": "HumanMessage", + "description": "Message from a human.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "human", + "enum": [ + "human" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "example": { + "title": "Example", + "default": false, + "type": "boolean" + } + }, + "required": [ + "content" + ] + }, + "ChatMessage": { + "title": "ChatMessage", + "description": "Message that can be assigned an arbitrary speaker (i.e. role).", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "chat", + "enum": [ + "chat" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "role": { + "title": "Role", + "type": "string" + } + }, + "required": [ + "content", + "role" + ] + }, + "SystemMessage": { + "title": "SystemMessage", + "description": "Message for priming AI behavior, usually passed in as the first of a sequence\nof input messages.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "system", + "enum": [ + "system" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content" + ] + }, + "FunctionMessage": { + "title": "FunctionMessage", + "description": "Message for passing the result of executing a function back to a model.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "function", + "enum": [ + "function" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "name" + ] + }, + "ToolMessage": { + "title": "ToolMessage", + "description": "Message for passing the result of executing a tool back to a model.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "default": "tool", + "enum": [ + "tool" + ], + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "tool_call_id": { + "title": "Tool Call Id", + "type": "string" + } + }, + "required": [ + "content", + "tool_call_id" + ] + } + } + } + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "test_pregel", + "FakeFuntionChatModel" + ], + "name": "FakeFuntionChatModel" + } + }, + { + "id": "action", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "call_tool" + } + }, + { + "id": "agent_should_continue", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_continue" + } + } + ], + "edges": [ + { + "source": "action", + "target": "agent", + "data": null + }, + { + "source": "agent", + "target": "agent_should_continue", + "data": null + }, + { + "source": "agent_should_continue", + "target": "action", + "data": "continue" + }, + { + "source": "agent_should_continue", + "target": "__end__", + "data": "end" + }, + { + "source": "__start__", + "target": "agent", + "data": null + } + ] + } + ''' +# --- +# name: test_message_graph.3 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + *+-------+* + ** *** + ** ** + ** ** + +-----------------------+ ** + | agent_should_continue | * + +-----------------------+ * + * ***** * + * **** * + * *** * + +---------+ +--------+ + | __end__ | | action | + +---------+ +--------+ + ''' +# --- +# name: test_prebuilt_chat + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' +# --- +# name: test_prebuilt_chat.1 + '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' +# --- +# name: test_prebuilt_chat.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput", + "$ref": "#/definitions/AgentState", + "definitions": { + "BaseMessage": { + "title": "BaseMessage", + "description": "Base abstract Message class.\n\nMessages are the inputs and outputs of ChatModels.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "type" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, + "required": [ + "messages" + ] + } + } + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput", + "$ref": "#/definitions/AgentState", + "definitions": { + "BaseMessage": { + "title": "BaseMessage", + "description": "Base abstract Message class.\n\nMessages are the inputs and outputs of ChatModels.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "type" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, + "required": [ + "messages" + ] + } + } + } + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "call_model" + } + }, + { + "id": "action", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "call_tool" + } + }, + { + "id": "agent_should_continue", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_continue" + } + } + ], + "edges": [ + { + "source": "action", + "target": "agent", + "data": null + }, + { + "source": "agent", + "target": "agent_should_continue", + "data": null + }, + { + "source": "agent_should_continue", + "target": "action", + "data": "continue" + }, + { + "source": "agent_should_continue", + "target": "__end__", + "data": "end" + }, + { + "source": "__start__", + "target": "agent", + "data": null + } + ] + } + ''' +# --- +# name: test_prebuilt_chat.3 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + *+-------+* + ** *** + ** ** + ** ** + +-----------------------+ ** + | agent_should_continue | * + +-----------------------+ * + * ***** * + * **** * + * *** * + +---------+ +--------+ + | __end__ | | action | + +---------+ +--------+ + ''' +# --- +# name: test_prebuilt_tool_chat + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' +# --- +# name: test_prebuilt_tool_chat.1 + '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' +# --- +# name: test_prebuilt_tool_chat.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput", + "$ref": "#/definitions/AgentState", + "definitions": { + "BaseMessage": { + "title": "BaseMessage", + "description": "Base abstract Message class.\n\nMessages are the inputs and outputs of ChatModels.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "type" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, + "required": [ + "messages" + ] + } + } + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput", + "$ref": "#/definitions/AgentState", + "definitions": { + "BaseMessage": { + "title": "BaseMessage", + "description": "Base abstract Message class.\n\nMessages are the inputs and outputs of ChatModels.", + "type": "object", + "properties": { + "content": { + "title": "Content", + "anyOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "object" + } + ] + } + } + ] + }, + "additional_kwargs": { + "title": "Additional Kwargs", + "type": "object" + }, + "type": { + "title": "Type", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "content", + "type" + ] + }, + "AgentState": { + "title": "AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, + "required": [ + "messages" + ] + } + } + } + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "call_model" + } + }, + { + "id": "action", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "call_tool" + } + }, + { + "id": "agent_should_continue", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_continue" + } + } + ], + "edges": [ + { + "source": "action", + "target": "agent", + "data": null + }, + { + "source": "agent", + "target": "agent_should_continue", + "data": null + }, + { + "source": "agent_should_continue", + "target": "action", + "data": "continue" + }, + { + "source": "agent_should_continue", + "target": "__end__", + "data": "end" + }, + { + "source": "__start__", + "target": "agent", + "data": null + } + ] + } + ''' +# --- +# name: test_prebuilt_tool_chat.3 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + *+-------+* + ** *** + ** ** + ** ** + +-----------------------+ ** + | agent_should_continue | * + +-----------------------+ * + * ***** * + * **** * + * *** * + +---------+ +--------+ + | __end__ | | action | + +---------+ +--------+ + ''' +# --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index c79d2cf3d..33788d9d5 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1074,7 +1074,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: ] -def test_prebuilt_tool_chat() -> None: +def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.messages import AIMessage, HumanMessage, ToolMessage @@ -1137,6 +1137,11 @@ def test_prebuilt_tool_chat() -> None: tools, ) + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.invoke( {"messages": [HumanMessage(content="what is weather in sf")]} ) == { @@ -1318,7 +1323,7 @@ def test_prebuilt_tool_chat() -> None: ] -def test_prebuilt_chat() -> None: +def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage @@ -1361,6 +1366,11 @@ def test_prebuilt_chat() -> None: tools, ) + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.invoke( {"messages": [HumanMessage(content="what is weather in sf")]} ) == { @@ -1463,7 +1473,7 @@ def test_prebuilt_chat() -> None: ] -def test_message_graph() -> None: +def test_message_graph(snapshot: SnapshotAssertion) -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.agents import AgentAction @@ -1574,6 +1584,11 @@ def test_message_graph() -> None: # meaning you can use it as you would any other runnable app = workflow.compile() + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.invoke(HumanMessage(content="what is weather in sf")) == [ HumanMessage(content="what is weather in sf"), AIMessage(