Updating agents notebook

Author: ajosh0504

labs/ai-agents-lab.ipynb

@@ -21,7 +21,15 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "import os\nimport sys\nfrom pymongo import MongoClient\n\n# Add parent directory to path to import from utils\nsys.path.append(os.path.join(os.path.dirname(os.getcwd())))\nfrom utils import track_progress"
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "from pymongo import MongoClient\n",
+    "\n",
+    "# Add parent directory to path to import from utils\n",
+    "sys.path.append(os.path.join(os.path.dirname(os.getcwd())))\n",
+    "from utils import track_progress, set_env"
+   ]
   },
   {
    "cell_type": "code",
@@ -48,10 +56,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "### **Pick an LLM provider of your choice below**"
+    "# Set the LLM provider and passkey provided by your workshop instructor\n",
+    "# NOTE: LLM_PROVIDER can be set to one of \"aws\"/ \"microsoft\" / \"google\"\n",
+    "LLM_PROVIDER = \"aws\"\n",
+    "PASSKEY = \"replace-with-passkey\""
    ]
   },
   {
@@ -60,9 +73,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "SERVERLESS_URL = os.environ.get(\"SERVERLESS_URL\")\n",
-    "# Can be one of \"aws\", \"google\" or \"microsoft\"\n",
-    "LLM_PROVIDER = \"aws\""
+    "# Obtain API keys from our AI model proxy and set them as an environment variables-- DO NOT CHANGE\n",
+    "set_env([LLM_PROVIDER,\"voyageai\"], PASSKEY)"
    ]
   },
   {
@@ -192,29 +204,22 @@
     "            {\n",
     "                \"type\": \"vector\",\n",
     "                \"path\": \"embedding\",\n",
-    "                \"numDimensions\": 384,\n",
+    "                \"numDimensions\": 1024,\n",
     "                \"similarity\": \"cosine\",\n",
     "            }\n",
     "        ]\n",
     "    },\n",
     "}"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "📚 Refer to the `utils.py` script under `notebooks/utils`"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Use the `create_index` function from the `utils` module to create a vector search index with the above definition for the `vs_collection` collection\n",
-    "<CODE_BLOCK_1>"
+    "create_index(vs_collection, VS_INDEX_NAME, model)"
    ]
   },
   {
@@ -252,9 +257,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# You may see a warning upon running this cell. You can ignore it.\n",
     "from langchain.agents import tool\n",
-    "from sentence_transformers import SentenceTransformer\n",
+    "import voyageai\n",
     "from typing import List"
    ]
   },
@@ -271,15 +275,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Load the `gte-small` model using the Sentence Transformers library\n",
-    "embedding_model = SentenceTransformer(\"thenlper/gte-small\")"
+    "# Initialize the Voyage AI client\n",
+    "vo = voyageai.Client()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "📚 https://huggingface.co/thenlper/gte-small#usage (See \"Use with sentence-transformers\" under Usage)"
+    "📚 https://docs.voyageai.com/docs/contextualized-chunk-embeddings#approach-2-contextualized-chunk-embeddings"
    ]
   },
   {
@@ -288,20 +292,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Define a function that takes a piece of text (`text`) as input, embeds it using the `embedding_model` instantiated above and returns the embedding as a list\n",
-    "# An array can be converted to a list using the `tolist()` method\n",
-    "def get_embedding(text: str) -> List[float]:\n",
+    "def get_embeddings(query: str) -> List[float]:\n",
     "    \"\"\"\n",
-    "    Generate the embedding for a piece of text.\n",
+    "    Get embeddings for an input query.\n",
     "\n",
     "    Args:\n",
-    "        text (str): Text to embed.\n",
+    "        query (str): Query string\n",
     "\n",
     "    Returns:\n",
-    "        List[float]: Embedding of the text as a list.\n",
+    "        List[float]: Embedding of the query string\n",
     "    \"\"\"\n",
-    "    embedding = <CODE_BLOCK_2>\n",
-    "    return embedding.tolist()"
+    "    # Use the `contextualized_embed` method of the Voyage AI API to embed the user query with the following arguments:\n",
+    "    # inputs: `query` wrapped in a list of lists\n",
+    "    # model: `voyage-context-3`\n",
+    "    # input_type: \"query\"\n",
+    "    embds_obj = <CODE_BLOCK_1>\n",
+    "    # Extract embeddings from the embeddings object\n",
+    "    embeddings = <CODE_BLOCK_2>\n",
+    "    return embeddings"
    ]
   },
   {
@@ -330,8 +338,8 @@
     "    str: The retrieved information formatted as a string.\n",
     "    \"\"\"\n",
     "\n",
-    "    # Generate embedding for the `user_query` using the `get_embedding` function defined above\n",
-    "    query_embedding = <CODE_BLOCK_33>\n",
+    "    # Generate embeddings for the `user_query` using the `get_embeddings` function defined above\n",
+    "    query_embedding = <CODE_BLOCK_3>\n",
     "\n",
     "    # Define an aggregation pipeline consisting of a $vectorSearch stage, followed by a $project stage\n",
     "    # Set the number of candidates to 150 and only return the top 5 documents from the vector search\n",
@@ -481,7 +489,7 @@
    "source": [
     "from langchain_core.load import load\n",
     "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
-    "import requests"
+    "from utils import get_llm"
    ]
   },
   {
@@ -490,11 +498,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Obtain the Langchain LLM object from our serverless endpoint\n",
-    "llm_dict = requests.post(\n",
-    "    url=SERVERLESS_URL, json={\"task\": \"get_llm\", \"data\": LLM_PROVIDER}\n",
-    ").json()\n",
-    "llm = load(llm_dict[\"llm\"], secrets_map=llm_dict[\"secrets_map\"])"
+    "# Obtain the Langchain LLM object using the `get_llm` function from the `utils`` module.\n",
+    "llm = get_llm(LLM_PROVIDER)"
    ]
   },
   {
@@ -560,7 +565,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Chain the `prompt` with the tool-bound llm using the `|` operator\n",
+    "# Chain the `prompt` with the tool-augmented llm using the `|` operator\n",
     "llm_with_tools = <CODE_BLOCK_10>"
    ]
   },
