{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ddfa9ae6-69fe-444a-b994-8c4c5970a7ec",
   "metadata": {},
   "source": [
    "# Student Contribution\n",
    "\n",
    "An awesome variation that includes a tool to make a booking! Thank you! -- Ed\n",
    "\n",
    "# Project - Airline AI Assistant with Booking Tool\n",
    "\n",
    "We'll now bring together what we've learned to make an AI Customer Support assistant for an Airline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b50bbe2-c0b1-49c3-9a5c-1ba7efa2bcb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "747e8786-9da8-4342-b6c9-f5f69c2e22ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialization\n",
    "\n",
    "load_dotenv()\n",
    "os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')\n",
    "MODEL = \"gpt-4o-mini\"\n",
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a521d84-d07c-49ab-a0df-d6451499ed97",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"You are a helpful assistant for an Airline called FlightAI. \"\n",
    "system_message += \"Give short, courteous answers, no more than 1 sentence. \"\n",
    "system_message += \"Always be accurate. If you don't know the answer, say so.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36bedabf-a0a7-4985-ad8e-07ed6a55a3a4",
   "metadata": {},
   "source": [
    "## Tools\n",
    "\n",
    "* Price and Booking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0696acb1-0b05-4dc2-80d5-771be04f1fb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's start by making a useful function\n",
    "\n",
    "ticket_prices = {\"london\": \"$799\", \"paris\": \"$899\", \"tokyo\": \"$1400\", \"berlin\": \"$499\"}\n",
    "\n",
    "def get_ticket_price(destination_city):\n",
    "    print(f\"Tool get_ticket_price called for {destination_city}\")\n",
    "    city = destination_city.lower()\n",
    "    return ticket_prices.get(city, \"Unknown\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4afceded-7178-4c05-8fa6-9f2085e6a344",
   "metadata": {},
   "outputs": [],
   "source": [
    "# There's a particular dictionary structure that's required to describe our function:\n",
    "\n",
    "price_function = {\n",
    "    \"name\": \"get_ticket_price\",\n",
    "    \"description\": \"Get the price of a return ticket to the destination city. Call this whenever you need to know the ticket price, for example when a customer asks 'How much is a ticket to this city'\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"destination_city\": {\n",
    "                \"type\": \"string\",\n",
    "                \"description\": \"The city that the customer wants to travel to\",\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"destination_city\"],\n",
    "        \"additionalProperties\": False\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6996ab93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"required\": [\"destination_city\", \"num_tickets\", \"ticket_class\", \"mail_address\"]\n",
    "\n",
    "booking_function = {\n",
    "        \"name\": \"book_ticket\",\n",
    "        \"description\": \"Book a ticket to a destination city.\",\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"destination_city\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"The city the customer wants to travel to\",\n",
    "                },\n",
    "                \"num_tickets\": {\n",
    "                    \"type\": \"integer\",\n",
    "                    \"description\": \"The number of tickets to book\",\n",
    "                },\n",
    "                \"ticket_class\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"The class of the ticket (e.g., economy, business)\",\n",
    "                },\n",
    "\n",
    "                \"mail_address\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"Mail address to send the ticket to\",\n",
    "                },\n",
    "            },\n",
    "            \"required\": [\"destination_city\", \"num_tickets\", \"ticket_class\", \"mail_address\"],\n",
    "            \"additionalProperties\": False\n",
    "        }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d62733e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And this is included in a list of tools:\n",
    "\n",
    "tools = [{\"type\": \"function\", \"function\": price_function},\n",
    "         {\"type\": \"function\", \"function\": booking_function}]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d3554f-b4e3-4ce7-af6f-68faa6dd2340",
   "metadata": {},
   "source": [
    "## Creating Application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f671dbb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating  ID for each booking\n",
    "# Produces IDs like 'BK-000001', 'BK-000002'\n",
    "\n",
    "booking_counter = 0\n",
    "\n",
    "def generate_unique_booking_id():\n",
    "    global booking_counter\n",
    "    booking_counter += 1\n",
    "    return f\"BK-{booking_counter:06d}\"  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ae8c75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Temporary storage for bookings\n",
    "bookingDB = pd.DataFrame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0992986-ea09-4912-a076-8e5603ee631f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have to write that function handle_tool_call:\n",
    "\n",
    "def handle_tool_call(message):\n",
    "    global bookingDB\n",
    "    tool_call = message.tool_calls[0]\n",
    "    if tool_call.function.name == \"get_ticket_price\":\n",
    "        arguments = json.loads(tool_call.function.arguments)\n",
    "        destination_city = arguments.get('destination_city')\n",
    "        price = get_ticket_price(destination_city)\n",
    "        response = {\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": json.dumps({\"destination_city\": destination_city,\"price\": price}),\n",
    "            \"tool_call_id\": message.tool_calls[0].id\n",
    "        }\n",
    "        return response\n",
    "    \n",
    "    elif tool_call.function.name == \"book_ticket\":\n",
    "        arguments = json.loads(tool_call.function.arguments)\n",
    "        destination_city = arguments.get('destination_city')\n",
    "        num_tickets = arguments.get('num_tickets')\n",
    "        ticket_class = arguments.get('ticket_class')\n",
    "        mail_address = arguments.get('mail_address')\n",
    "        response = {\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": json.dumps({\"destination_city\": destination_city,\"num_tickets\": num_tickets, \"ticket_class\": ticket_class}),\n",
    "            \"tool_call_id\": message.tool_calls[0].id\n",
    "        } \n",
    "        # Get the ticket price\n",
    "        price = int(get_ticket_price(destination_city).replace(\"$\", \"\"))\n",
    "        total_price = price * num_tickets\n",
    "        # Generate a unique booking ID\n",
    "        booked_ID = generate_unique_booking_id()\n",
    "        # Define the booking data\n",
    "        data = {\n",
    "            \"booking_id\": [booked_ID],\n",
    "            \"mail_address\": [mail_address],\n",
    "            \"destination_city\": [destination_city],\n",
    "            \"num_tickets\": [num_tickets],\n",
    "            \"ticket_class\": [ticket_class],\n",
    "            \"total_price\": [total_price],\n",
    "        }\n",
    "        booking_temp = pd.DataFrame(data)\n",
    "        bookingDB = bookingDB._append(booking_temp)\n",
    "        # Can update with real booking system\n",
    "        bookingDB.to_csv('bookingDB.csv', index=False)\n",
    "        return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9b0744-9c78-408d-b9df-9f6fd9ed78cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(message, history):\n",
    "    messages = [{\"role\": \"system\", \"content\": system_message}]\n",
    "    for human, assistant in history:\n",
    "        messages.append({\"role\": \"user\", \"content\": human})\n",
    "        messages.append({\"role\": \"assistant\", \"content\": assistant})\n",
    "    messages.append({\"role\": \"user\", \"content\": message})\n",
    "    response = openai.chat.completions.create(model=MODEL, messages=messages, tools=tools)\n",
    "\n",
    "    if response.choices[0].finish_reason==\"tool_calls\":\n",
    "        message = response.choices[0].message\n",
    "        response = handle_tool_call(message)\n",
    "        messages.append(message)\n",
    "        messages.append(response)\n",
    "        response = openai.chat.completions.create(model=MODEL, messages=messages)\n",
    "    \n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4be8a71-b19e-4c2f-80df-f59ff2661f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "gr.ChatInterface(fn=chat).launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11c9da69-d0cf-4cf2-a49e-e5669deec47b",
   "metadata": {},
   "source": [
    "## Enjoy, Happy Learning Thanks to Ed!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
