feat(graph): Add Custom Retrievers for Spanner Graph RAG. (#122)

* Add Spanner Graph QA Chain

* Formatted notebook. Added copyright message to prompts file.

* Add missing imports for random graph name

* Make input table name randomized in integration tests to avoid name collision for tests running parallely from different python environments

* Provide timeout to graph cleanup

* Make default timeout of 300 secs for ddl application

* Increase timeout of integration test

* Change integration test timeout

* Minor formatting fixes

* Make the ddl operations test fixture scoped for the module

* Addressed review comments

* Addressed a few other review comments.

* Remove unused function

* fix type check errors

* Addressed review comments

* Addressed review comments

* Clear default project id from notebook

* Add import statement for SpanerGraphQAChain to notebook

* Add retrievers for Spanner Graph RAG

* Add licence headers

* Fix DATABASE name key

* Fix lint error on import ordering

* Fix lint errors

* Few minor changes to the SpannerGraphNodeVectorRetriever

* Fix lint error

* Add an option to expand context graph by hops

* Fix lint error

* Addressed review comments

* Remove expansion query options

* Add backticks to property names

* Change copyright year

* Address review comments

* Rename the retrievers. Merge the semantic retriever with the gql retriever.

* Fixed lint errors

* Change vertex ai versionto latest

* Fix lint errors

* Add documentation. Fixes the case where expands_by_hops is 0

* Add unit test for expand_by_hops=0

* Fix formatting for documentation

* Addressed review comments

---------

Co-authored-by: Amarnath Mullick <[email protected]>

README.rst

@@ -179,6 +179,53 @@ See the full `Spanner Graph QA Chain`_ tutorial.
 
 .. _`Spanner Graph QA Chain`: https://github.com/googleapis/langchain-google-spanner-python/blob/main/docs/graph_qa_chain.ipynb
 
+Spanner Graph Retrievers Usage
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use ``SpannerGraphTextToGQLRetriever`` to translate natural language question to GQL and query SpannerGraphStore.
+
+.. code:: python
+
+    from langchain_google_spanner import SpannerGraphStore, SpannerGraphTextToGQLRetriever
+    from langchain_google_vertexai import ChatVertexAI
+
+
+    graph = SpannerGraphStore(
+        instance_id="my-instance",
+        database_id="my-database",
+        graph_name="my_graph",
+    )
+    llm = ChatVertexAI()
+    retriever = SpannerGraphTextToGQLRetriever.from_params(
+        graph_store=graph,
+        llm=llm
+    )
+    retriever.invoke("Where does Elias Thorne's sibling live?")
+
+Use ``SpannerGraphVectorContextRetriever`` to perform vector search on embeddings that are stored in the nodes in a SpannerGraphStore. If expand_by_hops is provided, the nodes and edges at a distance upto the expand_by_hops from the nodes found in the vector search will also be returned.
+
+.. code:: python
+
+    from langchain_google_spanner import SpannerGraphStore, SpannerGraphVectorContextRetriever
+    from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings
+
+
+    graph = SpannerGraphStore(
+        instance_id="my-instance",
+        database_id="my-database",
+        graph_name="my_graph",
+    )
+    embedding_service = VertexAIEmbeddings(model_name="text-embedding-004")
+    retriever = SpannerGraphVectorContextRetriever.from_params(
+            graph_store=graph,
+            embedding_service=embedding_service,
+            label_expr="Person",
+            embeddings_column="embeddings",
+            top_k=1,
+            expand_by_hops=1,
+        )
+    retriever.invoke("Who lives in desert?")
+
 
 Contributions
 ~~~~~~~~~~~~~

docs/graph_qa_chain.ipynb

@@ -150,7 +150,7 @@
    "source": [
     "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n",
     "\n",
-    "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+    "PROJECT_ID = \"my-project-id\"  # @param {type:\"string\"}\n",
     "\n",
     "# Set the project id\n",
     "!gcloud config set project {PROJECT_ID}\n",

src/langchain_google_spanner/__init__.py

@@ -14,6 +14,10 @@
 
 from langchain_google_spanner.chat_message_history import SpannerChatMessageHistory
 from langchain_google_spanner.graph_qa import SpannerGraphQAChain
+from langchain_google_spanner.graph_retriever import (
+    SpannerGraphTextToGQLRetriever,
+    SpannerGraphVectorContextRetriever,
+)
 from langchain_google_spanner.graph_store import SpannerGraphStore
 from langchain_google_spanner.vector_store import (
     DistanceStrategy,
@@ -38,4 +42,6 @@
     "SecondaryIndex",
     "QueryParameters",
     "DistanceStrategy",
+    "SpannerGraphTextToGQLRetriever",
+    "SpannerGraphVectorContextRetriever",
 ]

src/langchain_google_spanner/graph_qa.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import re
 from typing import Any, Dict, List, Optional
 
 from langchain.chains.base import Chain
@@ -28,6 +27,7 @@
 
 from langchain_google_spanner.graph_store import SpannerGraphStore
 
+from .graph_utils import extract_gql, fix_gql_syntax
 from .prompts import (
     DEFAULT_GQL_FIX_TEMPLATE,
     DEFAULT_GQL_TEMPLATE,
@@ -71,50 +71,6 @@ class VerifyGqlOutput(BaseModel):
 INTERMEDIATE_STEPS_KEY = "intermediate_steps"
 
 
-def fix_gql_syntax(query: str) -> str:
-    """Fixes the syntax of a GQL query.
-    Example 1:
-        Input:
-            MATCH (p:paper {id: 0})-[c:cites*8]->(p2:paper)
-        Output:
-            MATCH (p:paper {id: 0})-[c:cites]->{8}(p2:paper)
-    Example 2:
-        Input:
-            MATCH (p:paper {id: 0})-[c:cites*1..8]->(p2:paper)
-        Output:
-            MATCH (p:paper {id: 0})-[c:cites]->{1:8}(p2:paper)
-
-    Args:
-        query: The input GQL query.
-
-    Returns:
-        Possibly modified GQL query.
-    """
-
-    query = re.sub(r"-\[(.*?):(\w+)\*(\d+)\.\.(\d+)\]->", r"-[\1:\2]->{\3,\4}", query)
-    query = re.sub(r"-\[(.*?):(\w+)\*(\d+)\]->", r"-[\1:\2]->{\3}", query)
-    query = re.sub(r"<-\[(.*?):(\w+)\*(\d+)\.\.(\d+)\]-", r"<-[\1:\2]-{\3,\4}", query)
-    query = re.sub(r"<-\[(.*?):(\w+)\*(\d+)\]-", r"<-[\1:\2]-{\3}", query)
-    query = re.sub(r"-\[(.*?):(\w+)\*(\d+)\.\.(\d+)\]-", r"-[\1:\2]-{\3,\4}", query)
-    query = re.sub(r"-\[(.*?):(\w+)\*(\d+)\]-", r"-[\1:\2]-{\3}", query)
-    return query
-
-
-def extract_gql(text: str) -> str:
-    """Extract GQL query from a text.
-
-    Args:
-        text: Text to extract GQL query from.
-
-    Returns:
-        GQL query extracted from the text.
-    """
-    pattern = r"```(.*?)```"
-    matches = re.findall(pattern, text, re.DOTALL)
-    query = matches[0] if matches else text
-    return fix_gql_syntax(query)
-
-
 class SpannerGraphQAChain(Chain):
     """Chain for question-answering against a Spanner Graph database by
         generating GQL statements from natural language questions.