ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/prompt-engineering/tau_bench/envs/airline/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Sierra
+
+from tau_bench.envs.airline.env import MockAirlineDomainEnv as MockAirlineDomainEnv

chapter2/prompt-engineering/tau_bench/envs/airline/data/__init__.py

@@ -0,0 +1,21 @@
+# Copyright Sierra
+
+import json
+import os
+from typing import Any
+
+FOLDER_PATH = os.path.dirname(__file__)
+
+
+def load_data() -> dict[str, Any]:
+    with open(os.path.join(FOLDER_PATH, "flights.json")) as f:
+        flight_data = json.load(f)
+    with open(os.path.join(FOLDER_PATH, "reservations.json")) as f:
+        reservation_data = json.load(f)
+    with open(os.path.join(FOLDER_PATH, "users.json")) as f:
+        user_data = json.load(f)
+    return {
+        "flights": flight_data,
+        "reservations": reservation_data,
+        "users": user_data,
+    }

chapter2/prompt-engineering/tau_bench/envs/airline/env.py

@@ -0,0 +1,39 @@
+# Copyright Sierra
+
+from tau_bench.envs.airline.data import load_data
+from tau_bench.envs.airline.rules import RULES
+from tau_bench.envs.airline.tools import ALL_TOOLS
+from tau_bench.envs.airline.wiki import WIKI
+from tau_bench.envs.base import Env
+from typing import Optional, Union
+from tau_bench.envs.user import UserStrategy
+
+
+class MockAirlineDomainEnv(Env):
+    def __init__(
+        self,
+        user_strategy: Union[str, UserStrategy] = UserStrategy.LLM,
+        user_model: str = "gpt-4o",
+        user_provider: Optional[str] = None,
+        task_split: str = "test",
+        task_index: Optional[int] = None,
+        user_seed: Optional[int] = None,
+    ):
+        match task_split:
+            case "test":
+                from tau_bench.envs.airline.tasks_test import TASKS as tasks
+            case _:
+                raise ValueError(f"Unknown task split: {task_split}")
+        super().__init__(
+            data_load_func=load_data,
+            tools=ALL_TOOLS,
+            tasks=tasks,
+            wiki=WIKI,
+            rules=RULES,
+            user_strategy=user_strategy,
+            user_model=user_model,
+            user_provider=user_provider,
+            task_index=task_index,
+            user_seed=user_seed,
+        )
+        self.terminate_tools = ["transfer_to_human_agents"]

chapter2/prompt-engineering/tau_bench/envs/airline/rules.py

@@ -0,0 +1,3 @@
+# Copyright Sierra
+
+RULES = []

chapter2/prompt-engineering/tau_bench/envs/airline/tools/__init__.py

@@ -0,0 +1,33 @@
+# Copyright Sierra
+
+from .book_reservation import BookReservation
+from .calculate import Calculate
+from .cancel_reservation import CancelReservation
+from .get_reservation_details import GetReservationDetails
+from .get_user_details import GetUserDetails
+from .list_all_airports import ListAllAirports
+from .search_direct_flight import SearchDirectFlight
+from .search_onestop_flight import SearchOnestopFlight
+from .send_certificate import SendCertificate
+from .think import Think
+from .transfer_to_human_agents import TransferToHumanAgents
+from .update_reservation_baggages import UpdateReservationBaggages
+from .update_reservation_flights import UpdateReservationFlights
+from .update_reservation_passengers import UpdateReservationPassengers
+
+ALL_TOOLS = [
+    BookReservation,
+    Calculate,
+    CancelReservation,
+    GetReservationDetails,
+    GetUserDetails,
+    ListAllAirports,
+    SearchDirectFlight,
+    SearchOnestopFlight,
+    SendCertificate,
+    Think,
+    TransferToHumanAgents,
+    UpdateReservationBaggages,
+    UpdateReservationFlights,
+    UpdateReservationPassengers,
+]

chapter2/prompt-engineering/tau_bench/envs/airline/tools/book_reservation.py

