Feat - New Endpoint - Similarity Search

requirements.txt

@@ -5,4 +5,5 @@ uvicorn
 fastapi
 openai
 pandas
-datasets
+datasets
+scikit-learn

src/api/database.py

@@ -110,31 +110,6 @@ class Database:
             with self.lock:
                 self.pool.append(conn)
 
-    # async def fetch(self, query: str, *args) -> List[Dict]:
-    #     """
-    #     Execute a SELECT query and return the results as a list of dictionaries.
-
-    #     Args:
-    #         query (str): The SQL query to execute.
-    #         *args: Query parameters.
-
-    #     Returns:
-    #         List[Dict]: A list of dictionaries where keys are column names and values are column values.
-
-    #     Raises:
-    #         QueryExecutionError: If the query execution fails.
-    #     """
-    #     try:
-    #         async with self.get_connection() as conn:
-    #             cursor = conn.cursor()
-    #             cursor.execute(query, args)
-    #             rows = cursor.fetchall()
-    #             columns = [desc[0] for desc in cursor.description]
-    #             return [dict(zip(columns, row)) for row in rows]
-    #     except Pg8000DatabaseError as e:
-    #         logger.error(f"Query execution failed: {e}")
-    #         raise QueryExecutionError(f"Failed to execute query: {query}") from e
-
     async def fetch(self, query: str, *args) -> Dict[str, List]:
         """
         Execute a SELECT query and return the results as a dictionary of lists.

src/api/models/embedding_models.py

@@ -48,3 +48,11 @@ class EmbedRequest(BaseModel):
     output_column: str = (
         "embedding"  # Column to store embeddings (default: "embeddings")
     )
+
+
+class SearchEmbeddingRequest(BaseModel):
+    texts: List[str]  # List of texts to search for
+    target_column: str  # Column to return in the results
+    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

src/api/services/embedding_service.py

@@ -1,137 +1,10 @@
-# from openai import AsyncOpenAI
-# import logging
-# from typing import List, Dict, Union
-# import pandas as pd
-# import asyncio
-# from src.api.exceptions import OpenAIError
-
-# # Set up structured logging
-# logging.basicConfig(
-#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-# )
-# logger = logging.getLogger(__name__)
-
-
-# class EmbeddingService:
-#     def __init__(
-#         self,
-#         openai_api_key: str,
-#         model: str = "text-embedding-3-small",
-#         batch_size: int = 10,
-#         max_concurrent_requests: int = 10,  # Limit to 10 concurrent requests
-#     ):
-#         self.client = AsyncOpenAI(api_key=openai_api_key)
-#         self.model = model
-#         self.batch_size = batch_size
-#         self.semaphore = asyncio.Semaphore(max_concurrent_requests)  # Rate limiter
-#         self.total_requests = 0  # Total number of requests to process
-#         self.completed_requests = 0  # Number of completed requests
-
-#     async def get_embedding(self, text: str) -> List[float]:
-#         """Generate embeddings for the given text using OpenAI."""
-#         text = text.replace("\n", " ")
-#         try:
-#             async with self.semaphore:  # Acquire a semaphore slot
-#                 response = await self.client.embeddings.create(
-#                     input=[text], model=self.model
-#                 )
-#                 self.completed_requests += 1  # Increment completed requests
-#                 self._log_progress()  # Log progress
-#                 return response.data[0].embedding
-#         except Exception as e:
-#             logger.error(f"Failed to generate embedding: {e}")
-#             raise OpenAIError(f"OpenAI API error: {e}")
-
-#     async def create_embeddings(
-#         self,
-#         data: Union[pd.DataFrame, List[str]],
-#         target_column: str = None,
-#         output_column: str = "embeddings",
-#     ) -> Union[pd.DataFrame, List[List[float]]]:
-#         """
-#         Create embeddings for either a DataFrame or a list of strings.
-
-#         Args:
-#             data: Either a DataFrame or a list of strings.
-#             target_column: The column in the DataFrame to generate embeddings for (required if data is a DataFrame).
-#             output_column: The column to store embeddings in the DataFrame (default: "embeddings").
-
-#         Returns:
-#             If data is a DataFrame, returns the DataFrame with the embeddings column.
-#             If data is a list of strings, returns a list of embeddings.
-#         """
-#         if isinstance(data, pd.DataFrame):
-#             if not target_column:
-#                 raise ValueError("target_column is required when data is a DataFrame.")
-#             return await self._create_embeddings_for_dataframe(
-#                 data, target_column, output_column
-#             )
-#         elif isinstance(data, list):
-#             return await self._create_embeddings_for_texts(data)
-#         else:
-#             raise TypeError(
-#                 "data must be either a pandas DataFrame or a list of strings."
-#             )
-
-#     async def _create_embeddings_for_dataframe(
-#         self, df: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Create embeddings for the target column in the DataFrame."""
-#         logger.info("Generating embeddings for DataFrame...")
-#         self.total_requests = len(df)  # Set total number of requests
-#         self.completed_requests = 0  # Reset completed requests counter
-
-#         batches = [
-#             df[i : i + self.batch_size] for i in range(0, len(df), self.batch_size)
-#         ]
-#         processed_batches = await asyncio.gather(
-#             *[
-#                 self._process_batch(batch, target_column, output_column)
-#                 for batch in batches
-#             ]
-#         )
-#         return pd.concat(processed_batches)
-
-#     async def _create_embeddings_for_texts(self, texts: List[str]) -> List[List[float]]:
-#         """Create embeddings for a list of strings."""
-#         logger.info("Generating embeddings for list of texts...")
-#         self.total_requests = len(texts)  # Set total number of requests
-#         self.completed_requests = 0  # Reset completed requests counter
-
-#         batches = [
-#             texts[i : i + self.batch_size]
-#             for i in range(0, len(texts), self.batch_size)
-#         ]
-#         embeddings = []
-#         for batch in batches:
-#             batch_embeddings = await asyncio.gather(
-#                 *[self.get_embedding(text) for text in batch]
-#             )
-#             embeddings.extend(batch_embeddings)
-#         return embeddings
-
-#     async def _process_batch(
-#         self, df_batch: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Process a batch of rows to generate embeddings."""
-#         embeddings = await asyncio.gather(
-#             *[self.get_embedding(row[target_column]) for _, row in df_batch.iterrows()]
-#         )
-#         df_batch[output_column] = embeddings
-#         return df_batch
-
-#     def _log_progress(self):
-#         """Log the progress of embedding generation."""
-#         progress = (self.completed_requests / self.total_requests) * 100
-#         logger.info(
-#             f"Progress: {self.completed_requests}/{self.total_requests} ({progress:.2f}%)"
-#         )
-
 from openai import AsyncOpenAI
 import logging
 from typing import List, Dict, Union
 from datasets import Dataset
 import asyncio
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
 from src.api.exceptions import OpenAIError
 
 # Set up structured logging
