AGENTOPS_API_KEY=your_agentops_api_key_here
OPENAI_API_KEY=your_openai_api_key_here

# Required imports - Note: agentops must be imported before agents
import asyncio
import os
import uuid
from typing import Any, Literal
from dataclasses import dataclass
from dotenv import load_dotenv
from pydantic import BaseModel, Field

# Import agentops FIRST
import agentops

# Then import agents
from agents import (
    Agent,
    ItemHelpers,
    MessageOutputItem,
    Runner,
    trace,
    TResponseInputItem,
    FunctionToolResult,
    ModelSettings,
    RunContextWrapper,
    ToolsToFinalOutputFunction,
    ToolsToFinalOutputResult,
    function_tool,
    GuardrailFunctionOutput,
    InputGuardrailTripwireTriggered,
    OutputGuardrailTripwireTriggered,
    input_guardrail,
    output_guardrail,
)

from openai.types.responses import ResponseTextDeltaEvent

# Load environment variables and set API keys
load_dotenv()
os.environ["AGENTOPS_API_KEY"] = os.getenv("AGENTOPS_API_KEY", "your_api_key_here")
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_openai_api_key_here")

# Initialize AgentOps
agentops.init(auto_start_session=False)
# Note: tracer will be defined in each section's cell for clarity, using the specific tags for that pattern.

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Agents as Tools Pattern", tags=["agentops-example", "openai-agents"])

# Define specialized translation agents
spanish_agent = Agent(
    name="spanish_agent",
    instructions="You translate the user's message to Spanish",
    handoff_description="An english to spanish translator",
)

french_agent = Agent(
    name="french_agent",
    instructions="You translate the user's message to French",
    handoff_description="An english to french translator",
)

italian_agent = Agent(
    name="italian_agent",
    instructions="You translate the user's message to Italian",
    handoff_description="An english to italian translator",
)

# Orchestrator agent that uses other agents as tools
orchestrator_agent = Agent(
    name="orchestrator_agent",
    instructions=(
        "You are a translation agent. You use the tools given to you to translate."
        "If asked for multiple translations, you call the relevant tools in order."
        "You never translate on your own, you always use the provided tools."
    ),
    tools=[
        spanish_agent.as_tool(
            tool_name="translate_to_spanish",
            tool_description="Translate the user's message to Spanish",
        ),
        french_agent.as_tool(
            tool_name="translate_to_french",
            tool_description="Translate the user's message to French",
        ),
        italian_agent.as_tool(
            tool_name="translate_to_italian",
            tool_description="Translate the user's message to Italian",
        ),
    ],
)

synthesizer_agent = Agent(
    name="synthesizer_agent",
    instructions="You inspect translations, correct them if needed, and produce a final concatenated response.",
)


async def run_agents_as_tools_demo():
    msg = "Hello, how are you today?"
    print(f"Input: {msg}")

    with trace("Orchestrator evaluator"):
        orchestrator_result = await Runner.run(orchestrator_agent, msg)

        for item in orchestrator_result.new_items:
            if isinstance(item, MessageOutputItem):
                text = ItemHelpers.text_message_output(item)
                if text:
                    print(f"  - Translation step: {text}")

        synthesizer_result = await Runner.run(synthesizer_agent, orchestrator_result.to_input_list())

    print(f"Final response: {synthesizer_result.final_output}")


# Run the demo
await run_agents_as_tools_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Deterministic Flow Pattern", tags=["agentops-example", "openai-agents"])

# Define the story generation workflow
story_outline_agent = Agent(
    name="story_outline_agent",
    instructions="Generate a very short story outline based on the user's input.",
)


class OutlineCheckerOutput(BaseModel):
    good_quality: bool
    is_scifi: bool


outline_checker_agent = Agent(
    name="outline_checker_agent",
    instructions="Read the given story outline, and judge the quality. Also, determine if it is a scifi story.",
    output_type=OutlineCheckerOutput,
)

story_agent = Agent(
    name="story_agent",
    instructions="Write a short story based on the given outline.",
    output_type=str,
)