@@ -0,0 +1,226 @@
+# Copyright Sierra
+
+import json
+from copy import deepcopy
+from typing import Any, Dict, List
+from tau_bench.envs.tool import Tool
+
+
+class BookReservation(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        user_id: str,
+        origin: str,
+        destination: str,
+        flight_type: str,
+        cabin: str,
+        flights: List[Dict[str, Any]],
+        passengers: List[Dict[str, Any]],
+        payment_methods: List[Dict[str, Any]],
+        total_baggages: int,
+        nonfree_baggages: int,
+        insurance: str,
+    ) -> str:
+        reservations, users = data["reservations"], data["users"]
+        if user_id not in users:
+            return "Error: user not found"
+        user = users[user_id]
+
+        # assume each task makes at most 3 reservations
+        reservation_id = "HATHAT"
+        if reservation_id in reservations:
+            reservation_id = "HATHAU"
+            if reservation_id in reservations:
+                reservation_id = "HATHAV"
+
+        reservation = {
+            "reservation_id": reservation_id,
+            "user_id": user_id,
+            "origin": origin,
+            "destination": destination,
+            "flight_type": flight_type,
+            "cabin": cabin,
+            "flights": deepcopy(flights),
+            "passengers": passengers,
+            "payment_history": payment_methods,
+            "created_at": "2024-05-15T15:00:00",
+            "total_baggages": total_baggages,
+            "nonfree_baggages": nonfree_baggages,
+            "insurance": insurance,
+        }
+
+        # update flights and calculate price
+        total_price = 0
+        for flight in reservation["flights"]:
+            flight_number = flight["flight_number"]
+            if flight_number not in data["flights"]:
+                return f"Error: flight {flight_number} not found"
+            flight_data = data["flights"][flight_number]
+            if flight["date"] not in flight_data["dates"]:
+                return (
+                    f"Error: flight {flight_number} not found on date {flight['date']}"
+                )
+            flight_date_data = flight_data["dates"][flight["date"]]
+            if flight_date_data["status"] != "available":
+                return f"Error: flight {flight_number} not available on date {flight['date']}"
+            if flight_date_data["available_seats"][cabin] < len(passengers):
+                return f"Error: not enough seats on flight {flight_number}"
+            flight["price"] = flight_date_data["prices"][cabin]
+            flight["origin"] = flight_data["origin"]
+            flight["destination"] = flight_data["destination"]
+            total_price += flight["price"] * len(passengers)
+
+        if insurance == "yes":
+            total_price += 30 * len(passengers)
+
+        total_price += 50 * nonfree_baggages
+
+        for payment_method in payment_methods:
+            payment_id = payment_method["payment_id"]
+            amount = payment_method["amount"]
+            if payment_id not in user["payment_methods"]:
+                return f"Error: payment method {payment_id} not found"
+            if user["payment_methods"][payment_id]["source"] in [
+                "gift_card",
+                "certificate",
+            ]:
+                if user["payment_methods"][payment_id]["amount"] < amount:
+                    return f"Error: not enough balance in payment method {payment_id}"
+        if sum(payment["amount"] for payment in payment_methods) != total_price:
+            return f"Error: payment amount does not add up, total price is {total_price}, but paid {sum(payment['amount'] for payment in payment_methods)}"
+
+        # if checks pass, deduct payment and update seats
+        for payment_method in payment_methods:
+            payment_id = payment_method["payment_id"]
+            amount = payment_method["amount"]
+            if user["payment_methods"][payment_id]["source"] == "gift_card":
+                user["payment_methods"][payment_id]["amount"] -= amount
+            elif user["payment_methods"][payment_id]["source"] == "certificate":
+                del user["payment_methods"][payment_id]
+
+        reservations[reservation_id] = reservation
+        user["reservations"].append(reservation_id)
+        return json.dumps(reservation)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "book_reservation",
+                "description": "Book a reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "user_id": {
+                            "type": "string",
+                            "description": "The ID of the user to book the reservation, such as 'sara_doe_496'.",
+                        },
+                        "origin": {
+                            "type": "string",
+                            "description": "The IATA code for the origin city, such as 'SFO'.",
+                        },
+                        "destination": {
+                            "type": "string",
+                            "description": "The IATA code for the destination city, such as 'JFK'.",
+                        },
+                        "flight_type": {
+                            "type": "string",
+                            "enum": ["one_way", "round_trip"],
+                        },
+                        "cabin": {
+                            "type": "string",
+                            "enum": [
+                                "basic_economy",
+                                "economy",
+                                "business",
+                            ],
+                        },
+                        "flights": {
+                            "type": "array",
+                            "description": "An array of objects containing details about each piece of flight.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "flight_number": {
+                                        "type": "string",
+                                        "description": "Flight number, such as 'HAT001'.",
+                                    },
+                                    "date": {
+                                        "type": "string",
+                                        "description": "The date for the flight in the format 'YYYY-MM-DD', such as '2024-05-01'.",
+                                    },
+                                },
+                                "required": ["flight_number", "date"],
+                            },
+                        },
+                        "passengers": {
+                            "type": "array",
+                            "description": "An array of objects containing details about each passenger.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "first_name": {
+                                        "type": "string",
+                                        "description": "The first name of the passenger, such as 'Noah'.",
+                                    },
+                                    "last_name": {
+                                        "type": "string",
+                                        "description": "The last name of the passenger, such as 'Brown'.",
+                                    },
+                                    "dob": {
+                                        "type": "string",
+                                        "description": "The date of birth of the passenger in the format 'YYYY-MM-DD', such as '1990-01-01'.",
+                                    },
+                                },
+                                "required": ["first_name", "last_name", "dob"],
+                            },
+                        },
+                        "payment_methods": {
+                            "type": "array",
+                            "description": "An array of objects containing details about each payment method.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "payment_id": {
+                                        "type": "string",
+                                        "description": "The payment id stored in user profile, such as 'credit_card_7815826', 'gift_card_7815826', 'certificate_7815826'.",
+                                    },
+                                    "amount": {
+                                        "type": "number",
+                                        "description": "The amount to be paid.",
+                                    },
+                                },
+                                "required": ["payment_id", "amount"],
+                            },
+                        },
+                        "total_baggages": {
+                            "type": "integer",
+                            "description": "The total number of baggage items included in the reservation.",
+                        },
+                        "nonfree_baggages": {
+                            "type": "integer",
+                            "description": "The number of non-free baggage items included in the reservation.",
+                        },
+                        "insurance": {
+                            "type": "string",
+                            "enum": ["yes", "no"],
+                        },
+                    },
+                    "required": [
+                        "user_id",
+                        "origin",
+                        "destination",
+                        "flight_type",
+                        "cabin",
+                        "flights",
+                        "passengers",
+                        "payment_methods",
+                        "total_baggages",
+                        "nonfree_baggages",
+                        "insurance",
+                    ],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/calculate.py