@@ -245,3 +118,45 @@ class EmbeddingService:
         logger.info(
             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,
+    ) -> List[Dict]:
+        """
+        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 list of dictionaries 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)
+
+        # Get the top-k results for each query
+        results = []
+        for i, query_similarities in enumerate(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)
+
+        return results

src/api/services/huggingface_service.py

@@ -1,106 +1,3 @@
-# from datasets import Dataset, load_dataset, concatenate_datasets
-# from huggingface_hub import HfApi, HfFolder
-# import logging
-# import os
-# from typing import Optional, Dict, List
-# import pandas as pd
-# from src.api.services.embedding_service import EmbeddingService
-# from src.api.exceptions import (
-#     DatasetNotFoundError,
-#     DatasetPushError,
-#     DatasetDeleteError,
-# )
-
-# # Set up structured logging
-# logging.basicConfig(
-#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-# )
-# logger = logging.getLogger(__name__)
-
-
-# class HuggingFaceService:
-#     def __init__(self, hf_token: Optional[str] = None):
-#         """Initialize the HuggingFaceService with an optional token."""
-#         self.hf_api = HfApi()
-#         if hf_token:
-#             HfFolder.save_token(hf_token)  # Save the token for authentication
-
-#     async def push_to_hub(self, df: pd.DataFrame, dataset_name: str) -> None:
-#         """Push the dataset to Hugging Face Hub."""
-#         try:
-#             logger.info(f"Creating Hugging Face Dataset: {dataset_name}...")
-#             ds = Dataset.from_pandas(df)
-#             ds.push_to_hub(dataset_name)
-#             logger.info(f"Dataset pushed to Hugging Face Hub: {dataset_name}")
-#         except Exception as e:
-#             logger.error(f"Failed to push dataset to Hugging Face Hub: {e}")
-#             raise DatasetPushError(f"Failed to push dataset: {e}")
-
-#     async def read_dataset(self, dataset_name: str) -> Optional[pd.DataFrame]:
-#         """Read a dataset from Hugging Face Hub."""
-#         try:
-#             logger.info(f"Loading dataset from Hugging Face Hub: {dataset_name}...")
-#             ds = load_dataset(dataset_name)
-#             df = ds["train"].to_dict()
-#             return df
-#         except Exception as e:
-#             logger.error(f"Failed to read dataset: {e}")
-#             raise DatasetNotFoundError(f"Dataset not found: {e}")
-
-#     async def update_dataset(
-#         self,
-#         dataset_name: str,
-#         updates: Dict[str, List],
-#         target_column: str,
-#         output_column: str = "embeddings",
-#     ) -> Optional[pd.DataFrame]:
-#         """Update a dataset on Hugging Face Hub by generating embeddings for new data and concatenating it with the existing dataset."""
-#         try:
-#             # Step 1: Load the existing dataset from Hugging Face Hub
-#             logger.info(
-#                 f"Loading existing dataset from Hugging Face Hub: {dataset_name}..."
-#             )
-#             existing_ds = await self.read_dataset(dataset_name)
-#             existing_df = pd.DataFrame(existing_ds)
-
-#             # Step 2: Convert the new updates into a DataFrame
-#             logger.info("Converting updates to DataFrame...")
-#             new_df = pd.DataFrame(updates)
-
-#             # Step 3: Generate embeddings for the new data
-#             logger.info("Generating embeddings for the new data...")
-#             embedding_service = EmbeddingService(
-#                 openai_api_key=os.getenv("OPENAI_API_KEY")
-#             )  # Get the embedding service
-#             new_df = await embedding_service.create_embeddings(
-#                 new_df, target_column, output_column
-#             )
-
-#             # Step 4: Concatenate the existing DataFrame with the new DataFrame
-#             logger.info("Concatenating existing dataset with new data...")
-#             updated_df = pd.concat([existing_df, new_df], ignore_index=True)
-
-#             # Step 5: Push the updated dataset back to Hugging Face Hub
-#             logger.info(
-#                 f"Pushing updated dataset to Hugging Face Hub: {dataset_name}..."
-#             )
-#             await self.push_to_hub(updated_df, dataset_name)
-
-#             return updated_df
-#         except Exception as e:
-#             logger.error(f"Failed to update dataset: {e}")
-#             raise DatasetPushError(f"Failed to update dataset: {e}")
-
-#     async def delete_dataset(self, dataset_name: str) -> None:
-#         """Delete a dataset from Hugging Face Hub."""
-#         try:
-#             logger.info(f"Deleting dataset from Hugging Face Hub: {dataset_name}...")
-#             self.hf_api.delete_repo(repo_id=dataset_name, repo_type="dataset")
-#             logger.info(f"Dataset deleted from Hugging Face Hub: {dataset_name}")
-#         except Exception as e:
-#             logger.error(f"Failed to delete dataset: {e}")
-#             raise DatasetDeleteError(f"Failed to delete dataset: {e}")
-
 from datasets import Dataset, load_dataset, concatenate_datasets
 from huggingface_hub import HfApi, HfFolder
 import logging
@@ -141,7 +38,11 @@ class HuggingFaceService:
         """Read a dataset from Hugging Face Hub."""
         try:
             logger.info(f"Loading dataset from Hugging Face Hub: {dataset_name}...")
-            dataset = load_dataset(dataset_name)
+            dataset = load_dataset(
+                dataset_name,
+                keep_in_memory=True,
+                download_mode="force_redownload",
+            )
             return dataset["train"]
         except Exception as e:
             logger.error(f"Failed to read dataset: {e}")

src/main.py

@@ -1,252 +1,3 @@
-# import os
-# from fastapi import FastAPI, Depends, HTTPException
-# from fastapi.responses import JSONResponse, RedirectResponse
-# from fastapi.middleware.gzip import GZipMiddleware
-# from pydantic import BaseModel
-# from typing import List, Dict
-# from src.api.models.embedding_models import (
-#     CreateEmbeddingRequest,
-#     ReadEmbeddingRequest,
-#     UpdateEmbeddingRequest,
-#     DeleteEmbeddingRequest,
-#     EmbedRequest,
-# )
-# from src.api.database import get_db, Database, QueryExecutionError, HealthCheckError
-# from src.api.services.embedding_service import EmbeddingService
-# from src.api.services.huggingface_service import HuggingFaceService
-# from src.api.exceptions import DatasetNotFoundError, DatasetPushError, OpenAIError
-
-# # from src.api.dependency import get_embedding_service, get_huggingface_service
-# import pandas as pd
-# import logging
-# from dotenv import load_dotenv
-
-# # Load environment variables
-# load_dotenv()
-
-# # Set up structured logging
-# logging.basicConfig(
-#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-# )
-# logger = logging.getLogger(__name__)
-
-# description = """A FastAPI application for similarity search with PostgreSQL and OpenAI embeddings.
-
-# Direct/API URL:
-# https://re-mind-similarity-search.hf.space
-# """
-
-# # Initialize FastAPI app
-# app = FastAPI(
-#     title="Similarity Search API",
-#     description=description,
-#     version="1.0.0",
-# )
-
-# app.add_middleware(GZipMiddleware, minimum_size=1000)
-
-
-# # Dependency to get EmbeddingService
-# def get_embedding_service() -> EmbeddingService:
-#     return EmbeddingService(openai_api_key=os.getenv("OPENAI_API_KEY"))
-
-
-# # Dependency to get HuggingFaceService
-# def get_huggingface_service() -> HuggingFaceService:
-#     return HuggingFaceService()
-
-
-# # Root endpoint redirects to /docs
-# @app.get("/")
-# async def root():
-#     return RedirectResponse(url="/docs")
-
-
-# # Health check endpoint
-# @app.get("/health")
-# async def health_check(db: Database = Depends(get_db)):
-#     try:
-#         is_healthy = await db.health_check()
-#         if not is_healthy:
-#             raise HTTPException(status_code=500, detail="Database is unhealthy")
-#         return {"status": "healthy"}
-#     except HealthCheckError as e:
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# # Endpoint to generate embeddings for a list of strings
-# @app.post("/embed")
-# async def embed(
-#     request: EmbedRequest,
-#     embedding_service: EmbeddingService = Depends(get_embedding_service),
-# ):
-#     """
-#     Generate embeddings for a list of strings and return them in the response.
-#     """
-#     try:
-#         # Step 1: Generate embeddings
-#         logger.info("Generating embeddings for list of texts...")
-#         embeddings = await embedding_service.create_embeddings(request.texts)
-
-#         return JSONResponse(
-#             content={
-#                 "message": "Embeddings generated successfully.",
-#                 "embeddings": embeddings,
-#                 "num_texts": len(request.texts),
-#             }
-#         )
-#     except OpenAIError as e:
-#         logger.error(f"OpenAI API error: {e}")
-#         raise HTTPException(status_code=500, detail=f"OpenAI API error: {e}")
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
-# # Endpoint to create embeddings from a database query
-# @app.post("/create_embedding")
-# async def create_embedding(
-#     request: CreateEmbeddingRequest,
-#     db: Database = Depends(get_db),
-#     embedding_service: EmbeddingService = Depends(get_embedding_service),
-#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
-# ):
-#     """
-#     Create embeddings for the target column in the dataset.
-#     """
-#     try:
-#         # Step 1: Query the database
-#         logger.info("Fetching data from the database...")
-#         result = await db.fetch(request.query)
-#         df = pd.DataFrame(result)
-
-#         # Step 2: Generate embeddings
-#         df = await embedding_service.create_embeddings(
-#             df, request.target_column, request.output_column
-#         )
-
-#         # Step 3: Push to Hugging Face Hub
-#         await huggingface_service.push_to_hub(df, request.dataset_name)
-
-#         return JSONResponse(
-#             content={
-#                 "message": "Embeddings created and pushed to Hugging Face Hub.",
-#                 "dataset_name": request.dataset_name,
-#                 "num_rows": len(df),
-#             }
-#         )
-#     except QueryExecutionError as e:
-#         logger.error(f"Database query failed: {e}")
-#         raise HTTPException(status_code=500, detail=f"Database query failed: {e}")
-#     except OpenAIError as e:
-#         logger.error(f"OpenAI API error: {e}")
-#         raise HTTPException(status_code=500, detail=f"OpenAI API error: {e}")
-#     except DatasetPushError as e:
-#         logger.error(f"Failed to push dataset: {e}")
-#         raise HTTPException(status_code=500, detail=f"Failed to push dataset: {e}")
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
-# # Endpoint to read embeddings
-# @app.post("/read_embeddings")
-# async def read_embeddings(
-#     request: ReadEmbeddingRequest,
-#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
-# ):
-#     """
-#     Read embeddings from a Hugging Face dataset.
-#     """
-#     try:
-#         df = await huggingface_service.read_dataset(request.dataset_name)
-#         return df
-#     except DatasetNotFoundError as e:
-#         logger.error(f"Dataset not found: {e}")
-#         raise HTTPException(status_code=404, detail=f"Dataset not found: {e}")
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
-# # Endpoint to update embeddings
-# # @app.post("/update_embeddings")
-# # async def update_embeddings(
-# #     request: UpdateEmbeddingRequest,
-# #     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
-# # ):
-# #     """
-# #     Update embeddings in a Hugging Face dataset.
-# #     """
-# #     try:
-# #         df = await huggingface_service.update_dataset(
-# #             request.dataset_name, request.updates
-# #         )
-# #         return {
-# #             "message": "Embeddings updated successfully.",
-# #             "dataset_name": request.dataset_name,
-# #         }
-# #     except DatasetPushError as e:
-# #         logger.error(f"Failed to update dataset: {e}")
-# #         raise HTTPException(status_code=500, detail=f"Failed to update dataset: {e}")
-# #     except Exception as e:
-# #         logger.error(f"An error occurred: {e}")
-# #         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
-# @app.post("/update_embeddings")
-# async def update_embeddings(
-#     request: UpdateEmbeddingRequest,
-#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
-# ):
-#     """
-#     Update embeddings in a Hugging Face dataset by generating embeddings for new data and concatenating it with the existing dataset.
-#     """
-#     try:
-#         # Call the update_dataset method to generate embeddings, concatenate, and push the updated dataset
-#         updated_df = await huggingface_service.update_dataset(
-#             request.dataset_name,
-#             request.updates,
-#             request.target_column,
-#             request.output_column,
-#         )
-
-#         return {
-#             "message": "Embeddings updated successfully.",
-#             "dataset_name": request.dataset_name,
-#             "num_rows": len(updated_df),
-#         }
-#     except DatasetPushError as e:
-#         logger.error(f"Failed to update dataset: {e}")
-#         raise HTTPException(status_code=500, detail=f"Failed to update dataset: {e}")
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
-# # Endpoint to delete embeddings
-# @app.post("/delete_embeddings")
-# async def delete_embeddings(
-#     request: DeleteEmbeddingRequest,
-#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
-# ):
-#     """
-#     Delete embeddings from a Hugging Face dataset.
-#     """
-#     try:
-#         await huggingface_service.delete_dataset(request.dataset_name)
-#         return {
-#             "message": "Embeddings deleted successfully.",
-#             "dataset_name": request.dataset_name,
-#         }
-#     except DatasetPushError as e:
-#         logger.error(f"Failed to delete columns: {e}")
-#         raise HTTPException(status_code=500, detail=f"Failed to delete columns: {e}")
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
-
-
 import os
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.responses import JSONResponse, RedirectResponse
@@ -260,6 +11,7 @@ from src.api.models.embedding_models import (
     UpdateEmbeddingRequest,
     DeleteEmbeddingRequest,
     EmbedRequest,
+    SearchEmbeddingRequest,
 )
 from src.api.database import get_db, Database, QueryExecutionError, HealthCheckError
 from src.api.services.embedding_service import EmbeddingService
