reset embeddings effective

src/api/models/embedding_models.py

@@ -68,4 +68,4 @@ class SearchEmbeddingRequest(BaseModel):
 
 
 class ResetEmbeddingsRequest(BaseModel):
-    dataset_name: str
+    dataset_name: str = "re-mind/product_type_embedding"

src/main.py

@@ -2,6 +2,7 @@ import os
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.responses import JSONResponse, RedirectResponse
 from fastapi.middleware.gzip import GZipMiddleware
+from pg8000 import DatabaseError
 from pydantic import BaseModel
 from typing import List, Dict
 from datasets import Dataset
@@ -18,6 +19,7 @@ from src.api.models.embedding_models import (
 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.services.postgresql_service import PostgresqlService
 from src.api.exceptions import DatasetNotFoundError, DatasetPushError, OpenAIError
 
 import logging
@@ -309,16 +311,32 @@ async def search_embedding(
 @app.post("/reset_embeddings")
 async def reset_embeddings(
     request: ResetEmbeddingsRequest,
-    db: Database = Depends(get_db_from_url)
+    db: Database = Depends(get_db_from_url),
+    embedding_service: EmbeddingService = Depends(get_embedding_service),
+    huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
 ):
     """
     Reset embeddings from a Hugging Face dataset by deleting them, then reloading them
     using the actual database
     """
+    postgresql_service = PostgresqlService(db)
+
     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:
+        # List of rows from database
+        results = await postgresql_service.get_db_rows_from_dataset_name(request.dataset_name)
+
+        # Generation of embeddings for each row
+        dataset = Dataset.from_dict(results)
+        target_column = "type" if request.dataset_name == "re-mind/product_type_embedding" else "name"
+        dataset_embedded = await embedding_service.create_embeddings(dataset, target_column, "embedding")
+        # Embeddings up-to-date with database will overwrite old dataset
+        await huggingface_service.push_to_hub(dataset_embedded, request.dataset_name)
+
+        return {
+            "message": "Dataset updated succesfully with up-to-date rows from database",
+            "dataset_name": request.dataset_name,
+            "num_rows": len(dataset_embedded)
+        }
+
+    except DatabaseError as e:
         raise HTTPException(status_code=500, detail=str(e))