gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: fixing changed lancedb search + pruning

Browse source
This commit is contained in:
hajdul88 2025-01-16 17:32:44 +01:00
parent 1c4a605eb7
commit 935763b08d
1 changed files with 12 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py
View file

@ -152,7 +152,9 @@ class LanceDBAdapter(VectorDBInterface):
connection = await self.get_connection() connection = await self.get_connection()
collection = await connection.open_table(collection_name) collection = await connection.open_table(collection_name)
results = await collection.vector_search(query_vector).to_pandas() collection_size = await collection.count_rows()
results = await collection.vector_search(query_vector).limit(collection_size).to_pandas()
result_values = list(results.to_dict("index").values()) result_values = list(results.to_dict("index").values())
@ -250,9 +252,16 @@ class LanceDBAdapter(VectorDBInterface):
) )
async def prune(self): async def prune(self):
# Clean up the database if it was set up as temporary connection = await self.get_connection()
collection_names = await connection.table_names()
for collection_name in collection_names:
collection = await connection.open_table(collection_name)
await collection.delete("id IS NOT NULL")
await connection.drop_table(collection_name)
if self.url.startswith("/"): if self.url.startswith("/"):
LocalStorage.remove_all(self.url) # Remove the temporary directory and files inside LocalStorage.remove_all(self.url)
def get_data_point_schema(self, model_type): def get_data_point_schema(self, model_type):
return copy_model( return copy_model(