fix: Index store adds interface methods (#118)

Author: vishwarajanand

src/llama_index_cloud_sql_pg/async_index_store.py

@@ -150,6 +150,20 @@ async def aadd_index_struct(self, index_struct: IndexStruct) -> None:
  query = insert_query + values_statement + upsert_statement
  await self.__aexecute_query(query, index_row)
 
+ async def async_index_structs(self) -> list[IndexStruct]:
+ """Get all index structs.
+ Returns:
+ list[IndexStruct]: index structs
+ """
+ return await self.aindex_structs()
+
+ async def async_add_index_struct(self, index_struct: IndexStruct) -> None:
+ """Add an index struct.
+ Args:
+ index_struct (IndexStruct): index struct
+ """
+ await self.aadd_index_struct(index_struct)
+
  async def adelete_index_struct(self, key: str) -> None:
  """Delete an index struct.
 

src/llama_index_cloud_sql_pg/index_store.py

@@ -96,6 +96,17 @@ def create_sync(
  index_store = engine._run_as_sync(coro)
  return cls(cls.__create_key, engine, index_store)
 
+ def add_index_struct(self, index_struct: IndexStruct) -> None:
+ """Add an index struct.
+
+ Args:
+ index_struct (IndexStruct): index struct
+
+ """
+ return self._engine._run_as_sync(
+ self.__index_store.aadd_index_struct(index_struct)
+ )
+
  async def aindex_structs(self) -> list[IndexStruct]:
  """Get all index structs.
 
@@ -125,16 +136,19 @@ async def aadd_index_struct(self, index_struct: IndexStruct) -> None:
  self.__index_store.aadd_index_struct(index_struct)
  )
 
- def add_index_struct(self, index_struct: IndexStruct) -> None:
- """Add an index struct.
+ async def async_index_structs(self) -> list[IndexStruct]:
+ """Get all index structs.
+ Returns:
+ list[IndexStruct]: index structs
+ """
+ return await self.aindex_structs()
 
+ async def async_add_index_struct(self, index_struct: IndexStruct) -> None:
+ """Add an index struct.
  Args:
  index_struct (IndexStruct): index struct
-
  """
- return self._engine._run_as_sync(
- self.__index_store.aadd_index_struct(index_struct)
- )
+ await self.aadd_index_struct(index_struct)
 
  async def adelete_index_struct(self, key: str) -> None:
  """Delete an index struct.

tests/test_async_index_store.py

@@ -147,7 +147,7 @@ async def test_aindex_structs(self, index_store):
  index_graph_struct = IndexGraph()
 
  await index_store.aadd_index_struct(index_dict_struct)
- await index_store.aadd_index_struct(index_graph_struct)
+ await index_store.async_add_index_struct(index_graph_struct)
  await index_store.aadd_index_struct(index_list_struct)
 
  indexes = await index_store.aindex_structs()

tests/test_index_store.py

@@ -149,14 +149,16 @@ async def test_aindex_structs(self, index_store):
  index_graph_struct = IndexGraph()
 
  await index_store.aadd_index_struct(index_dict_struct)
- await index_store.aadd_index_struct(index_graph_struct)
+ await index_store.async_add_index_struct(index_graph_struct)
  await index_store.aadd_index_struct(index_list_struct)
 
  indexes = await index_store.aindex_structs()
+ indexes_with_async = await index_store.async_index_structs()
 
- index_store.add_index_struct(index_dict_struct)
- index_store.add_index_struct(index_graph_struct)
- index_store.add_index_struct(index_list_struct)
+ assert indexes == indexes_with_async
+ assert index_dict_struct in indexes
+ assert index_list_struct in indexes
+ assert index_graph_struct in indexes
 
  async def test_warning(self, index_store):
  index_dict_struct = IndexDict()
@@ -270,9 +272,9 @@ async def test_aindex_structs(self, index_store):
 
  indexes = index_store.index_structs()
 
- index_store.add_index_struct(index_dict_struct)
- index_store.add_index_struct(index_graph_struct)
- index_store.add_index_struct(index_list_struct)
+ assert index_dict_struct in indexes
+ assert index_list_struct in indexes
+ assert index_graph_struct in indexes
 
  async def test_warning(self, index_store):
  index_dict_struct = IndexDict()