@@ -0,0 +1,35 @@
+# Copyright Sierra
+
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class Calculate(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], expression: str) -> str:
+        if not all(char in "0123456789+-*/(). " for char in expression):
+            return "Error: invalid characters in expression"
+        try:
+            return str(round(float(eval(expression, {"__builtins__": None}, {})), 2))
+        except Exception as e:
+            return f"Error: {e}"
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "calculate",
+                "description": "Calculate the result of a mathematical expression.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "expression": {
+                            "type": "string",
+                            "description": "The mathematical expression to calculate, such as '2 + 2'. The expression can contain numbers, operators (+, -, *, /), parentheses, and spaces.",
+                        },
+                    },
+                    "required": ["expression"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/cancel_reservation.py

@@ -0,0 +1,50 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class CancelReservation(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        reservation_id: str,
+    ) -> str:
+        reservations = data["reservations"]
+        if reservation_id not in reservations:
+            return "Error: reservation not found"
+        reservation = reservations[reservation_id]
+
+        # reverse the payment
+        refunds = []
+        for payment in reservation["payment_history"]:
+            refunds.append(
+                {
+                    "payment_id": payment["payment_id"],
+                    "amount": -payment["amount"],
+                }
+            )
+        reservation["payment_history"].extend(refunds)
+        reservation["status"] = "cancelled"
+        return json.dumps(reservation)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "cancel_reservation",
+                "description": "Cancel the whole reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "reservation_id": {
+                            "type": "string",
+                            "description": "The reservation ID, such as 'ZFA04Y'.",
+                        },
+                    },
+                    "required": ["reservation_id"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/get_reservation_details.py

@@ -0,0 +1,34 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class GetReservationDetails(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], reservation_id: str) -> str:
+        reservations = data["reservations"]
+        if reservation_id in reservations:
+            return json.dumps(reservations[reservation_id])
+        return "Error: user not found"
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "get_reservation_details",
+                "description": "Get the details of a reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "reservation_id": {
+                            "type": "string",
+                            "description": "The reservation id, such as '8JX2WO'.",
+                        },
+                    },
+                    "required": ["reservation_id"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/get_user_details.py

@@ -0,0 +1,34 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class GetUserDetails(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], user_id: str) -> str:
+        users = data["users"]
+        if user_id in users:
+            return json.dumps(users[user_id])
+        return "Error: user not found"
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "get_user_details",
+                "description": "Get the details of an user, including their reservations.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "user_id": {
+                            "type": "string",
+                            "description": "The user id, such as 'sara_doe_496'.",
+                        },
+                    },
+                    "required": ["user_id"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/list_all_airports.py

@@ -0,0 +1,70 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class ListAllAirports(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any]) -> str:
+        airports = [
+            "SFO",
+            "JFK",
+            "LAX",
+            "ORD",
+            "DFW",
+            "DEN",
+            "SEA",
+            "ATL",
+            "MIA",
+            "BOS",
+            "PHX",
+            "IAH",
+            "LAS",
+            "MCO",
+            "EWR",
+            "CLT",
+            "MSP",
+            "DTW",
+            "PHL",
+            "LGA",
+        ]
+        cities = [
+            "San Francisco",
+            "New York",
+            "Los Angeles",
+            "Chicago",
+            "Dallas",
+            "Denver",
+            "Seattle",
+            "Atlanta",
+            "Miami",
+            "Boston",
+            "Phoenix",
+            "Houston",
+            "Las Vegas",
+            "Orlando",
+            "Newark",
+            "Charlotte",
+            "Minneapolis",
+            "Detroit",
+            "Philadelphia",
+            "LaGuardia",
+        ]
+        return json.dumps({airport: city for airport, city in zip(airports, cities)})
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "list_all_airports",
+                "description": "List all airports and their cities.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {},
+                    "required": [],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/search_direct_flight.py

