From 557f85d7fdc2c0abbf56fc7ba8f955ae024472f8 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Fri, 22 Nov 2024 16:21:07 -0500
Subject: [PATCH] x

---
 .../runnables/__snapshots__/test_graph.ambr       | 15 +++++++++++++++
 .../core/tests/unit_tests/runnables/test_graph.py | 14 ++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr b/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr
index 208bdbac45e..5a3b7126d99 100644
--- a/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr
+++ b/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr
@@ -468,6 +468,7 @@
                   'const': 'ai',
                   'default': 'ai',
                   'title': 'Type',
+                  'type': 'string',
                 }),
                 'usage_metadata': dict({
                   'anyOf': list([
@@ -580,6 +581,7 @@
                   'const': 'AIMessageChunk',
                   'default': 'AIMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
                 'usage_metadata': dict({
                   'anyOf': list([
@@ -664,6 +666,7 @@
                   'const': 'chat',
                   'default': 'chat',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -738,6 +741,7 @@
                   'const': 'ChatMessageChunk',
                   'default': 'ChatMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -809,6 +813,7 @@
                   'const': 'function',
                   'default': 'function',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -871,6 +876,7 @@
                   'const': 'FunctionMessageChunk',
                   'default': 'FunctionMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -969,6 +975,7 @@
                   'const': 'human',
                   'default': 'human',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1043,6 +1050,7 @@
                   'const': 'HumanMessageChunk',
                   'default': 'HumanMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1141,6 +1149,7 @@
                 'type': dict({
                   'const': 'invalid_tool_call',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1266,6 +1275,7 @@
                   'const': 'system',
                   'default': 'system',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1335,6 +1345,7 @@
                   'const': 'SystemMessageChunk',
                   'default': 'SystemMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1383,6 +1394,7 @@
                 'type': dict({
                   'const': 'tool_call',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1461,6 +1473,7 @@
                 'type': dict({
                   'const': 'tool_call_chunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1583,6 +1596,7 @@
                   'const': 'tool',
                   'default': 'tool',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
@@ -1664,6 +1678,7 @@
                   'const': 'ToolMessageChunk',
                   'default': 'ToolMessageChunk',
                   'title': 'Type',
+                  'type': 'string',
                 }),
               }),
               'required': list([
diff --git a/libs/core/tests/unit_tests/runnables/test_graph.py b/libs/core/tests/unit_tests/runnables/test_graph.py
index 39f2f2871a8..789875db623 100644
--- a/libs/core/tests/unit_tests/runnables/test_graph.py
+++ b/libs/core/tests/unit_tests/runnables/test_graph.py
@@ -12,6 +12,7 @@ from langchain_core.prompts.prompt import PromptTemplate
 from langchain_core.runnables.base import Runnable, RunnableConfig
 from langchain_core.runnables.graph import Edge, Graph, Node
 from langchain_core.runnables.graph_mermaid import _escape_node_label
+from langchain_core.utils.pydantic import PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION
 from tests.unit_tests.pydantic_utils import _normalize_schema
 
 
@@ -210,10 +211,15 @@ def test_graph_sequence_map(snapshot: SnapshotAssertion) -> None:
         }
     )
     graph = sequence.get_graph()
-    assert _normalize_schema(graph.to_json(with_schemas=True)) == snapshot(
-        name="graph_with_schema"
-    )
-    assert _normalize_schema(graph.to_json()) == snapshot(name="graph_no_schemas")
+
+    if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) >= (2, 10):
+        assert _normalize_schema(graph.to_json(with_schemas=True)) == snapshot(
+            name="graph_with_schema"
+        )
+
+    if (PYDANTIC_MAJOR_VERSION, PYDANTIC_MINOR_VERSION) >= (2, 10):
+        assert _normalize_schema(graph.to_json()) == snapshot(name="graph_no_schemas")
+
     assert graph.draw_ascii() == snapshot(name="ascii")
     assert graph.draw_mermaid() == snapshot(name="mermaid")
     assert graph.draw_mermaid(with_styles=False) == snapshot(name="mermaid-simple")