RAG application based on Qualcomm Cloud AI Playground API

playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "27756c1d-6c89-4996-8340-ccfde1893365",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install openai chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c96f690f-04ac-4e08-8ed2-8ba7423af5af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from openai import OpenAI \n",
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d0a01deb-9575-48c9-a367-f22de7badd44",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "api_key = os.getenv('CIRRASCALE_API_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b181019d-52ae-4c77-baf4-ffee64cfe8bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_client = OpenAI(api_key=api_key, base_url=\"https://cloudai.cirrascale.com/apis/v2\")\n",
    "chroma_client = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "COLLECTION_NAME = \"baggage_policy\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7c09ff4f-f7c9-4629-aa6c-571e377da929",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_function = embedding_functions.OpenAIEmbeddingFunction(\n",
    "    api_key=api_key,\n",
    "    model_name=\"BAAI/bge-large-en-v1.5\",\n",
    "    api_base=\"https://cloudai.cirrascale.com/apis/v2\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "991a0e4b-93d5-45a6-9e61-34a163eb9ac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index():\n",
    "    baggage_rules = [\n",
    "        \"Emirates Airlines offers a generous baggage policy that varies based on route, fare type, and cabin class.\",\n",
    "        \"For carry-on luggage, Economy passengers are allowed one piece weighing up to 7 kg with dimensions not exceeding 55 x 38 x 20 cm.\",\n",
    "        \"First and Business Class passengers can bring two pieces of carry-on baggage: one briefcase and either a handbag or garment bag.\",\n",
    "        \"Emirates uses either a 'weight concept' or 'piece concept' for checked baggage, depending on the route.\",\n",
    "        \"Economy Special fares typically allow 20 kg of checked baggage, while Economy Flex Plus offers up to 35 kg.\",\n",
    "        \"Each checked bag must not exceed 32 kg, with total dimensions not exceeding 300 cm (118 inches).\",\n",
    "        \"Passengers traveling with infants can bring a carrycot or collapsible stroller on board or check it for free.\",\n",
    "        \"Liquid items in carry-on luggage must be in containers of 100ml or less, placed in a single transparent, resealable 1-liter plastic bag.\",\n",
    "        \"Emirates Skywards members may be eligible for additional baggage allowance on routes using the weight concept.\",\n",
    "        \"For flights within the Americas and between the US and Europe, Emirates uses a 'piece concept' with different baggage allowances.\"\n",
    "    ]\n",
    "\n",
    "    collection = chroma_client.get_or_create_collection(\n",
    "        name=COLLECTION_NAME,\n",
    "        embedding_function=embedding_function,\n",
    "        metadata={\"hnsw:space\": \"cosine\"}\n",
    "    )\n",
    "    \n",
    "    if len(collection.get()['ids']) == 0:\n",
    "        for idx, rule in enumerate(baggage_rules):\n",
    "            collection.add(documents=[rule], ids=[f\"baggage_rule_{idx}\"])\n",
    "        print(f\"Stored {len(baggage_rules)} baggage rules in ChromaDB.\")\n",
    "\n",
    "    return collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "818f4c52-0543-4406-bafd-24471461a6bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve(query: str):\n",
    "    collection = chroma_client.get_collection(name=COLLECTION_NAME, embedding_function=embedding_function)\n",
    "    results = collection.query(query_texts=[query], n_results=3)\n",
    "    \n",
    "    if results and results['documents']:\n",
    "        return \" \".join(results['documents'][0])\n",
    "    return \"No relevant baggage information found.\"    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5463b657-c337-4598-bcca-82f5841111c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(query: str) -> str:\n",
    "    context=retrieve(query)\n",
    "    sys_prompt=\"\"\"\n",
    "    You are a helpful airline assistant. Answer the user query based on the context provided.\n",
    "    If you don't find the answer within the context, say I don't know. Give accurate and precise answers.\n",
    "    \"\"\"\n",
    "    usr_prompt=f\"Query: {query}\\nContext: {context}\"\n",
    "    \n",
    "    response = openai_client.chat.completions.create(\n",
    "        model=\"Llama-3.1-8B\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": sys_prompt},\n",
    "            {\"role\": \"user\", \"content\": usr_prompt},\n",
    "        ],\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ff563454-7dcd-4a5f-b89d-5ca505dae76d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Collection(name=baggage_policy)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bd3ed5b0-d68c-4462-9e1d-4dda735877f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q: What's the size limit for cabin baggage in economy class?\n",
      "A: The size limit for cabin baggage in economy class is 55 x 38 x 20 cm.\n"
     ]
    }
   ],
   "source": [
    "query=\"What's the size limit for cabin baggage in economy class?\"\n",
    "answer=generate(query)\n",
    "print(f\"Q: {query}\")\n",
    "print(f\"A: {answer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "78183629-cbac-4924-8fc5-cc989be5ca7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q: What's the weight of checked baggage for Economy Flex?\n",
      "A: For Economy Flex, the weight of checked baggage is up to 35 kg.\n"
     ]
    }
   ],
   "source": [
    "query=\"What's the weight of checked baggage for Economy Flex?\"\n",
    "answer=generate(query)\n",
    "print(f\"Q: {query}\")\n",
    "print(f\"A: {answer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1b70b21-8fba-4be7-af74-b169c30fd0e1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}