DOCSP-47958 -- langchain + langgraph natural language to mql notebook

ai-integrations/langchain-natural-language.ipynb

@@ -0,0 +1,182 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f9696293",
+   "metadata": {},
+   "source": [
+    "# Query Atlas with Natural Language Using LangChain and LangGraph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e696dea0",
+   "metadata": {},
+   "source": [
+    "This notebook is a companion to the [Query Atlas with Natural Language Using LangChain and LangGraph](https://www.mongodb.com/docs/atlas/atlas-vector-search/ai-integrations/langchain/natural-language-to-mql/) tutorial. Refer to the page for set-up instructions and detailed explanations.\n",
+    "\n",
+    "This notebook demonstrates how to query an Atlas cluster with a natural language prompt using an AI agent built with the [LangChain MongoDB Toolkit](https://langchain-mongodb.readthedocs.io/en/latest/langchain_mongodb/agent_toolkit/langchain_mongodb.agent_toolkit.toolkit.MongoDBDatabaseToolkit.html#langchain_mongodb.agent_toolkit.toolkit.MongoDBDatabaseToolkit) and the [LangGraph ReAct Agent Framework](https://langchain-ai.github.io/langgraph/agents/agents/).\n",
+    "\n",
+    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/mongodb/docs-notebooks/blob/main/ai-integrations/langchain-natural-language.ipynb\">\n",
+    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+    "</a>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f106dda9",
+   "metadata": {
+    "vscode": {
+     "languageId": "shellscript"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "pip install --quiet --upgrade langchain-mongodb langchain-openai langgraph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "998157e0",
+   "metadata": {},
+   "source": [
+    "## Set up your environment\n",
+    "\n",
+    "Before you begin, make sure you have the following:\n",
+    "\n",
+    "- An Atlas cluster up and running (you'll need the [connection string](https://www.mongodb.com/docs/guides/atlas/connection-string/))\n",
+    "- An API key to access an LLM (This tutorial uses a model from OpenAI, but you can use any model [supported by LangChain](https://python.langchain.com/docs/integrations/chat/))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "694ccd64",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"OPENAI_API_KEY\"] = '<api-key>'\n",
+    "ATLAS_CONNECTION_STRING = '<atlas-connection-string>'\n",
+    "ATLAS_DB_NAME = 'sample_restaurants'\n",
+    "NATURAL_LANGUAGE_QUERY = 'Find all restaurants that serve hamburgers.'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c764c565",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "from langgraph.prebuilt import create_react_agent\n",
+    "from langchain_mongodb.agent_toolkit import (\n",
+    "    MONGODB_AGENT_SYSTEM_PROMPT,\n",
+    "    MongoDBDatabase,\n",
+    "    MongoDBDatabaseToolkit,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5a6b006c",
+   "metadata": {},
+   "source": [
+    "## Build the agent\n",
+    "\n",
+    "Next, define the `NaturalLanguageToMQL` Python class.\n",
+    "\n",
+    "#### Key Points\n",
+    "\n",
+    "- `self.toolkit`, the tools that the agent can use, is an instance of the [MongoDB Toolkit](https://langchain-mongodb.readthedocs.io/en/latest/langchain_mongodb/agent_toolkit/langchain_mongodb.agent_toolkit.toolkit.MongoDBDatabaseToolkit.html#langchain_mongodb.agent_toolkit.toolkit.MongoDBDatabaseToolkit). \n",
+    "\n",
+    "- `self.agent`, the agent itself, is an instance of the [ReAct Agent framework](https://langchain-ai.github.io/langgraph/agents/agents/), which takes `self.toolkit` as a parameter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b45185db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class NaturalLanguageToMQL:\n",
+    "    def __init__(self):\n",
+    "        self.llm = ChatOpenAI(model=\"gpt-4o-mini\", timeout=60)\n",
+    "        self.system_message = MONGODB_AGENT_SYSTEM_PROMPT.format(top_k=5)\n",
+    "        self.db_wrapper = MongoDBDatabase.from_connection_string(\n",
+    "                            ATLAS_CONNECTION_STRING, \n",
+    "                            database=ATLAS_DB_NAME)\n",
+    "        self.toolkit = MongoDBDatabaseToolkit(db=self.db_wrapper, llm=self.llm)\n",
+    "        self.agent = create_react_agent(\n",
+    "                        self.llm, \n",
+    "                        self.toolkit.get_tools(), \n",
+    "                        state_modifier=self.system_message)\n",
+    "        self.messages = []\n",
+    "\n",
+    "    def convert_to_mql_and_execute_query(self, query):\n",
+    "        # Start the agent with the agent.stream() method\n",
+    "        events = self.agent.stream(\n",
+    "            {'messages': [('user', query)]},\n",
+    "            stream_mode='values',\n",
+    "        )\n",
+    "        # Add output (events) from the agent to the self.messages list\n",
+    "        for event in events:\n",
+    "            self.messages.extend(event['messages'])\n",
+    "    \n",
+    "    def print_results(self):\n",
+    "        # Print the the end-user's expected output from \n",
+    "        # the final message produced by the agent.\n",
+    "        print(self.messages[-1].content)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c90825eb",
+   "metadata": {},
+   "source": [
+    "## Run a sample query\n",
+    "\n",
+    "And finally, instantiate the `NaturalLanguageToMQL` class and run a sample query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7284c63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def main():\n",
+    "    converter = NaturalLanguageToMQL()\n",
+    "    converter.convert_to_mql_and_execute_query(NATURAL_LANGUAGE_QUERY)\n",
+    "    converter.print_results()\n",
+    "\n",
+    "main()"
+   ]
+  }
+ ],
+"metadata": {
+   "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.12"
+   }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+ }