CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-07 03:17:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Handle streaming with openai for guardrails. Only check guardrails on the last stream chunk.

Browse Source
This commit is contained in:
Hemang
2025-03-17 14:59:04 +01:00
committed by Hemang Sarkar
parent ca3c43ef76
commit 66973dbf30
1 changed files with 36 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/routes/open_ai.py
+36 -5
View File
@@ -131,9 +131,6 @@ async def stream_response(
if not chunk_text:
continue
# Yield chunk immediately to the client
yield chunk
# Process the chunk
# This will update merged_response with the data from the chunk
process_chunk_text(
@@ -143,6 +140,40 @@ async def stream_response(
tool_call_mapping_by_index,
)
# Check guardrails on the last chunk.
if (
chunk_text == "data: [DONE]"
and context.config
and context.config.guardrails
):
# Block on the guardrails check
guardrails_execution_result = await get_guardrails_check_result(
context, merged_response
)
if guardrails_execution_result.get("errors", []):
error_chunk = json.dumps(
{
"error": {
"message": "[Invariant] The response did not pass the guardrails",
"details": guardrails_execution_result,
}
}
)
# Push annotated trace to the explorer - don't block on its response
if context.dataset_name:
asyncio.create_task(
push_to_explorer(
context,
merged_response,
guardrails_execution_result,
)
)
yield f"data: {error_chunk}\n\n".encode()
return
# Yield chunk to the client
yield chunk
# Send full merged response to the explorer
# Don't block on the response from explorer
if context.dataset_name:
@@ -359,8 +390,8 @@ async def handle_non_streaming_response(
if guardrails_execution_result.get("errors", []):
response_string = json.dumps(
{
"error": "The request did not pass the guardrails",
"guadrails_check_result": guardrails_execution_result,
"error": "[Invariant] The response did not pass the guardrails",
"details": guardrails_execution_result,
}
)
response_code = 400