@@ -0,0 +1,50 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class SearchDirectFlight(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], origin: str, destination: str, date: str) -> str:
+        flights = data["flights"]
+        results = []
+        for flight in flights.values():
+            if flight["origin"] == origin and flight["destination"] == destination:
+                if (
+                    date in flight["dates"]
+                    and flight["dates"][date]["status"] == "available"
+                ):
+                    # results add flight except dates, but add flight["datas"][date]
+                    results.append({k: v for k, v in flight.items() if k != "dates"})
+                    results[-1].update(flight["dates"][date])
+        return json.dumps(results)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "search_direct_flight",
+                "description": "Search direct flights between two cities on a specific date.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "origin": {
+                            "type": "string",
+                            "description": "The origin city airport in three letters, such as 'JFK'.",
+                        },
+                        "destination": {
+                            "type": "string",
+                            "description": "The destination city airport in three letters, such as 'LAX'.",
+                        },
+                        "date": {
+                            "type": "string",
+                            "description": "The date of the flight in the format 'YYYY-MM-DD', such as '2024-01-01'.",
+                        },
+                    },
+                    "required": ["origin", "destination", "date"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/search_onestop_flight.py

@@ -0,0 +1,74 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class SearchOnestopFlight(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], origin: str, destination: str, date: str) -> str:
+        flights = data["flights"]
+        results = []
+        for flight1 in flights.values():
+            if flight1["origin"] == origin:
+                for flight2 in flights.values():
+                    if (
+                        flight2["destination"] == destination
+                        and flight1["destination"] == flight2["origin"]
+                    ):
+                        date2 = (
+                            f"2024-05-{int(date[-2:])+1}"
+                            if "+1" in flight1["scheduled_arrival_time_est"]
+                            else date
+                        )
+                        if (
+                            flight1["scheduled_arrival_time_est"]
+                            > flight2["scheduled_departure_time_est"]
+                        ):
+                            continue
+                        if date in flight1["dates"] and date2 in flight2["dates"]:
+                            if (
+                                flight1["dates"][date]["status"] == "available"
+                                and flight2["dates"][date2]["status"] == "available"
+                            ):
+                                result1 = {
+                                    k: v for k, v in flight1.items() if k != "dates"
+                                }
+                                result1.update(flight1["dates"][date])
+                                result1["date"] = date
+                                result2 = {
+                                    k: v for k, v in flight2.items() if k != "dates"
+                                }
+                                result2.update(flight2["dates"][date])
+                                result2["date"] = date2
+                                results.append([result1, result2])
+        return json.dumps(results)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "search_onestop_flight",
+                "description": "Search direct flights between two cities on a specific date.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "origin": {
+                            "type": "string",
+                            "description": "The origin city airport in three letters, such as 'JFK'.",
+                        },
+                        "destination": {
+                            "type": "string",
+                            "description": "The destination city airport in three letters, such as 'LAX'.",
+                        },
+                        "date": {
+                            "type": "string",
+                            "description": "The date of the flight in the format 'YYYY-MM-DD', such as '2024-05-01'.",
+                        },
+                    },
+                    "required": ["origin", "destination", "date"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/send_certificate.py

@@ -0,0 +1,52 @@
+# Copyright Sierra
+
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class SendCertificate(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        user_id: str,
+        amount: int,
+    ) -> str:
+        users = data["users"]
+        if user_id not in users:
+            return "Error: user not found"
+        user = users[user_id]
+
+        # add a certificate, assume at most 3 cases per task
+        for id in [3221322, 3221323, 3221324]:
+            payment_id = f"certificate_{id}"
+            if payment_id not in user["payment_methods"]:
+                user["payment_methods"][payment_id] = {
+                    "source": "certificate",
+                    "amount": amount,
+                    "id": payment_id,
+                }
+                return f"Certificate {payment_id} added to user {user_id} with amount {amount}."
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "send_certificate",
+                "description": "Send a certificate to a user. Be careful!",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "user_id": {
+                            "type": "string",
+                            "description": "The ID of the user to book the reservation, such as 'sara_doe_496'.",
+                        },
+                        "amount": {
+                            "type": "number",
+                            "description": "Certificate amount to send.",
+                        },
+                    },
+                    "required": ["user_id", "amount"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/think.py

@@ -0,0 +1,30 @@
+# Copyright Sierra
+
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class Think(Tool):
+    @staticmethod
+    def invoke(data: Dict[str, Any], thought: str) -> str:
+        return ""
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "think",
+                "description": "Use the tool to think about something. It will not obtain new information or change the database, but just append the thought to the log. Use it when complex reasoning is needed.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "thought": {
+                            "type": "string",
+                            "description": "A thought to think about.",
+                        },
+                    },
+                    "required": ["thought"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/transfer_to_human_agents.py

@@ -0,0 +1,35 @@
+# Copyright Sierra
+
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class TransferToHumanAgents(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        summary: str,
+    ) -> str:
+        return "Transfer successful"
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "transfer_to_human_agents",
+                "description": "Transfer the user to a human agent, with a summary of the user's issue. Only transfer if the user explicitly asks for a human agent, or if the user's issue cannot be resolved by the agent with the available tools.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "summary": {
+                            "type": "string",
+                            "description": "A summary of the user's issue.",
+                        },
+                    },
+                    "required": [
+                        "summary",
+                    ],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/update_reservation_baggages.py

