vcore/Labs/lab_4_langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import pymongo\n",
    "from typing import List\n",
    "from dotenv import load_dotenv\n",
    "from langchain.chat_models import AzureChatOpenAI\n",
    "from langchain.embeddings import AzureOpenAIEmbeddings\n",
    "from langchain.vectorstores import AzureCosmosDBVectorSearch\n",
    "from langchain_core.vectorstores import VectorStoreRetriever\n",
    "from langchain.schema.document import Document\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.agents import Tool\n",
    "from langchain.agents.agent_toolkits import create_conversational_retrieval_agent\n",
    "from langchain_core.messages import SystemMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load settings for the notebook\n",
    "load_dotenv()\n",
    "CONNECTION_STRING = os.environ.get(\"DB_CONNECTION_STRING\")\n",
    "EMBEDDINGS_DEPLOYMENT_NAME = \"embeddings\"\n",
    "COMPLETIONS_DEPLOYMENT_NAME = \"completions\"\n",
    "AOAI_ENDPOINT = os.environ.get(\"AOAI_ENDPOINT\")\n",
    "AOAI_KEY = os.environ.get(\"AOAI_KEY\")\n",
    "AOAI_API_VERSION = \"2023-09-01-preview\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Establish Azure OpenAI connectivity\n",
    "llm = AzureChatOpenAI(            \n",
    "        temperature = 0,\n",
    "        openai_api_version = AOAI_API_VERSION,\n",
    "        azure_endpoint = AOAI_ENDPOINT,\n",
    "        openai_api_key = AOAI_KEY,         \n",
    "        azure_deployment = \"completions\"\n",
    ")\n",
    "embedding_model = AzureOpenAIEmbeddings(\n",
    "    openai_api_version = AOAI_API_VERSION,\n",
    "    azure_endpoint = AOAI_ENDPOINT,\n",
    "    openai_api_key = AOAI_KEY,   \n",
    "    azure_deployment = \"embeddings\",\n",
    "    chunk_size=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector search with LangChain\n",
    "\n",
    "In the previous lab, the `pymongo` library was used to perform a vector search through a db command to find product documents that were most similar to the user's input. In this lab, you will use the `langchain` library to perform the same search. LangChain has a vector store class named **AzureCosmosDBVectorSearch**, a community contribution, that supports vector search in vCore-based Azure Cosmos DB for MongoDB.\n",
    "\n",
    "When establishing the connection to the vector store (vCore-based Azure Cosmos DB for MongoDB), recall that in previous labs the products collection was populated and a contentVector field added that contains the vectorized embeddings of the document itself. Finally, a vector index was also created on the contentVector field to enable vector search. The vector index in each collection is named `VectorSearchIndex`.\n",
    "\n",
    "The return value of a vector search in LangChain is a list of `Document` objects. The LangChain `Document` class contains two properties: `page_content`, that represents the textual content that is typically used to augment the prompt, and `metadata` that contains all other attributes of the document. In the cell below, we'll use the `_id` field as the page_content, and the rest of the fields are returned as metadata.\n",
    "\n",
    "The next two cells initiate a connection to the vector store and performs a vector search. Notice how much more concise the code is compared to the previous lab with the addition of LangChain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reference the existing vector store\n",
    "vector_store = AzureCosmosDBVectorSearch.from_connection_string(\n",
    "    connection_string = CONNECTION_STRING,\n",
    "    namespace = \"cosmic_works.products\",\n",
    "    embedding = embedding_model,\n",
    "    index_name = \"VectorSearchIndex\",    \n",
    "    embedding_key = \"contentVector\",\n",
    "    text_key = \"_id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What yellow products are there?\"\n",
    "vector_store.similarity_search(query, k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG with LangChain\n",
    "\n",
    "In this section, we'll implement the RAG pattern using LangChain. In LangChain, a **retriever** is used to augment the prompt with contextual data. In this case, the already established vector store will be used as the retriever. By default, the prompt is augmented with the `page_content` field of the retrieved document that customarily contains the text content of the embedded vector. In our case, the document itself serves as the textual content, so we'll have to do some pre-processing to format the text of the product list that is expected in our system prompt (JSON string) - see the **format_documents** function below for this implementation.\n",
    "\n",
    "We'll also define a reusable RAG [chain](https://python.langchain.com/docs/modules/chains/) to control the flow and behavior of the call into the LLM. This chain is defined using the LCEL syntax (LangChain Expression Language)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A system prompt describes the responsibilities, instructions, and persona of the AI.\n",
    "# Note the addition of the templated variable/placeholder for the list of products and the incoming question.\n",
    "system_prompt = \"\"\"\n",
    "You are a helpful, fun and friendly sales assistant for Cosmic Works, a bicycle and bicycle accessories store. \n",
    "Your name is Cosmo.\n",
    "You are designed to answer questions about the products that Cosmic Works sells.\n",
    "\n",
    "Only answer questions related to the information provided in the list of products below that are represented\n",
    "in JSON format.\n",
    "\n",
    "If you are asked a question that is not in the list, respond with \"I don't know.\"\n",
    "\n",
    "Only answer questions related to Cosmic Works products, customers, and sales orders.\n",
    "\n",
    "If a question is not related to Cosmic Works products, customers, or sales orders,\n",
    "respond with \"I only answer questions about Cosmic Works\"\n",
    "\n",
    "List of products:\n",
    "{products}\n",
    "\n",
    "Question:\n",
    "{question}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remember that each Document contains a page_content property\n",
    "# that is populated with the _id field of the document\n",
    "# all other document fields are located in the metadata property\n",
    "def format_docs(docs:List[Document]) -> str:\n",
    "        \"\"\"\n",
    "        Prepares the product list for the system prompt.\n",
    "        \"\"\"\n",
    "        str_docs = []\n",
    "        for doc in docs:\n",
    "                # Build the product document without the contentVector\n",
    "                doc_dict = {\"_id\": doc.page_content}\n",
    "                doc_dict.update(doc.metadata)\n",
    "                if \"contentVector\" in doc_dict:  \n",
    "                        del doc_dict[\"contentVector\"]\n",
    "                str_docs.append(json.dumps(doc_dict, default=str))                  \n",
    "        # Return a single string containing each product JSON representation\n",
    "        # separated by two newlines\n",
    "        return \"\\n\\n\".join(str_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a retriever from the vector store\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "# Create the prompt template from the system_prompt text\n",
    "llm_prompt = PromptTemplate.from_template(system_prompt)\n",
    "\n",
    "rag_chain = (\n",
    "    # populate the tokens/placeholders in the llm_prompt \n",
    "    # products takes the results of the vector store and formats the documents\n",
    "    # question is a passthrough that takes the incoming question\n",
    "    { \"products\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | llm_prompt\n",
    "    # pass the populated prompt to the language model\n",
    "    | llm\n",
    "    # return the string ouptut from the language model\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"What are the names and skus of yellow products? Output the answer as a bulleted list.\"\n",
    "response = rag_chain.invoke(question)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangChain Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create retrievers\n",
    "\n",
    "A separate retriever is required for each vector index. The following cell creates a VectorStoreRetriever for the products, customers, and sales collections and associated vector index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_cosmic_works_vector_store_retriever(collection_name: str, top_k: int = 3):\n",
    "    vector_store =  AzureCosmosDBVectorSearch.from_connection_string(\n",
    "        connection_string = CONNECTION_STRING,\n",
    "        namespace = f\"cosmic_works.{collection_name}\",\n",
    "        embedding = embedding_model,\n",
    "        index_name = \"VectorSearchIndex\",    \n",
    "        embedding_key = \"contentVector\",\n",
    "        text_key = \"_id\"\n",
    "    )\n",
    "    return vector_store.as_retriever(search_kwargs={\"k\": top_k})\n",
    "\n",
    "\n",
    "products_retriever = create_cosmic_works_vector_store_retriever(\"products\")\n",
    "customers_retriever = create_cosmic_works_vector_store_retriever(\"customers\")\n",
    "sales_retriever = create_cosmic_works_vector_store_retriever(\"sales\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Create Agent Tools\n",
    " \n",
    " LangChain does have a built-in [`create_retriever_tool`](https://python.langchain.com/docs/use_cases/question_answering/conversational_retrieval_agents#retriever-tool) that wraps a vector store retriever, however, because we are storing the embeddings in the `contentVector` field of the document, we must do some pre-processing of the retrieved documents to remove this field so that we don't needlessly expend the model's token quota. \n",
    " \n",
    " Instead, we'll create a RAG chain as our tool implementation that does the pre-processing through the `format_docs` function we defined above to return each document in its JSON representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create tools that will use vector search in vCore-based Azure Cosmos DB for MongoDB collections\n",
    "\n",
    "# create a chain on the retriever to format the documents as JSON\n",
    "products_retriever_chain = products_retriever | format_docs\n",
    "customers_retriever_chain = customers_retriever | format_docs\n",
    "sales_retriever_chain = sales_retriever | format_docs\n",
    "\n",
    "tools = [\n",
    "    Tool(\n",
    "        name = \"vector_search_products\", \n",
    "        func = products_retriever_chain.invoke,\n",
    "        description = \"Searches Cosmic Works product information for similar products based on the question. Returns the product information in JSON format.\"\n",
    "    ),\n",
    "    Tool(\n",
    "        name = \"vector_search_customers\", \n",
    "        func = customers_retriever_chain.invoke,\n",
    "        description = \"Searches Cosmic Works customer information and retrieves similar customers based on the question. Returns the customer information in JSON format.\"\n",
    "    ),\n",
    "    Tool(\n",
    "        name = \"vector_search_sales\", \n",
    "        func = sales_retriever_chain.invoke,\n",
    "        description = \"Searches Cosmic Works customer sales information and retrieves sales order details based on the question. Returns the sales order information in JSON format.\"\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tools part 2\n",
    "\n",
    "Certain properties do not have semantic meaning (such as the GUID _id fields) and attempting to use vector search on these fields will not yield meaningful results. We need a tool to retrieve specific documents based on popular searches criteria.\n",
    "\n",
    "The following tool definitions is not an exhaustive list of what may be needed but serves as an example to provide concrete lookups of a document in the Cosmic Works database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = pymongo.MongoClient(CONNECTION_STRING).cosmic_works\n",
    "\n",
    "def get_product_by_id(product_id: str) -> str:\n",
    "    \"\"\"\n",
    "    Retrieves a product by its ID.    \n",
    "    \"\"\"\n",
    "    doc = db.products.find_one({\"_id\": product_id})    \n",
    "    if \"contentVector\" in doc:\n",
    "        del doc[\"contentVector\"]\n",
    "    return json.dumps(doc)\n",
    "\n",
    "def get_product_by_sku(sku: str) -> str:\n",
    "    \"\"\"\n",
    "    Retrieves a product by its sku.\n",
    "    \"\"\"\n",
    "    doc = db.products.find_one({\"sku\": sku})\n",
    "    if \"contentVector\" in doc:\n",
    "        del doc[\"contentVector\"]\n",
    "    return json.dumps(doc, default=str)\n",
    "\n",
    "def get_sales_by_id(sales_id: str) -> str:\n",
    "    \"\"\"\n",
    "    Retrieves a sales order by its ID.\n",
    "    \"\"\"\n",
    "    doc = db.sales.find_one({\"_id\": sales_id})\n",
    "    if \"contentVector\" in doc:\n",
    "        del doc[\"contentVector\"]\n",
    "    return json.dumps(doc, default=str)    \n",
    "\n",
    "from langchain.tools import StructuredTool\n",
    "\n",
    "tools.extend([\n",
    "    StructuredTool.from_function(get_product_by_id),\n",
    "    StructuredTool.from_function(get_product_by_sku),\n",
    "    StructuredTool.from_function(get_sales_by_id)\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the agent\n",
    "\n",
    "The [`create_conversational_retrieval_agent`](https://python.langchain.com/docs/use_cases/question_answering/conversational_retrieval_agents#agent-constructor) is a built-in agent that includes conversational history as well uses the [OpenAIFunctionsAgent](https://python.langchain.com/docs/modules/agents/agent_types/openai_functions_agent#using-openaifunctionsagent) as its underlying implementation.\n",
    "\n",
    "Remember that an agent leverages the LLM to assess the incoming request with the current context to decide what steps or actions need to be executed and in what order. LangChain agents can leverage tools. A tool can be an integration into an external system, custom code, or even another chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = SystemMessage(\n",
    "    content = \"\"\"\n",
    "        You are a helpful, fun and friendly sales assistant for Cosmic Works, a bicycle and bicycle accessories store.\n",
    "\n",
    "        Your name is Cosmo.\n",
    "\n",
    "        You are designed to answer questions about the products that Cosmic Works sells, the customers that buy them, and the sales orders that are placed by customers.\n",
    "\n",
    "        If you don't know the answer to a question, respond with \"I don't know.\"\n",
    "        \n",
    "        Only answer questions related to Cosmic Works products, customers, and sales orders.\n",
    "        \n",
    "        If a question is not related to Cosmic Works products, customers, or sales orders,\n",
    "        respond with \"I only answer questions about Cosmic Works\"\n",
    "    \"\"\"    \n",
    ")\n",
    "agent_executor = create_conversational_retrieval_agent(llm, tools, system_message = system_message, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = agent_executor({\"input\": \"What products do you have that are yellow?\"})\n",
    "print(\"***********************************************************\")\n",
    "print(result['output'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = agent_executor({\"input\": \"What products were purchased for sales order '06FE91D2-B350-471A-AD29-906BF4EB97C4' ?\"})\n",
    "print(\"***********************************************************\")\n",
    "print(result['output'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = agent_executor({\"input\": \"What was the sales order total for sales order '93436616-4C8A-407D-9FDA-908707EFA2C5' ?\"})\n",
    "print(\"***********************************************************\")\n",
    "print(result['output'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = agent_executor({\"input\": \"What was the price of the product with sku `FR-R92B-58` ?\"})\n",
    "print(\"***********************************************************\")\n",
    "print(result['output'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}