neo4j-langchain-vertex-ai-reasoning-engine.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "name": "vertexAIreasoning.ipynb"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# GraphRAG with Neo4j and LangChain and Gemini on VertexAI Reasoning Engine\n",
        "\n",
        "This is a demonstration of a GeNAI API with advanced RAG patterns combining vector and graph search.\n",
        "\n",
        "It is deployed on Vertex AI Reasoning Engine (Preview) as scalable infrastructure and can then be integrated with GenAI applications on Cloud Run via REST APIs.\n",
        "\n",
        "## Dataset\n",
        "\n",
        "The dataset is a graph about companies, associated industries, and people and articles that report on those companies.\n",
        "\n",
        "![Graph Model](https://i.imgur.com/lWJZSEe.png)\n",
        "\n",
        "The articles are chunked and the chunks also stored in the graph.\n",
        "\n",
        "Embeddings are computed for each of the text chunks with `textembedding-gecko` (786 dim) and stored on each chunk node.\n",
        "A Neo4j vector index `news_google` and a fulltext index `news_fulltext` (for hybrid search) were created.\n",
        "\n",
        "The database is publicly available with a readonly user:\n",
        "\n",
        "https://demo.neo4jlabs.com:7473/browser/\n",
        "\n",
        "* URI: neo4j+s://demo.neo4jlabs.com\n",
        "* User: companies\n",
        "* Password: companies\n",
        "* Companies: companies\n",
        "\n",
        "We utilize the Neo4jVector LangChain integration, which allows for advanced RAG patterns.\n",
        "We will utilize both hybrid search as well as parent-child retrievers and GraphRAG (extract relevant context).\n",
        "\n",
        "In our configuration we provide both the vector and fulltext index as well as a retrieval query that fetches the following additional information for each chunk\n",
        "\n",
        "* Parent `Article` of the `Chunk` (aggregate all chunks for a single article)\n",
        "* `Organization`(s) mentioned\n",
        "* `IndustryCategory`(ies) for the organization\n",
        "* `Person`(s) connected to the organization and their roles (e.g. investor, chairman, ceo)\n",
        "\n",
        "We will retrieve the top-k = 5 results from the vector index.\n",
        "\n",
        "As LLM we will utilize Vertex AI *Gemini Pro 1.0*\n",
        "\n",
        "We use a temperature of 0.1, top-k=40, top-p=0.8\n",
        "\n",
        "Our `LangchainCode` class contains the methods for initialization which can only hold serializable information (strings and numbers).\n",
        "\n",
        "In `set_up()` Gemini as LLM, VertexAI Embeddings and the `Neo4jVector` retriever are combined into a LangChain chain.\n",
        "\n",
        "Which is then used in `query`  with `chain.invoke()`.\n",
        "\n",
        "The class is deployed as ReasoningEngine with the Google Vertex AI Python SDK.\n",
        "For the deployment you provide the instance of the class which captures relevant environment variables and configuration and the dependencies, in our case `google-cloud-vertexai, langchain, langchain_google_vertexai, neo4j`.\n",
        "\n",
        "And after successful deploymnet we can use the resulting object via the `query` method, passing in our user question."
      ],
      "metadata": {
        "id": "6u4maG9UQ685"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "PROJECT_ID = \"vertex-ai-neo4j-extension\"\n",
        "REGION = \"us-central1\"\n",
        "STAGING_BUCKET = \"gs://neo4j-vertex-ai-extension\"\n"
      ],
      "metadata": {
        "id": "3mIXWnf_7Vtz",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716387442760,
          "user_tz": -120,
          "elapsed": 118,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "!gcloud config set project vertex-ai-neo4j-extension"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DW7SGmhwRyoJ",
        "outputId": "5def019d-fa0d-4790-bcb4-de3df28908f5",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716387446921,
          "user_tz": -120,
          "elapsed": 2825,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Updated property [core/project].\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --quiet neo4j==5.19.0\n",
        "!pip install --quiet langchain_google_vertexai==1.0.4\n",
        "!pip install --quiet --force-reinstall langchain==0.2.0 langchain_community==0.2.0\n"
      ],
      "metadata": {
        "id": "wxp09CthdF89"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --quiet google-cloud-aiplatform==1.51.0\n",
        "!pip install --quiet  google-cloud-resource-manager==1.12.3"
      ],
      "metadata": {
        "id": "9rpCl9j4Hpmh",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716386739685,
          "user_tz": -120,
          "elapsed": 19630,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import vertexai\n",
        "from vertexai.preview import reasoning_engines\n",
        "\n",
        "vertexai.init(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")"
      ],
      "metadata": {
        "id": "guEu_Yh6Rv8n",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716393654282,
          "user_tz": -120,
          "elapsed": 226,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate\n",
        "from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings\n",
        "from langchain_community.vectorstores import Neo4jVector\n",
        "from langchain_core.output_parsers import StrOutputParser\n",
        "from langchain_core.runnables import RunnableParallel, RunnablePassthrough\n"
      ],
      "metadata": {
        "id": "mgNDGcXF-7zV",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716393656437,
          "user_tz": -120,
          "elapsed": 195,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you installed any packages above, you can restart the runtime to pick them up:\n",
        "\n",
        "Click on the \"Runtime\" button at the top of Colab.\n",
        "Select \"Restart session\".\n",
        "You would not need to re-run the cell above this one (to reinstall the packages)."
      ],
      "metadata": {
        "id": "HEljoz48bw6E"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "URI = os.getenv('NEO4J_URI', 'neo4j+s://demo.neo4jlabs.com')\n",
        "USER = os.getenv('NEO4J_USERNAME','companies')\n",
        "PASSWORD = os.getenv('NEO4J_PASSWORD','companies')\n",
        "DATABASE = os.getenv('NEO4J_DATABASE','companies')\n",
        "\n",
        "class LangchainCode:\n",
        "    def __init__(self):\n",
        "        self.model_name = \"gemini-1.5-pro-preview-0409\" #\"gemini-pro\"\n",
        "        self.max_output_tokens = 1024\n",
        "        self.temperature = 0.1\n",
        "        self.top_p = 0.8\n",
        "        self.top_k = 40\n",
        "        self.project_id = PROJECT_ID\n",
        "        self.location = REGION\n",
        "        self.uri = URI\n",
        "        self.username = USER\n",
        "        self.password = PASSWORD\n",
        "        self.database = DATABASE\n",
        "        self.prompt_input_variables = [\"query\"]\n",
        "        self.prompt_template=\"\"\"\n",
        "            You are a venture capital assistant that provides useful answers about companies, their boards, financing etc.\n",
        "            only using the information from a company database already provided in the context.\n",
        "            Prefer higher rated information in your context and add source links in your answers.\n",
        "            Context: {context}\"\"\"\n",
        "\n",
        "    def configure_qa_rag_chain(self, llm, embeddings):\n",
        "        qa_prompt = ChatPromptTemplate.from_messages([\n",
        "            SystemMessagePromptTemplate.from_template(self.prompt_template),\n",
        "            HumanMessagePromptTemplate.from_template(\"Question: {question}\"\n",
        "                                                      \"\\nWhat else can you tell me about it?\"),\n",
        "        ])\n",
        "\n",
        "        # Vector + Knowledge Graph response\n",
        "        kg = Neo4jVector.from_existing_index(\n",
        "            embedding=embeddings,\n",
        "            url=self.uri, username=self.username, password=self.password,database=self.database,\n",
        "            search_type=\"hybrid\",\n",
        "            keyword_index_name=\"news_fulltext\",\n",
        "            index_name=\"news_google\",\n",
        "            retrieval_query=\"\"\"\n",
        "              WITH node as c,score\n",
        "              MATCH (c)<-[:HAS_CHUNK]-(article:Article)\n",
        "\n",
        "              WITH article, collect(distinct c.text) as texts, avg(score) as score\n",
        "              RETURN article {.title, .sentiment, .siteName, .summary,\n",
        "                    organizations: [ (article)-[:MENTIONS]->(org:Organization) |\n",
        "                          org { .name, .revenue, .nbrEmployees, .isPublic, .motto, .summary,\n",
        "                          orgCategories: [ (org)-[:HAS_CATEGORY]->(i) | i.name],\n",
        "                          people: [ (org)-[rel]->(p:Person) | p { .name, .summary, role: replace(type(rel),\"HAS_\",\"\") }]}],\n",
        "                    texts: texts} as text,\n",
        "              score, {source: article.siteName} as metadata\n",
        "            \"\"\",\n",
        "        )\n",
        "        retriever = kg.as_retriever(search_kwargs={\"k\": 5})\n",
        "\n",
        "        def format_docs(docs):\n",
        "          return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
        "\n",
        "        chain = (\n",
        "            {\"context\": retriever | format_docs , \"question\": RunnablePassthrough()}\n",
        "            | qa_prompt\n",
        "            | llm\n",
        "            | StrOutputParser()\n",
        "        )\n",
        "        return chain\n",
        "\n",
        "    def set_up(self):\n",
        "        from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate\n",
        "        from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI\n",
        "        from langchain_community.vectorstores import Neo4jVector\n",
        "        from langchain_core.output_parsers import StrOutputParser\n",
        "        from langchain_core.runnables import RunnableParallel, RunnablePassthrough\n",
        "\n",
        "        llm = ChatVertexAI(\n",
        "            model_name=self.model_name,\n",
        "            max_output_tokens=self.max_output_tokens,\n",
        "            max_input_tokens=32000,\n",
        "            temperature=self.temperature,\n",
        "            top_p=self.top_p,\n",
        "            top_k=self.top_k,\n",
        "            project = self.project_id,\n",
        "            location = self.location,\n",
        "            # convert_system_message_to_human=True,\n",
        "            response_validation=False,\n",
        "            verbose=True\n",
        "        )\n",
        "        embeddings = VertexAIEmbeddings(\"textembedding-gecko@001\")\n",
        "\n",
        "        self.qa_chain = self.configure_qa_rag_chain(llm, embeddings)\n",
        "\n",
        "    def query(self, query):\n",
        "        return self.qa_chain.invoke(query)"
      ],
      "metadata": {
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716393704640,
          "user_tz": -120,
          "elapsed": 114,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "id": "Xbq3bR_Knf8k"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.globals import set_debug\n",
        "set_debug(False)\n"
      ],
      "metadata": {
        "id": "PfUcIghUElzt",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716303465137,
          "user_tz": -120,
          "elapsed": 149,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# testing locally\n",
        "lc = LangchainCode()\n",
        "lc.set_up()\n",
        "\n",
        "response = lc.query('What are the news about IBM and its acquisitions and who are the people involved?')\n",
        "print(response)"
      ],
      "metadata": {
        "id": "W4h6EQ89HlKR",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716396587908,
          "user_tz": -120,
          "elapsed": 12077,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "14b0c809-5d65-4d32-ca5b-a49a186b0322"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "IBM acquired several companies, including Ascential Software, Cognos, SPSS, Netezza, OpenPages, Algorithmics, CrossIdeas, Trusteer, Guardium, BigFix, Watchfire, Internet Security Systems, Q1 Labs, Turbonomic Inc., myInvenio and Bluetab Solutions Group. \n",
            "\n",
            "IBM's CEO at the time of these acquisitions was Arvind Krishna. Other key people involved in these acquisitions include Jim Whitehurst, former CEO of Red Hat, and Jim Kavanaugh, CFO of IBM. [https://newslanes.com/ibm-cloud-demand-lifts-sales-to-best-quarter-in-three-years-yahoo-finance/](https://newslanes.com/ibm-cloud-demand-lifts-sales-to-best-quarter-in-three-years-yahoo-finance/)\n",
            "\n",
            "IBM has been a multinational company for more than 90 years and has a long history of acquiring companies. IBM typically acquires companies that are experts in specific areas that are core to IBM's business. For example, IBM acquired Cognos and SPSS to augment its value proposition for customers in the area of data analytics. IBM has also divested parts of its company that it no longer considers to be core business areas. For example, IBM divested its PC business to Lenovo in 2005 and its printing business to Ricoh. [https://china.dailynews.com.cn/business/2012-11/01/content_15970464.htm](https://china.dailynews.com.cn/business/2012-11/01/content_15970464.htm)\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "remote_app = reasoning_engines.ReasoningEngine.create(\n",
        "    LangchainCode(),\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform==1.57.0\",\n",
        "        \"langchain_google_vertexai==1.0.6\",\n",
        "        \"langchain==0.2.6\",\n",
        "        \"langchain_community==0.2.6\",\n",
        "        \"neo4j==5.20.0\"\n",
        "    ],\n",
        "    display_name=\"Neo4j Vertex AI RE Companies\",\n",
        "    description=\"Neo4j Vertex AI RE Companies\",\n",
        "    sys_version=\"3.10\",\n",
        "    extra_packages=[]\n",
        ")"
      ],
      "metadata": {
        "id": "x6b4zSDDsUuw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "87f17308-3ea3-48a8-e6a4-4d49df785096",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716393943730,
          "user_tz": -120,
          "elapsed": 223162,
          "user": {
            "displayName": "",
            "userId": ""
          }
        }
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:vertexai.reasoning_engines._reasoning_engines:Using bucket neo4j-vertex-ai-extension\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://neo4j-vertex-ai-extension/reasoning_engine/reasoning_engine.pkl\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://neo4j-vertex-ai-extension/reasoning_engine/requirements.txt\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Creating in-memory tarfile of extra_packages\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://neo4j-vertex-ai-extension/reasoning_engine/dependencies.tar.gz\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Creating ReasoningEngine\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:Create ReasoningEngine backing LRO: projects/990868019953/locations/us-central1/reasoningEngines/7289005628154970112/operations/3531746006863446016\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:ReasoningEngine created. Resource name: projects/990868019953/locations/us-central1/reasoningEngines/7289005628154970112\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:To use this ReasoningEngine in another session:\n",
            "INFO:vertexai.reasoning_engines._reasoning_engines:reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/990868019953/locations/us-central1/reasoningEngines/7289005628154970112')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "response = remote_app.query(query=\"Who is on the board of Siemens?\")\n",
        "print(response)"
      ],
      "metadata": {
        "id": "hF95dtx-sXpT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1716393964595,
          "user_tz": -120,
          "elapsed": 20866,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "e8370694-6fd9-461e-8aa3-2af1a8003ca1"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The board of Siemens is composed of:\n",
            "\n",
            "* Jim Hagemann Snabe (Manager and chairman) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Dominika Bettman (CEO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Alejandro Preinfalk (CEO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Miguel Angel Lopez Borrego (CEO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Hanna Hennig (CIO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Barbara Humpton (CEO at Siemens Bank) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Matthias Rebellius (CEO at Siemens Corporate Technology) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Cedrik Neike (CEO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Ralf P. Thomas (CFO at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Michael Sigmund (Chairman at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Horst J. Kayser (Chairman at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Michael Diekmann (Business person) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Norbert Reithofer (German businessman) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Werner Brandt (German manager) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "* Birgit Steinborn (Chairman at Siemens) [source](https://automation.com/story/184274/siemens-brings-real-time-supply-chain-intelligence-to-siemens-xcelerator-and-the-digital-twin)\n",
            "\n",
            "Siemens is a German multinational conglomerate. \n",
            "\n"
          ]
        }
      ]
    }
  ]
}