@@ -363,6 +115,10 @@ async def create_embedding(
     Create embeddings for the target column in the dataset.
     """
     try:
+        embedding_service.model = request.model
+        embedding_service.batch_size = request.batch_size
+        # embedding_service.max_concurrent_requests = request.max_concurrent_requests
+
         # Step 1: Query the database
         logger.info("Fetching data from the database...")
         result = await db.fetch(request.query)
@@ -371,8 +127,6 @@ async def create_embedding(
 
         dataset = Dataset.from_dict(result)
 
-        embedding_service.batch_size = request.batch_size
-
         # Step 2: Generate embeddings
         dataset = await embedding_service.create_embeddings(
             dataset, request.target_column, request.output_column
@@ -474,3 +228,45 @@ async def delete_embeddings(
     except Exception as e:
         logger.error(f"An error occurred: {e}")
         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
+@app.post("/search_embedding")
+async def search_embedding(
+    request: SearchEmbeddingRequest,
+    embedding_service: EmbeddingService = Depends(get_embedding_service),
+    huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+):
+    """
+    Search for similar texts in a dataset using embeddings.
+    """
+    try:
+        # Step 1: Generate embeddings for the query texts
+        logger.info("Generating embeddings for query texts...")
+        query_embeddings = await embedding_service.create_embeddings(request.texts)
+
+        # Step 2: Load the dataset from Hugging Face Hub
+        logger.info(f"Loading dataset from Hugging Face Hub: {request.dataset_name}...")
+        dataset = await huggingface_service.read_dataset(request.dataset_name)
+
+        # Step 3: Perform cosine similarity search
+        logger.info("Performing cosine similarity search...")
+        results = await embedding_service.search_embeddings(
+            query_embeddings,
+            dataset,
+            request.embedding_column,
+            request.target_column,
+            request.num_results,
+        )
+
+        return JSONResponse(
+            content={
+                "message": "Search completed successfully.",
+                "results": results,
+            }
+        )
+    except DatasetNotFoundError as e:
+        logger.error(f"Dataset not found: {e}")
+        raise HTTPException(status_code=404, detail=f"Dataset not found: {e}")
+    except Exception as e:
+        logger.error(f"An error occurred: {e}")
+        raise HTTPException(status_code=500, detail=f"An error occurred: {e}")