feat: lint workflow + isort fix

Author: gRedHeadphone

.github/workflows/lint.yml

@@ -0,0 +1,56 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Lint
+on:
+  pull_request:
+    paths-ignore: # Changes to the paths list need to be reflected in lint_fallback.yml
+      - "*.md"
+      - ".kokoro/**"
+      - ".github/**"
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.ref }}
+      cancel-in-progress: true
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@f43a0e5ff2bd294095638e18286ca9a3d1956744 # v3
+
+      - name: Setup Python
+        uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5.1.0
+        with:
+          python-version: "3.11"
+
+      - name: Install requirements
+        run: pip install -r requirements.txt
+
+      - name: Install module (and test requirements)
+        run: pip install -e .[test] 
+
+      - name: Install linters and type-check
+        run: pip install black black[jupyter] isort mypy
+
+      - name: Run linters
+        run: |
+          black --check .
+          isort --check .
+
+      # - name: Run type-check
+      #   run: mypy --install-types --non-interactive .

requirements.txt

@@ -0,0 +1,2 @@
+google-cloud-spanner==3.51.0
+llama-index-core==0.12.0

src/llama_index_spanner/__init__.py

@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .version import __version__
+from .graph_retriever import (SpannerGraphCustomRetriever,
+                              SpannerGraphTextToGQLRetriever)
 from .property_graph_store import SpannerPropertyGraphStore
