diff --git a/guardrails_api/blueprints/guards.py b/guardrails_api/blueprints/guards.py index 4c8c115..dd8db10 100644 --- a/guardrails_api/blueprints/guards.py +++ b/guardrails_api/blueprints/guards.py @@ -178,6 +178,15 @@ def openai_v1_chat_completions(guard_name: str): decoded_guard_name = unquote_plus(guard_name) guard_struct = guard_client.get_guard(decoded_guard_name) guard = guard_struct + if guard_struct is None: + raise HttpError( + 404, + "NotFound", + "A Guard with the name {guard_name} does not exist!".format( + guard_name=decoded_guard_name + ), + ) + if not isinstance(guard_struct, Guard): guard: Guard = Guard.from_dict(guard_struct.to_dict()) stream = payload.get("stream", False) diff --git a/guardrails_api/utils/handle_error.py b/guardrails_api/utils/handle_error.py index 6dc0ad2..4fcf231 100644 --- a/guardrails_api/utils/handle_error.py +++ b/guardrails_api/utils/handle_error.py @@ -17,7 +17,7 @@ def decorator(*args, **kwargs): return str(validation_error), 400 except HttpError as http_error: logger.error(http_error) - traceback.print_exception(http_error) + traceback.print_exception(type(http_error), http_error, http_error.__traceback__) return http_error.to_dict(), http_error.status except HTTPException as http_exception: logger.error(http_exception) diff --git a/tests/blueprints/test_guards.py b/tests/blueprints/test_guards.py index 768017c..79167f0 100644 --- a/tests/blueprints/test_guards.py +++ b/tests/blueprints/test_guards.py @@ -614,6 +614,35 @@ def test_validate__call_throws_validation_error(mocker): del os.environ["PGHOST"] +def test_openai_v1_chat_completions__raises_404(mocker): + from guardrails_api.blueprints.guards import openai_v1_chat_completions + os.environ["PGHOST"] = "localhost" + mock_guard = None + + mock_request = MockRequest( + "POST", + json={ + "messages": [{"role":"user", "content":"Hello world!"}], + }, + headers={"x-openai-api-key": "mock-key"}, + ) + + mocker.patch("flask.Blueprint", new=MockBlueprint) + mocker.patch("guardrails_api.blueprints.guards.request", mock_request) + mock_get_guard = mocker.patch( + "guardrails_api.blueprints.guards.guard_client.get_guard", + return_value=mock_guard, + ) + mocker.patch("guardrails_api.blueprints.guards.CacheClient.set") + # expect 404 HttpError + with pytest.raises(Exception) as e: + response = openai_v1_chat_completions("My%20Guard's%20Name") + assert str(e.value) == '404 Not Found: Guard not found' + + mock_get_guard.assert_called_once_with("My Guard's Name") + + del os.environ["PGHOST"] + def test_openai_v1_chat_completions__call(mocker): from guardrails_api.blueprints.guards import openai_v1_chat_completions os.environ["PGHOST"] = "localhost"