Intial Deployment

.gitignore

@@ -0,0 +1,2 @@
+*pycache*
+*.env*

Dockerfile

@@ -0,0 +1,51 @@
+# Stage 1: Build stage
+FROM python:3.12-slim as builder
+
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+
+# Set the working directory
+WORKDIR /app
+
+# Copy only the requirements file first to leverage Docker cache
+COPY --chown=user ./requirements.txt /app/requirements.txt
+
+# Install dependencies in a virtual environment
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application code
+COPY --chown=user . /app
+
+# Stage 2: Runtime stage
+FROM python:3.12-slim
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+USER user
+
+# Copy the virtual environment from the builder stage
+COPY --from=builder /opt/venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# Set the working directory
+WORKDIR /app
+
+# Copy only the necessary files from the builder stage
+COPY --from=builder --chown=user /app /app
+
+# Expose the port the app runs on
+EXPOSE 7860
+
+# Health check to ensure the application is running
+HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:7860/health || exit 1
+
+# Command to run the application with hot reloading
+CMD ["uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "7860", "--reload"]

README.md

@@ -6,6 +6,9 @@ colorTo: gray
 sdk: docker
 pinned: false
 license: mit
+python_version: 3.12
+app_port: 7860
+app_file: src/main.py
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

docker-compose.yml

@@ -0,0 +1,43 @@
+version: "3.9"
+
+services:
+  app:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    container_name: similarity-search-app
+    ports:
+      - "7860:7860"
+    volumes:
+      - ./src:/app/src  # Mount the local src directory for hot reloading
+    environment:
+      - PYTHONUNBUFFERED=1
+    restart: unless-stopped
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:7860/health"]
+      interval: 30s
+      timeout: 10s
+      retries: 3
+#     depends_on:
+#       - db  # If you have a database service, add it here
+
+#   # Example database service (optional)
+#   db:
+#     image: postgres:latest
+#     container_name: similarity-search-db
+#     environment:
+#       POSTGRES_USER: user
+#       POSTGRES_PASSWORD: password
+#       POSTGRES_DB: mydatabase
+#     ports:
+#       - "5432:5432"
+#     volumes:
+#       - postgres_data:/var/lib/postgresql/data
+#     healthcheck:
+#       test: ["CMD-SHELL", "pg_isready -U user -d mydatabase"]
+#       interval: 5s
+#       timeout: 5s
+#       retries: 5
+
+# volumes:
+#   postgres_data:

requirements.txt

@@ -0,0 +1,8 @@
+pg8000
+pydantic
+pydantic-settings
+uvicorn
+fastapi
+openai
+pandas
+datasets

src/api/database.py

@@ -0,0 +1,596 @@
+# import logging
+# from typing import Dict, List, Optional, AsyncGenerator
+# from pydantic import BaseSettings, PostgresDsn
+# import pg8000
+# from pg8000 import Connection, Cursor
+# from pg8000.exceptions import DatabaseError
+# import asyncio
+# from contextlib import asynccontextmanager
+# from dataclasses import dataclass
+# from threading import Lock
+
+# # Set up structured logging
+# logging.basicConfig(
+#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+# )
+# logger = logging.getLogger(__name__)
+
+
+# class DatabaseSettings(BaseSettings):
+#     db_url: PostgresDsn
+#     pool_size: int = 5
+
+#     class Config:
+#         env_file = ".env"
+
+
+# @dataclass
+# class DatabaseConfig:
+#     username: str
+#     password: str
+#     hostname: str
+#     port: int
+#     database: str
+
+
+# class DatabaseError(Exception):
+#     """Custom exception for database errors."""
+
+#     pass
+
+
+# class Database:
+#     def __init__(self, db_url: str, pool_size: int):
+#         self.db_url = db_url
+#         self.pool_size = pool_size
+#         self.pool: List[Connection] = []
+#         self.lock = Lock()
+#         self.config = self._parse_db_url()
+
+#     def _parse_db_url(self) -> DatabaseConfig:
+#         """Parse the database URL into components."""
+#         result = urlparse(self.db_url)
+#         return DatabaseConfig(
+#             username=result.username,
+#             password=result.password,
+#             hostname=result.hostname,
+#             port=result.port or 5432,
+#             database=result.path.lstrip("/"),
+#         )
+
+#     async def connect(self) -> None:
+#         """Create a connection pool."""
+#         try:
+#             for _ in range(self.pool_size):
+#                 conn = await self._create_connection()
+#                 self.pool.append(conn)
+#             logger.info(
+#                 f"Database connection pool created with {self.pool_size} connections."
+#             )
+#         except DatabaseError as e:
+#             logger.error(f"Failed to create database connection pool: {e}")
+#             raise
+
+#     async def _create_connection(self) -> Connection:
+#         """Create a single database connection."""
+#         try:
+#             conn = pg8000.connect(
+#                 user=self.config.username,
+#                 password=self.config.password,
+#                 host=self.config.hostname,
+#                 port=self.config.port,
+#                 database=self.config.database,
+#             )
+#             return conn
+#         except DatabaseError as e:
+#             logger.error(f"Failed to create database connection: {e}")
+#             raise DatabaseError("Failed to create database connection.")
+
+#     async def disconnect(self) -> None:
+#         """Close all connections in the pool."""
+#         with self.lock:
+#             for conn in self.pool:
+#                 conn.close()
+#             self.pool.clear()
+#             logger.info("Database connection pool closed.")
+
+#     @asynccontextmanager
+#     async def get_connection(self) -> AsyncGenerator[Connection, None]:
+#         """Acquire a connection from the pool."""
+#         with self.lock:
+#             if not self.pool:
+#                 raise DatabaseError("Database connection pool is empty.")
+#             conn = self.pool.pop()
+#         try:
+#             yield conn
+#         finally:
+#             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.
+#         """
+#         try:
+#             async with self.get_connection() as conn:
+#                 cursor: 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 DatabaseError as e:
+#             logger.error(f"Error executing query: {query}. Error: {e}")
+#             raise DatabaseError(f"Failed to execute query: {query}")
+
+#     async def execute(self, query: str, *args) -> None:
+#         """
+#         Execute an INSERT, UPDATE, or DELETE query.
+
+#         Args:
+#             query (str): The SQL query to execute.
+#             *args: Query parameters.
+#         """
+#         try:
+#             async with self.get_connection() as conn:
+#                 cursor: Cursor = conn.cursor()
+#                 cursor.execute(query, args)
+#                 conn.commit()
+#         except DatabaseError as e:
+#             logger.error(f"Error executing query: {query}. Error: {e}")
+#             raise DatabaseError(f"Failed to execute query: {query}")
+
+
+# # Dependency to get the database instance
+# async def get_db() -> AsyncGenerator[Database, None]:
+#     settings = DatabaseSettings()
+#     db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+#     await db.connect()
+#     try:
+#         yield db
+#     finally:
+#         await db.disconnect()
+
+
+# # Example usage
+# if __name__ == "__main__":
+
+#     async def main():
+#         settings = DatabaseSettings()
+#         db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+#         await db.connect()
+
+#         try:
+#             # Example query
+#             query = """
+# SELECT
+#     ppt.type AS product_type,
+#     pc.name AS product_category
+# FROM
+#     product_producttype ppt
+# INNER JOIN
+#     product_category pc
+# ON
+#     ppt.category_id = pc.id
+# """
+#             result = await db.fetch(query)
+#             print(result)
+#         finally:
+#             await db.disconnect()
+
+#     asyncio.run(main())
+
+# import logging
+# from urllib.parse import urlparse
+# from typing import Dict, List, Optional, AsyncGenerator
+# from pydantic_settings import BaseSettings
+# from pydantic import PostgresDsn
+# import pg8000
+# from pg8000 import Connection, Cursor
+# from pg8000.exceptions import DatabaseError
+# import asyncio
+# from contextlib import asynccontextmanager
+# from dataclasses import dataclass
+# from threading import Lock
+
+# # Set up structured logging
+# logging.basicConfig(
+#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+# )
+# logger = logging.getLogger(__name__)
+
+
+# class DatabaseSettings(BaseSettings):
+#     db_url: PostgresDsn
+#     pool_size: int = 5
+
+#     class Config:
+#         env_file = ".env"
+
+
+# @dataclass
+# class DatabaseConfig:
+#     username: str
+#     password: str
+#     hostname: str
+#     port: int
+#     database: str
+
+
+# class DatabaseError(Exception):
+#     """Custom exception for database errors."""
+
+#     pass
+
+
+# class Database:
+#     def __init__(self, db_url: str, pool_size: int):
+#         self.db_url = db_url
+#         self.pool_size = pool_size
+#         self.pool: List[Connection] = []
+#         self.lock = Lock()
+#         self.config = self._parse_db_url()
+
+#     def _parse_db_url(self) -> DatabaseConfig:
+#         """Parse the database URL into components."""
+#         # Convert PostgresDsn to a string
+#         db_url_str = str(self.db_url)
+#         result = urlparse(db_url_str)
+#         return DatabaseConfig(
+#             username=result.username,
+#             password=result.password,
+#             hostname=result.hostname,
+#             port=result.port or 5432,
+#             database=result.path.lstrip("/"),
+#         )
+
+#     async def connect(self) -> None:
+#         """Create a connection pool."""
+#         try:
+#             for _ in range(self.pool_size):
+#                 conn = await self._create_connection()
+#                 self.pool.append(conn)
+#             logger.info(
+#                 f"Database connection pool created with {self.pool_size} connections."
+#             )
+#         except DatabaseError as e:
+#             logger.error(f"Failed to create database connection pool: {e}")
+#             raise
+
+#     async def _create_connection(self) -> Connection:
+#         """Create a single database connection."""
+#         try:
+#             conn = pg8000.connect(
+#                 user=self.config.username,
+#                 password=self.config.password,
+#                 host=self.config.hostname,
+#                 port=self.config.port,
+#                 database=self.config.database,
+#             )
+#             return conn
+#         except DatabaseError as e:
+#             logger.error(f"Failed to create database connection: {e}")
+#             raise DatabaseError("Failed to create database connection.")
+
+#     async def disconnect(self) -> None:
+#         """Close all connections in the pool."""
+#         with self.lock:
+#             for conn in self.pool:
+#                 conn.close()
+#             self.pool.clear()
+#             logger.info("Database connection pool closed.")
+
+#     @asynccontextmanager
+#     async def get_connection(self) -> AsyncGenerator[Connection, None]:
+#         """Acquire a connection from the pool."""
+#         with self.lock:
+#             if not self.pool:
+#                 raise DatabaseError("Database connection pool is empty.")
+#             conn = self.pool.pop()
+#         try:
+#             yield conn
+#         finally:
+#             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.
+#         """
+#         try:
+#             async with self.get_connection() as conn:
+#                 cursor: 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 DatabaseError as e:
+#             logger.error(f"Error executing query: {query}. Error: {e}")
+#             raise DatabaseError(f"Failed to execute query: {query}")
+
+#     async def execute(self, query: str, *args) -> None:
+#         """
+#         Execute an INSERT, UPDATE, or DELETE query.
+
+#         Args:
+#             query (str): The SQL query to execute.
+#             *args: Query parameters.
+#         """
+#         try:
+#             async with self.get_connection() as conn:
+#                 cursor: Cursor = conn.cursor()
+#                 cursor.execute(query, args)
+#                 conn.commit()
+#         except DatabaseError as e:
+#             logger.error(f"Error executing query: {query}. Error: {e}")
+#             raise DatabaseError(f"Failed to execute query: {query}")
+
+
+# # Dependency to get the database instance
+# async def get_db() -> AsyncGenerator[Database, None]:
+#     settings = DatabaseSettings()
+#     db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+#     await db.connect()
+#     try:
+#         yield db
+#     finally:
+#         await db.disconnect()
+
+
+# # Example usage
+# if __name__ == "__main__":
+
+#     async def main():
+#         settings = DatabaseSettings()
+#         db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+#         await db.connect()
+
+#         try:
+#             # Example query
+#             query = "SELECT * FROM your_table LIMIT 10"
+#             query = """
+# SELECT
+#     ppt.type AS product_type,
+#     pc.name AS product_category
+# FROM
+#     product_producttype ppt
+# INNER JOIN
+#     product_category pc
+# ON
+#     ppt.category_id = pc.id
+# """
+#             result = await db.fetch(query)
+#             print(result)
+#         finally:
+#             await db.disconnect()
+
+#     asyncio.run(main())
+
+import logging
+from typing import AsyncGenerator, List, Optional, Dict
+from pydantic_settings import BaseSettings
+from pydantic import PostgresDsn
+import pg8000
+from pg8000 import Connection
+from pg8000.exceptions import DatabaseError as Pg8000DatabaseError
+import asyncio
+from contextlib import asynccontextmanager
+from threading import Lock
+from urllib.parse import urlparse
+
+# Set up structured logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+class DatabaseSettings(BaseSettings):
+    db_url: PostgresDsn
+    pool_size: int = 5  # Default pool size is 5
+
+    class Config:
+        env_file = ".env"
+
+
+# Custom database errors
+class DatabaseError(Exception):
+    """Base exception for database errors."""
+
+    pass
+
+
+class ConnectionError(DatabaseError):
+    """Exception raised when a database connection fails."""
+
+    pass
+
+
+class PoolExhaustedError(DatabaseError):
+    """Exception raised when the connection pool is exhausted."""
+
+    pass
+
+
+class QueryExecutionError(DatabaseError):
+    """Exception raised when a query execution fails."""
+
+    pass
+
+
+class HealthCheckError(DatabaseError):
+    """Exception raised when a health check fails."""
+
+    pass
+
+
+class Database:
+    def __init__(self, db_url: PostgresDsn, pool_size: int):
+        self.db_url = db_url
+        self.pool_size = pool_size
+        self.pool: List[Connection] = []
+        self.lock = Lock()
+
+    async def connect(self) -> None:
+        """Create a connection pool."""
+        try:
+            # Convert PostgresDsn to a string
+            db_url_str = str(self.db_url)
+            result = urlparse(db_url_str)
+            for _ in range(self.pool_size):
+                conn = pg8000.connect(
+                    user=result.username,
+                    password=result.password,
+                    host=result.hostname,
+                    port=result.port or 5432,
+                    database=result.path.lstrip("/"),
+                )
+                self.pool.append(conn)
+            logger.info(
+                f"Database connection pool created with {self.pool_size} connections."
+            )
+        except Pg8000DatabaseError as e:
+            logger.error(f"Failed to create database connection pool: {e}")
+            raise ConnectionError("Failed to create database connection pool.") from e
+
+    async def disconnect(self) -> None:
+        """Close all connections in the pool."""
+        with self.lock:
+            for conn in self.pool:
+                conn.close()
+            self.pool.clear()
+            logger.info("Database connection pool closed.")
+
+    @asynccontextmanager
+    async def get_connection(self) -> AsyncGenerator[Connection, None]:
+        """Acquire a connection from the pool."""
+        with self.lock:
+            if not self.pool:
+                logger.error("Connection pool is exhausted.")
+                raise PoolExhaustedError("No available connections in the pool.")
+            conn = self.pool.pop()
+        try:
+            yield conn
+        except Pg8000DatabaseError as e:
+            logger.error(f"Connection error: {e}")
+            raise ConnectionError("Failed to use database connection.") from e
+        finally:
+            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 execute(self, query: str, *args) -> None:
+        """
+        Execute an INSERT, UPDATE, or DELETE query.
+
+        Args:
+            query (str): The SQL query to execute.
+            *args: Query parameters.
+
+        Raises:
+            QueryExecutionError: If the query execution fails.
+        """
+        try:
+            async with self.get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(query, args)
+                conn.commit()
+        except Pg8000DatabaseError as e:
+            logger.error(f"Query execution failed: {e}")
+            raise QueryExecutionError(f"Failed to execute query: {query}") from e
+
+    async def health_check(self) -> bool:
+        """
+        Perform a health check by executing a simple query (e.g., SELECT 1).
+
+        Returns:
+            bool: True if the database is healthy, False otherwise.
+
+        Raises:
+            HealthCheckError: If the health check fails.
+        """
+        try:
+            async with self.get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("SELECT 1")
+                result = cursor.fetchone()
+                cursor.close()
+
+                # Check if the result is as expected
+                if result and result[0] == 1:
+                    logger.info("Database health check succeeded.")
+                    return True
+                else:
+                    logger.error("Database health check failed: Unexpected result.")
+                    raise HealthCheckError("Unexpected result from health check query.")
+        except Pg8000DatabaseError as e:
+            logger.error(f"Health check failed: {e}")
+            raise HealthCheckError("Failed to perform health check.") from e
+
+
+# Dependency to get the database instance
+async def get_db() -> AsyncGenerator[Database, None]:
+    settings = DatabaseSettings()
+    db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+    await db.connect()
+    try:
+        yield db
+    finally:
+        await db.disconnect()
+
+
+# Example usage
+if __name__ == "__main__":
+
+    async def main():
+        settings = DatabaseSettings()
+        db = Database(db_url=settings.db_url, pool_size=settings.pool_size)
+        await db.connect()
+
+        try:
+            # Perform a health check
+            is_healthy = await db.health_check()
+            print(f"Database health check: {'Success' if is_healthy else 'Failure'}")
+        except HealthCheckError as e:
+            print(f"Health check failed: {e}")
+        finally:
+            await db.disconnect()
+
+    asyncio.run(main())

src/api/exceptions.py

@@ -0,0 +1,40 @@
+class DatabaseError(Exception):
+    """Base exception for database errors."""
+
+    pass
+
+
+class QueryExecutionError(DatabaseError):
+    """Exception raised when a database query fails."""
+
+    pass
+
+
+class EmbeddingError(Exception):
+    """Base exception for embedding-related errors."""
+
+    pass
+
+
+class OpenAIError(EmbeddingError):
+    """Exception raised when OpenAI API fails."""
+
+    pass
+
+
+class HuggingFaceError(Exception):
+    """Base exception for Hugging Face-related errors."""
+
+    pass
+
+
+class DatasetNotFoundError(HuggingFaceError):
+    """Exception raised when a dataset is not found."""
+
+    pass
+
+
+class DatasetPushError(HuggingFaceError):
+    """Exception raised when pushing a dataset to Hugging Face Hub fails."""
+
+    pass

src/api/models/embedding_models.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+from typing import List, Dict
+
+
+# Pydantic models for request validation
+class CreateEmbeddingRequest(BaseModel):
+    query: str
+    target_column: str = "product_type"
+    output_column: str = "embedding"
+    model: str = "text-embedding-3-small"
+    batch_size: int = 100
+    dataset_name: str = "re-mind/product_type_embedding"
+
+
+class UpdateEmbeddingRequest(BaseModel):
+    dataset_name: str
+    updates: Dict[str, List]  # Column name -> List of values
+
+
+class DeleteEmbeddingRequest(BaseModel):
+    dataset_name: str
+    columns: List[str]  # List of columns to delete

src/api/services/embedding_service.py

@@ -0,0 +1,62 @@
+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

src/api/services/huggingface_service.py

@@ -0,0 +1,69 @@
+from datasets import Dataset, load_dataset
+import logging
+from typing import Optional, Dict, List
+import pandas as pd
+from src.api.exceptions import DatasetNotFoundError, DatasetPushError
+
+# Set up structured logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceService:
+    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).remove_columns("__index_level_0__")
+            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_pandas()
+            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]
+    ) -> Optional[pd.DataFrame]:
+        """Update a dataset on Hugging Face Hub."""
+        try:
+            df = await self.read_dataset(dataset_name)
+            for column, values in updates.items():
+                if column in df.columns:
+                    df[column] = values
+                else:
+                    logger.warning(f"Column '{column}' not found in dataset.")
+            await self.push_to_hub(df, dataset_name)
+            return df
+        except Exception as e:
+            logger.error(f"Failed to update dataset: {e}")
+            raise DatasetPushError(f"Failed to update dataset: {e}")
+
+    async def delete_columns(
+        self, dataset_name: str, columns: List[str]
+    ) -> Optional[pd.DataFrame]:
+        """Delete columns from a dataset on Hugging Face Hub."""
+        try:
+            df = await self.read_dataset(dataset_name)
+            for column in columns:
+                if column in df.columns:
+                    df.drop(column, axis=1, inplace=True)
+                else:
+                    logger.warning(f"Column '{column}' not found in dataset.")
+            await self.push_to_hub(df, dataset_name)
+            return df
+        except Exception as e:
+            logger.error(f"Failed to delete columns: {e}")
+            raise DatasetPushError(f"Failed to delete columns: {e}")

src/main.py

@@ -0,0 +1,177 @@
+import os
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.responses import JSONResponse, RedirectResponse
+from pydantic import BaseModel
+from typing import List, Dict
+from src.api.models.embedding_models import (
+    CreateEmbeddingRequest,
+    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__)
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Similarity Search API",
+    description="A FastAPI application for similarity search with PostgreSQL and OpenAI embeddings.",
+    version="1.0.0",
+)
+
+
+# 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}")
+async def read_embeddings(
+    dataset_name: str,
+    huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+):
+    """
+    Read embeddings from a Hugging Face dataset.
+    """
+    try:
+        df = await huggingface_service.read_dataset(dataset_name)
+        return df.to_dict(orient="records")
+    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:
+        df = await huggingface_service.delete_columns(
+            request.dataset_name, request.columns
+        )
+        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}")