From c6f490354c4553ff095d71e6a5540192f83ba41e Mon Sep 17 00:00:00 2001 From: Joan Martinez Date: Thu, 7 Mar 2024 09:11:40 +0100 Subject: [PATCH 1/4] chore: free pydantic version --- extra-requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extra-requirements.txt b/extra-requirements.txt index 089c12007756e..fc1b16a07433f 100644 --- a/extra-requirements.txt +++ b/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 From ecdf2b645e71628491c7d4ed74fec7e713464c5b Mon Sep 17 00:00:00 2001 From: Joan Martinez Date: Thu, 7 Mar 2024 09:56:47 +0100 Subject: [PATCH 2/4] feat: adapt pydantic v2 --- jina/_docarray.py | 6 ++ .../runtimes/gateway/graph/topology_graph.py | 14 +++-- jina/serve/runtimes/head/request_handling.py | 11 ++-- jina/serve/runtimes/helper.py | 4 +- .../serve/runtimes/worker/request_handling.py | 11 ++-- tests/integration/docarray_v2/test_v2.py | 35 ++++++------ tests/unit/serve/runtimes/test_helper.py | 55 ++++++++++--------- 7 files changed, 80 insertions(+), 56 deletions(-) diff --git a/jina/_docarray.py b/jina/_docarray.py index b9cdfe14aa3fd..96ae3dd1d352d 100644 --- a/jina/_docarray.py +++ b/jina/_docarray.py @@ -8,3 +8,9 @@ from docarray import Document, DocumentArray docarray_v2 = False + + +import pydantic + +is_pydantic_v2 = pydantic.__version__.startswith('2.') + diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py index 89ad937c4698a..2a6c01ea9b633 100644 --- a/jina/serve/runtimes/gateway/graph/topology_graph.py +++ b/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 @@ -20,7 +20,11 @@ from docarray import DocList from docarray.documents.legacy import LegacyDocument - 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() @@ -239,7 +243,7 @@ async def task(): input_model = LegacyDocument else: input_model = ( - _create_pydantic_model_from_schema( + create_base_doc_from_schema( input_model_schema, input_model_name, models_created_by_name, @@ -269,7 +273,7 @@ async def task(): output_model = LegacyDocument 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 +310,7 @@ async def task(): 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, diff --git a/jina/serve/runtimes/head/request_handling.py b/jina/serve/runtimes/head/request_handling.py index 6891c68c02d6a..417c7a865ac6d 100644 --- a/jina/serve/runtimes/head/request_handling.py +++ b/jina/serve/runtimes/head/request_handling.py @@ -16,10 +16,13 @@ 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 @@ -359,7 +362,7 @@ async def task(): LegacyDocument ) 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 @@ -369,7 +372,7 @@ async def task(): LegacyDocument ) 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 diff --git a/jina/serve/runtimes/helper.py b/jina/serve/runtimes/helper.py index 70bb75a485c1b..6444f8fa5a03a 100644 --- a/jina/serve/runtimes/helper.py +++ b/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 diff --git a/jina/serve/runtimes/worker/request_handling.py b/jina/serve/runtimes/worker/request_handling.py index 2e095cb26da50..65472cd6d406f 100644 --- a/jina/serve/runtimes/worker/request_handling.py +++ b/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 @@ -1013,21 +1013,24 @@ async def endpoint_discovery(self, empty, context) -> jina_pb2.EndpointsProto: if docarray_v2: from docarray.documents.legacy import LegacyDocument - from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list + 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 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() diff --git a/tests/integration/docarray_v2/test_v2.py b/tests/integration/docarray_v2/test_v2.py index eebbafd1a572f..7643be4951af8 100644 --- a/tests/integration/docarray_v2/test_v2.py +++ b/tests/integration/docarray_v2/test_v2.py @@ -14,6 +14,7 @@ import numpy as np import pytest +from jina._docarray import is_pydantic_v2 from docarray import BaseDoc, DocList from docarray.documents import ImageDoc, TextDoc from docarray.documents.legacy import LegacyDocument @@ -302,10 +303,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: from jina.proto import jina_pb2 from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub from jina.serve.executors import __dry_run_endpoint__ - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + 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_pure_python_type_model, create_base_doc_from_schema channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}') stub = JinaDiscoverEndpointsRPCStub(channel) @@ -320,16 +322,16 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: v = schema_map['/bar'] assert ( v['input'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Input1).schema(), + == create_base_doc_from_schema( + create_pure_python_type_model(Input1).schema(), 'Input1', {}, ).schema() ) assert ( v['output'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Output2).schema(), + == create_base_doc_from_schema( + create_pure_python_type_model(Output2).schema(), 'Output2', {}, ).schema() @@ -390,10 +392,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: from jina.proto import jina_pb2 from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub from jina.serve.executors import __default_endpoint__, __dry_run_endpoint__ - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + 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_pure_python_type_model, create_base_doc_from_schema channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}') stub = JinaDiscoverEndpointsRPCStub(channel) @@ -411,14 +414,14 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]: v = schema_map[__default_endpoint__] assert ( v['input'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Input1).schema(), 'Input1', {} + == create_base_doc_from_schema( + create_pure_python_type_model(Input1).schema(), 'Input1', {} ).schema() ) assert ( v['output'] - == _create_pydantic_model_from_schema( - _create_aux_model_doc_list_to_list(Output2).schema(), 'Output2', {} + == create_base_doc_from_schema( + create_pure_python_type_model(Output2).schema(), 'Output2', {} ).schema() ) diff --git a/tests/unit/serve/runtimes/test_helper.py b/tests/unit/serve/runtimes/test_helper.py index 323d0cdbc8b02..ac5dec8c593cb 100644 --- a/tests/unit/serve/runtimes/test_helper.py +++ b/tests/unit/serve/runtimes/test_helper.py @@ -3,7 +3,7 @@ import pytest -from jina._docarray import docarray_v2 +from jina._docarray import docarray_v2, is_pydantic_v2 from jina.serve.helper import get_default_grpc_options from jina.serve.runtimes.helper import ( _get_name_from_replicas_name, @@ -96,10 +96,11 @@ def test_create_pydantic_model_from_schema(transformation): from docarray.documents import TextDoc from docarray.typing import AnyTensor, ImageUrl - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + 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_pure_python_type_model, create_base_doc_from_schema class Nested2Doc(BaseDoc): value: str @@ -124,8 +125,8 @@ class CustomDoc(BaseDoc): nested: Nested1Doc classvar: ClassVar[str] = 'classvar' - CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc) - new_custom_doc_model = _create_pydantic_model_from_schema( + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( CustomDocCopy.schema(), 'CustomDoc', {} ) @@ -199,8 +200,8 @@ class CustomDoc(BaseDoc): class TextDocWithId(BaseDoc): ia: str - TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId) - new_textdoc_with_id_model = _create_pydantic_model_from_schema( + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( TextDocWithIdCopy.schema(), 'TextDocWithId', {} ) @@ -229,8 +230,8 @@ class TextDocWithId(BaseDoc): class ResultTestDoc(BaseDoc): matches: DocList[TextDocWithId] - ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc) - new_result_test_doc_with_id_model = _create_pydantic_model_from_schema( + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( ResultTestDocCopy.schema(), 'ResultTestDoc', {} ) result_test_docs = DocList[ResultTestDoc]( @@ -268,10 +269,11 @@ def test_create_empty_doc_list_from_schema(transformation): from docarray.documents import TextDoc from docarray.typing import AnyTensor, ImageUrl - from jina.serve.runtimes.helper import ( - _create_aux_model_doc_list_to_list, - _create_pydantic_model_from_schema, - ) + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + 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_pure_python_type_model, create_base_doc_from_schema class CustomDoc(BaseDoc): tensor: Optional[AnyTensor] @@ -288,8 +290,8 @@ class CustomDoc(BaseDoc): tags: Optional[Dict[str, Any]] = None lf: List[float] = [3.0, 4.1] - CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc) - new_custom_doc_model = _create_pydantic_model_from_schema( + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( CustomDocCopy.schema(), 'CustomDoc', {} ) @@ -313,8 +315,8 @@ class CustomDoc(BaseDoc): class TextDocWithId(BaseDoc): ia: str - TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId) - new_textdoc_with_id_model = _create_pydantic_model_from_schema( + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( TextDocWithIdCopy.schema(), 'TextDocWithId', {} ) @@ -336,8 +338,8 @@ class TextDocWithId(BaseDoc): class ResultTestDoc(BaseDoc): matches: DocList[TextDocWithId] - ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc) - new_result_test_doc_with_id_model = _create_pydantic_model_from_schema( + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( ResultTestDocCopy.schema(), 'ResultTestDoc', {} ) result_test_docs = DocList[ResultTestDoc]() @@ -360,8 +362,11 @@ class ResultTestDoc(BaseDoc): @pytest.mark.skipif(not docarray_v2, reason='Test only working with docarray v2') def test_dynamic_class_creation_multiple_doclist_nested(): from docarray import BaseDoc, DocList - from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list - from jina.serve.runtimes.helper import _create_pydantic_model_from_schema + if not is_pydantic_v2: + from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model + 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_pure_python_type_model, create_base_doc_from_schema class MyTextDoc(BaseDoc): text: str @@ -374,8 +379,8 @@ class SearchResult(BaseDoc): textlist = DocList[MyTextDoc]([MyTextDoc(text='hey')]) models_created_by_name = {} - SearchResult_aux = _create_aux_model_doc_list_to_list(SearchResult) - _ = _create_pydantic_model_from_schema( + SearchResult_aux = create_pure_python_type_model(SearchResult) + _ = create_base_doc_from_schema( SearchResult_aux.schema(), 'SearchResult', models_created_by_name ) QuoteFile_reconstructed_in_gateway_from_Search_results = models_created_by_name[ From 9000f653f04e4291c9e1c33eab5c885086789987 Mon Sep 17 00:00:00 2001 From: Joan Martinez Date: Thu, 7 Mar 2024 14:30:18 +0100 Subject: [PATCH 3/4] fix: fix inherit config --- .../runtimes/gateway/http_fastapi_app_docarrayv2.py | 10 ++++++++-- jina/serve/runtimes/worker/http_fastapi_app.py | 11 ++++++++--- jina/serve/runtimes/worker/http_sagemaker_app.py | 12 +++++++++--- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py b/jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py index c4153ec3480fc..d826e6d0e0629 100644 --- a/jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py +++ b/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 @@ async def _shutdown(): 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 @@ async def event_generator(): 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', diff --git a/jina/serve/runtimes/worker/http_fastapi_app.py b/jina/serve/runtimes/worker/http_fastapi_app.py index 47006dd4be329..d20d1a68e83e5 100644 --- a/jina/serve/runtimes/worker/http_fastapi_app.py +++ b/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 @@ def get_fastapi_app( 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 @@ async def streaming_get(request: Request = None, body: input_doc_model = None): ) 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__ diff --git a/jina/serve/runtimes/worker/http_sagemaker_app.py b/jina/serve/runtimes/worker/http_sagemaker_app.py index e44082afc57f2..a0c93568b5054 100644 --- a/jina/serve/runtimes/worker/http_sagemaker_app.py +++ b/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 @@ def get_fastapi_app( 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 @@ def construct_model_from_line( ... 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], ...), From 4cc9e41f53bdcb2cfcfb132019aa34f4da7c65d0 Mon Sep 17 00:00:00 2001 From: Joan Martinez Date: Thu, 7 Mar 2024 17:54:28 +0100 Subject: [PATCH 4/4] fix: hack for LegacyDocument --- jina/_docarray.py | 2 + jina/_docarray_legacy.py | 49 +++++++++++++++++++ jina/serve/executors/__init__.py | 4 +- .../runtimes/gateway/graph/topology_graph.py | 10 ++-- jina/serve/runtimes/head/request_handling.py | 11 +++-- .../serve/runtimes/worker/request_handling.py | 7 +-- 6 files changed, 67 insertions(+), 16 deletions(-) create mode 100644 jina/_docarray_legacy.py diff --git a/jina/_docarray.py b/jina/_docarray.py index 96ae3dd1d352d..499fdb67e6f40 100644 --- a/jina/_docarray.py +++ b/jina/_docarray.py @@ -4,6 +4,8 @@ docarray_v2 = True + from jina._docarray_legacy import LegacyDocumentJina + except ImportError: from docarray import Document, DocumentArray diff --git a/jina/_docarray_legacy.py b/jina/_docarray_legacy.py new file mode 100644 index 0000000000000..61a2347f852b7 --- /dev/null +++ b/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 diff --git a/jina/serve/executors/__init__.py b/jina/serve/executors/__init__.py index c5c8f72a8e6c1..3f3c83e5a2083 100644 --- a/jina/serve/executors/__init__.py +++ b/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 @@ def get_function_with_schema(fn: Callable) -> T: 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 diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py index 2a6c01ea9b633..2c20e803bb73b 100644 --- a/jina/serve/runtimes/gateway/graph/topology_graph.py +++ b/jina/serve/runtimes/gateway/graph/topology_graph.py @@ -18,7 +18,7 @@ if docarray_v2: from docarray import DocList - from docarray.documents.legacy import LegacyDocument + from jina._docarray import LegacyDocumentJina if not is_pydantic_v2: from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema @@ -26,7 +26,7 @@ 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: @@ -222,8 +222,6 @@ async def task(): 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 = {} @@ -240,7 +238,7 @@ async def task(): 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_base_doc_from_schema( @@ -270,7 +268,7 @@ async def task(): 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_base_doc_from_schema( diff --git a/jina/serve/runtimes/head/request_handling.py b/jina/serve/runtimes/head/request_handling.py index 417c7a865ac6d..e883b901a55ae 100644 --- a/jina/serve/runtimes/head/request_handling.py +++ b/jina/serve/runtimes/head/request_handling.py @@ -26,6 +26,10 @@ 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 @@ -333,9 +337,6 @@ def _get_endpoints_from_workers( 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( @@ -359,7 +360,7 @@ async def task(): 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_base_doc_from_schema( @@ -369,7 +370,7 @@ async def task(): 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_base_doc_from_schema( diff --git a/jina/serve/runtimes/worker/request_handling.py b/jina/serve/runtimes/worker/request_handling.py index 65472cd6d406f..16b79d1e0e047 100644 --- a/jina/serve/runtimes/worker/request_handling.py +++ b/jina/serve/runtimes/worker/request_handling.py @@ -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 @@ -1011,14 +1014,12 @@ async def endpoint_discovery(self, empty, context) -> jina_pb2.EndpointsProto: 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 - 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