@@ -0,0 +1,84 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict
+from tau_bench.envs.tool import Tool
+
+
+class UpdateReservationBaggages(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        reservation_id: str,
+        total_baggages: int,
+        nonfree_baggages: int,
+        payment_id: str,
+    ) -> str:
+        users, reservations = data["users"], data["reservations"]
+        if reservation_id not in reservations:
+            return "Error: reservation not found"
+        reservation = reservations[reservation_id]
+
+        total_price = 50 * max(0, nonfree_baggages - reservation["nonfree_baggages"])
+        if payment_id not in users[reservation["user_id"]]["payment_methods"]:
+            return "Error: payment method not found"
+        payment_method = users[reservation["user_id"]]["payment_methods"][payment_id]
+        if payment_method["source"] == "certificate":
+            return "Error: certificate cannot be used to update reservation"
+        elif (
+            payment_method["source"] == "gift_card"
+            and payment_method["amount"] < total_price
+        ):
+            return "Error: gift card balance is not enough"
+
+        reservation["total_baggages"] = total_baggages
+        reservation["nonfree_baggages"] = nonfree_baggages
+        if payment_method["source"] == "gift_card":
+            payment_method["amount"] -= total_price
+
+        if total_price != 0:
+            reservation["payment_history"].append(
+                {
+                    "payment_id": payment_id,
+                    "amount": total_price,
+                }
+            )
+
+        return json.dumps(reservation)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "update_reservation_baggages",
+                "description": "Update the baggage information of a reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "reservation_id": {
+                            "type": "string",
+                            "description": "The reservation ID, such as 'ZFA04Y'.",
+                        },
+                        "total_baggages": {
+                            "type": "integer",
+                            "description": "The updated total number of baggage items included in the reservation.",
+                        },
+                        "nonfree_baggages": {
+                            "type": "integer",
+                            "description": "The updated number of non-free baggage items included in the reservation.",
+                        },
+                        "payment_id": {
+                            "type": "string",
+                            "description": "The payment id stored in user profile, such as 'credit_card_7815826', 'gift_card_7815826', 'certificate_7815826'.",
+                        },
+                    },
+                    "required": [
+                        "reservation_id",
+                        "total_baggages",
+                        "nonfree_baggages",
+                        "payment_id",
+                    ],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/update_reservation_flights.py

