Properly return error response in vertex_generate HTTP endpoint (#5956)
This commit is contained in:
@@ -675,6 +675,8 @@ async def vertex_generate(vertex_req: VertexGenerateReqInput, raw_request: Reque
|
||||
**(vertex_req.parameters or {}),
|
||||
)
|
||||
ret = await generate_request(req, raw_request)
|
||||
if isinstance(ret, Response):
|
||||
return ret
|
||||
return ORJSONResponse({"predictions": ret})
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ python3 -m unittest test_vertex_endpoint.TestVertexEndpoint.test_vertex_generate
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from http import HTTPStatus
|
||||
|
||||
import requests
|
||||
|
||||
@@ -49,6 +50,15 @@ class TestVertexEndpoint(CustomTestCase):
|
||||
for parameters in [None, {"sampling_params": {"max_new_tokens": 4}}]:
|
||||
self.run_generate(parameters)
|
||||
|
||||
def test_vertex_generate_fail(self):
|
||||
data = {
|
||||
"instances": [
|
||||
{"prompt": "The capital of France is"},
|
||||
],
|
||||
}
|
||||
response = requests.post(self.base_url + "/vertex_generate", json=data)
|
||||
assert response.status_code == HTTPStatus.BAD_REQUEST
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user