cosmos-demos.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG Pattern Application - Simple Implementation\n",
    "\n",
    "In this sample, we'll demonstrate how to build a RAG Pattern application using a subset of the Movie Lens dataset. This sample will leverage the SDK's for MongoDB to perform vector search and cache the results. And Azure OpenAI to generate embeddings and LLM completions.\n",
    "\n",
    "There are two implementations in this project. One using LangChain and this simple implementation. The simple implementation connects directly to Azure Cosmos DB for MongoDB to perform vector search, and cache responses. It also connects directly to Azure OpenAI to generate embeddings and completions. This version requires a user to define and build the LLM payloads for LLM generation and also define the RAG Pattern request pipeline. Cache must be manually consulted in the pipelin and responses must also be manually cached.\n",
    "\n",
    "The vector search will be done using Azure Cosmos DB for MongoDB's [vector similarity search](https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search) functionality to do vector search over the vectorized movie data as well as the conversation history which is also used as a cache.\n",
    "\n",
    "At the end we will create a simple UX using Gradio to allow users to type in questions and display responses generated by a GPT model or served from the cache. The resopnses will also display an elapsed time so you can see the impact caching has on performance versus generating a response."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preliminaries <a class=\"anchor\" id=\"preliminaries\"></a>\n",
    "First, let's start by installing the packages that we'll need later. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install json\n",
    "! pip install python-dotenv\n",
    "\n",
    "! pip install pymongo\n",
    "! pip install openai\n",
    "\n",
    "! pip install gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import json\n",
    "from dotenv import dotenv_values\n",
    "\n",
    "import pymongo\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "import gradio as gr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please use the example.env as a template to provide the necessary keys and endpoints in your own .env file.\n",
    "Make sure to modify the env_name accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify the name of the .env file name \n",
    "env_name = \"fabcondemo.env\" # following example.env template change to your own .env file name\n",
    "config = dotenv_values(env_name)\n",
    "\n",
    "\n",
    "cosmos_conn = config['cosmos_connection_string']\n",
    "cosmos_database = config['cosmos_database_name']\n",
    "cosmos_collection = config['cosmos_collection_name']\n",
    "cosmos_vector_property = config['cosmos_vector_property_name']\n",
    "cosmos_cache = config['cosmos_cache_collection_name']\n",
    "# Create the Azure Cosmos DB for MongoDB client\n",
    "cosmos_client = pymongo.MongoClient(cosmos_conn)\n",
    "\n",
    "\n",
    "openai_endpoint = config['openai_endpoint']\n",
    "openai_key = config['openai_key']\n",
    "openai_version = config['openai_version']\n",
    "openai_embeddings_deployment = config['openai_embeddings_deployment']\n",
    "openai_embeddings_model = config['openai_embeddings_model']\n",
    "openai_embeddings_dimensions = int(config['openai_embeddings_dimensions'])\n",
    "openai_completions_deployment = config['openai_completions_deployment']\n",
    "openai_completions_model = config['openai_completions_model']\n",
    "# Create the OpenAI client\n",
    "openai_client = AzureOpenAI(azure_endpoint=openai_endpoint, api_key=openai_key, api_version=openai_version)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Database and collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the database\n",
    "database = cosmos_client[cosmos_database]\n",
    "\n",
    "# Get the movie collection\n",
    "movies = database[cosmos_collection]\n",
    "\n",
    "# Get the cache collection\n",
    "cache = database[cosmos_cache]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate embeddings from Azure OpenAI\n",
    "\n",
    "This is used to vectorize the user input for the vector search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_embeddings(text):\n",
    "    \n",
    "    response = openai_client.embeddings.create(\n",
    "        input = text, \n",
    "        model = openai_embeddings_deployment, \n",
    "        dimensions = openai_embeddings_dimensions)\n",
    "\n",
    "    embeddings = response.model_dump()\n",
    "    return embeddings['data'][0]['embedding']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vector Search in Azure Cosmos DB for MongoDB\n",
    "\n",
    "This defines a function for performing a vector search over any collection in Azure Cosmos DB for MongoDB. Function takes a collection reference, array of vectors, and optional similarity score to filter for top matches and number of results to return to filter further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vector_search(collection, vectors, similarity_score=0.02, num_results=3):\n",
    "    \n",
    "    pipeline = [\n",
    "        {\n",
    "        '$search': {\n",
    "            \"cosmosSearch\": {\n",
    "                \"vector\": vectors,\n",
    "                \"path\": cosmos_vector_property,\n",
    "                \"k\": num_results,\n",
    "                \"efsearch\": 40 # optional for HNSW only \n",
    "            },\n",
    "            \"returnStoredSource\": True }},\n",
    "            { '$project': { 'similarityScore': { '$meta': 'searchScore' }, 'document' : '$$ROOT' } },\n",
    "            { '$match': { \"similarityScore\": { '$gt': similarity_score } } \n",
    "        }   \n",
    "    ]\n",
    "\n",
    "    results = list(collection.aggregate(pipeline))\n",
    "\n",
    "    # Exclude the 'vector' to reduce payload size to LLM and _id properties to avoid serialization issues \n",
    "    for result in results:\n",
    "        del result['document']['vector']\n",
    "        del result['_id']\n",
    "        del result['document']['_id']\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get recent chat history\n",
    "\n",
    "This function provides conversational context to the LLM, allowing it to better have a conversation with the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab chat history to as part of the payload to GPT model for completion.\n",
    "def get_chat_history(completions=3):\n",
    "\n",
    "    # Sort by _id in descending order and limit the results to the completions value passed in\n",
    "    results = cache.find({}, {\"prompt\": 1, \"completion\": 1}).sort([(\"_id\", -1)]).limit(completions)\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Completion Function\n",
    "\n",
    "This function assembles all of the required data as a payload to send to a GPT model to generate a completion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_completion(user_prompt, vector_search_results, chat_history):\n",
    "    \n",
    "    system_prompt = '''\n",
    "    You are an intelligent assistant for the Movie Lens Expert AI Assistant.\n",
    "    You are designed to provide helpful answers to user questions about movies in your database.\n",
    "    You are friendly, helpful, and informative and wee bit cheeky.\n",
    "        - Only answer questions related to the information provided below.\n",
    "        - Write two lines of whitespace between each answer in the list.\n",
    "        - If you're unsure of an answer, you can say \"\"I don't know\"\" or \"\"I'm not sure\"\" and recommend users search themselves.\"\n",
    "    '''\n",
    "\n",
    "    # Create a list of messages as a payload to send to the OpenAI Completions API\n",
    "\n",
    "    # system prompt\n",
    "    messages = [{'role': 'system', 'content': system_prompt}]\n",
    "\n",
    "    #chat history\n",
    "    for chat in chat_history:\n",
    "        messages.append({'role': 'user', 'content': chat['prompt'] + \" \" + chat['completion']})\n",
    "    \n",
    "    #user prompt\n",
    "    messages.append({'role': 'user', 'content': user_prompt})\n",
    "\n",
    "    #vector search results\n",
    "    for result in vector_search_results:\n",
    "        messages.append({'role': 'system', 'content': json.dumps(result['document'])})\n",
    "\n",
    "    # Create the completion\n",
    "        response = openai_client.chat.completions.create(\n",
    "            model = openai_completions_deployment,\n",
    "            messages = messages \n",
    "        )\n",
    "    \n",
    "    \n",
    "    return response.model_dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cache Generated Responses\n",
    "\n",
    "Save the user prompts and generated completions in a conversation. Used to answer the same questions from other users. This is cheaper and faster than generating results each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cache_response(user_prompt, prompt_vectors, response):\n",
    "\n",
    "    chat = [\n",
    "        {\n",
    "            'prompt': user_prompt,\n",
    "            'completion': response['choices'][0]['message']['content'],\n",
    "            'completionTokens': str(response['usage']['completion_tokens']),\n",
    "            'promptTokens': str(response['usage']['prompt_tokens']),\n",
    "            'totalTokens': str(response['usage']['total_tokens']),\n",
    "            'model': response['model'],\n",
    "            cosmos_vector_property: prompt_vectors\n",
    "         }\n",
    "    ]\n",
    "\n",
    "    cache.insert_one(chat[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLM Pipeline function\n",
    "\n",
    "This function defines the pipeline for our RAG Pattern application. When user submits a question, the cache is consulted first for an exact match. If no match then a vector search is made, chat history gathered, the LLM generates a response, which is then cached before returning to the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat_completion(user_input):\n",
    "\n",
    "    # Generate embeddings from the user input\n",
    "    user_embeddings = generate_embeddings(user_input)\n",
    "\n",
    "    # Query the chat history cache first to see if this question has been asked before\n",
    "    # Similarity score set to 0.99, will only return exact matches. Limit to 1 result.\n",
    "    cache_results = vector_search(cache, user_embeddings, similarity_score=0.99, num_results=1)\n",
    "\n",
    "    if len(cache_results) > 0:\n",
    "        \n",
    "        return cache_results[0]['document']['completion']\n",
    "        \n",
    "    else:\n",
    "        \n",
    "        #perform vector search on the movie collection\n",
    "        search_results = vector_search(movies, user_embeddings)\n",
    "\n",
    "        #chat history\n",
    "        chat_history = get_chat_history(3)\n",
    "\n",
    "        #generate the completion\n",
    "        completions_results = generate_completion(user_input, search_results, chat_history)\n",
    "\n",
    "        #cache the response\n",
    "        cache_response(user_input, user_embeddings, completions_results)\n",
    "\n",
    "        # Return the generated LLM completion\n",
    "        return completions_results['choices'][0]['message']['content'] \n",
    "       \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a simple UX in Gradio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_history = []\n",
    "with gr.Blocks() as demo:\n",
    "    chatbot = gr.Chatbot()\n",
    "    msg = gr.Textbox(label=\"Ask me anything about movies!\")\n",
    "    clear = gr.Button(\"Clear\")\n",
    "\n",
    "    def user(user_message, chat_history):\n",
    "\n",
    "        # Create a timer to measure the time it takes to complete the request\n",
    "        start_time = time.time()\n",
    "        \n",
    "        # Get LLM completion\n",
    "        response_payload = chat_completion(user_message)\n",
    "\n",
    "        # Stop the timer\n",
    "        end_time = time.time()\n",
    "\n",
    "        elapsed_time = round((end_time - start_time) * 1000, 2)\n",
    "\n",
    "        response = response_payload\n",
    "        \n",
    "        # Append user message and response to chat history\n",
    "        chat_history.append([user_message, response_payload + f\"\\n (Time: {elapsed_time}ms)\"])\n",
    "        \n",
    "        return gr.update(value=\"\"), chat_history\n",
    "    \n",
    "    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False)\n",
    "    \n",
    "    clear.click(lambda: None, None, chatbot, queue=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# launch the gradio interface\n",
    "demo.launch(debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# be sure to run this cell to close or restart the gradio demo\n",
    "demo.close()"
   ]
  }
 ],
 "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.8"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}