Pydantic V2 upgrade PR2 (#1196)

Author: papa99do

src/marqo/base_model.py

@@ -1,3 +1,5 @@
+import pydantic
+from pydantic import ConfigDict
 from pydantic.v1 import BaseModel
 
 
@@ -19,4 +21,8 @@ class Config(MarqoBaseModel.Config):
 
 class ImmutableStrictBaseModel(StrictBaseModel, ImmutableBaseModel):
     class Config(StrictBaseModel.Config, ImmutableBaseModel.Config):
-        pass
+        pass
+
+
+class MarqoBaseModelV2(pydantic.BaseModel):
+    model_config = ConfigDict(validate_by_name=True, validate_assignment=True)

src/marqo/core/models/marqo_index.py

@@ -279,6 +279,7 @@ class MarqoIndex(ImmutableBaseModel, ABC):
     marqo_version: str
     created_at: int = pydantic.Field(gt=0)
     updated_at: int = pydantic.Field(gt=0)
+    # TODO After upgraded to pydantic v2, _cache can be removed. We can use @cached_property instead
     _cache: Dict[str, Any] = PrivateAttr()
     version: Optional[int] = pydantic.Field(default=None)
 
@@ -624,7 +625,9 @@ def index_supports_partial_updates(self) -> bool:
         """
         Check if the index supports partial updates.
         """
-        return self.parsed_marqo_version() >= self._PARTIAL_UPDATE_SUPPORTED_VERSION
+        return self._cache_or_get(
+            'index_supports_partial_updates',
+            lambda: self.parsed_marqo_version() >= self._PARTIAL_UPDATE_SUPPORTED_VERSION)
 
 
 _PROTECTED_FIELD_NAMES = ['_id', '_tensor_facets', '_highlights', '_score', '_found']

src/marqo/core/semi_structured_vespa_index/semi_structured_document.py

@@ -46,7 +46,7 @@ def to_marqo_document(self, vespa_document: Dict[str, Any], return_highlights: b
         vespa_doc = SemiStructuredVespaDocument.from_vespa_document(vespa_document, marqo_index=self.get_marqo_index())
         marqo_doc = vespa_doc.to_marqo_document(marqo_index=self.get_marqo_index())
 
-        if return_highlights and vespa_doc.match_features:
+        if return_highlights and vespa_doc.fixed_fields.match_features:
             # Since tensor fields are stored in each individual field, we need to use same logic in structured
             # index to extract highlights
             marqo_doc[MARQO_DOC_HIGHLIGHTS] = StructuredVespaIndex._extract_highlights(

src/marqo/core/semi_structured_vespa_index/semi_structured_vespa_index.py

@@ -207,8 +207,28 @@ async def api_validation_exception_handler(request: Request, exc: RequestValidat
     )
 
 
+# For validation error raised from PydanticV1 model classes
 @app.exception_handler(pydantic.v1.ValidationError)
 async def validation_exception_handler(request, exc: pydantic.v1.ValidationError) -> JSONResponse:
+    """Catch pydantic v1 validation errors and rewrite as an InvalidArgError whilst keeping error messages from the ValidationError."""
+    error_messages = [{
+        'loc': error.get('loc', ''),
+        'msg': error.get('msg', ''),
+        'type': error.get('type', '')
+    } for error in exc.errors()]
+
+    body = {
+        "message": json.dumps(error_messages),
+        "code": InvalidArgError.code,
+        "type": InvalidArgError.error_type,
+        "link": InvalidArgError.link
+    }
+    return JSONResponse(content=body, status_code=InvalidArgError.status_code)
+
+
+# For validation error raised from PydanticV2 model classes
+@app.exception_handler(pydantic.ValidationError)
+async def validation_exception_handler(request, exc: pydantic.ValidationError) -> JSONResponse:
     """Catch pydantic validation errors and rewrite as an InvalidArgError whilst keeping error messages from the ValidationError."""
     error_messages = [{
         'loc': error.get('loc', ''),

src/marqo/tensor_search/api.py

@@ -1,6 +1,6 @@
 from typing import List, Dict, Any, Optional
 
-from pydantic.v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 
 
 # See https://docs.vespa.ai/en/reference/default-result-format.html

src/marqo/vespa/models/query_result.py

@@ -9,15 +9,19 @@
 import pydantic
 from fastapi.exceptions import RequestValidationError
 from fastapi.testclient import TestClient
+from pydantic.v1.error_wrappers import ErrorWrapper
+from pydantic_core import InitErrorDetails, PydanticCustomError
 
 import marqo.tensor_search.api as api
 from integ_tests.marqo_test import MarqoTestCase
 from marqo import exceptions as base_exceptions
+from marqo.api.exceptions import InvalidArgError
 from marqo.core import exceptions as core_exceptions
 from marqo.core.models.marqo_add_documents_response import MarqoAddDocumentsResponse, MarqoAddDocumentsItem
 from marqo.core.models.marqo_index import FieldType
 from marqo.core.models.marqo_index_request import FieldRequest
 from marqo.tensor_search.enums import EnvVars
+from marqo.tensor_search.models.api_models import SearchQuery
 from marqo.vespa import exceptions as vespa_exceptions
 
 
@@ -609,4 +613,37 @@ class PydanticV1Model(pydantic.v1.BaseModel):
 
         self.assertIn('field required', str(context.exception.errors()))
 
+    def test_handle_pydantic_v1_validation_errors(self):
+        """Test pydantic v1 ValidationError is correctly handled and converted to error response"""
+        error = pydantic.v1.ValidationError(errors=[ErrorWrapper(ValueError("some message"), loc="doc")],
+                                            model=SearchQuery)
+        with patch("marqo.tensor_search.tensor_search.search", side_effect=error):
+            response = self.client.post("/indexes/" + self.structured_index.name + "/search?device=cpu", json={
+                "q": "test",
+                "filter": ""
+            })
+
+            self.assertEqual(response.status_code, 400)
+            self.assertEqual(response.json()["code"], InvalidArgError.code)
+            self.assertEqual(response.json()["type"], InvalidArgError.error_type)
+            assert "some message" in response.json()["message"]
+
+    def test_handle_pydantic_v2_validation_errors(self):
+        """Test pydantic v2 ValidationError is correctly handled and converted to error response"""
+        error = pydantic.ValidationError.from_exception_data(
+            title='SearchQuery',
+            line_errors=[InitErrorDetails(
+                type=PydanticCustomError('type1', 'some message'), loc=('doc',), input=...)]
+        )
+        with patch("marqo.tensor_search.tensor_search.search", side_effect=error):
+            response = self.client.post("/indexes/" + self.structured_index.name + "/search?device=cpu", json={
+                "q": "test",
+                "filter": ""
+            })
+
+            self.assertEqual(response.status_code, 400)
+            self.assertEqual(response.json()["code"], InvalidArgError.code)
+            self.assertEqual(response.json()["type"], InvalidArgError.error_type)
+            assert "some message" in response.json()["message"]
+
     # TODO: Test how marqo handles generic exceptions, including Exception, RunTimeError, ValueError, etc.