-from .graph_retriever import (
-    SpannerGraphCustomRetriever,
-    SpannerGraphTextToGQLRetriever,
-)
+from .version import __version__
 
 __all__ = [
     __version__,

src/llama_index_spanner/graph_retriever.py

@@ -21,24 +21,17 @@
 from llama_index.core.output_parsers import PydanticOutputParser
 from llama_index.core.postprocessor.llm_rerank import LLMRerank
 from llama_index.core.prompts import PromptType
-from llama_index.core.retrievers import (
-    CustomPGRetriever,
-    VectorContextRetriever,
-)
+from llama_index.core.retrievers import (CustomPGRetriever,
+                                         VectorContextRetriever)
 from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
 from llama_index.core.vector_stores.types import VectorStore
 from pydantic import BaseModel
 
 from .graph_utils import extract_gql, fix_gql_syntax
-from .prompts import (
-    DEFAULT_GQL_FIX_TEMPLATE,
-    DEFAULT_GQL_VERIFY_TEMPLATE,
-    DEFAULT_SPANNER_GQL_TEMPLATE,
-    DEFAULT_SUMMARY_TEMPLATE,
-)
+from .prompts import (DEFAULT_GQL_FIX_TEMPLATE, DEFAULT_GQL_VERIFY_TEMPLATE,
+                      DEFAULT_SPANNER_GQL_TEMPLATE, DEFAULT_SUMMARY_TEMPLATE)
 from .property_graph_store import SpannerPropertyGraphStore
 
-
 GQL_GENERATION_PROMPT = PromptTemplate(
     template=DEFAULT_SPANNER_GQL_TEMPLATE,
     prompt_type=PromptType.TEXT_TO_GRAPH_QUERY,

src/llama_index_spanner/property_graph_store.py

@@ -17,25 +17,16 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from google.cloud import spanner
-from llama_index.core.graph_stores.types import (
-    ChunkNode,
-    EntityNode,
-    LabelledNode,
-    PropertyGraphStore,
-    Relation,
-    Triplet,
-)
+from llama_index.core.graph_stores.types import (ChunkNode, EntityNode,
+                                                 LabelledNode,
+                                                 PropertyGraphStore, Relation,
+                                                 Triplet)
 from llama_index.core.prompts import PromptTemplate, PromptType
 from llama_index.core.vector_stores.types import VectorStoreQuery
 
 from .prompts import DEFAULT_SPANNER_GQL_TEMPLATE
-from .schema import (
-    ElementSchema,
-    GraphDocumentUtility,
-    SpannerGraphSchema,
-    group_edges,
-    group_nodes,
-)
+from .schema import (ElementSchema, GraphDocumentUtility, SpannerGraphSchema,
+                     group_edges, group_nodes)
 from .spanner import SpannerImpl, SpannerInterface, client_with_user_agent
 from .version import __version__
 

src/llama_index_spanner/schema.py

@@ -14,16 +14,13 @@
 
 from __future__ import annotations
 
-import json, re
+import json
+import re
 from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple
 
 from google.cloud.spanner_v1 import JsonObject, param_types
-from llama_index.core.graph_stores.types import (
-    ChunkNode,
-    EntityNode,
-    LabelledNode,
-    Relation,
-)
+from llama_index.core.graph_stores.types import (ChunkNode, EntityNode,
+                                                 LabelledNode, Relation)
 from requests.structures import CaseInsensitiveDict
 
 from .type_utils import TypeUtility

src/llama_index_spanner/spanner.py

@@ -14,7 +14,9 @@
 
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple
+
 from google.cloud import spanner
+
 from .type_utils import TypeUtility
 
 MUTATION_BATCH_SIZE = 1000

tests/test_graph_retriever.py

@@ -12,21 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from llama_index.core import PropertyGraphIndex
-from llama_index.core import Settings
-from llama_index.core.graph_stores.types import (
-    ChunkNode,
-    EntityNode,
-    Relation,
-)
+import pytest
+from llama_index.core import PropertyGraphIndex, Settings
+from llama_index.core.graph_stores.types import ChunkNode, EntityNode, Relation
 from llama_index.core.indices.property_graph import SchemaLLMPathExtractor
 from llama_index.core.query_engine import RetrieverQueryEngine
 from llama_index.readers.wikipedia import WikipediaReader
-import pytest
+
 from llama_index_spanner.graph_retriever import (
-    SpannerGraphCustomRetriever,
-    SpannerGraphTextToGQLRetriever,
-)
+    SpannerGraphCustomRetriever, SpannerGraphTextToGQLRetriever)
 from tests.utils import get_resources
 
 

tests/test_spanner_property_graph_store.py

@@ -13,30 +13,21 @@
 # limitations under the License.
 
 from typing import Generator
-from llama_index.core.graph_stores.types import (
-    ChunkNode,
-    EntityNode,
-    Relation,
-)
-from llama_index.core.vector_stores.types import (
-    ExactMatchFilter,
-    FilterCondition,
-    FilterOperator,
-    MetadataFilter,
-    MetadataFilters,
-    VectorStoreQuery,
-)
+
 import pytest
+from llama_index.core.graph_stores.types import ChunkNode, EntityNode, Relation
+from llama_index.core.vector_stores.types import (ExactMatchFilter,
+                                                  FilterCondition,
+                                                  FilterOperator,
+                                                  MetadataFilter,
+                                                  MetadataFilters,
+                                                  VectorStoreQuery)
+
 from llama_index_spanner.property_graph_store import SpannerPropertyGraphStore
 from llama_index_spanner.schema import ElementSchema
-from tests.utils import (
-    get_random_suffix,
-    get_spanner_property_graph_store,
-    spanner_database_id,
-    spanner_graph_name,
-    spanner_instance_id,
-)
-
+from tests.utils import (get_random_suffix, get_spanner_property_graph_store,
+                         spanner_database_id, spanner_graph_name,
+                         spanner_instance_id)
 
 pytestmark = pytest.mark.skipif(
     (not spanner_instance_id or not spanner_database_id or not spanner_graph_name),

tests/utils.py

@@ -18,6 +18,7 @@
 from llama_index.core.storage import StorageContext
 from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
 from llama_index.llms.google_genai import GoogleGenAI
+
 from llama_index_spanner import SpannerPropertyGraphStore
 
 spanner_instance_id = os.environ.get("SPANNER_INSTANCE_ID") or "graphdb-spanner-llama"