Merge branch 'main' into issue_183

Author: averikitsch

.github/workflows/docs.yml

@@ -10,9 +10,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Checkout
-      uses: actions/checkout@a5ac7e51b41094c92402da3b24376905380afc29 # v4
+      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
     - name: Setup Python
-      uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5
+      uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
       with:
         python-version: "3.10"
     - name: Install nox
@@ -26,9 +26,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Checkout
-      uses: actions/checkout@a5ac7e51b41094c92402da3b24376905380afc29 # v4
+      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
     - name: Setup Python
-      uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5
+      uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
       with:
         python-version: "3.10"
     - name: Install nox

.github/workflows/lint.yml

@@ -34,7 +34,7 @@ jobs:
         uses: actions/checkout@f43a0e5ff2bd294095638e18286ca9a3d1956744 # v3
 
       - name: Setup Python
-        uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5.1.0
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: "3.11"
 

integration.cloudbuild.yaml

@@ -36,8 +36,8 @@ steps:
 timeout: "7200s"
 substitutions:
   _INSTANCE_ID: test-instance
-  _GOOGLE_DATABASE: test-google-db
-  _PG_DATABASE: test-pg-db
+  _GOOGLE_DATABASE: test-gsql-db
+  _PG_DATABASE: test-pgsql-db
   _VERSION: "3.9"
 
 options:

src/langchain_google_spanner/graph_qa.py

@@ -71,6 +71,12 @@ class VerifyGqlOutput(BaseModel):
 INTERMEDIATE_STEPS_KEY = "intermediate_steps"
 
 
+class InvalidGQLGenerationError(ValueError):
+    def __init__(self, message, intermediate_steps=None):
+        self.intermediate_steps = intermediate_steps
+        super().__init__(message)
+
+
 class SpannerGraphQAChain(Chain):
     """Chain for question-answering against a Spanner Graph database by
         generating GQL statements from natural language questions.
@@ -268,7 +274,9 @@ def execute_with_retry(
             finally:
                 retries += 1
 
-        raise ValueError("The generated gql query is invalid")
+        raise InvalidGQLGenerationError(
+            "The generated gql query is invalid", intermediate_steps
+        )
 
     def log_invalid_query(
         self,
@@ -309,6 +317,7 @@ def _call(
                 }
             )
             if "verified_gql" in verify_response:
+                intermediate_steps.append({"raw_generated_gql": generated_gql})
                 verified_gql = fix_gql_syntax(verify_response["verified_gql"])
                 intermediate_steps.append({"verified_gql": verified_gql})
             else:
@@ -322,7 +331,9 @@ def _call(
                 _run_manager, intermediate_steps, question, verified_gql
             )
             if not final_gql:
-                raise ValueError("No GQL was generated.")
+                raise InvalidGQLGenerationError(
+                    "No GQL was generated.", intermediate_steps
+                )
             _run_manager.on_text("Full Context:", end="\n", verbose=self.verbose)
             _run_manager.on_text(
                 str(context), color="green", end="\n", verbose=self.verbose

src/langchain_google_spanner/graph_store.py

@@ -958,8 +958,8 @@ def __repr__(self) -> str:
                 "Node properties per node label": {
                     label: [
                         {
-                            "property name": name,
-                            "property type": properties[name],
+                            "name": name,
+                            "type": properties[name],
                         }
                         for name in sorted(self.labels[label].prop_names)
                     ]
@@ -968,8 +968,8 @@ def __repr__(self) -> str:
                 "Edge properties per edge label": {
                     label: [
                         {
-                            "property name": name,
-                            "property type": properties[name],
+                            "name": name,
+                            "type": properties[name],
                         }
                         for name in sorted(self.labels[label].prop_names)
                     ]
@@ -1230,6 +1230,7 @@ def __init__(
         instance_id: str,
         database_id: str,
         client: Optional[spanner.Client] = None,
+        timeout: Optional[float] = None,
     ):
         """Initializes the Spanner implementation.
 
