add new endpoint

src/api/database.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import AsyncGenerator, List, Optional, Dict
 from pydantic_settings import BaseSettings
 from pydantic import PostgresDsn
@@ -191,6 +192,15 @@ class Database:
             raise HealthCheckError("Failed to perform health check.") from e
 
 
+async def get_db_from_url() -> AsyncGenerator[Database, None]:
+    db = Database(db_url=os.getenv("DB_URL"), pool_size=5)
+    await db.connect()
+    try:
+        yield db
+    finally:
+        await db.disconnect()
+
+
 # Dependency to get the database instance
 async def get_db() -> AsyncGenerator[Database, None]:
     settings = DatabaseSettings()

src/api/models/embedding_models.py

@@ -65,3 +65,7 @@ class SearchEmbeddingRequest(BaseModel):
     additional_columns: Optional[List[str]] = (
         None  # Optional list of additional columns to include in the results
     )
+
+
+class ResetEmbeddingsRequest(BaseModel):
+    dataset_name: str

src/main.py

@@ -13,8 +13,9 @@ from src.api.models.embedding_models import (
     DeleteEmbeddingRequest,
     EmbedRequest,
     SearchEmbeddingRequest,
+    ResetEmbeddingsRequest,
 )
-from src.api.database import get_db, Database, QueryExecutionError, HealthCheckError
+from src.api.database import get_db, Database, QueryExecutionError, HealthCheckError, get_db_from_url
 from src.api.services.embedding_service import EmbeddingService
 from src.api.services.huggingface_service import HuggingFaceService
 from src.api.exceptions import DatasetNotFoundError, DatasetPushError, OpenAIError
@@ -46,12 +47,19 @@ app = FastAPI(
 
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
+# def get_database_service() -> Database:
+#     return Database(db_url=os.getenv("DB_URL"))
+
 
 # Dependency to get EmbeddingService
 def get_embedding_service() -> EmbeddingService:
     return EmbeddingService(openai_api_key=os.getenv("OPENAI_API_KEY"))
 
 
+# def get_db_from_env():
+#     return get_db_from_url(os.getenv("DB_URL"))
+
+
 # Dependency to get HuggingFaceService
 def get_huggingface_service() -> HuggingFaceService:
     return HuggingFaceService()
@@ -296,3 +304,21 @@ async def search_embedding(
     except Exception as e:
         logger.error(f"An error occurred: {e}")
         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
+@app.post("/reset_embeddings")
+async def reset_embeddings(
+    request: ResetEmbeddingsRequest,
+    db: Database = Depends(get_db_from_url)
+):
+    """
+    Reset embeddings from a Hugging Face dataset by deleting them, then reloading them
+    using the actual database
+    """
+    try:
+        is_healthy = 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))