fix: fix caching models from all endpoints, inputs and outputs (#6005)

jina-ai · Aug 2, 2023 · 103ddb3 · 103ddb3
1 parent 6fd64da
commit 103ddb3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py
@@ -185,7 +185,7 @@ async def task():
                                         else:
                                             input_model = _create_pydantic_model_from_schema(input_model_schema,
                                                                                              input_model_name,
-                                                                                             {})
+                                                                                             models_created_by_name)
                                         models_created_by_name[input_model_name] = input_model
                                     input_model = models_created_by_name[input_model_name]
                                     models_schema_list.append(input_model_schema)
@@ -203,7 +203,7 @@ async def task():
                                         else:
                                             output_model = _create_pydantic_model_from_schema(output_model_schema,
                                                                                               output_model_name,
-                                                                                              {})
+                                                                                              models_created_by_name)
                                         models_created_by_name[output_model_name] = output_model
                                     output_model = models_created_by_name[output_model_name]
                                     models_schema_list.append(output_model)
@@ -221,7 +221,7 @@ async def task():
                                             from pydantic import BaseModel
                                             parameters_model = _create_pydantic_model_from_schema(parameters_model_schema,
                                                                                                   parameters_model_name,
-                                                                                                  {},
+                                                                                                  models_created_by_name,
                                                                                                   base_class=BaseModel)
                                             models_created_by_name[parameters_model_name] = parameters_model
                                     parameters_model = models_created_by_name[parameters_model_name]

diff --git a/tests/integration/docarray_v2/test_v2.py b/tests/integration/docarray_v2/test_v2.py
@@ -1518,3 +1518,22 @@ def foo(self, docs: DocList[MyDocWithExample], **kwargs) -> DocList[MyDocWithExa
         assert 'This test should be in description' in resp_str
         assert 'MyDocWithExampleTitle' in resp_str
         assert 'extra_key' in resp_str
+
+
+def test_issue_fastapi_multiple_models_same_name():
+    class MyRandomModel(BaseDoc):
+        a: str
+
+    class MyInputModel(BaseDoc):
+        b: Optional[MyRandomModel] = None
+
+
+    class MyFailingExecutor(Executor):
+        @requests(on='/generate')
+        def generate(self, docs: DocList[MyInputModel], **kwargs) -> DocList[MyRandomModel]:
+            return DocList[MyRandomModel]([doc.b for doc in docs])
+
+    with Flow(protocol='http').add(uses=MyFailingExecutor) as f:
+        input_doc = MyRandomModel(a='hello world')
+        res = f.post(on='/generate', inputs=[MyInputModel(b=MyRandomModel(a='hey'))], return_type=DocList[MyRandomModel])
+        assert res[0].a == 'hey'