multiple: pydantic 2 compatibility, v0.3 (#26443)

Signed-off-by: ChengZi <chen.zhang@zilliz.com>
Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
Co-authored-by: Dan O'Donovan <dan.odonovan@gmail.com>
Co-authored-by: Tom Daniel Grande <tomdgrande@gmail.com>
Co-authored-by: Grande <Tom.Daniel.Grande@statsbygg.no>
Co-authored-by: Bagatur <baskaryan@gmail.com>
Co-authored-by: ccurme <chester.curme@gmail.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
Co-authored-by: Tomaz Bratanic <bratanic.tomaz@gmail.com>
Co-authored-by: ZhangShenao <15201440436@163.com>
Co-authored-by: Friso H. Kingma <fhkingma@gmail.com>
Co-authored-by: ChengZi <chen.zhang@zilliz.com>
Co-authored-by: Nuno Campos <nuno@langchain.dev>
Co-authored-by: Morgante Pell <morgantep@google.com>

libs/community/langchain_community/llms/databricks.py

@@ -7,12 +7,12 @@ from typing import Any, Callable, Dict, List, Mapping, Optional
 import requests
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.language_models import LLM
-from langchain_core.pydantic_v1 import (
+from pydantic import (
     BaseModel,
+    ConfigDict,
     Field,
     PrivateAttr,
-    root_validator,
-    validator,
+    model_validator,
 )
 
 __all__ = ["Databricks"]
@@ -97,8 +97,9 @@ class _DatabricksServingEndpointClient(_DatabricksClientBase):
     def llm(self) -> bool:
         return self.task in ("llm/v1/chat", "llm/v1/completions", "llama2/chat")
 
-    @root_validator(pre=True)
-    def set_api_url(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    def set_api_url(cls, values: Dict[str, Any]) -> Any:
         if "api_url" not in values:
             host = values["host"]
             endpoint_name = values["endpoint_name"]
@@ -141,8 +142,9 @@ class _DatabricksClusterDriverProxyClient(_DatabricksClientBase):
     cluster_id: str
     cluster_driver_port: str
 
-    @root_validator(pre=True)
-    def set_api_url(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    def set_api_url(cls, values: Dict[str, Any]) -> Any:
         if "api_url" not in values:
             host = values["host"]
             cluster_id = values["cluster_id"]
@@ -395,9 +397,9 @@ class Databricks(LLM):
 
     _client: _DatabricksClientBase = PrivateAttr()
 
-    class Config:
-        extra = "forbid"
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     @property
     def _llm_params(self) -> Dict[str, Any]:
@@ -411,18 +413,21 @@ class Databricks(LLM):
             params["max_tokens"] = self.max_tokens
         return params
 
-    @validator("cluster_id", always=True)
-    def set_cluster_id(cls, v: Any, values: Dict[str, Any]) -> Optional[str]:
-        if v and values["endpoint_name"]:
+    @model_validator(mode="before")
+    @classmethod
+    def set_cluster_id(cls, values: Dict[str, Any]) -> dict:
+        cluster_id = values.get("cluster_id")
+        endpoint_name = values.get("endpoint_name")
+        if cluster_id and endpoint_name:
             raise ValueError("Cannot set both endpoint_name and cluster_id.")
-        elif values["endpoint_name"]:
-            return None
-        elif v:
-            return v
+        elif endpoint_name:
+            values["cluster_id"] = None
+        elif cluster_id:
+            pass
         else:
             try:
-                if v := get_repl_context().clusterId:
-                    return v
+                if context_cluster_id := get_repl_context().clusterId:
+                    values["cluster_id"] = context_cluster_id
                 raise ValueError("Context doesn't contain clusterId.")
             except Exception as e:
                 raise ValueError(
@@ -431,27 +436,28 @@ class Databricks(LLM):
                     f" error: {e}"
                 )
 
-    @validator("cluster_driver_port", always=True)
-    def set_cluster_driver_port(cls, v: Any, values: Dict[str, Any]) -> Optional[str]:
-        if v and values["endpoint_name"]:
+        cluster_driver_port = values.get("cluster_driver_port")
+        if cluster_driver_port and endpoint_name:
             raise ValueError("Cannot set both endpoint_name and cluster_driver_port.")
-        elif values["endpoint_name"]:
-            return None
-        elif v is None:
+        elif endpoint_name:
+            values["cluster_driver_port"] = None
+        elif cluster_driver_port is None:
             raise ValueError(
                 "Must set cluster_driver_port to connect to a cluster driver."
             )
-        elif int(v) <= 0:
-            raise ValueError(f"Invalid cluster_driver_port: {v}")
+        elif int(cluster_driver_port) <= 0:
+            raise ValueError(f"Invalid cluster_driver_port: {cluster_driver_port}")
         else:
-            return v
+            pass
 
-    @validator("model_kwargs", always=True)
-    def set_model_kwargs(cls, v: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
-        if v:
-            assert "prompt" not in v, "model_kwargs must not contain key 'prompt'"
-            assert "stop" not in v, "model_kwargs must not contain key 'stop'"
-        return v
+        if model_kwargs := values.get("model_kwargs"):
+            assert (
+                "prompt" not in model_kwargs
+            ), "model_kwargs must not contain key 'prompt'"
+            assert (
+                "stop" not in model_kwargs
+            ), "model_kwargs must not contain key 'stop'"
+        return values
 
     def __init__(self, **data: Any):
         if "transform_input_fn" in data and _is_hex_string(data["transform_input_fn"]):