@@ -1241,11 +1242,14 @@ def __init__(
         self.client = client or spanner.Client()
         self.instance = self.client.instance(instance_id)
         self.database = self.instance.database(database_id)
+        self.timeout = timeout
 
     def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
         param_types = {k: TypeUtility.value_to_param_type(v) for k, v in params.items()}
         with self.database.snapshot() as snapshot:
-            rows = snapshot.execute_sql(query, params=params, param_types=param_types)
+            rows = snapshot.execute_sql(
+                query, params=params, param_types=param_types, timeout=self.timeout
+            )
             return [
                 {
                     column: value
@@ -1286,6 +1290,7 @@ def __init__(
         static_node_properties: List[str] = [],
         static_edge_properties: List[str] = [],
         impl: Optional[SpannerInterface] = None,
+        timeout: Optional[float] = None,
     ):
         """Initializes SpannerGraphStore.
 
@@ -1300,11 +1305,13 @@ def __init__(
           properties as static;
           static_edge_properties: in flexible schema, treat these edge
           properties as static.
+          timeout (Optional[float]): The timeout for queries in seconds.
         """
         self.impl = impl or SpannerImpl(
             instance_id,
             database_id,
             client_with_user_agent(client, USER_AGENT_GRAPH_STORE),
+            timeout=timeout,
         )
         self.schema = SpannerGraphSchema(
             graph_name,

tests/integration/test_spanner_loader.py

@@ -35,9 +35,41 @@ def client() -> Client:
     return Client(project=project_id)
 
 
+@pytest.fixture(scope="class")
+def cleanupGSQL(client):
+    yield
+
+    print("\nPerforming GSQL cleanup after each test...")
+
+    database = client.instance(instance_id).database(google_database)
+    operation = database.update_ddl(
+        [
+            f"DROP TABLE IF EXISTS {table_name}",
+        ]
+    )
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    # Code to perform teardown after each test goes here
+    print("\nGSQL Cleanup complete.")
+
+
+@pytest.fixture(scope="class")
+def cleanupPGSQL(client):
+    yield
+
+    print("\nPerforming PGSQL cleanup after each test...")
+
+    database = client.instance(instance_id).database(pg_database)
+    operation = database.update_ddl([f"DROP TABLE IF EXISTS {table_name}"])
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    # Code to perform teardown after each test goes here
+    print("\n PGSQL Cleanup complete.")
+
+
 class TestSpannerDocumentLoaderGoogleSQL:
     @pytest.fixture(autouse=True, scope="class")
-    def setup_database(self, client):
+    def setup_database(self, client, cleanupGSQL):
         database = client.instance(instance_id).database(google_database)
         operation = database.update_ddl([f"DROP TABLE IF EXISTS {table_name}"])
         operation.result(OPERATION_TIMEOUT_SECONDS)
@@ -455,7 +487,7 @@ def test_loader_custom_json_metadata(self, client):
 
 class TestSpannerDocumentLoaderPostgreSQL:
     @pytest.fixture(autouse=True, scope="class")
-    def setup_database(self, client):
+    def setup_database(self, client, cleanupPGSQL):
         database = client.instance(instance_id).database(pg_database)
         operation = database.update_ddl([f"DROP TABLE IF EXISTS {table_name}"])
         operation.result(OPERATION_TIMEOUT_SECONDS)
@@ -872,15 +904,15 @@ def test_loader_custom_json_metadata(self, client):
 
 class TestSpannerDocumentSaver:
     @pytest.fixture(name="google_client")
-    def setup_google_client(self, client) -> Client:
+    def setup_google_client(self, client, cleanupGSQL) -> Client:
         database = client.instance(instance_id).database(google_database)
         operation = database.update_ddl([f"DROP TABLE IF EXISTS {table_name}"])
         print("table dropped")
         operation.result(OPERATION_TIMEOUT_SECONDS)
         yield client
 
     @pytest.fixture(name="pg_client")
-    def setup_pg_client(self, client) -> Client:
+    def setup_pg_client(self, client, cleanupPGSQL) -> Client:
         database = client.instance(instance_id).database(pg_database)
         operation = database.update_ddl([f"DROP TABLE IF EXISTS {table_name}"])
         operation.result(OPERATION_TIMEOUT_SECONDS)

tests/integration/test_spanner_vector_store.py

@@ -20,7 +20,7 @@
 
 import pytest
 from google.cloud.spanner import Client  # type: ignore
-from langchain_community.document_loaders import HNLoader
+from langchain_community.document_loaders import RecursiveUrlLoader
 from langchain_community.embeddings import FakeEmbeddings
 
 from langchain_google_spanner.vector_store import (  # type: ignore
@@ -245,11 +245,13 @@ def setup_database(self, client):
             id_column="row_id",
             metadata_columns=[
                 TableColumn(name="metadata", type="JSON", is_null=True),
-                TableColumn(name="title", type="STRING(MAX)", is_null=False),
+                TableColumn(name="title", type="STRING(MAX)"),
             ],
         )
 
-        loader = HNLoader("https://news.ycombinator.com/item?id=34817881")
+        loader = RecursiveUrlLoader(
+            "https://news.ycombinator.com/item?id=1", max_depth=1
+        )
 
         embeddings = FakeEmbeddings(size=3)
 
@@ -327,7 +329,7 @@ def test_spanner_vector_delete_data(self, setup_database):
 
         docs = loader.load()
 
-        deleted = db.delete(documents=[docs[0], docs[1]])
+        deleted = db.delete(documents=docs)
 
         assert deleted
 
@@ -459,7 +461,9 @@ def setup_database(self, client):
                 ],
             )
 
-        loader = HNLoader("https://news.ycombinator.com/item?id=34817881")
+        loader = RecursiveUrlLoader(
+            "https://news.ycombinator.com/item?id=1", max_depth=1
+        )
         embeddings = FakeEmbeddings(size=title_vector_size)
 
         def cleanup_db():
@@ -552,7 +556,7 @@ def test_delete(self, setup_database):
         )
 
         docs = loader.load()
-        deleted = db.delete(documents=[docs[0], docs[1]])
+        deleted = db.delete(documents=docs)
 
         assert deleted
 
@@ -677,8 +681,9 @@ def setup_database(self, client):
             ],
         )
 
-        loader = HNLoader("https://news.ycombinator.com/item?id=34817881")
-
+        loader = RecursiveUrlLoader(
+            "https://news.ycombinator.com/item?id=1", max_depth=1
+        )
         embeddings = FakeEmbeddings(size=3)
 
         yield loader, embeddings
@@ -755,7 +760,7 @@ def test_spanner_vector_delete_data(self, setup_database):
 
         docs = loader.load()
 
-        deleted = db.delete(documents=[docs[0], docs[1]])
+        deleted = db.delete(documents=docs)
 
         assert deleted