update

langchain-ai · Sep 7, 2024 · aa695be · aa695be
1 parent 1fd8a93
commit aa695be
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py
@@ -145,17 +145,19 @@ def test_delete_all():
     assert result == expected_result
 
 
-class MessagesStatePydanticV1(BaseModelV1):
-    messages: Annotated[list[AnyMessage], add_messages]
+MESSAGES_STATE_SCHEMAS = [MessagesState]
+if IS_LANGCHAIN_CORE_030_OR_GREATER:
 
+    class MessagesStatePydantic(BaseModel):
+        messages: Annotated[list[AnyMessage], add_messages]
 
-class MessagesStatePydantic(BaseModel):
-    messages: Annotated[list[AnyMessage], add_messages]
+    MESSAGES_STATE_SCHEMAS.append(MessagesStatePydantic)
+else:
 
+    class MessagesStatePydanticV1(BaseModelV1):
+        messages: Annotated[list[AnyMessage], add_messages]
 
-MESSAGES_STATE_SCHEMAS = [MessagesState, MessagesStatePydanticV1]
-if IS_LANGCHAIN_CORE_030_OR_GREATER:
-    MESSAGES_STATE_SCHEMAS.append(MessagesStatePydantic)
+    MESSAGES_STATE_SCHEMAS.append(MessagesStatePydanticV1)
 
 
 @pytest.mark.parametrize("state_schema", MESSAGES_STATE_SCHEMAS)