cosmos-demos-langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this tutorial, we'll demonstrate how to leverage a sample dataset stored in Azure Cosmos DB for MongoDB vCore to ground OpenAI models. We'll do this taking advantage of Azure Cosmos DB for Mongo DB vCore's [vector similarity search](https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search) functionality. In the end, we'll create an interatice chat session with the GPT-3.5 completions model to answer questions about Azure services informed by our dataset. This process is known as Retrieval Augmented Generation, or RAG."
   ]
  },
  {
   "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 gradio\n",
    "# ! pip install langchain\n",
    "# ! pip install langchain_community\n",
    "# ! pip install langchain_openai\n",
    "# ! pip install openai\n",
    "# ! pip install pymongo\n",
    "# ! pip install python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import dotenv_values\n",
    "import gradio as gr\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain.globals import set_llm_cache\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_community.cache import AzureCosmosDBSemanticCache\n",
    "from langchain_community.chat_message_histories import MongoDBChatMessageHistory\n",
    "from langchain_community.vectorstores.azure_cosmos_db import (\n",
    "    AzureCosmosDBVectorSearch,\n",
    "    CosmosDBSimilarityType,\n",
    "    CosmosDBVectorSearchType)\n",
    "from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings\n",
    "import pymongo"
   ]
  },
  {
   "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": null,
   "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",
    "# Azure Cosmos DB connection details\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_semcache = config['cosmos_semcache_collection_name']\n",
    "cosmos_chat_history = config['cosmos_chathistory_collection_name']\n",
    "\n",
    "# Azure OpenAI connection details\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']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure Cosmos DB for MongoDB Connections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Establish connection to db\n",
    "cosmos_client = pymongo.MongoClient(cosmos_conn)\n",
    "\n",
    "# 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_semcache]\n",
    "\n",
    "# Get the chat history collection\n",
    "chathistory = database[cosmos_chat_history]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Embeddings with Azure OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "azure_openai_embeddings = AzureOpenAIEmbeddings(\n",
    "    azure_deployment = openai_embeddings_deployment,\n",
    "    api_key= openai_key,\n",
    "    azure_endpoint= openai_endpoint,\n",
    "    model= openai_embeddings_model,\n",
    "    dimensions= openai_embeddings_dimensions,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "azure_openai_embeddings.embed_query(\"Hello, this is a test!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Configure Vector Search w/ LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# similarity_algorithm = CosmosDBSimilarityType.COS\n",
    "# kind = CosmosDBVectorSearchType.VECTOR_HNSW\n",
    "# m=10\n",
    "# ef_construction = 64\n",
    "# ef_search = 40\n",
    "# score_threshold = 0.7\n",
    "# vectorstore.create_index(\n",
    "#      0, openai_embeddings_dimensions, similarity_algorithm, kind, m, ef_construction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdb = AzureCosmosDBVectorSearch(\n",
    "    collection= cosmos_collection,\n",
    "    embedding=azure_openai_embeddings)\n",
    "\n",
    "vectorstore = cdb.from_connection_string(\n",
    "    connection_string=cosmos_conn,\n",
    "    namespace = cosmos_database + \".\" + cosmos_collection,\n",
    "    embedding = azure_openai_embeddings,\n",
    "    embedding_key = cosmos_vector_property,\n",
    "    text_key = \"overview\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore.similarity_search_with_score(\"Buzz Lightyear\", k=5, score_threshold=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup RAG, Semantic Caching, and History with your LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "You are an upbeat AI assistant who is excited to help answer questions. \n",
    "You can use this context\n",
    "\n",
    "{context},\n",
    "\n",
    "or this chat history\n",
    "\n",
    "{chat_history},\n",
    "\n",
    "to answer this question. \n",
    "\n",
    "Question: {question}\n",
    "If you don't know the answer, just say that you don't know. Don't try to make up an answer.\n",
    "\"\"\"\n",
    "chatbot_prompt = PromptTemplate(\n",
    "    template = prompt_template, input_variables = [\"context\", \"question\", \"chat_history\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_chain():\n",
    "    llm = AzureChatOpenAI(\n",
    "            azure_endpoint = openai_endpoint,\n",
    "            api_key = openai_key,\n",
    "            api_version = openai_version,\n",
    "            azure_deployment = \"completions\", \n",
    "            cache = True,\n",
    "            n = 1)\n",
    "\n",
    "    retriever = vectorstore.as_retriever(\n",
    "    search_type = \"similarity\",\n",
    "    search_kwargs = {\"k\": 5, 'score_threshold': 0.2})\n",
    "\n",
    "    sem_qa = ConversationalRetrievalChain.from_llm(\n",
    "    llm = llm,\n",
    "    chain_type = \"stuff\",\n",
    "    retriever = retriever,\n",
    "    return_source_documents = True,\n",
    "    combine_docs_chain_kwargs = {\"prompt\": chatbot_prompt})\n",
    "\n",
    "    similarity_algorithm = CosmosDBSimilarityType.COS\n",
    "    kind = CosmosDBVectorSearchType.VECTOR_IVF\n",
    "    num_lists = 1\n",
    "    score_threshold = 0.99\n",
    "\n",
    "    sem_cache = AzureCosmosDBSemanticCache(\n",
    "            cosmosdb_connection_string = cosmos_conn,\n",
    "            cosmosdb_client = None,\n",
    "            embedding = azure_openai_embeddings,\n",
    "            database_name = cosmos_database, \n",
    "            collection_name = cosmos_semcache, \n",
    "            similarity = similarity_algorithm,\n",
    "            num_lists = num_lists,\n",
    "            kind = kind,\n",
    "            dimensions = openai_embeddings_dimensions, \n",
    "            score_threshold = score_threshold)\n",
    "\n",
    "    set_llm_cache(sem_cache)\n",
    "\n",
    "    return retriever, llm, sem_qa, sem_cache "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever, llm, chain, semantic_cache = prepare_chain()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clearing Semantic Cache inbetween testing\n",
    "cache.drop_indexes()\n",
    "database.drop_collection(cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain.invoke({'question': \"Tell me about movies with Buzz Lightyear.\", 'chat_history': []})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain.invoke({'question': \"Tell me something about films with Buzz Lightyear.\", 'chat_history': []})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gradio / UI integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_chain_retriever, chatbot_chain_llm, chatbot_chain, chatbot_chain_cache = prepare_chain()\n",
    "\n",
    "# Clearing Semantic Cache inbetween testing\n",
    "cache.drop_indexes()\n",
    "database.drop_collection(cache)\n",
    "database.drop_collection(chathistory)\n",
    "\n",
    "cosmos_message_history = MongoDBChatMessageHistory(\n",
    "    session_id = \"test_session\",\n",
    "    connection_string = cosmos_conn,\n",
    "    database_name = cosmos_database,\n",
    "    collection_name = cosmos_chat_history)\n",
    "\n",
    "conversational_memory = ConversationBufferMemory(\n",
    "    chat_memory=cosmos_message_history,\n",
    "    memory_key='chat_history',\n",
    "    return_messages=True)\n",
    "\n",
    "# Load history locally. Grab last \n",
    "hist = []\n",
    "with gr.Blocks() as demo:\n",
    "    chatbot = gr.Chatbot()\n",
    "    msg = gr.Textbox()\n",
    "    clear = gr.Button(\"Clear\")\n",
    "\n",
    "    def user(user_message, chat_history):\n",
    "        # Get response from QA chain\n",
    "        response = chatbot_chain.invoke({\"question\": user_message, \"chat_history\":conversational_memory.buffer_as_messages[-6:]},temperature=0.2)\n",
    "        # Append user message and response to chat history\n",
    "        hist.append([\"User: \"+user_message, \"Chatbot: \"+response['answer']])\n",
    "        cosmos_message_history.add_user_message(user_message)\n",
    "        cosmos_message_history.add_ai_message(response['answer'])\n",
    "        return gr.update(value=\"\"), hist\n",
    "\n",
    "    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False)\n",
    "    clear.click(lambda: None, None, chatbot, queue=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo.launch(debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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.0"
  },
  "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
}