Building Production grade RAG with GPT 3_5-turbo and MongoDB Atlas.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "### **Setup**"
      ],
      "metadata": {
        "id": "cYnRIDsf0OzK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Y6NJxs2ySTR",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "72d1e54c-75c3-4b79-a593-29d2cb496584"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m676.9/676.9 kB\u001b[0m \u001b[31m5.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m15.4/15.4 MB\u001b[0m \u001b[31m17.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m21.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.7/307.7 kB\u001b[0m \u001b[31m7.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m267.1/267.1 kB\u001b[0m \u001b[31m12.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m4.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m136.0/136.0 kB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m38.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.9/3.9 MB\u001b[0m \u001b[31m45.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m290.4/290.4 kB\u001b[0m \u001b[31m21.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.9/77.9 kB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m30.8/30.8 MB\u001b[0m \u001b[31m31.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m3.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install --quiet llama-index pymongo llama-index-vector-stores-mongodb"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "! wget https://huggingface.co/spaces/rasyosef/RAG-with-Phi-2-and-LangChain/raw/main/Oppenheimer-movie-wiki.txt -P ./data"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O-GTXWBB8xpE",
        "outputId": "61da3ff2-32e0-4f1a-8010-a5338b602c6d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-04-10 00:36:29--  https://huggingface.co/spaces/rasyosef/RAG-with-Phi-2-and-LangChain/raw/main/Oppenheimer-movie-wiki.txt\n",
            "Resolving huggingface.co (huggingface.co)... 18.164.174.17, 18.164.174.55, 18.164.174.23, ...\n",
            "Connecting to huggingface.co (huggingface.co)|18.164.174.17|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 51987 (51K) [text/plain]\n",
            "Saving to: ‘./data/Oppenheimer-movie-wiki.txt’\n",
            "\n",
            "\r          Oppenheim   0%[                    ]       0  --.-KB/s               \rOppenheimer-movie-w 100%[===================>]  50.77K  --.-KB/s    in 0.06s   \n",
            "\n",
            "2024-04-10 00:36:29 (785 KB/s) - ‘./data/Oppenheimer-movie-wiki.txt’ saved [51987/51987]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = userdata.get(\"openai-api-key\")\n",
        "MONGO_URI = userdata.get(\"MONGO_URI\")"
      ],
      "metadata": {
        "id": "ctu83fS1EstN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Load Data**"
      ],
      "metadata": {
        "id": "5lu0s-q129pS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core import SimpleDirectoryReader\n",
        "\n",
        "reader = SimpleDirectoryReader(\n",
        "    input_dir=\"./data\"\n",
        "  )\n",
        "documents = reader.load_data()\n",
        "print(\"Number of Documents:\", len(documents))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jL3nklJ60cLD",
        "outputId": "20aa93ad-fb2c-4f2f-e349-dfc3103de01c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of Documents: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "\n",
        "node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=32)\n",
        "nodes = node_parser.get_nodes_from_documents(documents)\n",
        "print(\"Number of Nodes:\", len(nodes))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KfkrcApkAWS0",
        "outputId": "f5b57394-e65a-4c8c-b28e-cc1942107b32"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of Nodes: 25\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Embeddings**"
      ],
      "metadata": {
        "id": "WIpMj9mdCS5a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.embeddings.openai import OpenAIEmbedding\n",
        "\n",
        "embed_model = OpenAIEmbedding(\n",
        "    model=\"text-embedding-3-small\",\n",
        "    embed_batch_size=16,\n",
        ")"
      ],
      "metadata": {
        "id": "7LYxig13CxvY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **MongoDB Atlas Vector Store**"
      ],
      "metadata": {
        "id": "tSl9ZUGCHoFl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pymongo\n",
        "from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch\n",
        "from llama_index.core.indices import VectorStoreIndex\n",
        "from llama_index.core import StorageContext\n",
        "\n",
        "mongodb_client = pymongo.MongoClient(MONGO_URI)\n",
        "\n",
        "db = mongodb_client[\"oppenheimer\"]\n",
        "collection = db[\"oppenheimer_wiki_chunks\"]\n",
        "\n",
        "# Delete any existing records\n",
        "collection.delete_many({})"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "itM_6cQsFNu6",
        "outputId": "4b453686-e29d-42c6-be85-58a41c3fa5e0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeleteResult({'n': 100, 'electionId': ObjectId('7fffffff0000000000000064'), 'opTime': {'ts': Timestamp(1712709399, 92), 't': 100}, 'ok': 1.0, '$clusterTime': {'clusterTime': Timestamp(1712709399, 101), 'signature': {'hash': b'O\\x0fe\\x7fv#\\xb12-T%\\x0b\\xa7VK\\xbfP\\x88-\\xdb', 'keyId': 7314014067373376376}}, 'operationTime': Timestamp(1712709399, 92)}, acknowledged=True)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "store = MongoDBAtlasVectorSearch(\n",
        "    mongodb_client=mongodb_client,\n",
        "    db_name=\"oppenheimer\",\n",
        "    collection_name=\"oppenheimer_wiki_chunks\",\n",
        "    index_name=\"vector_index\",\n",
        "    embedding_key=\"embedding\",\n",
        "  )\n",
        "storage_context = StorageContext.from_defaults(vector_store=store)"
      ],
      "metadata": {
        "id": "OumuArgpQeIS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index = VectorStoreIndex(\n",
        "    nodes=nodes,\n",
        "    storage_context=storage_context,\n",
        "    embed_model=embed_model\n",
        ")"
      ],
      "metadata": {
        "id": "iiwdnoxrJgUX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Query Engine**"
      ],
      "metadata": {
        "id": "q7x_GPQcOKA0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.llms.openai import OpenAI\n",
        "\n",
        "llm = OpenAI(\n",
        "    model=\"gpt-3.5-turbo-0125\",\n",
        "    temperature=0,\n",
        "    max_tokens=256\n",
        ")"
      ],
      "metadata": {
        "id": "7CzNjVngJsz2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query_engine = index.as_query_engine(llm=llm, streaming=True, similarity_top_k=3)"
      ],
      "metadata": {
        "id": "L_iGMe3GOkhj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"What's the name of the actor that played Lewis Strauss?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lnWLceLvOtLc",
        "outputId": "3482c16f-c95b-4666-f35a-9de3f70cb2af"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Robert Downey Jr."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"Which character did Matthias Schweighöfer play in the movie?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bRf8_n4fO208",
        "outputId": "07ef6244-5272-42e6-802a-1de5932b718a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Matthias Schweighöfer played the character Werner Heisenberg in the movie."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"Who is Boris Pash?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "id": "QO3jyLcRPZU-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3c44b429-5d92-48b2-b228-e80fe121a6ea"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Boris Pash is a U.S. Army military intelligence officer and commander of the Alsos Mission."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"Which character did Casey Affleck play in the movie?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Qi4xJX19csdV",
        "outputId": "8cebc823-5590-4bed-95f6-659d739fda08"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Boris Pash"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"How much money did the Oppenheimer movie make at the US and global box office?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HGiGkxQDc1zX",
        "outputId": "4255ddd1-0821-493f-9f74-67df865a1cca"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Oppenheimer made $326.4 million in the United States and Canada, and $628.9 million in other territories, resulting in a worldwide total of $955.3 million."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"What score did the Oppenheimer movie get on Rotten Tomatoes and Metacritic?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "44TXBx-edOOu",
        "outputId": "accf19ac-08db-4686-ba64-09b46a5cf906"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The Oppenheimer movie received a score of 93% on Rotten Tomatoes and a score of 89 out of 100 on Metacritic."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"Why did Lewis Strauss have a grudge against J. Robert Oppenheimer?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vOQDB-5RK2rj",
        "outputId": "7ca1a2b0-04f2-4caf-b194-678d3d0114f3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Lewis Strauss had a grudge against J. Robert Oppenheimer because Oppenheimer publicly humiliated him by dismissing his concerns about exporting radioisotopes and recommending negotiations with the Soviet Union after they successfully detonated their own bomb. Additionally, Strauss believed that Oppenheimer denigrated him during a conversation Oppenheimer had with Einstein in 1947."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "streaming_response = query_engine.query(\"What happened while Oppenheimer was a student at the University of Cambridge?\")\n",
        "streaming_response.print_response_stream()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G_p69rMMKFVJ",
        "outputId": "08917e46-7239-439f-e8de-f98bdcebac00"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "While Oppenheimer was a student at the University of Cambridge, he grappled with anxiety and homesickness while studying under experimental physicist Patrick Blackett. Upset with Blackett's attitude, Oppenheimer left him a poisoned apple but later retrieved it. Niels Bohr recommended that Oppenheimer study theoretical physics at the University of Göttingen."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sources = streaming_response.source_nodes\n",
        "for node in sources:\n",
        "  print(node.text)\n",
        "  print(\"\\n\\n------------------------------------------\\n\\n\")\n",
        "\n",
        "len(streaming_response.source_nodes)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j2d8n9ZUKN2h",
        "outputId": "9a6b652d-8bbc-447d-e951-25fd0882f93a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dane DeHaan as Maj Gen. Kenneth Nichols, a U.S. Army officer and the deputy district engineer of the Manhattan Project.\n",
            "Alden Ehrenreich as a Senate aide to Lewis Strauss, a fictional character who is an aide during Strauss's nomination for United States Secretary of Commerce.\n",
            "Tony Goldwyn as Gordon Gray, a government official and chairman of the committee deciding the revoking of Oppenheimer security clearance.\n",
            "Jefferson Hall as Haakon Chevalier, a Berkeley professor who became friends with Oppenheimer at university.\n",
            "David Krumholtz as Isidor Isaac Rabi, a Nobel Prize-winning physicist who worked as a consultant on the Manhattan Project.\n",
            "Matthew Modine as Vannevar Bush, head of the Office of Scientific Research and Development.\n",
            "Scott Grimes as Counsel to Lewis Strauss.\n",
            "Kurt Koehler as Thomas A. Morgan, an industrialist and former chairman of the board of the Sperry Corporation who was one of the panel members at Oppenheimer's security clearance hearing.\n",
            "John Gowans as Ward V. Evans, a chemist and academic who served as one of the panel members at Oppenheimer's security clearance hearing.\n",
            "Macon Blair as Lloyd K. Garrison, a lawyer who helped to represent Oppenheimer at his security clearance hearing.\n",
            "Gregory Jbara as Sen. Warren Magnuson, Chairman of Senate Commerce Committee.\n",
            "Harry Groener as Sen. Gale W. McGee.\n",
            "Tim DeKay as Sen. John Pastore.\n",
            "Matthias Schweighöfer as Werner Heisenberg, a German Nobel Prize-winning physicist who worked in the country's nuclear weapons program during World War II.\n",
            "Alex Wolff as Luis Walter Alvarez, a Nobel-winning physicist who worked on the Manhattan Project.\n",
            "Josh Zuckerman as Giovanni Rossi Lomanitz, a physicist who became Oppenheimer's protégé at Berkeley.\n",
            "Rory Keane as Hartland Snyder, a physicist, who collaborated with Oppenheimer to calculate the gravitational collapse of a dust particle sphere.\n",
            "Michael Angarano as Robert Serber, a physicist who worked on the Manhattan Project.\n",
            "Emma Dumont as Jackie Oppenheimer, Frank's wife and Robert's sister-in-law.\n",
            "Guy Burnet as George C. Eltenton, a chemical engineer in the U.S. with ties to the Soviet Union.\n",
            "\n",
            "\n",
            "------------------------------------------\n",
            "\n",
            "\n",
            "In 1963, President Lyndon B. Johnson presents Oppenheimer with the Enrico Fermi Award as a gesture of political rehabilitation. A flashback reveals Oppenheimer and Einstein's 1947 conversation never mentioned Strauss. Oppenheimer instead expressed his belief that they had indeed started a chain reaction—a nuclear arms race—that would one day destroy the world.\n",
            "\n",
            "Cast\n",
            "Cillian Murphy as J. Robert Oppenheimer, a theoretical physicist and director of the Los Alamos National Laboratory.\n",
            "Emily Blunt as Katherine \"Kitty\" Oppenheimer, Robert Oppenheimer's wife and a former Communist Party USA member.\n",
            "Matt Damon as Gen. Leslie Groves, a United States Army Corps of Engineers (USACE) officer and director of the Manhattan Project.\n",
            "Robert Downey Jr. as Rear Admiral Lewis Strauss, a retired Naval officer and high-ranking member of the U.S. Atomic Energy Commission (AEC).\n",
            "Florence Pugh as Jean Tatlock, a psychiatrist, Communist Party USA member, and Robert Oppenheimer's romantic interest.\n",
            "Josh Hartnett as Ernest Lawrence, a Nobel-winning nuclear physicist who worked with Oppenheimer at the University of California, Berkeley.\n",
            "Casey Affleck as Boris Pash, a U.S. Army military intelligence officer and commander of the Alsos Mission.\n",
            "Rami Malek as David L. Hill, a nuclear physicist at the Met Lab, who helped to create the Chicago Pile.\n",
            "Kenneth Branagh as Niels Bohr, a Nobel-winning Danish physicist, philosopher and Oppenheimer's personal idol.\n",
            "Benny Safdie as Edward Teller, a Hungarian theoretical physicist known for being the \"father of the hydrogen bomb\".\n",
            "Jason Clarke as Roger Robb, an attorney and future U.S. circuit judge who served as special counsel to the AEC at Oppenheimer's security hearing.\n",
            "Dylan Arnold as Frank Oppenheimer, Robert's younger brother and a particle physicist who worked on the Manhattan Project.\n",
            "Tom Conti as Albert Einstein, Nobel-winning German theoretical physicist known for developing the theory of relativity.\n",
            "James D'Arcy as Patrick Blackett, Oppenheimer's college professor and Nobel-winning physicist at Cambridge University.\n",
            "David Dastmalchian as William L. Borden, a lawyer and executive director of the JCAE.\n",
            "Dane DeHaan as Maj Gen. Kenneth Nichols, a U.S.\n",
            "\n",
            "\n",
            "------------------------------------------\n",
            "\n",
            "\n",
            "Guy Burnet as George C. Eltenton, a chemical engineer in the U.S. with ties to the Soviet Union.\n",
            "Louise Lombard as Ruth Tolman, a psychologist close to Oppenheimer during the development of the atomic bomb.\n",
            "Tom Jenkins as Richard C. Tolman, Ruth's husband and General Groves' chief scientific adviser on the Manhattan Project.\n",
            "Olli Haaskivi as Edward Condon, a nuclear physicist who helped with the development of radar and briefly took part in the Manhattan Project.\n",
            "David Rysdahl as Donald Hornig, a chemist who worked on the firing unit at Los Alamos.\n",
            "Josh Peck as Kenneth Bainbridge, a physicist who was the director of the Manhattan Project's Trinity nuclear test.\n",
            "Jack Quaid as Richard Feynman, theoretical physicist who worked in the Theoretical Division at Los Alamos.\n",
            "Gustaf Skarsgård as Hans Bethe, a German-American Nobel-winning theoretical physicist and the head of the Theoretical Division at Los Alamos.\n",
            "James Urbaniak as Kurt Gödel, a mathematician known for his theorems that revolutionized mathematics and had far-reaching implications for philosophy and computer science.\n",
            "Trond Fausa as George Kistiakowsky, a Harvard professor who took part in the Manhattan Project.\n",
            "Devon Bostick as Seth Neddermeyer, a physicist who discovered the muon and advocated for the implosion-type nuclear weapon used in the Trinity Test.\n",
            "Danny Deferrari as Enrico Fermi, an Italian Nobel-winning physicist and creator of the Chicago Pile.\n",
            "Christopher Denham as Klaus Fuchs, a German-born physicist who worked on the Manhattan Project and spied for the Soviet Union.\n",
            "Jessica Erin Martin as Charlotte Serber, head technical librarian at Los Alamos.\n",
            "Ronald Auguste as J. Ernest Wilkins Jr., an African American nuclear scientist, mechanical engineer and mathematician who worked with Oppenheimer on the Manhattan Project.\n",
            "Máté Haumann as Leo Szilard, a Hungarian physicist who conceived the idea of nuclear chain reaction in 1933, and later in July 1945 at the Chicago branch of the Manhattan Project circulated the petition to President Truman against unannounced use of atomic weapons on Japan.\n",
            "Olivia Thirlby as Lilli Hornig, a Czech-American scientist who worked on the Manhattan Project.\n",
            "\n",
            "\n",
            "------------------------------------------\n",
            "\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "3"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Vector Search**"
      ],
      "metadata": {
        "id": "ulQYZ1NNSipI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def query_atlas(query):\n",
        "  query_embedding = embed_model.get_query_embedding(query)\n",
        "  pipeline = [\n",
        "    {\n",
        "      \"$vectorSearch\": {\n",
        "        \"index\": \"vector_index\",\n",
        "        \"path\": \"embedding\",\n",
        "        \"queryVector\": query_embedding,\n",
        "        \"numCandidates\": 128,\n",
        "        \"limit\": 3 # Return top 3 matches\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "        \"$project\": {\n",
        "            \"_id\": 0,\n",
        "            \"text\": 1,\n",
        "            \"score\": {\"$meta\": \"vectorSearchScore\"}\n",
        "        }\n",
        "    }\n",
        "  ]\n",
        "\n",
        "  results = collection.aggregate(pipeline)\n",
        "  return list(results)"
      ],
      "metadata": {
        "id": "aEF-S9T9TZi7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query_atlas(\"Who is Boris Pash?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5PU5HZY8Ugpy",
        "outputId": "ee45d819-47b3-41ad-b5dc-7481c9b5cab8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'text': 'In 1963, President Lyndon B. Johnson presents Oppenheimer with the Enrico Fermi Award as a gesture of political rehabilitation. A flashback reveals Oppenheimer and Einstein\\'s 1947 conversation never mentioned Strauss. Oppenheimer instead expressed his belief that they had indeed started a chain reaction—a nuclear arms race—that would one day destroy the world.\\n\\nCast\\nCillian Murphy as J. Robert Oppenheimer, a theoretical physicist and director of the Los Alamos National Laboratory.\\nEmily Blunt as Katherine \"Kitty\" Oppenheimer, Robert Oppenheimer\\'s wife and a former Communist Party USA member.\\nMatt Damon as Gen. Leslie Groves, a United States Army Corps of Engineers (USACE) officer and director of the Manhattan Project.\\nRobert Downey Jr. as Rear Admiral Lewis Strauss, a retired Naval officer and high-ranking member of the U.S. Atomic Energy Commission (AEC).\\nFlorence Pugh as Jean Tatlock, a psychiatrist, Communist Party USA member, and Robert Oppenheimer\\'s romantic interest.\\nJosh Hartnett as Ernest Lawrence, a Nobel-winning nuclear physicist who worked with Oppenheimer at the University of California, Berkeley.\\nCasey Affleck as Boris Pash, a U.S. Army military intelligence officer and commander of the Alsos Mission.\\nRami Malek as David L. Hill, a nuclear physicist at the Met Lab, who helped to create the Chicago Pile.\\nKenneth Branagh as Niels Bohr, a Nobel-winning Danish physicist, philosopher and Oppenheimer\\'s personal idol.\\nBenny Safdie as Edward Teller, a Hungarian theoretical physicist known for being the \"father of the hydrogen bomb\".\\nJason Clarke as Roger Robb, an attorney and future U.S. circuit judge who served as special counsel to the AEC at Oppenheimer\\'s security hearing.\\nDylan Arnold as Frank Oppenheimer, Robert\\'s younger brother and a particle physicist who worked on the Manhattan Project.\\nTom Conti as Albert Einstein, Nobel-winning German theoretical physicist known for developing the theory of relativity.\\nJames D\\'Arcy as Patrick Blackett, Oppenheimer\\'s college professor and Nobel-winning physicist at Cambridge University.\\nDavid Dastmalchian as William L. Borden, a lawyer and executive director of the JCAE.\\nDane DeHaan as Maj Gen. Kenneth Nichols, a U.S.',\n",
              "  'score': 0.5782333612442017},\n",
              " {'text': \"Guy Burnet as George C. Eltenton, a chemical engineer in the U.S. with ties to the Soviet Union.\\nLouise Lombard as Ruth Tolman, a psychologist close to Oppenheimer during the development of the atomic bomb.\\nTom Jenkins as Richard C. Tolman, Ruth's husband and General Groves' chief scientific adviser on the Manhattan Project.\\nOlli Haaskivi as Edward Condon, a nuclear physicist who helped with the development of radar and briefly took part in the Manhattan Project.\\nDavid Rysdahl as Donald Hornig, a chemist who worked on the firing unit at Los Alamos.\\nJosh Peck as Kenneth Bainbridge, a physicist who was the director of the Manhattan Project's Trinity nuclear test.\\nJack Quaid as Richard Feynman, theoretical physicist who worked in the Theoretical Division at Los Alamos.\\nGustaf Skarsgård as Hans Bethe, a German-American Nobel-winning theoretical physicist and the head of the Theoretical Division at Los Alamos.\\nJames Urbaniak as Kurt Gödel, a mathematician known for his theorems that revolutionized mathematics and had far-reaching implications for philosophy and computer science.\\nTrond Fausa as George Kistiakowsky, a Harvard professor who took part in the Manhattan Project.\\nDevon Bostick as Seth Neddermeyer, a physicist who discovered the muon and advocated for the implosion-type nuclear weapon used in the Trinity Test.\\nDanny Deferrari as Enrico Fermi, an Italian Nobel-winning physicist and creator of the Chicago Pile.\\nChristopher Denham as Klaus Fuchs, a German-born physicist who worked on the Manhattan Project and spied for the Soviet Union.\\nJessica Erin Martin as Charlotte Serber, head technical librarian at Los Alamos.\\nRonald Auguste as J. Ernest Wilkins Jr., an African American nuclear scientist, mechanical engineer and mathematician who worked with Oppenheimer on the Manhattan Project.\\nMáté Haumann as Leo Szilard, a Hungarian physicist who conceived the idea of nuclear chain reaction in 1933, and later in July 1945 at the Chicago branch of the Manhattan Project circulated the petition to President Truman against unannounced use of atomic weapons on Japan.\\nOlivia Thirlby as Lilli Hornig, a Czech-American scientist who worked on the Manhattan Project.\",\n",
              "  'score': 0.5758520364761353},\n",
              " {'text': \"Dane DeHaan as Maj Gen. Kenneth Nichols, a U.S. Army officer and the deputy district engineer of the Manhattan Project.\\nAlden Ehrenreich as a Senate aide to Lewis Strauss, a fictional character who is an aide during Strauss's nomination for United States Secretary of Commerce.\\nTony Goldwyn as Gordon Gray, a government official and chairman of the committee deciding the revoking of Oppenheimer security clearance.\\nJefferson Hall as Haakon Chevalier, a Berkeley professor who became friends with Oppenheimer at university.\\nDavid Krumholtz as Isidor Isaac Rabi, a Nobel Prize-winning physicist who worked as a consultant on the Manhattan Project.\\nMatthew Modine as Vannevar Bush, head of the Office of Scientific Research and Development.\\nScott Grimes as Counsel to Lewis Strauss.\\nKurt Koehler as Thomas A. Morgan, an industrialist and former chairman of the board of the Sperry Corporation who was one of the panel members at Oppenheimer's security clearance hearing.\\nJohn Gowans as Ward V. Evans, a chemist and academic who served as one of the panel members at Oppenheimer's security clearance hearing.\\nMacon Blair as Lloyd K. Garrison, a lawyer who helped to represent Oppenheimer at his security clearance hearing.\\nGregory Jbara as Sen. Warren Magnuson, Chairman of Senate Commerce Committee.\\nHarry Groener as Sen. Gale W. McGee.\\nTim DeKay as Sen. John Pastore.\\nMatthias Schweighöfer as Werner Heisenberg, a German Nobel Prize-winning physicist who worked in the country's nuclear weapons program during World War II.\\nAlex Wolff as Luis Walter Alvarez, a Nobel-winning physicist who worked on the Manhattan Project.\\nJosh Zuckerman as Giovanni Rossi Lomanitz, a physicist who became Oppenheimer's protégé at Berkeley.\\nRory Keane as Hartland Snyder, a physicist, who collaborated with Oppenheimer to calculate the gravitational collapse of a dust particle sphere.\\nMichael Angarano as Robert Serber, a physicist who worked on the Manhattan Project.\\nEmma Dumont as Jackie Oppenheimer, Frank's wife and Robert's sister-in-law.\\nGuy Burnet as George C. Eltenton, a chemical engineer in the U.S. with ties to the Soviet Union.\",\n",
              "  'score': 0.567152738571167}]"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    }
  ]
}