reasoning engine.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "lJAbLZ1ZEaHF",
      "metadata": {
        "id": "lJAbLZ1ZEaHF"
      },
      "source": [
        "# RAG with MongoDB Atlas and VertexAI Reasoning Engine using Langchain"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d8095d1a",
      "metadata": {},
      "source": [
        "**Vertex AI Reasoning Engine with Langchain**  is a powerful duo for building and deploying generative AI applications. It is one of the managed services in the Vertex AI console that porvides secure scalable runtime environment for your workload. Langchain provides the tool to design your application logic while reasoning engine provides environment to run it. With its flexibility to connect external data sources we can connect the MongoDB Atlas to Google VertexAI reasoning engine. \n",
        "\n",
        "At the core lies **MongoDB Atlas Vector Search**. It excels at searching unstructured data using vector embeddings, allowing you to find similar information even if phrased differently. This empowers your AI to grasp the true meaning behind user queries. Langchain then steps in, providing a user-friendly framework to design your application logic. Here, you can leverage Langchain's flexibility to seamlessly integrate MongoDB Atlas Vector Search, enabling your AI to retrieve highly relevant data based on semantic similarity. Finally, Vertex AI Reasoning Engine provides a secure and scalable environment to run your creations. This trio simplifies development, offering a pre-built foundation and tools to focus on building innovative solutions. With MongoDB Atlas Vector Search's semantic understanding, your generative AI applications can deliver superior results and user experiences."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5REfyDmgEpzE",
      "metadata": {
        "id": "5REfyDmgEpzE"
      },
      "source": [
        "In this Notebook we will cover *How to build a RAG and deploy it as endpoints using Reasoning Engine, MongoDB Atlas and VertexAI*\n",
        "\n",
        "First we will install all thre required dependecies and restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "id": "ZKCZmNCsBK2W73bGDaZ61YMo",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "executionInfo": {
          "elapsed": 19528,
          "status": "ok",
          "timestamp": 1718715616217,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "ZKCZmNCsBK2W73bGDaZ61YMo",
        "outputId": "2580f891-89b4-4427-b0c2-3706a89d1cc1",
        "tags": []
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'status': 'ok', 'restart': True}"
            ]
          },
          "execution_count": 49,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "!pip install --upgrade --quiet \\\n",
        "    \"google-cloud-aiplatform[langchain,reasoningengine]\" \\\n",
        "    cloudpickle==3.0.0 \\\n",
        "    pydantic==2.7.4 \\\n",
        "    requests \\\n",
        "    datasets \\\n",
        "    pymongo \\\n",
        "    langchain \\\n",
        "    langchain-mongodb \\\n",
        "    langchain-google-vertexai \\\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "MYcbMsd2FDzr",
      "metadata": {
        "id": "MYcbMsd2FDzr"
      },
      "source": [
        "## Ingest data\n",
        "To begin with the setup we will load the dataset to MongoDB Atlas. For user convinience, we are using an existing Hugingface MongoDB embedding dataset. Run the below code to import the *MongoDB/subset_arxiv_papers_with_embeddings* dataset as ds and load to MongoDB Atlas.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "OP-uRUPMFkXV",
      "metadata": {
        "collapsed": true,
        "id": "OP-uRUPMFkXV"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "from bs4 import BeautifulSoup\n",
        "from google.cloud import aiplatform\n",
        "from pymongo import MongoClient\n",
        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
        "from pymongo import MongoClient\n",
        "import certifi\n",
        "\n",
        "# Scrape the website content\n",
        "def scrape_website(url):\n",
        "    response = requests.get(url)\n",
        "    soup = BeautifulSoup(response.text, 'html.parser')\n",
        "    content = ' '.join([p.text for p in soup.find_all('p')])\n",
        "    return content\n",
        "\n",
        "# Split the content into chunks of 1000 characters\n",
        "def split_into_chunks(text, chunk_size=1000):\n",
        "    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]\n",
        "\n",
        "def get_text_embeddings(chunks):\n",
        "    model = TextEmbeddingModel.from_pretrained(\"text-embedding-004\")\n",
        "    inputs = chunks[0]\n",
        "    embeddings = model.get_embeddings(chunks)\n",
        "    return [embedding.values for embedding in embeddings]\n",
        "\n",
        "def write_to_mongoDB(embeddings, chunks, db_name, coll_name):\n",
        "    client = MongoClient(\"URI\", tlsCAFile=certifi.where())\n",
        "    db = client[db_name]\n",
        "    collection = db[coll_name]\n",
        "    \n",
        "    for i in range(len(chunks)):\n",
        "        collection.insert_one({\n",
        "            \"chunk\": chunks[i],\n",
        "            \"embedding\": embeddings[i]\n",
        "        })\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "content = scrape_website(\"https://en.wikipedia.org/wiki/Star_Wars#The_Skywalker_Saga\")\n",
        "chunks = split_into_chunks(content)\n",
        "embeddings_starwars = get_text_embeddings(chunks)\n",
        "write_to_mongoDB(embeddings_starwars, chunks, \"embeddings\", \"sample_starwars_embeddings\")\n",
        "\n",
        "content = scrape_website(\"https://en.wikipedia.org/wiki/Star_Trek\")\n",
        "chunks = split_into_chunks(content)\n",
        "embeddings_starteck = get_text_embeddings(chunks)\n",
        "write_to_mongoDB(embeddings_starteck, chunks, \"embeddings\", \"sample_startrek_embeddings\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "FNmiWoYLKNAi",
      "metadata": {
        "id": "FNmiWoYLKNAi"
      },
      "source": [
        "## Create vector search index on the newly created MongoDb collection\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "RmuEkpybNLza",
      "metadata": {
        "id": "RmuEkpybNLza"
      },
      "source": [
        "// To do: add code for creating atlas vector search index on the collection"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9df5222f",
      "metadata": {},
      "source": [
        "### Initilize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/apis/enableflow?apiid=aiplatform.googleapis.com)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "0pifIeELue8b",
      "metadata": {
        "executionInfo": {
          "elapsed": 4361,
          "status": "ok",
          "timestamp": 1718715628153,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "0pifIeELue8b"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"PROJECT ID\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = \"gs://bucket_name\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e8b884ec",
      "metadata": {},
      "source": [
        "### Import reasoning engine library"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "BxpZ1MCBwkR8",
      "metadata": {
        "executionInfo": {
          "elapsed": 671,
          "status": "ok",
          "timestamp": 1718715632800,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "BxpZ1MCBwkR8"
      },
      "outputs": [],
      "source": [
        "from vertexai.preview import reasoning_engines"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "28fe84de",
      "metadata": {},
      "source": [
        "### 1. Define Model\n",
        "\n",
        "As you construct your reasoning engine agent from the bottom up, the first component deals with which generative model you want to use in your agent. We are using \"gemini-1.5-pro\" which is latest at the time of creation of this python notebook. This LLM model will be used to build the RAG itself."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "ouoPfB8OwsD1",
      "metadata": {
        "executionInfo": {
          "elapsed": 644,
          "status": "ok",
          "timestamp": 1718715635977,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "ouoPfB8OwsD1"
      },
      "outputs": [],
      "source": [
        "model = \"gemini-1.5-pro-001\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "543c6e04",
      "metadata": {},
      "source": [
        "### 2. Defile Function to read from MongoDB Atlas using langchain\n",
        "\n",
        "The second component of your agent includes tools and functions, which enable the generative model to interact with MongoDB Atlas. We use Langchain to interact and query vectors from MongoDB Atlas. The function takes \"query\" as input and is trasformed into embeddings using Googles \"textembedding-gecko@001\" model. The function returns the queried data from MongoDB that has most similarity with the queried data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "id": "rq1nNT_fcMbV",
      "metadata": {
        "executionInfo": {
          "elapsed": 657,
          "status": "ok",
          "timestamp": 1718716092270,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "rq1nNT_fcMbV"
      },
      "outputs": [],
      "source": [
        "def star_wars_query_tool(\n",
        "    query: str\n",
        "):\n",
        "    \"\"\"\n",
        "    Retrieves vectors from a MongoDB database and uses them to answer a question related to Star wars.\n",
        "\n",
        "    Args:\n",
        "        query: The question to be answered about star wars.\n",
        "\n",
        "    Returns:\n",
        "        A dictionary containing the response to the question.\n",
        "    \"\"\"\n",
        "    from langchain.chains import ConversationalRetrievalChain, RetrievalQA\n",
        "    from langchain_mongodb import MongoDBAtlasVectorSearch\n",
        "    from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI\n",
        "    from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory\n",
        "    from pymongo import MongoClient\n",
        "\n",
        "    from langchain.prompts import PromptTemplate\n",
        "\n",
        "\n",
        "    prompt_template = \"\"\"Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Do not return any answers from your own knowledge.\n",
        "\n",
        "\n",
        "    {context}\n",
        "\n",
        "\n",
        "    Question: {question}\n",
        "    \"\"\"\n",
        "    PROMPT = PromptTemplate(\n",
        "        template=prompt_template, input_variables=[\"context\", \"question\"]\n",
        "    )\n",
        "\n",
        "    # Add your connection string in srv format below in place of URI\n",
        "    client = MongoClient(\"mongodb+srv://venkatesh:ashwin123@freetier.kxcgwh2.mongodb.net/\")\n",
        "    db = client[\"embeddings\"]\n",
        "\n",
        "    embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko@003\")\n",
        "\n",
        "    vs = MongoDBAtlasVectorSearch(\n",
        "        collection=db[\"sample_starwars_embeddings\"],\n",
        "        embedding=embeddings,\n",
        "        index_name=\"vector_index\",\n",
        "        embedding_key=\"embedding\",\n",
        "        text_key=\"chunk\",\n",
        "    )\n",
        "\n",
        "    llm = ChatVertexAI(\n",
        "        model_name=\"gemini-pro\",\n",
        "        convert_system_message_to_human=True,\n",
        "        max_output_tokens=1000,\n",
        "    )\n",
        "    retriever = vs.as_retriever(\n",
        "        search_type=\"mmr\", search_kwargs={\"k\": 10, \"lambda_mult\": 0.25}\n",
        "    )\n",
        "    memory = ConversationBufferWindowMemory(\n",
        "        memory_key=\"chat_history\", k=5, return_messages=True\n",
        "    )\n",
        "    conversation_chain = ConversationalRetrievalChain.from_llm(\n",
        "        llm=llm,\n",
        "        retriever=retriever,\n",
        "        memory=memory,\n",
        "        combine_docs_chain_kwargs={\"prompt\": PROMPT},\n",
        "    )\n",
        "    response = conversation_chain({\"question\": query})\n",
        "\n",
        "    return response"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "44edbefe",
      "metadata": {},
      "outputs": [],
      "source": [
        "def star_trek_query_tool(\n",
        "    query: str\n",
        "):\n",
        "    \"\"\"\n",
        "    Retrieves vectors from a MongoDB database and uses them to answer a question related to star trek.\n",
        "\n",
        "    Args:\n",
        "        query: The question to be answered about star trek.\n",
        "\n",
        "    Returns:\n",
        "        A dictionary containing the response to the question.\n",
        "    \"\"\"\n",
        "    from langchain.chains import ConversationalRetrievalChain, RetrievalQA\n",
        "    from langchain_mongodb import MongoDBAtlasVectorSearch\n",
        "    from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI\n",
        "    from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory\n",
        "    from pymongo import MongoClient\n",
        "\n",
        "    from langchain.prompts import PromptTemplate\n",
        "\n",
        "\n",
        "    prompt_template = \"\"\"Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Do not return any answers from your own knowledge.\n",
        "\n",
        "\n",
        "    {context}\n",
        "\n",
        "\n",
        "    Question: {question}\n",
        "    \"\"\"\n",
        "    PROMPT = PromptTemplate(\n",
        "        template=prompt_template, input_variables=[\"context\", \"question\"]\n",
        "    )\n",
        "\n",
        "    # Add your connection string in srv format below in place of URI\n",
        "    client = MongoClient(\"mongodb+srv://venkatesh:ashwin123@freetier.kxcgwh2.mongodb.net/\")\n",
        "    db = client[\"embeddings\"]\n",
        "\n",
        "    embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko@003\")\n",
        "\n",
        "    vs = MongoDBAtlasVectorSearch(\n",
        "        collection=db[\"sample_startrek_embeddings\"],\n",
        "        embedding=embeddings,\n",
        "        index_name=\"vector_index\",\n",
        "        embedding_key=\"embedding\",\n",
        "        text_key=\"chunk\",\n",
        "    )\n",
        "\n",
        "    llm = ChatVertexAI(\n",
        "        model_name=\"gemini-pro\",\n",
        "        convert_system_message_to_human=True,\n",
        "        max_output_tokens=1000,\n",
        "    )\n",
        "    retriever = vs.as_retriever(\n",
        "        search_type=\"mmr\", search_kwargs={\"k\": 10, \"lambda_mult\": 0.25}\n",
        "    )\n",
        "    memory = ConversationBufferWindowMemory(\n",
        "        memory_key=\"chat_history\", k=5, return_messages=True\n",
        "    )\n",
        "    conversation_chain = ConversationalRetrievalChain.from_llm(\n",
        "        llm=llm,\n",
        "        retriever=retriever,\n",
        "        memory=memory,\n",
        "        combine_docs_chain_kwargs={\"prompt\": PROMPT},\n",
        "    )\n",
        "    response = conversation_chain({\"question\": query})\n",
        "\n",
        "    return response"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c9501830",
      "metadata": {},
      "source": [
        "### 3. Define agent for calling your function\n",
        "\n",
        "The third component of your agent involves adding a reasoning layer, which helps your agent use the tools that you provided to help the end user achieve a higher-level goal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "id": "B56uiBCPxDM0",
      "metadata": {
        "executionInfo": {
          "elapsed": 644,
          "status": "ok",
          "timestamp": 1718715514590,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "B56uiBCPxDM0"
      },
      "outputs": [],
      "source": [
        "agent = reasoning_engines.LangchainAgent(\n",
        "    model=model,\n",
        "    tools=[star_wars_query_tool, star_trek_query_tool],\n",
        "    agent_executor_kwargs={\"return_intermediate_steps\": True},\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8frYrUr8I8rF",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 4844,
          "status": "ok",
          "timestamp": 1718716151526,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -330
        },
        "id": "8frYrUr8I8rF",
        "outputId": "4f77fb6a-ab5e-4002-f2d6-ea2b04f0a39e"
      },
      "outputs": [],
      "source": [
        "agent.query(input=\"tell me about star wars?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bf896a3e",
      "metadata": {},
      "source": [
        "### 4. Deploy the agent\n",
        "\n",
        "Now that you've specified a model, tools, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8VyDfLbxQEU",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b8VyDfLbxQEU",
        "outputId": "78cf063f-96ab-433f-c428-6c94e9eb0f6b"
      },
      "outputs": [],
      "source": [
        "remote_agent = reasoning_engines.ReasoningEngine.create(\n",
        "    agent,\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[langchain,reasoningengine]\",\n",
        "        \"cloudpickle==3.0.0\",\n",
        "        \"pydantic==2.7.4\",\n",
        "        \"langchain-mongodb\",\n",
        "        \"pymongo\",\n",
        "        \"langchain-google-vertexai\",\n",
        "\n",
        "    ],\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "MYcbMsd2FDzr"
      ],
      "name": "reasoning engine",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}