chore: update imports in tests + update types in schema

Author: gRedHeadphone

src/llama_index_spanner/__init__.py

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

src/llama_index_spanner/schema.py

@@ -38,7 +38,7 @@ def remove_empty_values(input_dict):
 
 def group_nodes(
     nodes: List[LabelledNode],
-) -> dict[str, List[LabelledNode]]:
+) -> Dict[str, List[LabelledNode]]:
   """Groups nodes by their respective types.
 
   Args:
@@ -47,7 +47,7 @@ def group_nodes(
   Returns:
     A dictionary mapping node types to lists of LabelledNodes.
   """
-  nodes_group: CaseInsensitiveDict[dict[str, LabelledNode]] = (
+  nodes_group: CaseInsensitiveDict[Dict[str, LabelledNode]] = (
       CaseInsensitiveDict()
   )
   for node in nodes:
@@ -63,7 +63,7 @@ def group_nodes(
 
 def group_edges(
     edges: List[Tuple[Relation, str, str, str, str]],
-) -> dict[str, List[Tuple[Relation, str, str]]]:
+) -> Dict[str, List[Tuple[Relation, str, str]]]:
   """Groups edges by their respective types.
 
   Args:
@@ -79,7 +79,7 @@ def group_edges(
       - The target node table name.
   """
   edges_group: CaseInsensitiveDict[
-      dict[Tuple[str, str, str], Tuple[Relation, str, str]]
+      Dict[Tuple[str, str, str], Tuple[Relation, str, str]]
   ] = CaseInsensitiveDict()
   for edge, source_label, source_table, target_label, target_table in edges:
     edge_name = (
@@ -159,9 +159,9 @@ class ElementSchema(object):
   key_columns: List[str]
   base_table_name: str
   labels: List[str]
-  properties: Dict[str, str]
+  properties: CaseInsensitiveDict[str]
   # types: A dictionary where keys are property names (strings) and values are Spanner type definitions
-  types: Dict[str, param_types.Type]
+  types: CaseInsensitiveDict[param_types.Type]
   source: NodeReference
   target: NodeReference
 
@@ -181,7 +181,7 @@ def is_dynamic_schema(self) -> bool:
   def make_node_schema(
       node_label: str,
       graph_name: str,
-      property_types: Dict[str, param_types.Type],
+      property_types: CaseInsensitiveDict[param_types.Type],
   ) -> ElementSchema:
     """Creates a node schema for a given node type and label."""
     node = ElementSchema()
@@ -199,7 +199,7 @@ def make_edge_schema(
       edge_label: str,
       graph_schema: SpannerGraphSchema,
       key_columns: List[str],
-      property_types: Dict[str, param_types.Type],
+      property_types: CaseInsensitiveDict[param_types.Type],
       source_node_table: str,
       target_node_table: str,
   ) -> ElementSchema:
@@ -741,20 +741,20 @@ def __init__(
       )
 
     self.graph_name: str = graph_name
-    self.node_tables: Dict[str, ElementSchema] = CaseInsensitiveDict({})
-    self.edge_tables: Dict[str, ElementSchema] = CaseInsensitiveDict({})
-    self.labels: Dict[str, Label] = CaseInsensitiveDict({})
-    self.properties: Dict[str, param_types.Type] = CaseInsensitiveDict({})
-    self.node_properties: Dict[str, param_types.Type] = CaseInsensitiveDict({})
+    self.node_tables: CaseInsensitiveDict[ElementSchema] = CaseInsensitiveDict({})
+    self.edge_tables: CaseInsensitiveDict[ElementSchema] = CaseInsensitiveDict({})
+    self.labels: CaseInsensitiveDict[Label] = CaseInsensitiveDict({})
+    self.properties: CaseInsensitiveDict[param_types.Type] = CaseInsensitiveDict({})
+    self.node_properties: CaseInsensitiveDict[param_types.Type] = CaseInsensitiveDict({})
     self.use_flexible_schema = use_flexible_schema
     self.static_node_properties = set(static_node_properties or [])
     self.static_edge_properties = set(static_edge_properties or [])
     self.graph_exists = False
 
   def evolve_from_nodes(
       self,
-      nodes: dict[str, List[LabelledNode]],
-  ) -> Tuple[List[str], dict[str, ElementSchema]]:
+      nodes: Dict[str, List[LabelledNode]],
+  ) -> Tuple[List[str], Dict[str, ElementSchema]]:
     """Evolves the graph schema based on new nodes and edges.
 
     This method updates the internal schema representation by adding new
@@ -789,8 +789,8 @@ def evolve_from_nodes(
 
   def evolve_from_edges(
       self,
-      edges: dict[str, List[Tuple[Relation, str, str]]],
-  ) -> Tuple[List[str], dict[str, ElementSchema], dict[str, ElementSchema]]:
+      edges: Dict[str, List[Tuple[Relation, str, str]]],
+  ) -> Tuple[List[str], Dict[str, ElementSchema]]:
     """Evolves the graph schema based on new edges.
 
     This method updates the internal schema representation by adding new
@@ -902,7 +902,7 @@ def to_ddl(self) -> str:
 
     def construct_label_and_properties(
         target_label: str,
-        labels: Dict[str, Label],
+        labels: CaseInsensitiveDict[Label],
         element: ElementSchema,
     ) -> str:
       props = labels[target_label].prop_names
@@ -917,7 +917,7 @@ def construct_label_and_properties(
 
     def construct_label_and_properties_list(
         target_labels: List[str],
-        labels: Dict[str, Label],
+        labels: CaseInsensitiveDict[Label],
         element: ElementSchema,
     ) -> str:
       return "\n".join((
@@ -942,7 +942,7 @@ def construct_node_reference(
       )
 
     def construct_element_table(
-        element: ElementSchema, labels: Dict[str, Label]
+        element: ElementSchema, labels: CaseInsensitiveDict[Label]
     ) -> str:
       definition = [
           "{} AS {}".format(

tests/__init__.py

@@ -23,7 +23,7 @@
 from llama_index.core.query_engine import RetrieverQueryEngine
 from llama_index.readers.wikipedia import WikipediaReader
 import pytest
-from src.llama_index_spanner.graph_retriever import (
+from llama_index_spanner.graph_retriever import (
     SpannerGraphCustomRetriever,
     SpannerGraphTextToGQLRetriever,
 )

tests/test_graph_retriever.py

@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
 from typing import Generator
 from llama_index.core.graph_stores.types import (
     ChunkNode,
     EntityNode,
-    LabelledNode,
-    PropertyGraphStore,
     Relation,
 )
 from llama_index.core.vector_stores.types import (
@@ -30,8 +27,8 @@
     VectorStoreQuery,
 )
 import pytest
-from src.llama_index_spanner.property_graph_store import SpannerPropertyGraphStore
-from src.llama_index_spanner.schema import ElementSchema
+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,
@@ -55,7 +52,7 @@
 
 
 @pytest.fixture
-def property_graph_store_static() -> Generator[SpannerPropertyGraphStore]:
+def property_graph_store_static() -> Generator[SpannerPropertyGraphStore, None, None]:
   """Provides a fresh SpannerPropertyGraphStore for each test."""
   graph_store = get_spanner_property_graph_store(
       graph_name_suffix=get_random_suffix()
@@ -65,7 +62,7 @@ def property_graph_store_static() -> Generator[SpannerPropertyGraphStore]:
 
 
 @pytest.fixture
-def property_graph_store_dynamic() -> Generator[SpannerPropertyGraphStore]:
+def property_graph_store_dynamic() -> Generator[SpannerPropertyGraphStore, None, None]:
   """Provides a fresh SpannerPropertyGraphStore for each test."""
   graph_store_dynamic_schema = get_spanner_property_graph_store(
       use_flexible_schema=True,

tests/test_spanner_property_graph_store.py

@@ -18,7 +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 src.llama_index_spanner import SpannerGraphStore, SpannerPropertyGraphStore
+from llama_index_spanner import SpannerPropertyGraphStore
 
 spanner_instance_id = (
     os.environ.get("SPANNER_INSTANCE_ID") or "graphdb-spanner-llama"
@@ -29,27 +29,13 @@
 spanner_graph_name = os.environ.get("SPANNER_GRAPH_NAME") or "llama_index_graph"
 
 
-def get_spanner_graph_store(
-    graph_name_suffix: str = "", clean_up: bool = False
-) -> SpannerGraphStore:
-  """Get a SpannerGraphStore instance for testing."""
-  graph_name = spanner_graph_name
-  if graph_name_suffix:
-    graph_name += "_" + graph_name_suffix
-  return SpannerGraphStore(
-      instance_id=spanner_instance_id,
-      database_id=spanner_database_id,
-      graph_name=graph_name,
-      clean_up=clean_up,
-  )
-
 
 def get_spanner_property_graph_store(
     graph_name_suffix: str = "",
     use_flexible_schema: bool = False,
     clean_up: bool = False,
 ) -> SpannerPropertyGraphStore:
-  """Get a SpannerGraphStore instance for testing."""
+  """Get a SpannerPropertyGraphStore instance for testing."""
   graph_name = spanner_graph_name
   if graph_name_suffix:
     graph_name += "_" + graph_name_suffix