diff --git a/guardrails_api/api/guards.py b/guardrails_api/api/guards.py index 08bd529..a9f70d6 100644 --- a/guardrails_api/api/guards.py +++ b/guardrails_api/api/guards.py @@ -114,7 +114,7 @@ async def openai_v1_chat_completions(guard_name: str, request: Request): ) guard = ( - Guard.from_dict(guard_struct.to_dict()) + AsyncGuard.from_dict(guard_struct.to_dict()) if not isinstance(guard_struct, Guard) else guard_struct ) @@ -125,7 +125,7 @@ async def openai_v1_chat_completions(guard_name: str, request: Request): ) if not stream: - validation_outcome: ValidationOutcome = guard(num_reasks=0, **payload) + validation_outcome: ValidationOutcome = await guard(num_reasks=0, **payload) llm_response = guard.history.last.iterations.last.outputs.llm_response_info result = outcome_to_chat_completion( validation_outcome=validation_outcome, @@ -136,8 +136,8 @@ async def openai_v1_chat_completions(guard_name: str, request: Request): else: async def openai_streamer(): - guard_stream = guard(num_reasks=0, **payload) - for result in guard_stream: + guard_stream = await guard(num_reasks=0, **payload) + async for result in guard_stream: chunk = json.dumps( outcome_to_stream_response(validation_outcome=result) ) diff --git a/tests/api/test_guards.py b/tests/api/test_guards.py index 453a976..416baf0 100644 --- a/tests/api/test_guards.py +++ b/tests/api/test_guards.py @@ -14,6 +14,9 @@ from tests.mocks.mock_guard_client import MockGuardStruct from guardrails_api.api.guards import router as guards_router + +import asyncio + # TODO: Should we mock this somehow? # Right now it's just empty, but it technically does a file read register_config() @@ -344,7 +347,9 @@ def test_openai_v1_chat_completions__call(mocker): ) mock___call__ = mocker.patch.object(MockGuardStruct, "__call__") - mock___call__.return_value = mock_outcome + future = asyncio.Future() + future.set_result(mock_outcome) + mock___call__.return_value = future mock_from_dict = mocker.patch("guardrails_api.api.guards.Guard.from_dict") mock_from_dict.return_value = mock_guard