@@ -0,0 +1,138 @@
+# Copyright Sierra
+
+import json
+from copy import deepcopy
+from typing import Any, Dict, List
+from tau_bench.envs.tool import Tool
+
+
+class UpdateReservationFlights(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        reservation_id: str,
+        cabin: str,
+        flights: List[Dict[str, Any]],
+        payment_id: str,
+    ) -> str:
+        users, reservations = data["users"], data["reservations"]
+        if reservation_id not in reservations:
+            return "Error: reservation not found"
+        reservation = reservations[reservation_id]
+
+        # update flights and calculate price
+        total_price = 0
+        flights = deepcopy(flights)
+        for flight in flights:
+            # if existing flight, ignore
+            if _ := [
+                f
+                for f in reservation["flights"]
+                if f["flight_number"] == flight["flight_number"]
+                and f["date"] == flight["date"]
+                and cabin == reservation["cabin"]
+            ]:
+                total_price += _[0]["price"] * len(reservation["passengers"])
+                flight["price"] = _[0]["price"]
+                flight["origin"] = _[0]["origin"]
+                flight["destination"] = _[0]["destination"]
+                continue
+            flight_number = flight["flight_number"]
+            if flight_number not in data["flights"]:
+                return f"Error: flight {flight_number} not found"
+            flight_data = data["flights"][flight_number]
+            if flight["date"] not in flight_data["dates"]:
+                return (
+                    f"Error: flight {flight_number} not found on date {flight['date']}"
+                )
+            flight_date_data = flight_data["dates"][flight["date"]]
+            if flight_date_data["status"] != "available":
+                return f"Error: flight {flight_number} not available on date {flight['date']}"
+            if flight_date_data["available_seats"][cabin] < len(
+                reservation["passengers"]
+            ):
+                return f"Error: not enough seats on flight {flight_number}"
+            flight["price"] = flight_date_data["prices"][cabin]
+            flight["origin"] = flight_data["origin"]
+            flight["destination"] = flight_data["destination"]
+            total_price += flight["price"] * len(reservation["passengers"])
+
+        total_price -= sum(flight["price"] for flight in reservation["flights"]) * len(
+            reservation["passengers"]
+        )
+
+        # check payment
+        if payment_id not in users[reservation["user_id"]]["payment_methods"]:
+            return "Error: payment method not found"
+        payment_method = users[reservation["user_id"]]["payment_methods"][payment_id]
+        if payment_method["source"] == "certificate":
+            return "Error: certificate cannot be used to update reservation"
+        elif (
+            payment_method["source"] == "gift_card"
+            and payment_method["amount"] < total_price
+        ):
+            return "Error: gift card balance is not enough"
+
+        # if checks pass, deduct payment and update seats
+        if payment_method["source"] == "gift_card":
+            payment_method["amount"] -= total_price
+        reservation["flights"] = flights
+        if total_price != 0:
+            reservation["payment_history"].append(
+                {
+                    "payment_id": payment_id,
+                    "amount": total_price,
+                }
+            )
+        # do not make flight database update here, assume it takes time to be updated
+        return json.dumps(reservation)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "update_reservation_flights",
+                "description": "Update the flight information of a reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "reservation_id": {
+                            "type": "string",
+                            "description": "The reservation ID, such as 'ZFA04Y'.",
+                        },
+                        "cabin": {
+                            "type": "string",
+                            "enum": [
+                                "basic_economy",
+                                "economy",
+                                "business",
+                            ],
+                        },
+                        "flights": {
+                            "type": "array",
+                            "description": "An array of objects containing details about each piece of flight in the ENTIRE new reservation. Even if the a flight segment is not changed, it should still be included in the array.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "flight_number": {
+                                        "type": "string",
+                                        "description": "Flight number, such as 'HAT001'.",
+                                    },
+                                    "date": {
+                                        "type": "string",
+                                        "description": "The date for the flight in the format 'YYYY-MM-DD', such as '2024-05-01'.",
+                                    },
+                                },
+                                "required": ["flight_number", "date"],
+                            },
+                        },
+                        "payment_id": {
+                            "type": "string",
+                            "description": "The payment id stored in user profile, such as 'credit_card_7815826', 'gift_card_7815826', 'certificate_7815826'.",
+                        },
+                    },
+                    "required": ["reservation_id", "cabin", "flights", "payment_id"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/tools/update_reservation_passengers.py

@@ -0,0 +1,64 @@
+# Copyright Sierra
+
+import json
+from typing import Any, Dict, List
+from tau_bench.envs.tool import Tool
+
+
+class UpdateReservationPassengers(Tool):
+    @staticmethod
+    def invoke(
+        data: Dict[str, Any],
+        reservation_id: str,
+        passengers: List[Dict[str, Any]],
+    ) -> str:
+        reservations = data["reservations"]
+        if reservation_id not in reservations:
+            return "Error: reservation not found"
+        reservation = reservations[reservation_id]
+        if len(passengers) != len(reservation["passengers"]):
+            return "Error: number of passengers does not match"
+        reservation["passengers"] = passengers
+        return json.dumps(reservation)
+
+    @staticmethod
+    def get_info() -> Dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": "update_reservation_passengers",
+                "description": "Update the passenger information of a reservation.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "reservation_id": {
+                            "type": "string",
+                            "description": "The reservation ID, such as 'ZFA04Y'.",
+                        },
+                        "passengers": {
+                            "type": "array",
+                            "description": "An array of objects containing details about each passenger.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "first_name": {
+                                        "type": "string",
+                                        "description": "The first name of the passenger, such as 'Noah'.",
+                                    },
+                                    "last_name": {
+                                        "type": "string",
+                                        "description": "The last name of the passenger, such as 'Brown'.",
+                                    },
+                                    "dob": {
+                                        "type": "string",
+                                        "description": "The date of birth of the passenger in the format 'YYYY-MM-DD', such as '1990-01-01'.",
+                                    },
+                                },
+                                "required": ["first_name", "last_name", "dob"],
+                            },
+                        },
+                    },
+                    "required": ["reservation_id", "passengers"],
+                },
+            },
+        }

