Feat - Data Format

src/api/services/embedding_service.py

@@ -126,7 +126,7 @@ class EmbeddingService:
         embedding_column: str,
         target_column: str,
         num_results: int,
-    ) -> List[Dict]:
+    ) -> Dict[str, List]:
         """
         Perform a cosine similarity search between query embeddings and dataset embeddings.
 
@@ -138,7 +138,7 @@ class EmbeddingService:
             num_results: The number of results to return.
 
         Returns:
-            A list of dictionaries containing the target column values and their similarity scores.
+            A dictionary of lists containing the target column values and their similarity scores.
         """
         dataset_embeddings = np.array(dataset[embedding_column])
         query_embeddings = np.array(query_embeddings)
@@ -146,17 +146,17 @@ class EmbeddingService:
         # Compute cosine similarity
         similarities = cosine_similarity(query_embeddings, dataset_embeddings)
 
+        # Initialize the results dictionary
+        results = {
+            target_column: [],
+            "similarity": [],
+        }
+
         # Get the top-k results for each query
-        results = []
-        for i, query_similarities in enumerate(similarities):
+        for query_similarities in similarities:
             top_k_indices = np.argsort(query_similarities)[-num_results:][::-1]
-            top_k_results = [
-                {
-                    target_column: dataset[target_column][idx],
-                    "similarity": float(query_similarities[idx]),
-                }
-                for idx in top_k_indices
-            ]
-            results.append(top_k_results)
+            for idx in top_k_indices:
+                results[target_column].append(dataset[target_column][idx])
+                results["similarity"].append(float(query_similarities[idx]))
 
         return results