diff --git a/client.py b/client.py index 759337d..cb7bec7 100644 --- a/client.py +++ b/client.py @@ -14,7 +14,7 @@ openai_client = OpenAI( "Invariant-Authorization": "Bearer " + os.getenv("INVARIANT_API_KEY"), "Invariant-Guardrails": guardrails, }, - base_url="http://localhost:8000/api/v1/gateway/non-streaming/openai", + base_url="http://localhost:9999/api/v1/gateway/non-streaming/openai", ) response = openai_client.chat.completions.create( diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index c785dfd..c303ee8 100644 --- a/gateway/common/config_manager.py +++ b/gateway/common/config_manager.py @@ -33,7 +33,7 @@ class GatewayConfig: """Common configurations for the Gateway Server.""" def __init__(self): - self.guardrails = self._load_guardrails_from_file() + self.guardrails_from_file = self._load_guardrails_from_file() def _load_guardrails_from_file(self) -> str: """ diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index 5fbb8d8..eeed533 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -51,6 +51,18 @@ def create_annotations_from_guardrails_errors( extra_metadata={ "source": "guardrails-error", "guardrail-action": action, + # if included in error, also include the source of the guardrail + **( + {"guardrail_content": error.get("guardrail_content")} + if error.get("guardrail_content") + else {} + ), + # same for guardrail-action + **( + {"guardrail_action": error.get("guardrail_action")} + if error.get("guardrail_action") + else {} + ), }, ) ) diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index 2785788..ae6cf31 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -368,8 +368,16 @@ async def check_guardrails( guardrails_result = result.json() aggregated_errors = {"errors": []} - for res in guardrails_result.get("result", []): - aggregated_errors["errors"].extend(res.get("errors", [])) + for res, guardrail in zip(guardrails_result.get("result", []), guardrails): + for error in res.get("errors", []): + # aggregated_errors["errors"].extend(res.get("errors", [])) + aggregated_errors["errors"].append( + { + **error, + "guardrail_content": guardrail.content, + "guardrail_action": guardrail.action, + } + ) # check for any error_message if error_message := res.get("error_message"):