chapter2/prompt-engineering/tau_bench/envs/airline/wiki.md

@@ -0,0 +1,70 @@
+# Airline Agent Policy
+
+The current time is 2024-05-15 15:00:00 EST.
+
+As an airline agent, you can help users book, modify, or cancel flight reservations.
+
+- Before taking any actions that update the booking database (booking, modifying flights, editing baggage, upgrading cabin class, or updating passenger information), you must list the action details and obtain explicit user confirmation (yes) to proceed.
+
+- You should not provide any information, knowledge, or procedures not provided by the user or available tools, or give subjective recommendations or comments.
+
+- You should only make one tool call at a time, and if you make a tool call, you should not respond to the user simultaneously. If you respond to the user, you should not make a tool call at the same time.
+
+- You should deny user requests that are against this policy.
+
+- You should transfer the user to a human agent if and only if the request cannot be handled within the scope of your actions.
+
+## Domain Basic
+
+- Each user has a profile containing user id, email, addresses, date of birth, payment methods, reservation numbers, and membership tier.
+
+- Each reservation has an reservation id, user id, trip type (one way, round trip), flights, passengers, payment methods, created time, baggages, and travel insurance information.
+
+- Each flight has a flight number, an origin, destination, scheduled departure and arrival time (local time), and for each date:
+  - If the status is "available", the flight has not taken off, available seats and prices are listed.
+  - If the status is "delayed" or "on time", the flight has not taken off, cannot be booked.
+  - If the status is "flying", the flight has taken off but not landed, cannot be booked.
+
+## Book flight
+
+- The agent must first obtain the user id, then ask for the trip type, origin, destination.
+
+- Passengers: Each reservation can have at most five passengers. The agent needs to collect the first name, last name, and date of birth for each passenger. All passengers must fly the same flights in the same cabin.
+
+- Payment: each reservation can use at most one travel certificate, at most one credit card, and at most three gift cards. The remaining amount of a travel certificate is not refundable. All payment methods must already be in user profile for safety reasons.
+
+- Checked bag allowance: If the booking user is a regular member, 0 free checked bag for each basic economy passenger, 1 free checked bag for each economy passenger, and 2 free checked bags for each business passenger. If the booking user is a silver member, 1 free checked bag for each basic economy passenger, 2 free checked bag for each economy passenger, and 3 free checked bags for each business passenger. If the booking user is a gold member, 2 free checked bag for each basic economy passenger, 3 free checked bag for each economy passenger, and 3 free checked bags for each business passenger. Each extra baggage is 50 dollars.
+
+- Travel insurance: the agent should ask if the user wants to buy the travel insurance, which is 30 dollars per passenger and enables full refund if the user needs to cancel the flight given health or weather reasons.
+
+## Modify flight
+
+- The agent must first obtain the user id and the reservation id.
+
+- Change flights: Basic economy flights cannot be modified. Other reservations can be modified without changing the origin, destination, and trip type. Some flight segments can be kept, but their prices will not be updated based on the current price. The API does not check these for the agent, so the agent must make sure the rules apply before calling the API!
+
+- Change cabin: all reservations, including basic economy, can change cabin without changing the flights. Cabin changes require the user to pay for the difference between their current cabin and the new cabin class. Cabin class must be the same across all the flights in the same reservation; changing cabin for just one flight segment is not possible.
+
+- Change baggage and insurance: The user can add but not remove checked bags. The user cannot add insurance after initial booking.
+
+- Change passengers: The user can modify passengers but cannot modify the number of passengers. This is something that even a human agent cannot assist with.
+
+- Payment: If the flights are changed, the user needs to provide one gift card or credit card for payment or refund method. The agent should ask for the payment or refund method instead.
+
+## Cancel flight
+
+- The agent must first obtain the user id, the reservation id, and the reason for cancellation (change of plan, airline cancelled flight, or other reasons)
+
+- All reservations can be cancelled within 24 hours of booking, or if the airline cancelled the flight. Otherwise, basic economy or economy flights can be cancelled only if travel insurance is bought and the condition is met, and business flights can always be cancelled. The rules are strict regardless of the membership status. The API does not check these for the agent, so the agent must make sure the rules apply before calling the API!
+
+- The agent can only cancel the whole trip that is not flown. If any of the segments are already used, the agent cannot help and transfer is needed.
+
+- The refund will go to original payment methods in 5 to 7 business days.
+
+## Refund
+
+- If the user is silver/gold member or has travel insurance or flies business, and complains about cancelled flights in a reservation, the agent can offer a certificate as a gesture after confirming the facts, with the amount being $100 times the number of passengers.
+
+- If the user is silver/gold member or has travel insurance or flies business, and complains about delayed flights in a reservation and wants to change or cancel the reservation, the agent can offer a certificate as a gesture after confirming the facts and changing or cancelling the reservation, with the amount being $50 times the number of passengers.
+
+- Do not proactively offer these unless the user complains about the situation and explicitly asks for some compensation. Do not compensate if the user is regular member and has no travel insurance and flies (basic) economy.

chapter2/prompt-engineering/tau_bench/envs/airline/wiki.py

@@ -0,0 +1,8 @@
+# Copyright Sierra
+
+import os
+
+FOLDER_PATH = os.path.dirname(__file__)
+
+with open(os.path.join(FOLDER_PATH, "wiki.md"), "r") as f:
+    WIKI = f.read()