fix: Handle missing import for RunnableConfig when generating schema for LangChain templates

PiperOrigin-RevId: 627419996

Author: yeesian

setup.py

@@ -140,7 +140,7 @@
 
 reasoning_engine_extra_require = [
     "cloudpickle >= 2.2.1, < 3.0",
-    "pydantic < 3",
+    "pydantic >= 2.6.3, < 3",
 ]
 
 rapid_evaluation_extra_require = [
@@ -154,9 +154,13 @@
     "langchain-google-vertexai < 0.2",
 ]
 
-langchain_testing_extra_require = langchain_extra_require + [
-    "pytest-xdist",
-]
+langchain_testing_extra_require = list(
+    set(
+        langchain_extra_require
+        + reasoning_engine_extra_require
+        + ["absl-py", "pytest-xdist"]
+    )
+)
 
 full_extra_require = list(
     set(
@@ -174,7 +178,6 @@
         + autologging_extra_require
         + preview_extra_require
         + ray_extra_require
-        + reasoning_engine_extra_require
         + rapid_evaluation_extra_require
     )
 )

tests/unit/vertexai/test_reasoning_engines.py tests/unit/vertex_langchain/test_reasoning_engines.py

@@ -24,6 +24,15 @@
 from google.protobuf import struct_pb2
 from google.protobuf import json_format
 
+try:
+    # For LangChain templates, they might not import langchain_core and get
+    #   PydanticUserError: `query` is not fully defined; you should define
+    #   `RunnableConfig`, then call `query.model_rebuild()`.
+    import langchain_core.runnables.config
+    RunnableConfig = langchain_core.runnables.config.RunnableConfig
+except ImportError:
+    RunnableConfig = Any
+
 JsonDict = Dict[str, Any]
 
 
@@ -214,11 +223,10 @@ def _import_cloudpickle_or_raise() -> types.ModuleType:
 def _import_pydantic_or_raise() -> types.ModuleType:
     """Tries to import the pydantic module."""
     try:
-        # For compatibility across Pydantic V1 and V2
-        try:
-            from pydantic import v1 as pydantic
-        except ImportError:
-            import pydantic
+        import pydantic
+        _ = pydantic.Field
+    except AttributeError:
+        from pydantic import v1 as pydantic
     except ImportError as e:
         raise ImportError(
             "pydantic is not installed. Please call "