chore: free pydantic version

extra-requirements.txt

@@ -52,7 +52,7 @@ pathspec:                   standard,devel
 filelock:                   standard,devel
 requests:                   standard,devel
 websockets:                 standard,devel
-pydantic<2.0.0:             core
+pydantic:             core
 python-multipart:           standard,devel
 aiofiles:                   standard,devel
 aiohttp:                    standard,devel

jina/_docarray.py

@@ -4,7 +4,15 @@
 
     docarray_v2 = True
 
+    from jina._docarray_legacy import LegacyDocumentJina
+
 except ImportError:
     from docarray import Document, DocumentArray
 
     docarray_v2 = False
+
+
+import pydantic
+
+is_pydantic_v2 = pydantic.__version__.startswith('2.')
+

jina/_docarray_legacy.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from docarray import BaseDoc
+from docarray import DocList
+
+docarray_v2 = True
+
+from typing import Any, Dict, Optional, List, Union
+
+from docarray.typing import AnyEmbedding, AnyTensor
+
+
+class LegacyDocumentJina(BaseDoc):
+    """
+    This Document is the LegacyDocumentJina. It follows the same schema as in DocArray <=0.21.
+    It can be useful to start migrating a codebase from v1 to v2.
+
+    Nevertheless, the API is not totally compatible with DocArray <=0.21 `Document`.
+    Indeed, none of the method associated with `Document` are present. Only the schema
+    of the data is similar.
+
+    ```python
+    from docarray import DocList
+    from docarray.documents.legacy import LegacyDocument
+    import numpy as np
+
+    doc = LegacyDocument(text='hello')
+    doc.url = 'http://myimg.png'
+    doc.tensor = np.zeros((3, 224, 224))
+    doc.embedding = np.zeros((100, 1))
+
+    doc.tags['price'] = 10
+
+    doc.chunks = DocList[Document]([Document() for _ in range(10)])
+
+    doc.chunks = DocList[Document]([Document() for _ in range(10)])
+    ```
+
+    """
+
+    tensor: Optional[AnyTensor] = None
+    chunks: Optional[Union[DocList[LegacyDocumentJina], List[LegacyDocumentJina]]] = None
+    matches: Optional[Union[DocList[LegacyDocumentJina], List[LegacyDocumentJina]]] = None
+    blob: Optional[bytes] = None
+    text: Optional[str] = None
+    url: Optional[str] = None
+    embedding: Optional[AnyEmbedding] = None
+    tags: Dict[str, Any] = dict()
+    scores: Optional[Dict[str, Any]] = None

jina/serve/executors/__init__.py

@@ -52,7 +52,7 @@
 from jina.serve.instrumentation import MetricsTimer
 
 if docarray_v2:
-    from docarray.documents.legacy import LegacyDocument
+    from jina._docarray import LegacyDocumentJina
 
 if TYPE_CHECKING:  # pragma: no cover
     from opentelemetry.context.context import Context
@@ -257,7 +257,7 @@
             from docarray import BaseDoc, DocList
 
             default_annotations = (
-                DocList[LegacyDocument] if is_batch_docs else LegacyDocument
+                DocList[LegacyDocumentJina] if is_batch_docs else LegacyDocumentJina
             )
         else:
             from jina import Document, DocumentArray

jina/serve/runtimes/gateway/graph/topology_graph.py

@@ -7,7 +7,7 @@
 
 import grpc.aio
 
-from jina._docarray import DocumentArray, docarray_v2
+from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
 from jina.constants import __default_endpoint__
 from jina.excepts import InternalNetworkError
 from jina.logging.logger import JinaLogger
@@ -18,11 +18,15 @@
 
 if docarray_v2:
     from docarray import DocList
-    from docarray.documents.legacy import LegacyDocument
+    from jina._docarray import LegacyDocumentJina
 
-    from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
+    if not is_pydantic_v2:
+        from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
+    else:
+        from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema
 
-    legacy_doc_schema = LegacyDocument.schema()
+
+    legacy_doc_schema = LegacyDocumentJina.schema()
 
 
 class TopologyGraph:
@@ -218,8 +222,6 @@
                         endp, _ = endpoints_proto
                         self.endpoints = endp.endpoints
                         if docarray_v2:
-                            from docarray.documents.legacy import LegacyDocument
-
                             schemas = json_format.MessageToDict(endp.schemas)
                             self._pydantic_models_by_endpoint = {}
                             models_created_by_name = {}
@@ -236,10 +238,10 @@
                                 else:
                                     if input_model_name not in models_created_by_name:
                                         if input_model_schema == legacy_doc_schema:
