Skip to content

Commit

Permalink
add handling to have informative 404 when a guard doesnt exist on openai
Browse files Browse the repository at this point in the history
  • Loading branch information
dtam committed Aug 16, 2024
1 parent af02055 commit f02d38b
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 1 deletion.
9 changes: 9 additions & 0 deletions guardrails_api/blueprints/guards.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion guardrails_api/utils/handle_error.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
29 changes: 29 additions & 0 deletions tests/blueprints/test_guards.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit f02d38b

Please sign in to comment.