Skip to content

Commit ecdf2b6

Browse files
JoanFMJoanFM
committed
feat: adapt pydantic v2
1 parent c6f4903 commit ecdf2b6

File tree

7 files changed

+80
-56
lines changed

7 files changed

+80
-56
lines changed

‎jina/_docarray.py

+6
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,9 @@
88
from docarray import Document, DocumentArray
99

1010
docarray_v2 = False
11+
12+
13+
import pydantic
14+
15+
is_pydantic_v2 = pydantic.__version__.startswith('2.')
16+

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

+9-5
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
import grpc.aio
99

10-
from jina._docarray import DocumentArray, docarray_v2
10+
from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
1111
from jina.constants import __default_endpoint__
1212
from jina.excepts import InternalNetworkError
1313
from jina.logging.logger import JinaLogger
@@ -20,7 +20,11 @@
2020
from docarray import DocList
2121
from docarray.documents.legacy import LegacyDocument
2222

23-
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
23+
if not is_pydantic_v2:
24+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
25+
else:
26+
from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema
27+
2428

2529
legacy_doc_schema = LegacyDocument.schema()
2630

@@ -239,7 +243,7 @@ async def task():
239243
input_model = LegacyDocument
240244
else:
241245
input_model = (
242-
_create_pydantic_model_from_schema(
246+
create_base_doc_from_schema(
243247
input_model_schema,
244248
input_model_name,
245249
models_created_by_name,
@@ -269,7 +273,7 @@ async def task():
269273
output_model = LegacyDocument
270274
else:
271275
output_model = (
272-
_create_pydantic_model_from_schema(
276+
create_base_doc_from_schema(
273277
output_model_schema,
274278
output_model_name,
275279
models_created_by_name,
@@ -306,7 +310,7 @@ async def task():
306310
from pydantic import BaseModel
307311

308312
parameters_model = (
309-
_create_pydantic_model_from_schema(
313+
create_base_doc_from_schema(
310314
parameters_model_schema,
311315
parameters_model_name,
312316
models_created_by_name,

‎jina/serve/runtimes/head/request_handling.py

+7-4
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616
from jina.serve.runtimes.monitoring import MonitoringRequestMixin
1717
from jina.serve.runtimes.worker.request_handling import WorkerRequestHandler
1818
from jina.types.request.data import DataRequest, Response
19-
from jina._docarray import docarray_v2
19+
from jina._docarray import docarray_v2, is_pydantic_v2
2020

2121
if docarray_v2:
22-
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
22+
if not is_pydantic_v2:
23+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
24+
else:
25+
from docarray.utils.create_dynamic_doc_class import create_base_doc_from_schema
2326
from docarray import DocList
2427
from docarray.base_doc.any_doc import AnyDoc
2528

@@ -359,7 +362,7 @@ async def task():
359362
LegacyDocument
360363
)
361364
elif input_model_name not in models_created_by_name:
362-
input_model = _create_pydantic_model_from_schema(
365+
input_model = create_base_doc_from_schema(
363366
input_model_schema, input_model_name, {}
364367
)
365368
models_created_by_name[input_model_name] = input_model
@@ -369,7 +372,7 @@ async def task():
369372
LegacyDocument
370373
)
371374
elif output_model_name not in models_created_by_name:
372-
output_model = _create_pydantic_model_from_schema(
375+
output_model = create_base_doc_from_schema(
373376
output_model_schema, output_model_name, {}
374377
)
375378
models_created_by_name[output_model_name] = output_model

‎jina/serve/runtimes/helper.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import copy
22
from typing import Any, Dict, List, Optional, Tuple, Union
33

4-
from jina._docarray import docarray_v2
4+
from jina._docarray import docarray_v2, is_pydantic_v2
55

66
_SPECIFIC_EXECUTOR_SEPARATOR = '__'
77

@@ -79,7 +79,7 @@ def _parse_specific_params(parameters: Dict, executor_name: str):
7979
return parsed_params
8080

8181

82-
if docarray_v2:
82+
if docarray_v2 and not is_pydantic_v2:
8383
from docarray import BaseDoc, DocList
8484
from docarray.typing import AnyTensor
8585
from pydantic import create_model

‎jina/serve/runtimes/worker/request_handling.py

+7-4
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020

2121
from google.protobuf.struct_pb2 import Struct
2222

23-
from jina._docarray import DocumentArray, docarray_v2
23+
from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
2424
from jina.constants import __default_endpoint__
2525
from jina.excepts import BadConfigSource, RuntimeTerminated
2626
from jina.helper import get_full_version
@@ -1013,21 +1013,24 @@ async def endpoint_discovery(self, empty, context) -> jina_pb2.EndpointsProto:
10131013
if docarray_v2:
10141014
from docarray.documents.legacy import LegacyDocument
10151015

1016-
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list
1016+
if not is_pydantic_v2:
1017+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
1018+
else:
1019+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model
10171020

10181021
legacy_doc_schema = LegacyDocument.schema()
10191022
for endpoint_name, inner_dict in schemas.items():
10201023
if inner_dict['input']['model'].schema() == legacy_doc_schema:
10211024
inner_dict['input']['model'] = legacy_doc_schema
10221025
else:
1023-
inner_dict['input']['model'] = _create_aux_model_doc_list_to_list(
1026+
inner_dict['input']['model'] = create_pure_python_type_model(
10241027
inner_dict['input']['model']
10251028
).schema()
10261029

10271030
if inner_dict['output']['model'].schema() == legacy_doc_schema:
10281031
inner_dict['output']['model'] = legacy_doc_schema
10291032
else:
1030-
inner_dict['output']['model'] = _create_aux_model_doc_list_to_list(
1033+
inner_dict['output']['model'] = create_pure_python_type_model(
10311034
inner_dict['output']['model']
10321035
).schema()
10331036

‎tests/integration/docarray_v2/test_v2.py

+19-16
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import numpy as np
1616
import pytest
17+
from jina._docarray import is_pydantic_v2
1718
from docarray import BaseDoc, DocList
1819
from docarray.documents import ImageDoc, TextDoc
1920
from docarray.documents.legacy import LegacyDocument
@@ -302,10 +303,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]:
302303
from jina.proto import jina_pb2
303304
from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub
304305
from jina.serve.executors import __dry_run_endpoint__
305-
from jina.serve.runtimes.helper import (
306-
_create_aux_model_doc_list_to_list,
307-
_create_pydantic_model_from_schema,
308-
)
306+
if not is_pydantic_v2:
307+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
308+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
309+
else:
310+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema
309311

310312
channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}')
311313
stub = JinaDiscoverEndpointsRPCStub(channel)
@@ -320,16 +322,16 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]:
320322
v = schema_map['/bar']
321323
assert (
322324
v['input']
323-
== _create_pydantic_model_from_schema(
324-
_create_aux_model_doc_list_to_list(Input1).schema(),
325+
== create_base_doc_from_schema(
326+
create_pure_python_type_model(Input1).schema(),
325327
'Input1',
326328
{},
327329
).schema()
328330
)
329331
assert (
330332
v['output']
331-
== _create_pydantic_model_from_schema(
332-
_create_aux_model_doc_list_to_list(Output2).schema(),
333+
== create_base_doc_from_schema(
334+
create_pure_python_type_model(Output2).schema(),
333335
'Output2',
334336
{},
335337
).schema()
@@ -390,10 +392,11 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]:
390392
from jina.proto import jina_pb2
391393
from jina.proto.jina_pb2_grpc import JinaDiscoverEndpointsRPCStub
392394
from jina.serve.executors import __default_endpoint__, __dry_run_endpoint__
393-
from jina.serve.runtimes.helper import (
394-
_create_aux_model_doc_list_to_list,
395-
_create_pydantic_model_from_schema,
396-
)
395+
if not is_pydantic_v2:
396+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
397+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
398+
else:
399+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema
397400

398401
channel = grpc.insecure_channel(f'0.0.0.0:{ports[0]}')
399402
stub = JinaDiscoverEndpointsRPCStub(channel)
@@ -411,14 +414,14 @@ def bar(self, docs: DocList[Output1], **kwargs) -> DocList[Output2]:
411414
v = schema_map[__default_endpoint__]
412415
assert (
413416
v['input']
414-
== _create_pydantic_model_from_schema(
415-
_create_aux_model_doc_list_to_list(Input1).schema(), 'Input1', {}
417+
== create_base_doc_from_schema(
418+
create_pure_python_type_model(Input1).schema(), 'Input1', {}
416419
).schema()
417420
)
418421
assert (
419422
v['output']
420-
== _create_pydantic_model_from_schema(
421-
_create_aux_model_doc_list_to_list(Output2).schema(), 'Output2', {}
423+
== create_base_doc_from_schema(
424+
create_pure_python_type_model(Output2).schema(), 'Output2', {}
422425
).schema()
423426
)
424427

‎tests/unit/serve/runtimes/test_helper.py

+30-25
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
import pytest
55

6-
from jina._docarray import docarray_v2
6+
from jina._docarray import docarray_v2, is_pydantic_v2
77
from jina.serve.helper import get_default_grpc_options
88
from jina.serve.runtimes.helper import (
99
_get_name_from_replicas_name,
@@ -96,10 +96,11 @@ def test_create_pydantic_model_from_schema(transformation):
9696
from docarray.documents import TextDoc
9797
from docarray.typing import AnyTensor, ImageUrl
9898

99-
from jina.serve.runtimes.helper import (
100-
_create_aux_model_doc_list_to_list,
101-
_create_pydantic_model_from_schema,
102-
)
99+
if not is_pydantic_v2:
100+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
101+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
102+
else:
103+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema
103104

104105
class Nested2Doc(BaseDoc):
105106
value: str
@@ -124,8 +125,8 @@ class CustomDoc(BaseDoc):
124125
nested: Nested1Doc
125126
classvar: ClassVar[str] = 'classvar'
126127

127-
CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc)
128-
new_custom_doc_model = _create_pydantic_model_from_schema(
128+
CustomDocCopy = create_pure_python_type_model(CustomDoc)
129+
new_custom_doc_model = create_base_doc_from_schema(
129130
CustomDocCopy.schema(), 'CustomDoc', {}
130131
)
131132

@@ -199,8 +200,8 @@ class CustomDoc(BaseDoc):
199200
class TextDocWithId(BaseDoc):
200201
ia: str
201202

202-
TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId)
203-
new_textdoc_with_id_model = _create_pydantic_model_from_schema(
203+
TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId)
204+
new_textdoc_with_id_model = create_base_doc_from_schema(
204205
TextDocWithIdCopy.schema(), 'TextDocWithId', {}
205206
)
206207

@@ -229,8 +230,8 @@ class TextDocWithId(BaseDoc):
229230
class ResultTestDoc(BaseDoc):
230231
matches: DocList[TextDocWithId]
231232

232-
ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc)
233-
new_result_test_doc_with_id_model = _create_pydantic_model_from_schema(
233+
ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc)
234+
new_result_test_doc_with_id_model = create_base_doc_from_schema(
234235
ResultTestDocCopy.schema(), 'ResultTestDoc', {}
235236
)
236237
result_test_docs = DocList[ResultTestDoc](
@@ -268,10 +269,11 @@ def test_create_empty_doc_list_from_schema(transformation):
268269
from docarray.documents import TextDoc
269270
from docarray.typing import AnyTensor, ImageUrl
270271

271-
from jina.serve.runtimes.helper import (
272-
_create_aux_model_doc_list_to_list,
273-
_create_pydantic_model_from_schema,
274-
)
272+
if not is_pydantic_v2:
273+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
274+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
275+
else:
276+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema
275277

276278
class CustomDoc(BaseDoc):
277279
tensor: Optional[AnyTensor]
@@ -288,8 +290,8 @@ class CustomDoc(BaseDoc):
288290
tags: Optional[Dict[str, Any]] = None
289291
lf: List[float] = [3.0, 4.1]
290292

291-
CustomDocCopy = _create_aux_model_doc_list_to_list(CustomDoc)
292-
new_custom_doc_model = _create_pydantic_model_from_schema(
293+
CustomDocCopy = create_pure_python_type_model(CustomDoc)
294+
new_custom_doc_model = create_base_doc_from_schema(
293295
CustomDocCopy.schema(), 'CustomDoc', {}
294296
)
295297

@@ -313,8 +315,8 @@ class CustomDoc(BaseDoc):
313315
class TextDocWithId(BaseDoc):
314316
ia: str
315317

316-
TextDocWithIdCopy = _create_aux_model_doc_list_to_list(TextDocWithId)
317-
new_textdoc_with_id_model = _create_pydantic_model_from_schema(
318+
TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId)
319+
new_textdoc_with_id_model = create_base_doc_from_schema(
318320
TextDocWithIdCopy.schema(), 'TextDocWithId', {}
319321
)
320322

@@ -336,8 +338,8 @@ class TextDocWithId(BaseDoc):
336338
class ResultTestDoc(BaseDoc):
337339
matches: DocList[TextDocWithId]
338340

339-
ResultTestDocCopy = _create_aux_model_doc_list_to_list(ResultTestDoc)
340-
new_result_test_doc_with_id_model = _create_pydantic_model_from_schema(
341+
ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc)
342+
new_result_test_doc_with_id_model = create_base_doc_from_schema(
341343
ResultTestDocCopy.schema(), 'ResultTestDoc', {}
342344
)
343345
result_test_docs = DocList[ResultTestDoc]()
@@ -360,8 +362,11 @@ class ResultTestDoc(BaseDoc):
360362
@pytest.mark.skipif(not docarray_v2, reason='Test only working with docarray v2')
361363
def test_dynamic_class_creation_multiple_doclist_nested():
362364
from docarray import BaseDoc, DocList
363-
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list
364-
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
365+
if not is_pydantic_v2:
366+
from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
367+
from jina.serve.runtimes.helper import _create_pydantic_model_from_schema as create_base_doc_from_schema
368+
else:
369+
from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model, create_base_doc_from_schema
365370

366371
class MyTextDoc(BaseDoc):
367372
text: str
@@ -374,8 +379,8 @@ class SearchResult(BaseDoc):
374379

375380
textlist = DocList[MyTextDoc]([MyTextDoc(text='hey')])
376381
models_created_by_name = {}
377-
SearchResult_aux = _create_aux_model_doc_list_to_list(SearchResult)
378-
_ = _create_pydantic_model_from_schema(
382+
SearchResult_aux = create_pure_python_type_model(SearchResult)
383+
_ = create_base_doc_from_schema(
379384
SearchResult_aux.schema(), 'SearchResult', models_created_by_name
380385
)
381386
QuoteFile_reconstructed_in_gateway_from_Search_results = models_created_by_name[

0 commit comments

Comments
 (0)