Correct Pydantic warning(s) issed for llama-index-llms-ibm (#16141)

Fix Pydantic warnings in llam-index-llms-ibm
run-llama · Sep 21, 2024 · 9083c6d · 9083c6d
1 parent da8e444
commit 9083c6d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 18 deletions.
diff --git a/llama-index-integrations/llms/llama-index-llms-ibm/llama_index/llms/ibm/base.py b/llama-index-integrations/llms/llama-index-llms-ibm/llama_index/llms/ibm/base.py
@@ -14,6 +14,7 @@
 )
 from llama_index.core.constants import DEFAULT_CONTEXT_WINDOW
 from llama_index.core.bridge.pydantic import (
+    ConfigDict,
     Field,
     PrivateAttr,
 )
@@ -57,10 +58,10 @@ class WatsonxLLM(CustomLLM):
     """
 
     model_id: Optional[str] = Field(
-        default=None, description="Type of model to use.", allow_mutation=False
+        default=None, description="Type of model to use.", frozen=True
     )
     deployment_id: Optional[str] = Field(
-        default=None, description="Id of deployed model to use.", allow_mutation=False
+        default=None, description="Id of deployed model to use.", frozen=True
     )
 
     temperature: Optional[float] = Field(
@@ -79,43 +80,43 @@ class WatsonxLLM(CustomLLM):
     project_id: Optional[str] = Field(
         default=None,
         description="ID of the Watson Studio project.",
-        allow_mutation=False,
+        frozen=True,
     )
 
     space_id: Optional[str] = Field(
-        default=None, description="ID of the Watson Studio space.", allow_mutation=False
+        default=None, description="ID of the Watson Studio space.", frozen=True
     )
 
     url: Optional[SecretStr] = Field(
         default=None,
         description="Url to Watson Machine Learning or CPD instance",
-        allow_mutation=False,
+        frozen=True,
     )
 
     apikey: Optional[SecretStr] = Field(
         default=None,
         description="Apikey to Watson Machine Learning or CPD instance",
-        allow_mutation=False,
+        frozen=True,
     )
 
     token: Optional[SecretStr] = Field(
-        default=None, description="Token to CPD instance", allow_mutation=False
+        default=None, description="Token to CPD instance", frozen=True
     )
 
     password: Optional[SecretStr] = Field(
-        default=None, description="Password to CPD instance", allow_mutation=False
+        default=None, description="Password to CPD instance", frozen=True
     )
 
     username: Optional[SecretStr] = Field(
-        default=None, description="Username to CPD instance", allow_mutation=False
+        default=None, description="Username to CPD instance", frozen=True
     )
 
     instance_id: Optional[SecretStr] = Field(
-        default=None, description="Instance_id of CPD instance", allow_mutation=False
+        default=None, description="Instance_id of CPD instance", frozen=True
     )
 
     version: Optional[SecretStr] = Field(
-        default=None, description="Version of CPD instance", allow_mutation=False
+        default=None, description="Version of CPD instance", frozen=True
     )
 
     verify: Union[str, bool, None] = Field(
@@ -127,11 +128,11 @@ class WatsonxLLM(CustomLLM):
         True - default path to truststore will be taken
         False - no verification will be made
         """,
-        allow_mutation=False,
+        frozen=True,
     )
 
     validate_model: bool = Field(
-        default=True, description="Model id validation", allow_mutation=False
+        default=True, description="Model id validation", frozen=True
     )
 
     _model: ModelInference = PrivateAttr()
@@ -244,8 +245,7 @@ def __init__(
         self._model_info = None
         self._deployment_info = None
 
-    class Config:
-        validate_assignment = True
+    model_config = ConfigDict(protected_namespaces=(), validate_assignment=True)
 
     @property
     def model_info(self):

diff --git a/llama-index-integrations/llms/llama-index-llms-ibm/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-ibm/pyproject.toml
@@ -31,7 +31,7 @@ license = "MIT"
 name = "llama-index-llms-ibm"
 packages = [{include = "llama_index/"}]
 readme = "README.md"
-version = "0.2.0"
+version = "0.2.1"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"

diff --git a/llama-index-integrations/llms/llama-index-llms-ibm/tests/test_ibm.py b/llama-index-integrations/llms/llama-index-llms-ibm/tests/test_ibm.py
@@ -132,15 +132,15 @@ def test_initialization(self) -> None:
             ValueError, match=r"^Did not find 'apikey' or 'token',"
         ) as e_info:
             _ = WatsonxLLM(
-                model=self.TEST_MODEL,
+                model_id=self.TEST_MODEL,
                 url=self.TEST_URL,
                 project_id=self.TEST_PROJECT_ID,
             )
 
         # CPD scenario
         with pytest.raises(ValueError, match=r"^Did not find instance_id") as e_info:
             _ = WatsonxLLM(
-                model=self.TEST_MODEL,
+                model_id=self.TEST_MODEL,
                 token="123",
                 url="cpd-instance",
                 project_id=self.TEST_PROJECT_ID,