From f02d38b24b4f1958486712b056b1b76fa6063a72 Mon Sep 17 00:00:00 2001 From: David Tam Date: Fri, 16 Aug 2024 09:45:13 -0700 Subject: [PATCH 1/3] add handling to have informative 404 when a guard doesnt exist on openai --- guardrails_api/blueprints/guards.py | 9 +++++++++ guardrails_api/utils/handle_error.py | 2 +- tests/blueprints/test_guards.py | 29 ++++++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) 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" From b0c18b32be5df5b685a5611cc94c28a60ec50b14 Mon Sep 17 00:00:00 2001 From: David Tam Date: Fri, 16 Aug 2024 10:08:24 -0700 Subject: [PATCH 2/3] lint --- tests/blueprints/test_guards.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/blueprints/test_guards.py b/tests/blueprints/test_guards.py index 79167f0..062a458 100644 --- a/tests/blueprints/test_guards.py +++ b/tests/blueprints/test_guards.py @@ -636,7 +636,7 @@ def test_openai_v1_chat_completions__raises_404(mocker): 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") + 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") From ac8d8c071f1eb26d92f1160dfe6cec5dcd62b408 Mon Sep 17 00:00:00 2001 From: David Tam Date: Fri, 16 Aug 2024 10:30:33 -0700 Subject: [PATCH 3/3] fix tests --- tests/blueprints/test_guards.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/blueprints/test_guards.py b/tests/blueprints/test_guards.py index 062a458..3cce59d 100644 --- a/tests/blueprints/test_guards.py +++ b/tests/blueprints/test_guards.py @@ -634,10 +634,11 @@ def test_openai_v1_chat_completions__raises_404(mocker): return_value=mock_guard, ) mocker.patch("guardrails_api.blueprints.guards.CacheClient.set") - # expect 404 HttpError - with pytest.raises(Exception) as e: - openai_v1_chat_completions("My%20Guard's%20Name") - assert str(e.value) == '404 Not Found: Guard not found' + + response = openai_v1_chat_completions("My%20Guard's%20Name") + assert response[1] == 404 + assert response[0]["message"] == 'NotFound' + mock_get_guard.assert_called_once_with("My Guard's Name")