Feat - embed endpoint created

src/api/services/embedding_service.py

@@ -61,9 +61,89 @@
 #         df_batch[output_column] = embeddings
 #         return df_batch
 
+# from openai import AsyncOpenAI
+# import logging
+# from typing import List, Dict
+# 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, df: pd.DataFrame, target_column: str, output_column: str
+#     ) -> pd.DataFrame:
+#         """Create embeddings for the target column in the dataset."""
+#         logger.info("Generating embeddings...")
+#         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 _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
+from typing import List, Dict, Union
 import pandas as pd
 import asyncio
 from src.api.exceptions import OpenAIError
@@ -106,10 +186,41 @@ class EmbeddingService:
             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 dataset."""
-        logger.info("Generating embeddings...")
+        """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
 
@@ -124,6 +235,24 @@ class EmbeddingService:
         )
         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:

src/main.py

@@ -1,3 +1,191 @@
+# 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,
+# )
+# 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
+# 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)
+
+
+# # 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))
+
+
+# # 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()
+
+
+# # Endpoint to create embeddings
+# @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.get("/read_embeddings/{dataset_name}")
+# @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}")
+
+
+# # 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
@@ -71,7 +259,44 @@ def get_huggingface_service() -> HuggingFaceService:
     return HuggingFaceService()
 
 
-# Endpoint to create embeddings
+# Request model for the /embed endpoint
+class EmbedRequest(BaseModel):
+    texts: List[str]  # List of strings to generate embeddings for
+    output_column: str = (
+        "embeddings"  # Column to store embeddings (default: "embeddings")
+    )
+
+
+# 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,
@@ -118,7 +343,6 @@ async def create_embedding(
 
 
 # Endpoint to read embeddings
-# @app.get("/read_embeddings/{dataset_name}")
 @app.post("/read_embeddings")
 async def read_embeddings(
     request: ReadEmbeddingRequest,