-                                            input_model = LegacyDocument
+                                            input_model = LegacyDocumentJina
                                         else:
                                             input_model = (
-                                                _create_pydantic_model_from_schema(
+                                                create_base_doc_from_schema(
                                                     input_model_schema,
                                                     input_model_name,
                                                     models_created_by_name,
@@ -266,10 +268,10 @@
                                 else:
                                     if output_model_name not in models_created_by_name:
                                         if output_model_name == legacy_doc_schema:
-                                            output_model = LegacyDocument
+                                            output_model = LegacyDocumentJina
                                         else:
                                             output_model = (
-                                                _create_pydantic_model_from_schema(
+                                                create_base_doc_from_schema(
                                                     output_model_schema,
                                                     output_model_name,
                                                     models_created_by_name,
@@ -306,7 +308,7 @@
                                             from pydantic import BaseModel
 
                                             parameters_model = (
-                                                _create_pydantic_model_from_schema(
+                                                create_base_doc_from_schema(
                                                     parameters_model_schema,
                                                     parameters_model_name,
                                                     models_created_by_name,

jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py

@@ -6,6 +6,7 @@
 from jina.logging.logger import JinaLogger
 from jina.serve.networking.sse import EventSourceResponse
 from jina.types.request.data import DataRequest
+from jina._docarray import is_pydantic_v2
 
 if TYPE_CHECKING:  # pragma: no cover
     from opentelemetry import trace
@@ -80,7 +81,9 @@
     import os
 
     from pydantic import BaseModel
-    from pydantic.config import BaseConfig, inherit_config
+    from pydantic.config import BaseConfig
+    if not is_pydantic_v2:
+        from pydantic.config import inherit_config
 
     from jina.proto import jina_pb2
     from jina.serve.runtimes.gateway.models import (
@@ -275,7 +278,10 @@
             parameters_model = input_output_map['parameters'] or Optional[Dict]
             default_parameters = ... if input_output_map['parameters'] else None
 
-            _config = inherit_config(InnerConfig, BaseDoc.__config__)
+            if not is_pydantic_v2:
+                _config = inherit_config(InnerConfig, BaseDoc.__config__)
+            else:
+                _config = InnerConfig
 
             endpoint_input_model = pydantic.create_model(
                 f'{endpoint.strip("/")}_input_model',

jina/serve/runtimes/head/request_handling.py

@@ -16,13 +16,20 @@
 from jina.serve.runtimes.monitoring import MonitoringRequestMixin
 from jina.serve.runtimes.worker.request_handling import WorkerRequestHandler
 from jina.types.request.data import DataRequest, Response
-from jina._docarray import docarray_v2
+from jina._docarray import docarray_v2, is_pydantic_v2
 
 if docarray_v2:
-    from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
+    if not is_pydantic_v2:
+        from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
+    else:
+        from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema
     from docarray import DocList
     from docarray.base_doc.any_doc import AnyDoc
 
+    from jina._docarray import LegacyDocumentJina
+
+    legacy_doc_schema = LegacyDocumentJina.schema()
+
 
 if TYPE_CHECKING:  # pragma: no cover
     from prometheus_client import CollectorRegistry
@@ -330,9 +337,6 @@
         self, connection_pool: GrpcConnectionPool, name: str, retries: int, stop_event
     ):
         from google.protobuf import json_format
-        from docarray.documents.legacy import LegacyDocument
-
-        legacy_doc_schema = LegacyDocument.schema()
 
         async def task():
             self.logger.debug(
@@ -356,20 +360,20 @@
 
                             if input_model_schema == legacy_doc_schema:
                                 models_created_by_name[input_model_name] = (
-                                    LegacyDocument
+                                    LegacyDocumentJina
                                 )
                             elif input_model_name not in models_created_by_name:
-                                input_model = _create_pydantic_model_from_schema(
+                                input_model = create_base_doc_from_schema(
                                     input_model_schema, input_model_name, {}
                                 )
                                 models_created_by_name[input_model_name] = input_model
 
                             if output_model_name == legacy_doc_schema:
                                 models_created_by_name[output_model_name] = (
-                                    LegacyDocument
+                                    LegacyDocumentJina
                                 )
                             elif output_model_name not in models_created_by_name:
-                                output_model = _create_pydantic_model_from_schema(
+                                output_model = create_base_doc_from_schema(
                                     output_model_schema, output_model_name, {}
                                 )
                                 models_created_by_name[output_model_name] = output_model

jina/serve/runtimes/helper.py

@@ -1,7 +1,7 @@
 import copy
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-from jina._docarray import docarray_v2
+from jina._docarray import docarray_v2, is_pydantic_v2
 
 _SPECIFIC_EXECUTOR_SEPARATOR = '__'
 
@@ -79,7 +79,7 @@ def _parse_specific_params(parameters: Dict, executor_name: str):
     return parsed_params
 
 
-if docarray_v2:
+if docarray_v2 and not is_pydantic_v2:
     from docarray import BaseDoc, DocList
     from docarray.typing import AnyTensor
     from pydantic import create_model

jina/serve/runtimes/worker/http_fastapi_app.py

@@ -2,7 +2,7 @@
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 
 from jina import Document, DocumentArray
-from jina._docarray import docarray_v2
+from jina._docarray import docarray_v2, is_pydantic_v2
 from jina.importer import ImportExtensions
 from jina.serve.networking.sse import EventSourceResponse
 from jina.types.request.data import DataRequest
@@ -38,7 +38,9 @@
     import os
 
     from pydantic import BaseModel, Field
-    from pydantic.config import BaseConfig, inherit_config
+    from pydantic.config import BaseConfig
+    if not is_pydantic_v2:
+        from pydantic.config import inherit_config
 
     from jina.proto import jina_pb2
     from jina.serve.runtimes.gateway.models import _to_camel_case
@@ -168,7 +170,10 @@
             )
 
             if docarray_v2:
-                _config = inherit_config(InnerConfig, BaseDoc.__config__)
+                if not is_pydantic_v2:
+                    _config = inherit_config(InnerConfig, BaseDoc.__config__)
+                else:
+                    _config = InnerConfig
             else:
                 _config = input_doc_model.__config__
 

jina/serve/runtimes/worker/http_sagemaker_app.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 
-from jina._docarray import docarray_v2
+from jina._docarray import docarray_v2, is_pydantic_v2
 from jina.importer import ImportExtensions
 from jina.types.request.data import DataRequest
 
@@ -33,7 +33,9 @@
         from fastapi import FastAPI, HTTPException, Request
         from fastapi.middleware.cors import CORSMiddleware
         from pydantic import BaseModel, Field
-        from pydantic.config import BaseConfig, inherit_config
+        from pydantic.config import BaseConfig
+        if not is_pydantic_v2:
+            from pydantic.config import inherit_config
 
     import os
 
@@ -235,7 +237,11 @@
                 ... if input_output_map['parameters']['model'] else None
             )
 
-            _config = inherit_config(InnerConfig, BaseDoc.__config__)
+            if not is_pydantic_v2:
+                _config = inherit_config(InnerConfig, BaseDoc.__config__)
+            else:
+                _config = InnerConfig
+
             endpoint_input_model = pydantic.create_model(
                 f'{endpoint.strip("/")}_input_model',
                 data=(Union[List[input_doc_model], input_doc_model], ...),

jina/serve/runtimes/worker/request_handling.py

@@ -20,7 +20,7 @@
 
 from google.protobuf.struct_pb2 import Struct
 
-from jina._docarray import DocumentArray, docarray_v2
+from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
 from jina.constants import __default_endpoint__
 from jina.excepts import BadConfigSource, RuntimeTerminated
 from jina.helper import get_full_version
@@ -33,6 +33,9 @@
 
 if docarray_v2:
     from docarray import DocList
+    from jina._docarray import LegacyDocumentJina
+    legacy_doc_schema = LegacyDocumentJina.schema()
+
 
 if TYPE_CHECKING:  # pragma: no cover
     import grpc
@@ -1012,23 +1015,24 @@
         endpoints_proto.write_endpoints.extend(list(self._executor.write_endpoints))
         schemas = self._executor._get_endpoint_models_dict()
         if docarray_v2:
-            from docarray.documents.legacy import LegacyDocument
+            if not is_pydantic_v2:
+                from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
+            else:
+                from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model
 
-            from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list
 
-            legacy_doc_schema = LegacyDocument.schema()
             for endpoint_name, inner_dict in schemas.items():
                 if inner_dict['input']['model'].schema() == legacy_doc_schema:
                     inner_dict['input']['model'] = legacy_doc_schema
                 else:
-                    inner_dict['input']['model'] = _create_aux_model_doc_list_to_list(
+                    inner_dict['input']['model'] = create_pure_python_type_model(
                         inner_dict['input']['model']
                     ).schema()
 
                 if inner_dict['output']['model'].schema() == legacy_doc_schema:
                     inner_dict['output']['model'] = legacy_doc_schema
                 else:
-                    inner_dict['output']['model'] = _create_aux_model_doc_list_to_list(
+                    inner_dict['output']['model'] = create_pure_python_type_model(
                         inner_dict['output']['model']
                     ).schema()