@@ -830,11 +835,7 @@
    "outputs": [],
    "source": [
     "# Visualize the graph\n",
-    "try:\n",
-    "    display(Image(app.get_graph().draw_mermaid_png()))\n",
-    "except Exception:\n",
-    "    # This requires some extra dependencies and is optional\n",
-    "    pass"
+    "app"
    ]
   },
   {
@@ -850,24 +851,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Stream outputs from the graph as they pass through its nodes\n",
     "def execute_graph(user_input: str) -> None:\n",
     "    \"\"\"\n",
     "    Stream outputs from the graph\n",
     "\n",
     "    Args:\n",
     "        user_input (str): User query string\n",
     "    \"\"\"\n",
-    "    # Add user input to the messages attribute of the graph state\n",
-    "    # The role of the message should be \"user\" and content should be `user_input`\n",
-    "    input = {\"messages\": [(\"user\", user_input)]}\n",
-    "    # Pass input to the graph and stream the outputs\n",
-    "    for output in app.stream(input):\n",
-    "        for key, value in output.items():\n",
-    "            print(f\"Node {key}:\")\n",
-    "            print(value)\n",
-    "    print(\"---FINAL ANSWER---\")\n",
-    "    print(value[\"messages\"][-1].content)"
+    "    # Stream outputs from each step in the graph\n",
+    "    for step in app.stream(\n",
+    "        {\"messages\": [{\"role\": \"user\", \"content\": user_input}]},\n",
+    "        # Stream full value of the state after each step\n",
+    "        stream_mode=\"values\",\n",
+    "    ):\n",
+    "        # Print the latest message from the step\n",
+    "        step[\"messages\"][-1].pretty_print()"
    ]
   },
   {
@@ -939,26 +937,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def execute_graph(thread_id: str, user_input: str) -> None:\n",
+    "def execute_graph_with_memory(thread_id: str, user_input: str) -> None:\n",
     "    \"\"\"\n",
     "    Stream outputs from the graph\n",
     "\n",
     "    Args:\n",
     "        thread_id (str): Thread ID for the checkpointer\n",
     "        user_input (str): User query string\n",
     "    \"\"\"\n",
-    "    # Add user input to the messages attribute of the graph state\n",
-    "    # The role of the message should be \"user\" and content should be `user_input`\n",
-    "    input = {\"messages\": [(\"user\", user_input)]}\n",
-    "    # Define a config containing the thread ID\n",
+    "    # Create a runtime config for the thread ID `thread_id`\n",
     "    config = <CODE_BLOCK_19>\n",
-    "    # Pass `input` and `config` to the graph and stream outputs\n",
-    "    for output in app.stream(input, config):\n",
-    "        for key, value in output.items():\n",
-    "            print(f\"Node {key}:\")\n",
-    "            print(value)\n",
-    "    print(\"---FINAL ANSWER---\")\n",
-    "    print(value[\"messages\"][-1].content)"
+    "    # Stream outputs from each step in the graph\n",
+    "    for step in app.stream(\n",
+    "        {\"messages\": [{\"role\": \"user\", \"content\": user_input}]},\n",
+    "        # Pass the config as an additional parameter\n",
+    "        config,\n",
+    "        stream_mode=\"values\",\n",
+    "    ):\n",
+    "        # Print the latest message from the step\n",
+    "        step[\"messages\"][-1].pretty_print()"
    ]
   },
   {
@@ -968,7 +965,7 @@
    "outputs": [],
    "source": [
     "# Test graph execution with thread ID\n",
-    "execute_graph(\n",
+    "execute_graph_with_memory(\n",
     "    \"1\",\n",
     "    \"What are some best practices for data backups in MongoDB?\",\n",
     ")"
@@ -981,7 +978,7 @@
    "outputs": [],
    "source": [
     "# Follow-up question to ensure message history works\n",
-    "execute_graph(\n",
+    "execute_graph_with_memory(\n",
     "    \"1\",\n",
     "    \"What did I just ask you?\",\n",
     ")"
@@ -1014,4 +1011,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

labs/ai-rag-lab.ipynb

@@ -231,7 +231,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def get_embeddings(content: List[str], input_type: str) -> List[List[float]]:\n",
+    "def get_embeddings(content: List[str], input_type: str) -> List[float] | List[List[float]]:\n",
     "    \"\"\"\n",
     "    Get contextualized embeddings for each chunk.\n",
     "\n",
@@ -240,7 +240,7 @@
     "        input_type (str): Type of input, either \"document\" or \"query\" \n",
     "\n",
     "    Returns:\n",
-    "        List[List[float]]: Contextualized embeddings\n",
+    "        List[float] | List[List[float]]: Contextualized embeddings\n",
     "    \"\"\"\n",
     "    # Use the `contextualized_embed` method of the Voyage AI API to get contextualized embeddings for each chunk with the following arguments:\n",
     "    # inputs: `content` wrapped in another list\n",