|
| 1 | +from __future__ import annotations as _annotations |
| 2 | + |
| 3 | +import random |
| 4 | + |
| 5 | +from azure.durable_functions.openai_agents import DurableAIAgentContext |
| 6 | +from pydantic import BaseModel |
| 7 | + |
| 8 | +from agents import ( |
| 9 | + Agent, |
| 10 | + HandoffOutputItem, |
| 11 | + ItemHelpers, |
| 12 | + MessageOutputItem, |
| 13 | + RunContextWrapper, |
| 14 | + Runner, |
| 15 | + ToolCallItem, |
| 16 | + ToolCallOutputItem, |
| 17 | + TResponseInputItem, |
| 18 | + function_tool, |
| 19 | + handoff, |
| 20 | + trace, |
| 21 | +) |
| 22 | +from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX |
| 23 | + |
| 24 | +### CONTEXT |
| 25 | + |
| 26 | + |
| 27 | +class AirlineAgentContext(BaseModel): |
| 28 | + passenger_name: str | None = None |
| 29 | + confirmation_number: str | None = None |
| 30 | + seat_number: str | None = None |
| 31 | + flight_number: str | None = None |
| 32 | + |
| 33 | + |
| 34 | +### TOOLS |
| 35 | + |
| 36 | + |
| 37 | +@function_tool( |
| 38 | + name_override="faq_lookup_tool", description_override="Lookup frequently asked questions." |
| 39 | +) |
| 40 | +async def faq_lookup_tool(question: str) -> str: |
| 41 | + question_lower = question.lower() |
| 42 | + if any( |
| 43 | + keyword in question_lower |
| 44 | + for keyword in ["bag", "baggage", "luggage", "carry-on", "hand luggage", "hand carry"] |
| 45 | + ): |
| 46 | + return ( |
| 47 | + "You are allowed to bring one bag on the plane. " |
| 48 | + "It must be under 50 pounds and 22 inches x 14 inches x 9 inches." |
| 49 | + ) |
| 50 | + elif any(keyword in question_lower for keyword in ["seat", "seats", "seating", "plane"]): |
| 51 | + return ( |
| 52 | + "There are 120 seats on the plane. " |
| 53 | + "There are 22 business class seats and 98 economy seats. " |
| 54 | + "Exit rows are rows 4 and 16. " |
| 55 | + "Rows 5-8 are Economy Plus, with extra legroom. " |
| 56 | + ) |
| 57 | + elif any( |
| 58 | + keyword in question_lower |
| 59 | + for keyword in ["wifi", "internet", "wireless", "connectivity", "network", "online"] |
| 60 | + ): |
| 61 | + return "We have free wifi on the plane, join Airline-Wifi" |
| 62 | + return "I'm sorry, I don't know the answer to that question." |
| 63 | + |
| 64 | + |
| 65 | +@function_tool |
| 66 | +async def update_seat( |
| 67 | + context: RunContextWrapper[AirlineAgentContext], confirmation_number: str, new_seat: str |
| 68 | +) -> str: |
| 69 | + """ |
| 70 | + Update the seat for a given confirmation number. |
| 71 | +
|
| 72 | + Args: |
| 73 | + confirmation_number: The confirmation number for the flight. |
| 74 | + new_seat: The new seat to update to. |
| 75 | + """ |
| 76 | + # Update the context based on the customer's input |
| 77 | + context.context.confirmation_number = confirmation_number |
| 78 | + context.context.seat_number = new_seat |
| 79 | + # Ensure that the flight number has been set by the incoming handoff |
| 80 | + assert context.context.flight_number is not None, "Flight number is required" |
| 81 | + return f"Updated seat to {new_seat} for confirmation number {confirmation_number}" |
| 82 | + |
| 83 | + |
| 84 | +### HOOKS |
| 85 | + |
| 86 | + |
| 87 | +async def on_seat_booking_handoff(context: RunContextWrapper[AirlineAgentContext]) -> None: |
| 88 | + flight_number = f"FLT-{random.randint(100, 999)}" |
| 89 | + context.context.flight_number = flight_number |
| 90 | + |
| 91 | + |
| 92 | +### AGENTS |
| 93 | + |
| 94 | +faq_agent = Agent[AirlineAgentContext]( |
| 95 | + name="FAQ Agent", |
| 96 | + handoff_description="A helpful agent that can answer questions about the airline.", |
| 97 | + instructions=f"""{RECOMMENDED_PROMPT_PREFIX} |
| 98 | + You are an FAQ agent. If you are speaking to a customer, you probably were transferred to from the triage agent. |
| 99 | + Use the following routine to support the customer. |
| 100 | + # Routine |
| 101 | + 1. Identify the last question asked by the customer. |
| 102 | + 2. Use the faq lookup tool to answer the question. Do not rely on your own knowledge. |
| 103 | + 3. If you cannot answer the question, transfer back to the triage agent.""", |
| 104 | + tools=[faq_lookup_tool], |
| 105 | +) |
| 106 | + |
| 107 | +seat_booking_agent = Agent[AirlineAgentContext]( |
| 108 | + name="Seat Booking Agent", |
| 109 | + handoff_description="A helpful agent that can update a seat on a flight.", |
| 110 | + instructions=f"""{RECOMMENDED_PROMPT_PREFIX} |
| 111 | + You are a seat booking agent. If you are speaking to a customer, you probably were transferred to from the triage agent. |
| 112 | + Use the following routine to support the customer. |
| 113 | + # Routine |
| 114 | + 1. Ask for their confirmation number. |
| 115 | + 2. Ask the customer what their desired seat number is. |
| 116 | + 3. Use the update seat tool to update the seat on the flight. |
| 117 | + If the customer asks a question that is not related to the routine, transfer back to the triage agent. """, |
| 118 | + tools=[update_seat], |
| 119 | +) |
| 120 | + |
| 121 | +triage_agent = Agent[AirlineAgentContext]( |
| 122 | + name="Triage Agent", |
| 123 | + handoff_description="A triage agent that can delegate a customer's request to the appropriate agent.", |
| 124 | + instructions=( |
| 125 | + f"{RECOMMENDED_PROMPT_PREFIX} " |
| 126 | + "You are a helpful triaging agent. You can use your tools to delegate questions to other appropriate agents." |
| 127 | + ), |
| 128 | + handoffs=[ |
| 129 | + faq_agent, |
| 130 | + handoff(agent=seat_booking_agent, on_handoff=on_seat_booking_handoff), |
| 131 | + ], |
| 132 | +) |
| 133 | + |
| 134 | +faq_agent.handoffs.append(triage_agent) |
| 135 | +seat_booking_agent.handoffs.append(triage_agent) |
| 136 | + |
| 137 | + |
| 138 | +### RUN |
| 139 | + |
| 140 | + |
| 141 | +def main(context: DurableAIAgentContext): |
| 142 | + current_agent: Agent[AirlineAgentContext] = triage_agent |
| 143 | + input_items: list[TResponseInputItem] = [] |
| 144 | + airline_agent_context = AirlineAgentContext() |
| 145 | + |
| 146 | + conversation_id = context.instance_id |
| 147 | + |
| 148 | + context.set_custom_status("How can I help you today?") |
| 149 | + while True: |
| 150 | + user_input = yield context.wait_for_external_event("UserInput") |
| 151 | + if user_input is None or user_input.strip().lower() in ["exit", "quit", "bye"]: |
| 152 | + context.set_custom_status("Goodbye!") |
| 153 | + break |
| 154 | + context.set_custom_status("Thinking...") |
| 155 | + with trace("Customer service", group_id=conversation_id): |
| 156 | + input_items.append({"content": user_input, "role": "user"}) |
| 157 | + result = Runner.run_sync(current_agent, input_items, context=airline_agent_context) |
| 158 | + |
| 159 | + for new_item in result.new_items: |
| 160 | + agent_name = new_item.agent.name |
| 161 | + if isinstance(new_item, MessageOutputItem): |
| 162 | + print(f"{agent_name}: {ItemHelpers.text_message_output(new_item)}") |
| 163 | + elif isinstance(new_item, HandoffOutputItem): |
| 164 | + print( |
| 165 | + f"Handed off from {new_item.source_agent.name} to {new_item.target_agent.name}" |
| 166 | + ) |
| 167 | + elif isinstance(new_item, ToolCallItem): |
| 168 | + print(f"{agent_name}: Calling a tool") |
| 169 | + elif isinstance(new_item, ToolCallOutputItem): |
| 170 | + print(f"{agent_name}: Tool call output: {new_item.output}") |
| 171 | + else: |
| 172 | + print(f"{agent_name}: Skipping item: {new_item.__class__.__name__}") |
| 173 | + input_items = result.to_input_list() |
| 174 | + current_agent = result.last_agent |
| 175 | + |
| 176 | + context.set_custom_status(result.final_output) |
0 commit comments