feat(guardrail): Add guardrail decorator (#3521)

Author: nina-kollman

packages/sample-app/sample_app/agents/travel_agent_example.py

@@ -654,11 +654,13 @@ async def handle_runner_stream(runner: "Runner"):
     """Process runner events and display output."""
 
     tool_calls_made = []
+    response_text_parts = []
 
     async for event in runner.stream_events():
         if event.type == "raw_response_event":
             if isinstance(event.data, ResponseTextDeltaEvent):
                 print(event.data.delta, end="", flush=True)
+                response_text_parts.append(event.data.delta)
             elif isinstance(event.data, ResponseOutputItemAddedEvent):
                 if isinstance(event.data.item, ResponseFunctionToolCall):
                     tool_name = event.data.item.name
@@ -688,17 +690,29 @@ async def handle_runner_stream(runner: "Runner"):
                     for part in getattr(raw_item, "content", []):
                         if isinstance(part, ResponseOutputText):
                             content_parts.append(part.text)
+                            response_text_parts.append(part.text)
                         elif isinstance(part, ResponseOutputRefusal):
                             content_parts.append(part.refusal)
+                            response_text_parts.append(part.refusal)
                     if content_parts:
                         print("".join(content_parts), end="", flush=True)
 
     print()
-    return tool_calls_made
+    return tool_calls_made, "".join(response_text_parts)
 
 
-async def run_travel_query(query: str):
-    """Run a single travel planning query."""
+async def run_travel_query(query: str, return_response_text: bool = False):
+    """
+    Run a single travel planning query.
+
+    Args:
+        query: The travel planning query
+        return_response_text: If True, returns the response text.
+            If False, returns tool_calls (for backward compatibility)
+
+    Returns:
+        Either response_text (str) or tool_calls (list) depending on parameter
+    """
 
     print("=" * 80)
     print(f"Query: {query}")
@@ -710,13 +724,16 @@ async def run_travel_query(query: str):
 
     messages = [{"role": "user", "content": query}]
     runner = Runner().run_streamed(starting_agent=travel_agent, input=messages)
-    tool_calls = await handle_runner_stream(runner)
+    tool_calls, response_text = await handle_runner_stream(runner)
 
     print(f"\n{'='*80}")
     print(f"✅ Query completed! Tools used: {', '.join(tool_calls) if tool_calls else 'None'}")
     print(f"{'='*80}\n")
 
-    return tool_calls
+    if return_response_text:
+        return response_text
+    else:
+        return tool_calls  # Backward compatibility for existing code
 
 
 def generate_travel_queries(n: int = 10) -> List[str]:

packages/sample-app/sample_app/guardrail_medical_chat_example.py

@@ -0,0 +1,158 @@
+import asyncio
+import os
+from openai import AsyncOpenAI
+from traceloop.sdk import Traceloop
+from traceloop.sdk.guardrails.guardrails import guardrail
+from traceloop.sdk.evaluator import EvaluatorMadeByTraceloop
+
+
+Traceloop.init(
+    app_name="medical-chat-example"
+)
+
+api_key = os.getenv("OPENAI_API_KEY")
+if not api_key:
+    raise ValueError("OPENAI_API_KEY environment variable is required. Please set it before running this example.")
+
+client = AsyncOpenAI(api_key=api_key)
+
+
+# Custom callback function to handle evaluation results
+def handle_medical_evaluation(evaluator_result, original_result):
+    """
+    Custom handler for medical advice evaluation.
+
+    Args:
+        evaluator_result: The evaluation result with 'success' and 'reason' fields
+        original_result: The original AI response dict (e.g., {"text": "..."})
+
+    Returns:
+        Either the original result dict or a modified version
+    """
+    if not evaluator_result.success:
+        # Return a modified dict with error message
+        print(f"handle_medical_evaluation was activated - evaluator_result: {evaluator_result}")
+        return {
+            "text": "There is an issue with the request. Please try again."
+        }
+    return original_result
+
+
+@guardrail(
+    evaluator=EvaluatorMadeByTraceloop.pii_detector(probability_threshold=0.8),
+    on_evaluation_complete=handle_medical_evaluation
+)
+async def get_doctor_response_with_pii_check(patient_message: str) -> dict:
+    """Get a doctor's response with PII detection guardrail and custom callback."""
+
+    system_prompt = """You are a medical AI assistant. Provide helpful,
+      general medical information and advice while being clear about your limitations.
+      Always recommend consulting with qualified healthcare providers for proper diagnosis and treatment.
+      Be empathetic and professional in your responses."""
+    # This is the system prompt for the personal information case
+    personal_info_system_prompt = """You are a medical AI assistant that provides helpful, general medical information  # noqa: F841
+        tailored to the individual user.
+
+        When personal information is available (such as age, sex, symptoms, medical history,
+        lifestyle, medications, or concerns), actively incorporate it into your responses
+        to make guidance more relevant and personalized.
+
+        Adapt explanations, examples, and recommendations to the user’s context whenever possible.
+        If key personal details are missing, ask concise and relevant follow-up questions
+        before giving advice.
+
+        Be clear about your limitations as an AI and do not provide diagnoses or definitive
+        treatment plans. Always encourage consultation with qualified healthcare professionals
+        for diagnosis, treatment, or urgent concerns.
+
+        Maintain a professional, empathetic, and supportive tone.
+        Avoid assumptions, respect privacy, and clearly distinguish general information
+        from personalized considerations."""
+
+    response = await client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": patient_message}
+        ],
+        max_tokens=500,
+        temperature=0
+    )
+
+    return {
+        "text": response.choices[0].message.content
+    }
+
+
+# Main function using the simple example
+@guardrail(evaluator="medicaladvicegiven")
+async def get_doctor_response(patient_message: str) -> dict:
+    """Get a doctor's response to patient input using GPT-4o."""
+
+    system_prompt = """You are a medical AI assistant. Provide helpful,
+      general medical information and advice while being clear about your limitations.
+      Always recommend consulting with qualified healthcare providers for proper diagnosis and treatment.
+      Be empathetic and professional in your responses."""
+
+    response = await client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": patient_message}
+        ],
+        max_tokens=500,
+        temperature=0
+    )
+
+    # Return dict with 'text' field
+    return {
+        "text": response.choices[0].message.content
+    }
+
+
+async def medical_chat_session():
+    """Run an interactive medical chat session."""
+    print("🏥 Welcome to the Medical Chat")
+    print("=" * 50)
+    print("This example simulates a conversation between a patient and a doctor.")
+    print("The doctor's responses are processed through guardrails to ensure safety.")
+    print("Type 'quit' to exit the chat.\n")
+
+    while True:
+        try:
+            patient_input = input("Patient: ").strip()
+
+            if patient_input.lower() in ['quit', 'exit', 'q']:
+                print("\n👋 Thank you for using the medical chat. Take care!")
+                break
+
+            if not patient_input:
+                print("Please enter your symptoms or medical concern.")
+                continue
+
+            print("\n🤖 Processing your request through the medical AI system...\n")
+
+            # Get the doctor's response with guardrails applied
+            doctor_response = await get_doctor_response_with_pii_check(patient_input)
+
+            # Extract text from the response dict
+            response_text = doctor_response.get("text", str(doctor_response))
+            print(f"👨‍⚕️ Doctor response: {response_text}")
+
+            print("-" * 50)
+
+        except KeyboardInterrupt:
+            print("\n\n👋 Chat session interrupted. Goodbye!")
+            break
+        except Exception as e:
+            print(f"\n❌ An error occurred: {e}")
+            print("Please try again or type 'quit' to exit.")
+
+
+async def main():
+    """Main function to run the medical chat example."""
+    await medical_chat_session()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

