Added safety header oaig. Ruff formatting (#18)

* Added safety header oaig. Ruff formatting
* updated safety header name
* additional test for agents

Author: steven10a

examples/basic/azure_implementation.py

@@ -54,9 +54,7 @@
 }
 
 
-async def process_input(
-    guardrails_client: GuardrailsAsyncAzureOpenAI, user_input: str
-) -> None:
+async def process_input(guardrails_client: GuardrailsAsyncAzureOpenAI, user_input: str) -> None:
     """Process user input with complete response validation using GuardrailsClient."""
     try:
         # Use GuardrailsClient to handle all guardrail checks and LLM calls

examples/basic/custom_context.py

@@ -26,14 +26,15 @@
                     "system_prompt_details": "Check if the text contains any math problems.",
                 },
             },
-        ]
-    }
+        ],
+    },
 }
 
 
 async def main() -> None:
     # Use Ollama for guardrail LLM checks
     from openai import AsyncOpenAI
+
     guardrail_llm = AsyncOpenAI(
         base_url="http://127.0.0.1:11434/v1/",  # Ollama endpoint
         api_key="ollama",
@@ -50,10 +51,7 @@ async def main() -> None:
         while True:
             try:
                 user_input = input("Enter a message: ")
-                response = await client.chat.completions.create(
-                    model="gpt-4.1-nano",
-                    messages=[{"role": "user", "content": user_input}]
-                )
+                response = await client.chat.completions.create(model="gpt-4.1-nano", messages=[{"role": "user", "content": user_input}])
                 print("Assistant:", response.llm_response.choices[0].message.content)
             except EOFError:
                 break
@@ -65,5 +63,3 @@ async def main() -> None:
 
 if __name__ == "__main__":
     asyncio.run(main())
-
-

examples/basic/hello_world.py

@@ -49,15 +49,11 @@ async def process_input(
             previous_response_id=response_id,
         )
 
-        console.print(
-            f"\nAssistant output: {response.llm_response.output_text}", end="\n\n"
-        )
+        console.print(f"\nAssistant output: {response.llm_response.output_text}", end="\n\n")
 
         # Show guardrail results if any were run
         if response.guardrail_results.all_results:
-            console.print(
-                f"[dim]Guardrails checked: {len(response.guardrail_results.all_results)}[/dim]"
-            )
+            console.print(f"[dim]Guardrails checked: {len(response.guardrail_results.all_results)}[/dim]")
 
         return response.llm_response.id
 
@@ -76,16 +72,12 @@ async def main() -> None:
         while True:
             try:
                 user_input = input("Enter a message: ")
-                response_id = await process_input(
-                    guardrails_client, user_input, response_id
-                )
+                response_id = await process_input(guardrails_client, user_input, response_id)
             except EOFError:
                 break
             except GuardrailTripwireTriggered as exc:
                 stage_name = exc.guardrail_result.info.get("stage_name", "unknown")
-                console.print(
-                    f"\n🛑 [bold red]Guardrail triggered in stage '{stage_name}'![/bold red]"
-                )
+                console.print(f"\n🛑 [bold red]Guardrail triggered in stage '{stage_name}'![/bold red]")
                 console.print(
                     Panel(
                         str(exc.guardrail_result),

examples/basic/local_model.py

@@ -80,12 +80,8 @@ async def main() -> None:
                 break
             except GuardrailTripwireTriggered as exc:
                 stage_name = exc.guardrail_result.info.get("stage_name", "unknown")
-                guardrail_name = exc.guardrail_result.info.get(
-                    "guardrail_name", "unknown"
-                )
-                console.print(
-                    f"\n🛑 [bold red]Guardrail '{guardrail_name}' triggered in stage '{stage_name}'![/bold red]"
-                )
+                guardrail_name = exc.guardrail_result.info.get("guardrail_name", "unknown")
+                console.print(f"\n🛑 [bold red]Guardrail '{guardrail_name}' triggered in stage '{stage_name}'![/bold red]")
                 console.print(
                     Panel(
                         str(exc.guardrail_result),

examples/basic/multi_bundle.py

@@ -73,9 +73,7 @@ async def process_input(
 
             # Get the response ID from the final chunk
             response_id_to_return = None
-            if hasattr(chunk.llm_response, "response") and hasattr(
-                chunk.llm_response.response, "id"
-            ):
+            if hasattr(chunk.llm_response, "response") and hasattr(chunk.llm_response.response, "id"):
                 response_id_to_return = chunk.llm_response.response.id
 
             return response_id_to_return
@@ -98,16 +96,12 @@ async def main() -> None:
         while True:
             try:
                 prompt = input("Enter a message: ")
-                response_id = await process_input(
-                    guardrails_client, prompt, response_id
-                )
+                response_id = await process_input(guardrails_client, prompt, response_id)
             except (EOFError, KeyboardInterrupt):
                 break
             except GuardrailTripwireTriggered as exc:
                 stage_name = exc.guardrail_result.info.get("stage_name", "unknown")
-                guardrail_name = exc.guardrail_result.info.get(
-                    "guardrail_name", "unknown"
-                )
+                guardrail_name = exc.guardrail_result.info.get("guardrail_name", "unknown")
                 console.print(
                     f"🛑 Guardrail '{guardrail_name}' triggered in stage '{stage_name}'!",
                     style="bold red",

examples/basic/multiturn_chat_with_alignment.py

@@ -43,9 +43,7 @@ def get_weather(location: str, unit: str = "celsius") -> dict[str, str | int]:
     }
 
 
-def get_flights(
-    origin: str, destination: str, date: str
-) -> dict[str, list[dict[str, str]]]:
+def get_flights(origin: str, destination: str, date: str) -> dict[str, list[dict[str, str]]]:
     flights = [
         {"flight": "GA123", "depart": f"{date} 08:00", "arrive": f"{date} 12:30"},
         {"flight": "GA456", "depart": f"{date} 15:45", "arrive": f"{date} 20:10"},
@@ -160,9 +158,7 @@ def _stage_lines(stage_name: str, stage_results: Iterable) -> list[str]:
         # Header with status and confidence
         lines.append(f"[bold]{stage_name.upper()}[/bold] · {name} · {status}")
         if confidence != "N/A":
-            lines.append(
-                f"  📊 Confidence: {confidence} (threshold: {info.get('threshold', 'N/A')})"
-            )
+            lines.append(f"  📊 Confidence: {confidence} (threshold: {info.get('threshold', 'N/A')})")
 
         # Prompt injection detection-specific details
         if name == "Prompt Injection Detection":
@@ -176,9 +172,7 @@ def _stage_lines(stage_name: str, stage_results: Iterable) -> list[str]:
 
             # Add interpretation
             if r.tripwire_triggered:
-                lines.append(
-                    "  ⚠️  PROMPT INJECTION DETECTED: Action does not serve user's goal!"
-                )
+                lines.append("  ⚠️  PROMPT INJECTION DETECTED: Action does not serve user's goal!")
             else:
                 lines.append("  ✨ ALIGNED: Action serves user's goal")
         else:
@@ -235,9 +229,7 @@ async def main(malicious: bool = False) -> None:
             messages.append({"role": "user", "content": user_input})
 
             try:
-                resp = await client.chat.completions.create(
-                    model="gpt-4.1-nano", messages=messages, tools=tools
-                )
+                resp = await client.chat.completions.create(model="gpt-4.1-nano", messages=messages, tools=tools)
                 print_guardrail_results("initial", resp)
                 choice = resp.llm_response.choices[0]
                 message = choice.message
@@ -246,12 +238,12 @@ async def main(malicious: bool = False) -> None:
                 info = getattr(e, "guardrail_result", None)
                 info = info.info if info else {}
                 lines = [
-                    f"Guardrail: {info.get('guardrail_name','Unknown')}",
-                    f"Stage: {info.get('stage_name','unknown')}",
-                    f"User goal: {info.get('user_goal','N/A')}",
-                    f"Action: {info.get('action','N/A')}",
-                    f"Observation: {info.get('observation','N/A')}",
-                    f"Confidence: {info.get('confidence','N/A')}",
+                    f"Guardrail: {info.get('guardrail_name', 'Unknown')}",
+                    f"Stage: {info.get('stage_name', 'unknown')}",
+                    f"User goal: {info.get('user_goal', 'N/A')}",
+                    f"Action: {info.get('action', 'N/A')}",
+                    f"Observation: {info.get('observation', 'N/A')}",
+                    f"Confidence: {info.get('confidence', 'N/A')}",
                 ]
                 console.print(
                     Panel(
@@ -292,12 +284,8 @@ async def main(malicious: bool = False) -> None:
 
                         # Malicious injection test mode
                         if malicious:
-                            console.print(
-                                "[yellow]⚠️  MALICIOUS TEST: Injecting unrelated sensitive data into function output[/yellow]"
-                            )
-                            console.print(
-                                "[yellow]   This should trigger the Prompt Injection Detection guardrail as misaligned![/yellow]"
-                            )
+                            console.print("[yellow]⚠️  MALICIOUS TEST: Injecting unrelated sensitive data into function output[/yellow]")
+                            console.print("[yellow]   This should trigger the Prompt Injection Detection guardrail as misaligned![/yellow]")
                             result = {
                                 **result,
                                 "bank_account": "1234567890",
@@ -319,17 +307,13 @@ async def main(malicious: bool = False) -> None:
                                 "role": "tool",
                                 "tool_call_id": call.id,
                                 "name": fname,
-                                "content": json.dumps(
-                                    {"error": f"Unknown function: {fname}"}
-                                ),
+                                "content": json.dumps({"error": f"Unknown function: {fname}"}),
                             }
                         )
 
                 # Final call
                 try:
-                    resp = await client.chat.completions.create(
-                        model="gpt-4.1-nano", messages=messages, tools=tools
-                    )
+                    resp = await client.chat.completions.create(model="gpt-4.1-nano", messages=messages, tools=tools)
 
                     print_guardrail_results("final", resp)
                     final_message = resp.llm_response.choices[0].message
@@ -342,19 +326,17 @@ async def main(malicious: bool = False) -> None:
                     )
 
                     # Add final assistant response to conversation
-                    messages.append(
-                        {"role": "assistant", "content": final_message.content}
-                    )
+                    messages.append({"role": "assistant", "content": final_message.content})
                 except GuardrailTripwireTriggered as e:
                     info = getattr(e, "guardrail_result", None)
                     info = info.info if info else {}
                     lines = [
-                        f"Guardrail: {info.get('guardrail_name','Unknown')}",
-                        f"Stage: {info.get('stage_name','unknown')}",
-                        f"User goal: {info.get('user_goal','N/A')}",
-                        f"Action: {info.get('action','N/A')}",
-                        f"Observation: {info.get('observation','N/A')}",
-                        f"Confidence: {info.get('confidence','N/A')}",
+                        f"Guardrail: {info.get('guardrail_name', 'Unknown')}",
+                        f"Stage: {info.get('stage_name', 'unknown')}",
+                        f"User goal: {info.get('user_goal', 'N/A')}",
+                        f"Action: {info.get('action', 'N/A')}",
+                        f"Observation: {info.get('observation', 'N/A')}",
+                        f"Confidence: {info.get('confidence', 'N/A')}",
                     ]
                     console.print(
                         Panel(
@@ -380,9 +362,7 @@ async def main(malicious: bool = False) -> None:
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Chat Completions with Prompt Injection Detection guardrails"
-    )
+    parser = argparse.ArgumentParser(description="Chat Completions with Prompt Injection Detection guardrails")
     parser.add_argument(
         "--malicious",
         action="store_true",

examples/basic/pii_mask_example.py

@@ -42,9 +42,7 @@
     },
     "input": {
         "version": 1,
-        "guardrails": [
-            {"name": "Moderation", "config": {"categories": ["hate", "violence"]}}
-        ],
+        "guardrails": [{"name": "Moderation", "config": {"categories": ["hate", "violence"]}}],
         "config": {"concurrency": 5, "suppress_tripwire": False},
     },
     "output": {
@@ -98,9 +96,7 @@ async def process_input(
         # Show PII masking information if detected in pre-flight
         if response.guardrail_results.preflight:
             for result in response.guardrail_results.preflight:
-                if result.info.get(
-                    "guardrail_name"
-                ) == "Contains PII" and result.info.get("pii_detected", False):
+                if result.info.get("guardrail_name") == "Contains PII" and result.info.get("pii_detected", False):
                     detected_entities = result.info.get("detected_entities", {})
                     masked_text = result.info.get("checked_text", user_input)
 
@@ -118,9 +114,7 @@ async def process_input(
         # Show if PII was detected in output
         if response.guardrail_results.output:
             for result in response.guardrail_results.output:
-                if result.info.get(
-                    "guardrail_name"
-                ) == "Contains PII" and result.info.get("pii_detected", False):
+                if result.info.get("guardrail_name") == "Contains PII" and result.info.get("pii_detected", False):
                     detected_entities = result.info.get("detected_entities", {})
                     console.print(
                         Panel(
@@ -134,14 +128,8 @@ async def process_input(
     except GuardrailTripwireTriggered as exc:
         stage_name = exc.guardrail_result.info.get("stage_name", "unknown")
         guardrail_name = exc.guardrail_result.info.get("guardrail_name", "unknown")
-        console.print(
-            f"[bold red]Guardrail '{guardrail_name}' triggered in stage '{stage_name}'![/bold red]"
-        )
-        console.print(
-            Panel(
-                str(exc.guardrail_result), title="Guardrail Result", border_style="red"
-            )
-        )
+        console.print(f"[bold red]Guardrail '{guardrail_name}' triggered in stage '{stage_name}'![/bold red]")
+        console.print(Panel(str(exc.guardrail_result), title="Guardrail Result", border_style="red"))
         raise
 
 

examples/basic/structured_outputs_example.py

@@ -10,6 +10,7 @@
 # Define a simple Pydantic model for structured output
 class UserInfo(BaseModel):
     """User information extracted from text."""
+
     name: str = Field(description="Full name of the user")
     age: int = Field(description="Age of the user")
     email: str = Field(description="Email address of the user")
@@ -22,21 +23,18 @@ class UserInfo(BaseModel):
         "version": 1,
         "guardrails": [
             {"name": "Moderation", "config": {"categories": ["hate", "violence"]}},
-        ]
-    }
+        ],
+    },
 }
 
 
 async def extract_user_info(guardrails_client: GuardrailsAsyncOpenAI, text: str) -> UserInfo:
     """Extract user information using responses_parse with structured output."""
     try:
         response = await guardrails_client.responses.parse(
-            input=[
-                {"role": "system", "content": "Extract user information from the provided text."},
-                {"role": "user", "content": text}
-            ],
+            input=[{"role": "system", "content": "Extract user information from the provided text."}, {"role": "user", "content": text}],
             model="gpt-4.1-nano",
-            text_format=UserInfo
+            text_format=UserInfo,
         )
 
         # Access the parsed structured output

examples/basic/suppress_tripwire.py

@@ -55,9 +55,7 @@ async def process_input(
             for result in response.guardrail_results.all_results:
                 guardrail_name = result.info.get("guardrail_name", "Unknown Guardrail")
                 if result.tripwire_triggered:
-                    console.print(
-                        f"[bold yellow]Guardrail '{guardrail_name}' triggered![/bold yellow]"
-                    )
+                    console.print(f"[bold yellow]Guardrail '{guardrail_name}' triggered![/bold yellow]")
                     console.print(
                         Panel(
                             str(result),
@@ -66,15 +64,11 @@ async def process_input(
                         )
                     )
                 else:
-                    console.print(
-                        f"[bold green]Guardrail '{guardrail_name}' passed.[/bold green]"
-                    )
+                    console.print(f"[bold green]Guardrail '{guardrail_name}' passed.[/bold green]")
         else:
             console.print("[bold green]No guardrails triggered.[/bold green]")
 
-        console.print(
-            f"\n[bold blue]Assistant output:[/bold blue] {response.llm_response.output_text}\n"
-        )
+        console.print(f"\n[bold blue]Assistant output:[/bold blue] {response.llm_response.output_text}\n")
         return response.llm_response.id
 
     except Exception as e:
@@ -95,9 +89,7 @@ async def main() -> None:
                 user_input = input("Enter a message: ")
             except EOFError:
                 break
-            response_id = await process_input(
-                guardrails_client, user_input, response_id
-            )
+            response_id = await process_input(guardrails_client, user_input, response_id)
 
 
 if __name__ == "__main__":