chore: mypy fixes (#6)

* chore: mypy initial fixes

* chore: poetry config change + spanner import type ignore

* chore: isort fix

Author: gRedHeadphone

.gitignore

@@ -30,6 +30,8 @@ pip-log.txt
 .cache
 .pytest_cache
 
+# Mypy cache
+.mypy_cache
 
 # Mac
 .DS_Store

poetry.lock

@@ -30,4 +30,9 @@ build-backend = "poetry.core.masonry.api"
 target-version = ['py39']
 
 [tool.isort]
-profile = "black"
+profile = "black"
+
+[tool.mypy]
+python_version = 3.9
+warn_unused_configs = true
+ignore_missing_imports = true

pyproject.toml

@@ -17,8 +17,8 @@
 from .version import __version__
 
 __all__ = [
-    __version__,
-    SpannerPropertyGraphStore,
-    SpannerGraphTextToGQLRetriever,
-    SpannerGraphCustomRetriever,
+    "__version__",
+    "SpannerPropertyGraphStore",
+    "SpannerGraphTextToGQLRetriever",
+    "SpannerGraphCustomRetriever",
 ]

src/llama_index_spanner/__init__.py

@@ -14,9 +14,9 @@
 
 import itertools
 import json
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
-from google.cloud import spanner
+from google.cloud import spanner  # type: ignore
 from llama_index.core.graph_stores.types import (
     ChunkNode,
     EntityNode,
@@ -26,7 +26,7 @@
     Triplet,
 )
 from llama_index.core.prompts import PromptTemplate, PromptType
-from llama_index.core.vector_stores.types import VectorStoreQuery
+from llama_index.core.vector_stores.types import MetadataFilter, VectorStoreQuery
 
 from .prompts import DEFAULT_SPANNER_GQL_TEMPLATE
 from .schema import (
@@ -57,7 +57,7 @@ def node_from_json(label: str, json_node_properties: Dict[str, Any]) -> Labelled
     Returns:
       A LabelledNode.
     """
-    id_, name, text, properties, embedding = None, None, None, {}, None
+    id_, name, text, properties, embedding = "", "", "", {}, None
     for k, v in json_node_properties.items():
         if k == ElementSchema.NODE_KEY_COLUMN_NAME:
             id_ = v
@@ -93,7 +93,7 @@ def edge_from_json(json_edge_properties: Dict[str, Any]) -> Relation:
     Returns:
       A Relation.
     """
-    source_id, target_id, properties, label = None, None, {}, None
+    source_id, target_id, properties, label = "", "", {}, ""
     for k, v in json_edge_properties.items():
         if k == ElementSchema.NODE_KEY_COLUMN_NAME:
             source_id = v
@@ -116,7 +116,7 @@ def edge_from_json(json_edge_properties: Dict[str, Any]) -> Relation:
 def update_condition(
     cond: List[str],
     params: Dict[str, Any],
-    schema: SpannerGraphSchema = None,
+    schema: SpannerGraphSchema,
     ids: Optional[List[str]] = None,
     properties: Optional[Dict[str, Any]] = None,
     entity_names: Optional[List[str]] = None,
@@ -251,7 +251,7 @@ def __init__(
     def client(self):
         return self.impl
 
-    def upsert_nodes(self, nodes: List[LabelledNode]) -> None:
+    def upsert_nodes(self, nodes: Sequence[LabelledNode]) -> None:
         """Upserts nodes into the graph store.
 
         This method takes a list of LabelledNodes and upserts them into the
@@ -372,7 +372,7 @@ def get(
             else "labels(n)[0]"
         )
         cond = ["1 = 1"]
-        params = {}
+        params: Dict[str, Any] = {}
 
         if not update_condition(
             cond, params, self.schema, ids=ids, properties=properties
@@ -416,7 +416,7 @@ def get_triplets(
         label_field = ElementSchema.DYNAMIC_LABEL_COLUMN_NAME
 
         cond = ["1 = 1"]
-        params = {}
+        params: Dict[str, Any] = {}
 
         if not update_condition(
             cond,
@@ -550,8 +550,8 @@ def delete(
         node_key_field = ElementSchema.NODE_KEY_COLUMN_NAME
         target_node_key_field = ElementSchema.TARGET_NODE_KEY_COLUMN_NAME
 
-        cond = []
-        params = {}
+        cond: List[str] = []
+        params: Dict[str, Any] = {}
 
         if (
             update_condition(
@@ -580,7 +580,7 @@ def delete(
                 ):
                     self.impl.delete(
                         self.schema.labels[node_table_label].base_table_name,
-                        [node_data[1:] for node_data in nodes],
+                        [[node_id] for _, node_id in nodes],
                     )
 
         if relation_names and self.schema.edge_tables:
@@ -610,7 +610,10 @@ def delete(
                     ):
                         self.impl.delete(
                             self.schema.labels[edge_label].base_table_name,
-                            [edge_data[1:] for edge_data in edges],
+                            [
+                                [edge_id, edge_target_id, edge_label]
+                                for _, edge_id, edge_target_id, edge_label in edges
+                            ],
                         )
             else:
                 data = self.structured_query(
@@ -669,7 +672,7 @@ def vector_query(
           - A list of LabelledNodes, representing the nodes that were found.
           - A list of floats, representing the similarity scores of the nodes.
         """
-        if not self.schema.graph_exists:
+        if not self.schema.graph_exists or query.query_embedding is None:
             return ([], [])
 
         query_condition = "1 = 1"
@@ -680,6 +683,11 @@ def vector_query(
         if query.filters:
             cond = []
             for i, query_filter in enumerate(query.filters.filters):
+                if not isinstance(
+                    query_filter, MetadataFilter
+                ):  # doesn't support nested MetadataFilters
+                    continue
+
                 if (
                     query_filter.key not in self.schema.node_properties
                     and (property_prefix + query_filter.key)

src/llama_index_spanner/property_graph_store.py

@@ -16,7 +16,7 @@
 
 import json
 import re
-from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple
+from typing import Any, Dict, Generator, Iterable, List, Optional, Sequence, Tuple
 
 from google.cloud.spanner_v1 import JsonObject, param_types
 from llama_index.core.graph_stores.types import (
@@ -38,7 +38,7 @@ def remove_empty_values(input_dict):
 
 
 def group_nodes(
-    nodes: List[LabelledNode],
+    nodes: Sequence[LabelledNode],
 ) -> Dict[str, List[LabelledNode]]:
     """Groups nodes by their respective types.
 

src/llama_index_spanner/schema.py

@@ -15,7 +15,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple
 
-from google.cloud import spanner
+from google.cloud import spanner  # type: ignore
 
 from .type_utils import TypeUtility
 
@@ -65,7 +65,7 @@ def apply_ddls(self, ddls: List[str], options: Dict[str, Any] = {}) -> None:
 
     @abstractmethod
     def insert_or_update(
-        self, table: str, columns: Tuple[str], values: List[List[Any]]
+        self, table: str, columns: Tuple[str, ...], values: List[List[Any]]
     ) -> None:
         """Insert or update the table.
 
@@ -130,7 +130,7 @@ def apply_ddls(self, ddls: List[str], options: Dict[str, Any] = {}) -> None:
         return op.result(options.get("timeout", DEFAULT_DDL_TIMEOUT))
 
     def insert_or_update(
-        self, table: str, columns: Tuple[str], values: List[List[Any]]
+        self, table: str, columns: Tuple[str, ...], values: List[List[Any]]
     ) -> None:
         for i in range(0, len(values), MUTATION_BATCH_SIZE):
             value_batch = values[i : i + MUTATION_BATCH_SIZE]

src/llama_index_spanner/spanner.py

@@ -55,7 +55,7 @@ def get_resources(
     graph_store = get_spanner_property_graph_store(
         graph_name_suffix, use_flexible_schema, clean_up
     )
-    storage_context = StorageContext.from_defaults(graph_store=graph_store)
+    storage_context = StorageContext.from_defaults(property_graph_store=graph_store)
     llm = GoogleGenAI(
         model="gemini-2.0-flash",
     )