packages/sample-app/sample_app/guardrail_travel_agent_example.py

@@ -0,0 +1,166 @@
+"""
+Guardrail wrapper for the existing travel agent example.
+This demonstrates how to add PII detection guardrails to an existing agentic system.
+"""
+import asyncio
+import sys
+from pathlib import Path
+
+# Add the agents directory to the path
+agents_dir = Path(__file__).parent / "agents"
+sys.path.insert(0, str(agents_dir))
+
+from traceloop.sdk.guardrails.guardrails import guardrail  # noqa: E402
+from traceloop.sdk.evaluator import EvaluatorMadeByTraceloop  # noqa: E402
+
+# Import the travel agent function
+try:
+    from travel_agent_example import run_travel_query
+except ImportError:
+    print("Error: Could not import travel_agent_example.")
+    print(f"Make sure {agents_dir}/travel_agent_example.py exists")
+    sys.exit(1)
+
+
+# Custom callback to handle PII detection in travel agent responses
+def handle_pii_detection(evaluator_result, original_result):
+    """
+    Custom handler for PII detection in travel agent itineraries.
+
+    Args:
+        evaluator_result: The evaluation result with 'success' and 'reason' fields
+        original_result: The original travel itinerary response dict
+
+    Returns:
+        Either the original result dict or a sanitized warning
+    """
+    # Get captured stdout
+    captured_stdout = original_result.get("_captured_stdout", "")
+
+    if not evaluator_result.success:
+        # PII was detected - don't display the output, return warning
+        return {
+            "text": (
+                "⚠️ PRIVACY ALERT: The generated travel itinerary contains personally "
+                "identifiable information (PII) that could compromise your privacy.\n\n"
+                "For your security, we cannot display this itinerary. Please:\n"
+                "1. Contact our support team through secure channels\n"
+                "2. Request a generic itinerary without personal details\n"
+                "3. Avoid sharing sensitive information in your travel requests\n\n"
+                f"Detection reason: {evaluator_result.reason or 'PII detected in response'}"
+            )
+        }
+    else:
+        # Guardrail passed - now safe to display the captured output
+        if captured_stdout:
+            print(captured_stdout, end="")
+
+        # Remove internal field before returning
+        result = original_result.copy()
+        result.pop("_captured_stdout", None)
+        return result
+
+
+@guardrail(
+    evaluator=EvaluatorMadeByTraceloop.pii_detector(probability_threshold=0.7),
+    on_evaluation_complete=handle_pii_detection
+)
+async def guarded_travel_agent(query: str) -> dict:
+    """
+    Wrapper around the travel agent that adds PII detection guardrails.
+
+    This function:
+    1. Runs the full travel agent flow (tools, API calls, itinerary generation)
+    2. Gets the final response text from the agent
+    3. Runs PII detection on the complete output
+    4. Returns sanitized response if PII is detected
+
+    Args:
+        query: User's travel planning request
+
+    Returns:
+        Dict with 'text' field containing the travel itinerary or privacy warning
+    """
+    import io
+
+    # Capture stdout to prevent streaming output before guardrail check
+    old_stdout = sys.stdout
+    sys.stdout = captured_output = io.StringIO()
+
+    try:
+        # Run the travel agent and get the response text
+        # return_response_text=True makes it return the agent's text instead of tool_calls
+        response_text = await run_travel_query(query, return_response_text=True)
+
+        # Return dict with 'text' field as required by pii_detector
+        return {"text": response_text, "_captured_stdout": captured_output.getvalue()}
+
+    finally:
+        # Restore stdout
+        sys.stdout = old_stdout
+
+
+async def main():
+    """
+    Interactive travel agent with PII detection guardrails.
+    """
+    print("=" * 80)
+    print("🛡️  Travel Agent with PII Detection Guardrails")
+    print("=" * 80)
+    print("This travel agent uses PII detection to protect your privacy.")
+    print("The agent's output is hidden until the guardrail check completes.")
+    print("Type 'quit' or 'exit' to stop.\n")
+    print("💡 Example queries:")
+    print("  - Plan a 5-day trip to Paris for couples interested in food")
+    print("  - I want to visit Tokyo for 7 days with a moderate budget")
+    print("  - Create an itinerary for a family trip to Barcelona")
+    print("=" * 80)
+    print()
+
+    while True:
+        try:
+            # Get user input
+            user_query = input("\n✈️  Your travel request: ").strip()
+
+            if not user_query:
+                print("Please enter a travel planning request.")
+                continue
+
+            if user_query.lower() in ['quit', 'exit', 'q']:
+                print("\n👋 Thank you for using the Travel Agent. Safe travels!")
+                break
+
+            print("\n🔒 Running travel agent with PII guardrail check...")
+            print("(Agent output will appear after guardrail validation)\n")
+
+            # Run the guarded travel agent
+            result = await guarded_travel_agent(user_query)
+
+            # Display the result
+            print("\n" + "=" * 80)
+            print("📋 FINAL RESPONSE (after PII guardrail check):")
+            print("=" * 80)
+            response_text = result.get("text", "")
+            print(response_text)
+            print("=" * 80)
+
+            # Check if this was a warning (PII detected)
+            if "PRIVACY ALERT" in response_text:
+                print("❌ Response blocked due to PII detection")
+            else:
+                print("✅ Response approved by guardrail")
+
+            print("=" * 80)
+
+        except KeyboardInterrupt:
+            print("\n\n👋 Session interrupted. Goodbye!")
+            break
+        except Exception as e:
+            print(f"\n❌ Error: {e}")
+            import traceback
+            traceback.print_exc()
+            print("\nPlease try again or type 'quit' to exit.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())