|
| 1 | +from __future__ import annotations as _annotations |
| 2 | + |
| 3 | +import asyncio |
| 4 | + |
| 5 | +####################################### |
| 6 | +import os |
| 7 | +import random |
| 8 | +import uuid |
| 9 | + |
| 10 | +from agents import ( |
| 11 | + Agent, |
| 12 | + HandoffOutputItem, |
| 13 | + ItemHelpers, |
| 14 | + MessageOutputItem, |
| 15 | + RunContextWrapper, |
| 16 | + Runner, |
| 17 | + ToolCallItem, |
| 18 | + ToolCallOutputItem, |
| 19 | + TResponseInputItem, |
| 20 | + function_tool, |
| 21 | + handoff, |
| 22 | + set_default_openai_client, |
| 23 | + trace, |
| 24 | +) |
| 25 | +from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX |
| 26 | +from dotenv import load_dotenv |
| 27 | +from openai import AsyncAzureOpenAI |
| 28 | +from pydantic import BaseModel |
| 29 | + |
| 30 | +from opentelemetry.instrumentation.openai_agents import ( |
| 31 | + OpenAIAgentsInstrumentor, |
| 32 | +) |
| 33 | + |
| 34 | +load_dotenv() |
| 35 | + |
| 36 | + |
| 37 | +OpenAIAgentsInstrumentor().instrument() |
| 38 | + |
| 39 | +custom_client = AsyncAzureOpenAI( |
| 40 | + azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"], |
| 41 | + azure_deployment="gpt-4.1-mini", |
| 42 | + api_version=os.environ["AZURE_OPENAI_API_VERSION"], |
| 43 | + api_key=os.environ["AZURE_OPENAI_API_KEY"], |
| 44 | +) |
| 45 | + |
| 46 | +set_default_openai_client(custom_client, use_for_tracing=False) |
| 47 | + |
| 48 | +####################################### |
| 49 | + |
| 50 | + |
| 51 | +### CONTEXT |
| 52 | + |
| 53 | + |
| 54 | +class AirlineAgentContext(BaseModel): |
| 55 | + passenger_name: str | None = None |
| 56 | + confirmation_number: str | None = None |
| 57 | + seat_number: str | None = None |
| 58 | + flight_number: str | None = None |
| 59 | + |
| 60 | + |
| 61 | +### TOOLS |
| 62 | + |
| 63 | + |
| 64 | +@function_tool( |
| 65 | + name_override="faq_lookup_tool", |
| 66 | + description_override="Lookup frequently asked questions.", |
| 67 | +) |
| 68 | +async def faq_lookup_tool(question: str) -> str: |
| 69 | + if "bag" in question or "baggage" in question: |
| 70 | + return ( |
| 71 | + "You are allowed to bring one bag on the plane. " |
| 72 | + "It must be under 50 pounds and 22 inches x 14 inches x 9 inches." |
| 73 | + ) |
| 74 | + elif "seats" in question or "plane" in question: |
| 75 | + return ( |
| 76 | + "There are 120 seats on the plane. " |
| 77 | + "There are 22 business class seats and 98 economy seats. " |
| 78 | + "Exit rows are rows 4 and 16. " |
| 79 | + "Rows 5-8 are Economy Plus, with extra legroom. " |
| 80 | + ) |
| 81 | + elif "wifi" in question: |
| 82 | + return "We have free wifi on the plane, join Airline-Wifi" |
| 83 | + return "I'm sorry, I don't know the answer to that question." |
| 84 | + |
| 85 | + |
| 86 | +@function_tool |
| 87 | +async def update_seat( |
| 88 | + context: RunContextWrapper[AirlineAgentContext], |
| 89 | + confirmation_number: str, |
| 90 | + new_seat: str, |
| 91 | +) -> str: |
| 92 | + """ |
| 93 | + Update the seat for a given confirmation number. |
| 94 | +
|
| 95 | + Args: |
| 96 | + confirmation_number: The confirmation number for the flight. |
| 97 | + new_seat: The new seat to update to. |
| 98 | + """ |
| 99 | + # Update the context based on the customer's input |
| 100 | + context.context.confirmation_number = confirmation_number |
| 101 | + context.context.seat_number = new_seat |
| 102 | + # Ensure that the flight number has been set by the incoming handoff |
| 103 | + assert ( |
| 104 | + context.context.flight_number is not None |
| 105 | + ), "Flight number is required" |
| 106 | + return f"Updated seat to {new_seat} for confirmation number {confirmation_number}" |
| 107 | + |
| 108 | + |
| 109 | +### HOOKS |
| 110 | + |
| 111 | + |
| 112 | +async def on_seat_booking_handoff( |
| 113 | + context: RunContextWrapper[AirlineAgentContext], |
| 114 | +) -> None: |
| 115 | + flight_number = f"FLT-{random.randint(100, 999)}" |
| 116 | + context.context.flight_number = flight_number |
| 117 | + |
| 118 | + |
| 119 | +### AGENTS |
| 120 | + |
| 121 | +faq_agent = Agent[AirlineAgentContext]( |
| 122 | + name="FAQ Agent", |
| 123 | + handoff_description="A helpful agent that can answer questions about the airline.", |
| 124 | + instructions=f"""{RECOMMENDED_PROMPT_PREFIX} |
| 125 | + You are an FAQ agent. If you are speaking to a customer, you probably were transferred to from the triage agent. |
| 126 | + Use the following routine to support the customer. |
| 127 | + # Routine |
| 128 | + 1. Identify the last question asked by the customer. |
| 129 | + 2. Use the faq lookup tool to answer the question. Do not rely on your own knowledge. |
| 130 | + 3. If you cannot answer the question, transfer back to the triage agent.""", |
| 131 | + tools=[faq_lookup_tool], |
| 132 | +) |
| 133 | + |
| 134 | +seat_booking_agent = Agent[AirlineAgentContext]( |
| 135 | + name="Seat Booking Agent", |
| 136 | + handoff_description="A helpful agent that can update a seat on a flight.", |
| 137 | + instructions=f"""{RECOMMENDED_PROMPT_PREFIX} |
| 138 | + You are a seat booking agent. If you are speaking to a customer, you probably were transferred to from the triage agent. |
| 139 | + Use the following routine to support the customer. |
| 140 | + # Routine |
| 141 | + 1. Ask for their confirmation number. |
| 142 | + 2. Ask the customer what their desired seat number is. |
| 143 | + 3. Use the update seat tool to update the seat on the flight. |
| 144 | + If the customer asks a question that is not related to the routine, transfer back to the triage agent. """, |
| 145 | + tools=[update_seat], |
| 146 | +) |
| 147 | + |
| 148 | +triage_agent = Agent[AirlineAgentContext]( |
| 149 | + name="Triage Agent", |
| 150 | + handoff_description="A triage agent that can delegate a customer's request to the appropriate agent.", |
| 151 | + instructions=( |
| 152 | + f"{RECOMMENDED_PROMPT_PREFIX} " |
| 153 | + "You are a helpful triaging agent. You can use your tools to delegate questions to other appropriate agents." |
| 154 | + ), |
| 155 | + handoffs=[ |
| 156 | + faq_agent, |
| 157 | + handoff(agent=seat_booking_agent, on_handoff=on_seat_booking_handoff), |
| 158 | + ], |
| 159 | +) |
| 160 | + |
| 161 | +faq_agent.handoffs.append(triage_agent) |
| 162 | +seat_booking_agent.handoffs.append(triage_agent) |
| 163 | + |
| 164 | + |
| 165 | +### RUN |
| 166 | + |
| 167 | + |
| 168 | +async def main(): |
| 169 | + current_agent: Agent[AirlineAgentContext] = triage_agent |
| 170 | + input_items: list[TResponseInputItem] = [] |
| 171 | + context = AirlineAgentContext() |
| 172 | + |
| 173 | + # Normally, each input from the user would be an API request to your app, and you can wrap the request in a trace() |
| 174 | + # Here, we'll just use a random UUID for the conversation ID |
| 175 | + conversation_id = uuid.uuid4().hex[:16] |
| 176 | + |
| 177 | + while True: |
| 178 | + user_input = input("Enter your message: ") |
| 179 | + with trace("Customer service", group_id=conversation_id): |
| 180 | + input_items.append({"content": user_input, "role": "user"}) |
| 181 | + result = await Runner.run( |
| 182 | + current_agent, input_items, context=context |
| 183 | + ) |
| 184 | + |
| 185 | + for new_item in result.new_items: |
| 186 | + agent_name = new_item.agent.name |
| 187 | + if isinstance(new_item, MessageOutputItem): |
| 188 | + print( |
| 189 | + f"{agent_name}: {ItemHelpers.text_message_output(new_item)}" |
| 190 | + ) |
| 191 | + elif isinstance(new_item, HandoffOutputItem): |
| 192 | + print( |
| 193 | + f"Handed off from {new_item.source_agent.name} to {new_item.target_agent.name}" |
| 194 | + ) |
| 195 | + elif isinstance(new_item, ToolCallItem): |
| 196 | + print(f"{agent_name}: Calling a tool") |
| 197 | + elif isinstance(new_item, ToolCallOutputItem): |
| 198 | + print(f"{agent_name}: Tool call output: {new_item.output}") |
| 199 | + else: |
| 200 | + print( |
| 201 | + f"{agent_name}: Skipping item: {new_item.__class__.__name__}" |
| 202 | + ) |
| 203 | + input_items = result.to_input_list() |
| 204 | + current_agent = result.last_agent |
| 205 | + |
| 206 | + |
| 207 | +if __name__ == "__main__": |
| 208 | + asyncio.run(main()) |
0 commit comments