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

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 = \"google.com:cloud-spanner-demo\"  # @param {type:\"string\"}\n",
+    "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
     "\n",
     "# Set the project id\n",
     "!gcloud config set project {PROJECT_ID}\n",
@@ -494,6 +494,7 @@
    "cell_type": "code",
    "source": [
     "from google.cloud import spanner\n",
+    "from langchain_google_spanner import SpannerGraphQAChain\n",
     "from langchain_google_vertexai import ChatVertexAI\n",
     "from IPython.core.display import HTML\n",
     "\n",

src/langchain_google_spanner/__init__.py

@@ -14,6 +14,11 @@
 
 from langchain_google_spanner.chat_message_history import SpannerChatMessageHistory
 from langchain_google_spanner.graph_qa import SpannerGraphQAChain
+from langchain_google_spanner.graph_retriever import (
+    SpannerGraphGQLRetriever,
+    SpannerGraphNodeVectorRetriever,
+    SpannerGraphSemanticGQLRetriever,
+)
 from langchain_google_spanner.graph_store import SpannerGraphStore
 from langchain_google_spanner.vector_store import (
     DistanceStrategy,
@@ -38,4 +43,7 @@
     "SecondaryIndex",
     "QueryParameters",
     "DistanceStrategy",
+    "SpannerGraphGQLRetriever",
+    "SpannerGraphNodeVectorRetriever",
+    "SpannerGraphSemanticGQLRetriever",
 ]

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.