Feat - concurrent request update

src/api/models/embedding_models.py

@@ -8,7 +8,8 @@ class CreateEmbeddingRequest(BaseModel):
     target_column: str = "product_type"
     output_column: str = "embedding"
     model: str = "text-embedding-3-small"
-    batch_size: int = 100
+    batch_size: int = 10
+    max_concurrent_requests: int = 10
     dataset_name: str = "re-mind/product_type_embedding"
 
 

src/api/services/embedding_service.py

@@ -1,3 +1,66 @@
+# 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 = 100,
+#     ):
+#         self.client = AsyncOpenAI(api_key=openai_api_key)
+#         self.model = model
+#         self.batch_size = batch_size
+
+#     async def get_embedding(self, text: str) -> List[float]:
+#         """Generate embeddings for the given text using OpenAI."""
+#         text = text.replace("\n", " ")
+#         try:
+#             response = await self.client.embeddings.create(
+#                 input=[text], model=self.model
+#             )
+#             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...")
+#         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
+
 from openai import AsyncOpenAI
 import logging
 from typing import List, Dict
@@ -17,20 +80,23 @@ class EmbeddingService:
         self,
         openai_api_key: str,
         model: str = "text-embedding-3-small",
-        batch_size: int = 100,
+        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
 
     async def get_embedding(self, text: str) -> List[float]:
         """Generate embeddings for the given text using OpenAI."""
         text = text.replace("\n", " ")
         try:
-            response = await self.client.embeddings.create(
-                input=[text], model=self.model
-            )
-            return response.data[0].embedding
+            async with self.semaphore:  # Acquire a semaphore slot
+                response = await self.client.embeddings.create(
+                    input=[text], model=self.model
+                )
+                return response.data[0].embedding
         except Exception as e:
             logger.error(f"Failed to generate embedding: {e}")
             raise OpenAIError(f"OpenAI API error: {e}")