langchain/docs/docs/use_cases/graph/graph_gremlin_cosmosdb_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c94240f5",
   "metadata": {},
   "source": [
    "# Gremlin (with CosmosDB) QA chain\n",
    "\n",
    "This notebook shows how to use LLMs to provide a natural language interface to a graph database you can query with the Gremlin query language."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc0ee68",
   "metadata": {},
   "source": [
    "You will need to have a Azure CosmosDB Graph database instance. One option is to create a [free CosmosDB Graph database instance in Azure](https://learn.microsoft.com/en-us/azure/cosmos-db/free-tier). \n",
    "\n",
    "When you create your Cosmos DB account and Graph, use /type as partition key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62812aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "from langchain.chains.graph_qa import GremlinQAChain\n",
    "from langchain.schema import Document\n",
    "from langchain_community.graphs import GremlinGraph\n",
    "from langchain_community.graphs.graph_document import GraphDocument, Node, Relationship\n",
    "from langchain_openai import AzureChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0928915d",
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmosdb_name = \"mycosmosdb\"\n",
    "cosmosdb_db_id = \"graphtesting\"\n",
    "cosmosdb_db_graph_id = \"mygraph\"\n",
    "cosmosdb_access_Key = \"longstring==\"\n",
    "\n",
    "graph = GremlinGraph(\n",
    "    url=f\"=wss://{cosmosdb_name}.gremlin.cosmos.azure.com:443/\",\n",
    "    username=f\"/dbs/{cosmosdb_db_id}/colls/{cosmosdb_db_graph_id}\",\n",
    "    password=cosmosdb_access_Key,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "995ea9b9",
   "metadata": {},
   "source": [
    "## Seeding the database\n",
    "\n",
    "Assuming your database is empty, you can populate it using the GraphDocuments\n",
    "\n",
    "For Gremlin, always add property called 'label' for each Node.\n",
    "If no label is set, Node.type is used as a label.\n",
    "For cosmos using natural id's make sense, as they are visible in the graph explorer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fedd26b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_doc = Document(\n",
    "    page_content=\"Matrix is a movie where Keanu Reeves, Laurence Fishburne and Carrie-Anne Moss acted.\"\n",
    ")\n",
    "movie = Node(id=\"The Matrix\", properties={\"label\": \"movie\", \"title\": \"The Matrix\"})\n",
    "actor1 = Node(id=\"Keanu Reeves\", properties={\"label\": \"actor\", \"name\": \"Keanu Reeves\"})\n",
    "actor2 = Node(\n",
    "    id=\"Laurence Fishburne\", properties={\"label\": \"actor\", \"name\": \"Laurence Fishburne\"}\n",
    ")\n",
    "actor3 = Node(\n",
    "    id=\"Carrie-Anne Moss\", properties={\"label\": \"actor\", \"name\": \"Carrie-Anne Moss\"}\n",
    ")\n",
    "rel1 = Relationship(\n",
    "    id=5, type=\"ActedIn\", source=actor1, target=movie, properties={\"label\": \"ActedIn\"}\n",
    ")\n",
    "rel2 = Relationship(\n",
    "    id=6, type=\"ActedIn\", source=actor2, target=movie, properties={\"label\": \"ActedIn\"}\n",
    ")\n",
    "rel3 = Relationship(\n",
    "    id=7, type=\"ActedIn\", source=actor3, target=movie, properties={\"label\": \"ActedIn\"}\n",
    ")\n",
    "rel4 = Relationship(\n",
    "    id=8,\n",
    "    type=\"Starring\",\n",
    "    source=movie,\n",
    "    target=actor1,\n",
    "    properties={\"label\": \"Strarring\"},\n",
    ")\n",
    "rel5 = Relationship(\n",
    "    id=9,\n",
    "    type=\"Starring\",\n",
    "    source=movie,\n",
    "    target=actor2,\n",
    "    properties={\"label\": \"Strarring\"},\n",
    ")\n",
    "rel6 = Relationship(\n",
    "    id=10,\n",
    "    type=\"Straring\",\n",
    "    source=movie,\n",
    "    target=actor3,\n",
    "    properties={\"label\": \"Strarring\"},\n",
    ")\n",
    "graph_doc = GraphDocument(\n",
    "    nodes=[movie, actor1, actor2, actor3],\n",
    "    relationships=[rel1, rel2, rel3, rel4, rel5, rel6],\n",
    "    source=source_doc,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d18f77a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The underlying python-gremlin has a problem when running in notebook\n",
    "# The following line is a workaround to fix the problem\n",
    "nest_asyncio.apply()\n",
    "\n",
    "# Add the document to the CosmosDB graph.\n",
    "graph.add_graph_documents([graph_doc])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c1a8ea",
   "metadata": {},
   "source": [
    "## Refresh graph schema information\n",
    "If the schema of database changes (after updates), you can refresh the schema information.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3de44f",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.refresh_schema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe76ccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(graph.schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a3c677",
   "metadata": {},
   "source": [
    "## Querying the graph\n",
    "\n",
    "We can now use the gremlin QA chain to ask question of the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7476ce98",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = GremlinQAChain.from_llm(\n",
    "    AzureChatOpenAI(\n",
    "        temperature=0,\n",
    "        azure_deployment=\"gpt-4-turbo\",\n",
    "    ),\n",
    "    graph=graph,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8ee27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain.invoke(\"Who played in The Matrix?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c64027-cf42-493a-9c76-2d10ba753728",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain.run(\"How many people played in The Matrix?\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}