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

refactor: Remove broad exception handling from PGVectorAdapter

Browse source 
Removed broad exception handling for the search function in PGVectorAdapter

Refactor #COG-170
This commit is contained in:
Igor Ilic 2024-10-22 15:41:07 +02:00
parent dc46304a8d
commit 0c6f019838
1 changed files with 24 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
View file

@ -157,39 +157,34 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
# Use async session to connect to the database # Use async session to connect to the database
async with self.get_async_session() as session: async with self.get_async_session() as session:
try: # Get PGVectorDataPoint Table from database
# Get PGVectorDataPoint Table from database PGVectorDataPoint = await self.get_table(collection_name)
PGVectorDataPoint = await self.get_table(collection_name)
# Find closest vectors to query_vector # Find closest vectors to query_vector
closest_items = await session.execute( closest_items = await session.execute(
select( select(
PGVectorDataPoint, PGVectorDataPoint,
PGVectorDataPoint.c.vector.cosine_distance(query_vector).label( PGVectorDataPoint.c.vector.cosine_distance(query_vector).label(
"similarity" "similarity"
), ),
)
.order_by("similarity")
.limit(limit)
) )
.order_by("similarity")
.limit(limit)
)
vector_list = [] vector_list = []
# Extract distances and find min/max for normalization # Extract distances and find min/max for normalization
for vector in closest_items: for vector in closest_items:
# TODO: Add normalization of similarity score # TODO: Add normalization of similarity score
vector_list.append(vector) vector_list.append(vector)
# Create and return ScoredResult objects # Create and return ScoredResult objects
return [ return [
ScoredResult( ScoredResult(
id=str(row.id), payload=row.payload, score=row.similarity id=str(row.id), payload=row.payload, score=row.similarity
) )
for row in vector_list for row in vector_list
] ]
except Exception as e:
print(f"Error during search: {e}")
return []
async def batch_search( async def batch_search(
self, self,