Skip to content

fix: Index store adds interface methods #118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 21, 2025
Merged

fix: Index store adds interface methods #118

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/llama_index_cloud_sql_pg/async_index_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
26 changes: 20 additions & 6 deletions src/llama_index_cloud_sql_pg/index_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion tests/test_async_index_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
16 changes: 9 additions & 7 deletions tests/test_index_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down