async def run_deterministic_flow_demo():
    input_prompt = "A story about robots exploring space"
    print(f"Input: {input_prompt}")

    with trace("Deterministic story flow"):
        # 1. Generate an outline
        outline_result = await Runner.run(story_outline_agent, input_prompt)
        print("Outline generated")

        # 2. Check the outline
        outline_checker_result = await Runner.run(outline_checker_agent, outline_result.final_output)

        # 3. Add a gate to stop if the outline is not good quality or not a scifi story
        assert isinstance(outline_checker_result.final_output, OutlineCheckerOutput)
        if not outline_checker_result.final_output.good_quality:
            print("Outline is not good quality, so we stop here.")
            return

        if not outline_checker_result.final_output.is_scifi:
            print("Outline is not a scifi story, so we stop here.")
            return

        print("Outline is good quality and a scifi story, so we continue to write the story.")

        # 4. Write the story
        story_result = await Runner.run(story_agent, outline_result.final_output)
        print(f"Story: {story_result.final_output}")


# Run the demo
await run_deterministic_flow_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Forcing Tool Use Pattern", tags=["agentops-example", "openai-agents"])


# Define the weather tool and agent
class Weather(BaseModel):
    city: str
    temperature_range: str
    conditions: str


@function_tool
def get_weather(city: str) -> Weather:
    print("[debug] get_weather called")
    return Weather(city=city, temperature_range="14-20C", conditions="Sunny with wind")


async def custom_tool_use_behavior(
    context: RunContextWrapper[Any], results: list[FunctionToolResult]
) -> ToolsToFinalOutputResult:
    weather: Weather = results[0].output
    return ToolsToFinalOutputResult(is_final_output=True, final_output=f"{weather.city} is {weather.conditions}.")


# User can choose which behavior to test
print("Choose tool use behavior:")
print("1. default - Send tool output to LLM")
print("2. first_tool - Use first tool result as final output")
print("3. custom - Use custom tool behavior function")

choice = input("Enter choice (1, 2, or 3): ").strip()

if choice == "1":
    tool_use_behavior = "default"
elif choice == "2":
    tool_use_behavior = "first_tool"
elif choice == "3":
    tool_use_behavior = "custom"
else:
    tool_use_behavior = "default"
    print("Invalid choice, using default")


async def run_forcing_tool_use_demo(tool_use_behavior: str):
    print(f"Testing {tool_use_behavior} behavior:")

    if tool_use_behavior == "default":
        behavior: Literal["run_llm_again", "stop_on_first_tool"] | ToolsToFinalOutputFunction = "run_llm_again"
    elif tool_use_behavior == "first_tool":
        behavior = "stop_on_first_tool"
    elif tool_use_behavior == "custom":
        behavior = custom_tool_use_behavior

    agent = Agent(
        name="Weather agent",
        instructions="You are a helpful agent.",
        tools=[get_weather],
        tool_use_behavior=behavior,
        model_settings=ModelSettings(tool_choice="required" if tool_use_behavior != "default" else None),
    )

    result = await Runner.run(agent, input="What's the weather in Tokyo?")
    print(f"Result: {result.final_output}")


# Run the demo
await run_forcing_tool_use_demo(tool_use_behavior)

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Input Guardrails Pattern", tags=["agentops-example", "openai-agents"])


# Define the guardrail
class MathHomeworkOutput(BaseModel):
    reasoning: str
    is_math_homework: bool


guardrail_agent = Agent(
    name="Guardrail check",
    instructions="Check if the user is asking you to do their math homework.",
    output_type=MathHomeworkOutput,
)


@input_guardrail
async def math_guardrail(
    context: RunContextWrapper[None], agent: Agent, input: str | list[TResponseInputItem]
) -> GuardrailFunctionOutput:
    result = await Runner.run(guardrail_agent, input, context=context.context)
    final_output = result.final_output_as(MathHomeworkOutput)

    return GuardrailFunctionOutput(
        output_info=final_output,
        tripwire_triggered=final_output.is_math_homework,
    )


async def run_input_guardrails_demo():
    agent = Agent(
        name="Customer support agent",
        instructions="You are a customer support agent. You help customers with their questions.",
        input_guardrails=[math_guardrail],
    )

    # Test with different inputs
    test_inputs = ["What's the capital of California?", "Can you help me solve for x: 2x + 5 = 11"]

    for user_input in test_inputs:
        print(f"Testing input: {user_input}")
        try:
            result = await Runner.run(agent, user_input)
            print(f"Response: {result.final_output}")
        except InputGuardrailTripwireTriggered:
            message = "Sorry, I can't help you with your math homework."
            print(f"Guardrail triggered: {message}")
        print()


# Run the demo
await run_input_guardrails_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="LLM as a Judge Pattern", tags=["agentops-example", "openai-agents"])

