Skip to content

Commit

Permalink
fix: issue in gcp app
Browse files Browse the repository at this point in the history
  • Loading branch information
zac-li committed Jan 17, 2024
1 parent 8a48521 commit 17d5f90
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions jina/serve/runtimes/worker/http_gcp_app.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ def add_post_route(
input_model,
output_model,
input_doc_list_model=None,
output_doc_list_model=None,
):
from docarray.base_doc.docarray_response import DocArrayResponse

Expand Down Expand Up @@ -128,7 +129,7 @@ async def process(body) -> output_model:
if status.code == jina_pb2.StatusProto.ERROR:
raise HTTPException(status_code=499, detail=status.description)
else:
return {"predictions": resp.docs}
return VertexAIResponse(predictions=output_model(data=resp.docs, parameters=resp.parameters))

@app.api_route(**app_kwargs)
async def post(request: Request):
Expand All @@ -151,6 +152,7 @@ async def post(request: Request):
for endpoint, input_output_map in request_models_map.items():
if endpoint != '_jina_dry_run_':
input_doc_model = input_output_map['input']['model']
output_doc_model = input_output_map['output']['model']
parameters_model = input_output_map['parameters']['model'] or Optional[Dict]
default_parameters = (
... if input_output_map['parameters']['model'] else None
Expand All @@ -165,11 +167,19 @@ async def post(request: Request):
__config__=_config,
)

endpoint_output_model = pydantic.create_model(
f'{endpoint.strip("/")}_output_model',
data=(Union[List[output_doc_model], output_doc_model], ...),
parameters=(Optional[Dict], None),
__config__=_config,
)

add_post_route(
endpoint,
input_model=endpoint_input_model,
output_model=VertexAIResponse,
output_model=endpoint_output_model,
input_doc_list_model=input_doc_model,
output_doc_list_model=VertexAIResponse,
)

from jina.serve.runtimes.gateway.health_model import JinaHealthModel
Expand Down

0 comments on commit 17d5f90

Please sign in to comment.