From d588b97d3855a1c207cc2ffdf65aa4ba05e57c76 Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Fri, 16 May 2025 11:57:05 +0000 Subject: [PATCH] fix: Index store adds interface methods --- .../async_index_store.py | 14 ++++++++++ src/llama_index_cloud_sql_pg/index_store.py | 26 ++++++++++++++----- tests/test_async_index_store.py | 2 +- tests/test_index_store.py | 16 +++++++----- 4 files changed, 44 insertions(+), 14 deletions(-) diff --git a/src/llama_index_cloud_sql_pg/async_index_store.py b/src/llama_index_cloud_sql_pg/async_index_store.py index 76a5917..9a5ca09 100644 --- a/src/llama_index_cloud_sql_pg/async_index_store.py +++ b/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. diff --git a/src/llama_index_cloud_sql_pg/index_store.py b/src/llama_index_cloud_sql_pg/index_store.py index 3103d54..ea8eed2 100644 --- a/src/llama_index_cloud_sql_pg/index_store.py +++ b/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. diff --git a/tests/test_async_index_store.py b/tests/test_async_index_store.py index 5c2389e..b59975d 100644 --- a/tests/test_async_index_store.py +++ b/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() diff --git a/tests/test_index_store.py b/tests/test_index_store.py index 5f840e7..a2ee0fc 100644 --- a/tests/test_index_store.py +++ b/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()