# Define the story generation and evaluation agents
story_outline_generator = Agent(
    name="story_outline_generator",
    instructions=(
        "You generate a very short story outline based on the user's input."
        "If there is any feedback provided, use it to improve the outline."
    ),
)


@dataclass
class EvaluationFeedback:
    feedback: str
    score: Literal["pass", "needs_improvement", "fail"]


evaluator = Agent[None](
    name="evaluator",
    instructions=(
        "You evaluate a story outline and decide if it's good enough."
        "If it's not good enough, you provide feedback on what needs to be improved."
        "Never give it a pass on the first try."
    ),
    output_type=EvaluationFeedback,
)


async def run_llm_as_judge_demo():
    msg = "A story about time travel"
    print(f"Input: {msg}")
    input_items: list[TResponseInputItem] = [{"content": msg, "role": "user"}]

    latest_outline: str | None = None
    iteration = 0
    max_iterations = 3  # Limit iterations for demo

    with trace("LLM as a judge"):
        while iteration < max_iterations:
            iteration += 1
            print(f"Iteration {iteration}:")

            story_outline_result = await Runner.run(
                story_outline_generator,
                input_items,
            )

            input_items = story_outline_result.to_input_list()
            latest_outline = ItemHelpers.text_message_outputs(story_outline_result.new_items)
            print("Story outline generated")

            evaluator_result = await Runner.run(evaluator, input_items)
            result: EvaluationFeedback = evaluator_result.final_output

            print(f"Evaluator score: {result.score}")

            if result.score == "pass":
                print("Story outline is good enough, exiting.")
                break

            print("Re-running with feedback")
            input_items.append({"content": f"Feedback: {result.feedback}", "role": "user"})

    print(f"Final story outline: {latest_outline}")


# Run the demo
await run_llm_as_judge_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Output Guardrails Pattern", tags=["agentops-example", "openai-agents"])


# The agent's output type
class MessageOutput(BaseModel):
    reasoning: str = Field(description="Thoughts on how to respond to the user's message")
    response: str = Field(description="The response to the user's message")
    user_name: str | None = Field(description="The name of the user who sent the message, if known")


@output_guardrail
async def sensitive_data_check(
    context: RunContextWrapper, agent: Agent, output: MessageOutput
) -> GuardrailFunctionOutput:
    phone_number_in_response = "650" in output.response
    phone_number_in_reasoning = "650" in output.reasoning

    return GuardrailFunctionOutput(
        output_info={
            "phone_number_in_response": phone_number_in_response,
            "phone_number_in_reasoning": phone_number_in_reasoning,
        },
        tripwire_triggered=phone_number_in_response or phone_number_in_reasoning,
    )


output_guardrail_agent = Agent(
    name="Assistant",
    instructions="You are a helpful assistant.",
    output_type=MessageOutput,
    output_guardrails=[sensitive_data_check],
)


async def run_output_guardrails_demo():
    # Test with safe input
    print("Testing with safe input:")
    try:
        result = await Runner.run(output_guardrail_agent, "What's the capital of California?")
        print("Safe message passed")
        print(f"Response: {result.final_output.response}")
    except OutputGuardrailTripwireTriggered as e:
        print(f"Unexpected guardrail trigger: {e.guardrail_result.output.output_info}")

    print()
    # Test with potentially sensitive input
    print("Testing with potentially sensitive input:")
    try:
        result = await Runner.run(output_guardrail_agent, "My phone number is 650-123-4567. Where do you think I live?")
        print(f"Guardrail didn't trip - this is unexpected. Output: {result.final_output.response}")
    except OutputGuardrailTripwireTriggered as e:
        print(f"Guardrail tripped as expected. Info: {e.guardrail_result.output.output_info}")


# Run the demo
await run_output_guardrails_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Output Guardrails Pattern", tags=["agentops-example", "openai-agents"])

# Define agents for parallelization
spanish_translation_agent = Agent(
    name="spanish_agent",
    instructions="You translate the user's message to Spanish",
)

translation_picker = Agent(
    name="translation_picker",
    instructions="You pick the best Spanish translation from the given options.",
)


