refactor: migrate pydantic to v2 (#182)

This PR removes `pydantic` v1 compatibility of use `import pydantic.v1` and applies need changes to `pydantic` v2 such as: - `dict()`, `parse_obj()`, `parse_file()` methods replace with `model_dump()`, `model_validate()`, `model_validate_json()` - `validate_arguments` decorator renamed to `validate_call`, `root_validator` to `model_validator` - add dependency `pydantic-settings`, new repo where settings class moved such as `BaseSettings` and updated imports - add custom validator for unique items Signed-off-by: Tiago Santana <[email protected]>
DS4SD · Aug 5, 2024 · 5b172db · 5b172db
1 parent e731d24
commit 5b172db
Show file tree

Hide file tree

Showing 25 changed files with 337 additions and 164 deletions.
diff --git a/deepsearch/chemistry/queries/molecules.py b/deepsearch/chemistry/queries/molecules.py
@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import List, Union
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.client.queries import Query
 

diff --git a/deepsearch/core/client/config.py b/deepsearch/core/client/config.py
@@ -1,6 +1,6 @@
 from typing import Union
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 
 class DeepSearchBearerTokenAuth(BaseModel):

diff --git a/deepsearch/core/client/settings.py b/deepsearch/core/client/settings.py
@@ -4,17 +4,18 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from pydantic.v1 import BaseSettings, SecretStr
+from pydantic import SecretStr
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class DumpableSettings(BaseSettings):
     @classmethod
-    def get_env_var_name(cls, attr_name) -> str:
-        return cls.Config.env_prefix + attr_name.upper()
+    def get_env_var_name(cls, attr_name: str) -> str:
+        return cls.model_config["env_prefix"] + attr_name.upper()
 
     def _get_serializable_dict(self) -> Dict[str, str]:
         result = {}
-        model_dict = self.dict()
+        model_dict = self.model_dump()
         for k in model_dict:
             new_key = self.get_env_var_name(attr_name=k)
             if isinstance((old_val := model_dict[k]), SecretStr):
@@ -38,9 +39,7 @@ class ProfileSettings(DumpableSettings):
     username: str
     api_key: SecretStr
     verify_ssl: bool = True
-
-    class Config:
-        env_prefix = "DEEPSEARCH_"
+    model_config = SettingsConfigDict(env_prefix="DEEPSEARCH_")
 
     @classmethod
     def from_cli_prompt(cls) -> ProfileSettings:
@@ -54,13 +53,9 @@ def from_cli_prompt(cls) -> ProfileSettings:
 
 class MainSettings(DumpableSettings):
     profile: Optional[str] = None
-
-    class Config:
-        env_prefix = "DEEPSEARCH_"
+    model_config = SettingsConfigDict(env_prefix="DEEPSEARCH_")
 
 
 class CLISettings(DumpableSettings):
     show_cli_stack_traces: bool = False
-
-    class Config:
-        env_prefix = "DEEPSEARCH_"
+    model_config = SettingsConfigDict(env_prefix="DEEPSEARCH_")
diff --git a/deepsearch/core/client/settings_manager.py b/deepsearch/core/client/settings_manager.py
@@ -5,7 +5,7 @@
 from typing import Dict, List, Optional
 
 import platformdirs
-from pydantic.v1 import ValidationError
+from pydantic import ValidationError
 
 from deepsearch.core.cli.profile_utils import (
     MSG_AMBIGUOUS_SUCCESSOR,
@@ -111,7 +111,7 @@ def _migrate_legacy_config(self) -> None:
         if self._main_settings.profile is None:
             legacy_cfg_path = self.config_root_path / LEGACY_CFG_FILENAME
             if legacy_cfg_path.exists():
-                legacy_cfg = DeepSearchConfig.parse_file(legacy_cfg_path)
+                legacy_cfg = DeepSearchConfig.model_validate(legacy_cfg_path)
                 if isinstance(legacy_cfg.auth, DeepSearchKeyAuth):
                     new_cfg = ProfileSettings(
                         host=legacy_cfg.host,

diff --git a/deepsearch/cps/cli/data_indices_typer.py b/deepsearch/cps/cli/data_indices_typer.py
@@ -165,7 +165,7 @@ def upload_files(
 
     if conv_settings is not None:
         try:
-            final_conv_settings = ConversionSettings.parse_obj(
+            final_conv_settings = ConversionSettings.model_validate(
                 json.loads(conv_settings)
             )
         except json.JSONDecodeError:
@@ -177,7 +177,7 @@ def upload_files(
 
     if target_settings is not None:
         try:
-            final_target_settings = TargetSettings.parse_file(target_settings)
+            final_target_settings = TargetSettings.model_validate_json(target_settings)
         except Exception as e:
             raise e
     else:

diff --git a/deepsearch/cps/client/api.py b/deepsearch/cps/client/api.py
@@ -3,7 +3,7 @@
 from typing import Optional
 
 import requests
-from pydantic.v1 import ValidationError
+from pydantic import ValidationError
 
 import deepsearch.cps.apis.user
 from deepsearch.core.client import (

diff --git a/deepsearch/cps/client/components/data_indices.py b/deepsearch/cps/client/components/data_indices.py
@@ -8,7 +8,7 @@
 from urllib.parse import urlparse
 
 import requests
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.apis import public as sw_client
 from deepsearch.cps.apis.public.models.attachment_upload_data import (
@@ -40,7 +40,7 @@ def list(self, proj_key: str) -> List[DataIndex]:
 
         # filter out saved searchs index
         return [
-            DataIndex.parse_obj(item.to_dict())
+            DataIndex.model_validate(item.to_dict())
             for item in response
             if item.to_dict()["type"] != "View"
         ]
@@ -94,7 +94,7 @@ def create(
             self.sw_api.create_project_data_index(proj_key=proj_key, data=data)
         )
 
-        return DataIndex.parse_obj(response.to_dict())
+        return DataIndex.model_validate(response.to_dict())
 
     def delete(
         self,
@@ -146,7 +146,7 @@ def convert_from_page_urls(
         task: Task = self.sw_api.html_print_convert_upload(
             proj_key=coords.proj_key,
             index_key=coords.index_key,
-            body={"urls": [item.dict() for item in urls]},
+            body={"urls": [item.model_dump() for item in urls]},
         )
         return task
 

diff --git a/deepsearch/cps/client/components/documents.py b/deepsearch/cps/client/components/documents.py
@@ -5,7 +5,7 @@
 import urllib.parse
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.apis.public_v2 import SemanticApi
 from deepsearch.cps.apis.public_v2.models.cps_task import CpsTask

diff --git a/deepsearch/cps/client/components/elastic.py b/deepsearch/cps/client/components/elastic.py
@@ -3,7 +3,7 @@
 import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Union
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.apis import public as sw_client
 from deepsearch.cps.client.components.data_indices import (
@@ -31,7 +31,9 @@ def list(self, domain: str = "all") -> List[ElasticDataCollection]:
             index_type="all", index_domain=domain
         )
 
-        return [ElasticDataCollection.parse_obj(item.to_dict()) for item in response]
+        return [
+            ElasticDataCollection.model_validate(item.to_dict()) for item in response
+        ]
 
 
 class ElasticDataCollectionSource(BaseModel):

diff --git a/deepsearch/cps/client/components/projects.py b/deepsearch/cps/client/components/projects.py
@@ -4,7 +4,7 @@
 from enum import Enum
 from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 import deepsearch.cps.apis.user
 from deepsearch.cps.apis.user.models.token_response import TokenResponse

diff --git a/deepsearch/cps/client/components/uploader.py b/deepsearch/cps/client/components/uploader.py
@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Union
 
 import requests
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.apis import public as sw_client
 from deepsearch.cps.apis.public.models.temporary_upload_file_result import (

diff --git a/deepsearch/cps/data_indices/utils.py b/deepsearch/cps/data_indices/utils.py
@@ -168,7 +168,9 @@ def process_local_file(
             if conv_settings is not None:
                 payload["conversion_settings"] = conv_settings.to_ccs_spec()
             if target_settings is not None:
-                payload["target_settings"] = target_settings.dict(exclude_none=True)
+                payload["target_settings"] = target_settings.model_dump(
+                    exclude_none=True
+                )
 
             task_id = api.data_indices.upload_file(coords=coords, body=payload)
             task_ids.append(task_id)
@@ -203,7 +205,7 @@ def process_external_cos(
         bar_format=progressbar.bar_format,
     ) as progress:
         # upload using coordinates
-        payload = {"s3_source": {"coordinates": s3_coordinates.dict()}}
+        payload = {"s3_source": {"coordinates": s3_coordinates.model_dump()}}
         task_id = api.data_indices.upload_file(
             coords=coords,
             body=payload,

diff --git a/deepsearch/cps/queries/__init__.py b/deepsearch/cps/queries/__init__.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from pydantic.v1 import BaseModel, Field, validate_arguments
+from pydantic import BaseModel, Field, validate_call
+from pydantic_settings import SettingsConfigDict
 from typing_extensions import Annotated
 
 from deepsearch.cps.client.components.documents import (
@@ -114,8 +115,10 @@ class _APISemanticRagParameters(_APISemanticRetrievalParameters):
     chunk_refs: Optional[List[ChunkRef]] = None
     gen_timeout: Optional[float] = None
 
+    model_config = SettingsConfigDict(protected_namespaces=())
 
-@validate_arguments
+
+@validate_call
 def RAGQuery(
     question: str,
     *,
@@ -202,7 +205,7 @@ def RAGQuery(
     return query
 
 
-@validate_arguments
+@validate_call
 def SemanticQuery(
     question: str,
     *,
@@ -254,7 +257,7 @@ def SemanticQuery(
     task = query.add(
         task_id="QA",
         kind_or_task="SemanticRetrieval",
-        parameters=params.dict(),
+        parameters=params.model_dump(),
         coordinates=coords,
     )
     task.output("items").output_as("items")

diff --git a/deepsearch/cps/queries/results.py b/deepsearch/cps/queries/results.py
@@ -2,7 +2,7 @@
 
 from typing import List, Optional
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from deepsearch.cps.client.components.queries import RunQueryResult
 
@@ -71,7 +71,9 @@ def from_api_output(cls, data: RunQueryResult, raise_on_error=True):
                         grounding=RAGGroundingInfo(
                             retr_items=(
                                 [
-                                    SearchResultItem.parse_obj(search_result_items[i])
+                                    SearchResultItem.model_validate(
+                                        search_result_items[i]
+                                    )
                                     for i in retr_idxs
                                 ]
                                 if retr_idxs is not None and retrieval_part is not None

diff --git a/deepsearch/documents/core/input_process.py b/deepsearch/documents/core/input_process.py
@@ -114,7 +114,7 @@ def process_cos_input(
         cps_proj_key=cps_proj_key,
         source={
             "type": "s3",
-            "coordinates": source_cos.dict(),
+            "coordinates": source_cos.model_dump(),
         },
         target=target,
         conversion_settings=conversion_settings,

diff --git a/deepsearch/documents/core/lookup.py b/deepsearch/documents/core/lookup.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict, List
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 
 def _resolve_item(item, doc):