From 8a7c6b4fa7996a81f78462522d9820bfd7297f0c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 19 Aug 2024 12:05:36 -0700 Subject: [PATCH] Add xray test --- .../tests/__snapshots__/test_pregel.ambr | 27 +++++++++++++ libs/langgraph/tests/test_pregel.py | 39 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 004daa1e1..644c99b95 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -1356,6 +1356,33 @@ # name: test_state_graph_w_config_inherited_state_keys.2 '{"title": "LangGraphOutput", "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"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "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": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "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"]}}}' # --- +# name: test_xray_issue + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + p_one(p_one) + p_two___start__(__start__) + p_two_c_one(c_one) + p_two_c_two(c_two) + p_two___end__(__end__) + __end__([__end__]):::last + subgraph p_two + p_two___start__ --> p_two_c_one; + p_two_c_two --> p_two_c_one; + p_two_c_one -.  0  .-> p_two_c_two; + p_two_c_one -.  1  .-> p_two___end__; + end + __start__ --> p_one; + p_two___end__ --> p_one; + p_one -.  0  .-> p_two___start__; + p_one -.  1  .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_xray_lance dict({ 'edges': list([ diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 68bbc65ce..14a96b577 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5,6 +5,7 @@ import warnings from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from random import randrange from typing import ( Annotated, Any, @@ -9682,3 +9683,41 @@ def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) ) app.invoke({"input": 1, "ephemeral": "meow"}, config) assert checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1} + + +def test_xray_issue(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: Annotated[list, add_messages] + + def node(name): + def _node(state: State): + return {"messages": [("human", f"entered {name} node")]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + + child.add_edge("__start__", "c_one") + child.add_edge("c_two", "c_one") + + child.add_conditional_edges( + "c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"} + ) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + + parent.add_edge("__start__", "p_one") + parent.add_edge("p_two", "p_one") + + parent.add_conditional_edges( + "p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"} + ) + + app = parent.compile() + + assert app.get_graph(xray=True).draw_mermaid() == snapshot