async def run_parallelization_demo():
    msg = "Good morning, I hope you have a wonderful day!"
    print(f"Input: {msg}")

    with trace("Parallel translation"):
        # Run three translation agents in parallel
        res_1, res_2, res_3 = await asyncio.gather(
            Runner.run(spanish_translation_agent, msg),
            Runner.run(spanish_translation_agent, msg),
            Runner.run(spanish_translation_agent, msg),
        )

        outputs = [
            ItemHelpers.text_message_outputs(res_1.new_items),
            ItemHelpers.text_message_outputs(res_2.new_items),
            ItemHelpers.text_message_outputs(res_3.new_items),
        ]

        translations = "\n\n".join(outputs)
        print(f"\n\nTranslations:\n\n{translations}")

        best_translation = await Runner.run(
            translation_picker,
            f"Input: {msg}\n\nTranslations:\n{translations}",
        )

    print("\n\n-----")
    print(f"Best translation: {best_translation.final_output}")


# Run the demo
await run_parallelization_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Routing Pattern", tags=["agentops-example", "openai-agents"])

# Define language-specific agents
french_routing_agent = Agent(
    name="french_agent",
    instructions="You only speak French",
)

spanish_routing_agent = Agent(
    name="spanish_agent",
    instructions="You only speak Spanish",
)

english_routing_agent = Agent(
    name="english_agent",
    instructions="You only speak English",
)

triage_agent = Agent(
    name="triage_agent",
    instructions="Handoff to the appropriate agent based on the language of the request.",
    handoffs=[french_routing_agent, spanish_routing_agent, english_routing_agent],
)


async def run_routing_demo():
    # Create an ID for this conversation
    conversation_id = str(uuid.uuid4().hex[:16])

    test_messages = ["Hello, how can you help me?", "Bonjour, comment allez-vous?", "Hola, ¿cómo estás?"]

    for msg in test_messages:
        print(f"\nTesting message: {msg}")

        with trace("Routing example", group_id=conversation_id):
            inputs = [{"content": msg, "role": "user"}]
            result = await Runner.run(triage_agent, input=inputs)
            print(f"Response: {result.final_output}")


# Run the demo
await run_routing_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")

# Start the AgentOps trace session
tracer = agentops.start_trace(trace_name="Streaming Guardrails Pattern", tags=["agentops-example", "openai-agents"])

# Define streaming guardrail agent
streaming_agent = Agent(
    name="Assistant",
    instructions=(
        "You are a helpful assistant. You ALWAYS write long responses, making sure to be verbose and detailed."
    ),
)


class GuardrailOutput(BaseModel):
    reasoning: str = Field(description="Reasoning about whether the response could be understood by a ten year old.")
    is_readable_by_ten_year_old: bool = Field(description="Whether the response is understandable by a ten year old.")


guardrail_streaming_agent = Agent(
    name="Checker",
    instructions=(
        "You will be given a question and a response. Your goal is to judge whether the response "
        "is simple enough to be understood by a ten year old."
    ),
    output_type=GuardrailOutput,
    model="gpt-4o-mini",
)


async def check_guardrail(text: str) -> GuardrailOutput:
    result = await Runner.run(guardrail_streaming_agent, text)
    return result.final_output_as(GuardrailOutput)


async def run_streaming_guardrails_demo():
    question = "What is a black hole, and how does it behave?"
    print(f"Question: {question}")

    result = Runner.run_streamed(streaming_agent, question)
    current_text = ""

    # We will check the guardrail every N characters
    next_guardrail_check_len = 300
    guardrail_task = None

    async for event in result.stream_events():
        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
            print(event.data.delta, end="", flush=True)
            current_text += event.data.delta

            # Check if it's time to run the guardrail check
            if len(current_text) >= next_guardrail_check_len and not guardrail_task:
                print("\n[Running guardrail check]")
                guardrail_task = asyncio.create_task(check_guardrail(current_text))
                next_guardrail_check_len += 300

        # Every iteration of the loop, check if the guardrail has been triggered
        if guardrail_task and guardrail_task.done():
            guardrail_result = guardrail_task.result()
            if not guardrail_result.is_readable_by_ten_year_old:
                print("\n\n================\n\n")
                print(f"Guardrail triggered. Reasoning:\n{guardrail_result.reasoning}")
                break
            guardrail_task = None

    # Do one final check on the final output
    if current_text:
        guardrail_result = await check_guardrail(current_text)
        if not guardrail_result.is_readable_by_ten_year_old:
            print("\n\n================\n\n")
            print(f"Final guardrail triggered. Reasoning:\n{guardrail_result.reasoning}")


# Run the demo
await run_streaming_guardrails_demo()

# End the AgentOps trace session
agentops.end_trace(tracer, end_state="Success")