Feat - Additional Columns Returned

src/api/models/embedding_models.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from typing import List, Dict
+from typing import List, Dict, Optional
 
 
 # Pydantic models for request validation
@@ -56,3 +56,6 @@ class SearchEmbeddingRequest(BaseModel):
     embedding_column: str  # Column containing the embeddings to search against
     num_results: int  # Number of results to return
     dataset_name: str  # Name of the dataset to search in
+    additional_columns: Optional[List[str]] = (
+        None  # Optional list of additional columns to include in the results
+    )

src/api/services/embedding_service.py

@@ -1,6 +1,6 @@
 from openai import AsyncOpenAI
 import logging
-from typing import List, Dict, Union
+from typing import List, Dict, Union, Optional
 from datasets import Dataset
 import asyncio
 import numpy as np
@@ -119,6 +119,48 @@ class EmbeddingService:
             f"Progress: {self.completed_requests}/{self.total_requests} ({progress:.2f}%)"
         )
 
+    # async def search_embeddings(
+    #     self,
+    #     query_embeddings: List[List[float]],
+    #     dataset: Dataset,
+    #     embedding_column: str,
+    #     target_column: str,
+    #     num_results: int,
+    # ) -> Dict[str, List]:
+    #     """
+    #     Perform a cosine similarity search between query embeddings and dataset embeddings.
+
+    #     Args:
+    #         query_embeddings: List of embeddings for the query texts.
+    #         dataset: The dataset to search in.
+    #         embedding_column: The column in the dataset containing embeddings.
+    #         target_column: The column to return in the results.
+    #         num_results: The number of results to return.
+
+    #     Returns:
+    #         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)
+
+    #     # 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
+    #     for query_similarities in similarities:
+    #         top_k_indices = np.argsort(query_similarities)[-num_results:][::-1]
+    #         for idx in top_k_indices:
+    #             results[target_column].append(dataset[target_column][idx])
+    #             results["similarity"].append(float(query_similarities[idx]))
+
+    #     return results
+
     async def search_embeddings(
         self,
         query_embeddings: List[List[float]],
@@ -126,6 +168,7 @@ class EmbeddingService:
         embedding_column: str,
         target_column: str,
         num_results: int,
+        additional_columns: Optional[List[str]] = None,
     ) -> Dict[str, List]:
         """
         Perform a cosine similarity search between query embeddings and dataset embeddings.
@@ -136,9 +179,11 @@ class EmbeddingService:
             embedding_column: The column in the dataset containing embeddings.
             target_column: The column to return in the results.
             num_results: The number of results to return.
+            additional_columns: List of additional columns to include in the results.
 
         Returns:
-            A dictionary of lists containing the target column values and their similarity scores.
+            A dictionary of lists containing the target column values, their similarity scores,
+            and any additional columns specified.
         """
         dataset_embeddings = np.array(dataset[embedding_column])
         query_embeddings = np.array(query_embeddings)
@@ -152,11 +197,19 @@ class EmbeddingService:
             "similarity": [],
         }
 
+        # Add additional columns to the results dictionary
+        if additional_columns:
+            for column in additional_columns:
+                results[column] = []
+
         # Get the top-k results for each query
         for query_similarities in similarities:
             top_k_indices = np.argsort(query_similarities)[-num_results:][::-1]
             for idx in top_k_indices:
                 results[target_column].append(dataset[target_column][idx])
                 results["similarity"].append(float(query_similarities[idx]))
+                if additional_columns:
+                    for column in additional_columns:
+                        results[column].append(dataset[column][idx])
 
         return results

src/main.py

@@ -256,6 +256,7 @@ async def search_embedding(
             request.embedding_column,
             request.target_column,
             request.num_results,
+            request.additional_columns,
         )
 
         return JSONResponse(