From c0d1aa12160ed4f6402d8c29766c22135c30c832 Mon Sep 17 00:00:00 2001
From: Boris Arzentar <borisarzentar@gmail.com>
Date: Mon, 11 Nov 2024 17:54:00 +0100
Subject: [PATCH] fix: update entities collection name in cognee_demo notebook

---
 .../vector/pgvector/PGVectorAdapter.py        | 23 ++++++++++++-------
 notebooks/cognee_demo.ipynb                   |  2 +-
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
index 84a32e3e2..01691714b 100644
--- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
+++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
@@ -179,6 +179,8 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         # Get PGVectorDataPoint Table from database
         PGVectorDataPoint = await self.get_table(collection_name)
 
+        closest_items = []
+
         # Use async session to connect to the database
         async with self.get_async_session() as session:
             # Find closest vectors to query_vector
@@ -195,14 +197,19 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
 
         vector_list = []
 
-            # Create and return ScoredResult objects
-            return [
-                ScoredResult(
-                    id = UUID(str(row.id)),
-                    payload = row.payload,
-                    score = row.similarity
-                ) for row in vector_list
-            ]
+        # Extract distances and find min/max for normalization
+        for vector in closest_items:
+            # TODO: Add normalization of similarity score
+            vector_list.append(vector)
+
+        # Create and return ScoredResult objects
+        return [
+            ScoredResult(
+                id = UUID(str(row.id)),
+                payload = row.payload,
+                score = row.similarity
+            ) for row in vector_list
+        ]
 
     async def batch_search(
         self,
diff --git a/notebooks/cognee_demo.ipynb b/notebooks/cognee_demo.ipynb
index 396d7b980..5f4dfa227 100644
--- a/notebooks/cognee_demo.ipynb
+++ b/notebooks/cognee_demo.ipynb
@@ -758,7 +758,7 @@
     "from cognee.infrastructure.databases.vector import get_vector_engine\n",
     "\n",
     "vector_engine = get_vector_engine()\n",
-    "results = await search(vector_engine, \"entities\", \"sarah.nguyen@example.com\")\n",
+    "results = await search(vector_engine, \"Entity_name\", \"sarah.nguyen@example.com\")\n",
     "for result in results:\n",
     "    print(result)"
    ]