x

libs/partners/robocorp/tests/unit_tests/test_toolkits.py

@@ -12,6 +12,10 @@ from langchain_core.utils.function_calling import (
 from langchain_robocorp.toolkits import ActionServerToolkit
 
 from ._fixtures import FakeChatLLMT
+from pydantic import __version__
+
+# Get the version tuple (it may contain a suffix like 'rc1')
+VERSION_TUPLE = tuple(int(x) if x.isdigit() else x for x in __version__.split("."))
 
 
 def test_initialization() -> None:
@@ -91,6 +95,8 @@ Strictly adhere to the schema."""
         assert set(params["properties"].keys()) == {"sheet", "rows_to_add"}
 
         desc = "The columns that make up the row"
+
+        if VERSION_TUPLE < (2, 9):
             expected = {
                 "description": "the rows to be added to the sheet",
                 "allOf": [
@@ -121,6 +127,28 @@ Strictly adhere to the schema."""
                     }
                 ],
             }
+        else:
+            expected = {
+                "properties": {
+                    "rows": {
+                        "description": "The rows that need to be added",
+                        "items": {
+                            "properties": {
+                                "columns": {
+                                    "description": "The columns that make up the row",
+                                    "items": {"type": "string"},
+                                    "type": "array",
+                                }
+                            },
+                            "required": ["columns"],
+                            "type": "object",
+                        },
+                        "type": "array",
+                    }
+                },
+                "required": ["rows"],
+                "type": "object",
+            }
         assert params["properties"]["rows_to_add"] == expected