From 209341e546c3712a160840c13754629ea124fd48 Mon Sep 17 00:00:00 2001
From: Emanuele Ferrelli <emanuele.ferrelli@accenture.com>
Date: Wed, 17 Dec 2025 10:25:07 +0100
Subject: [PATCH] Update

---
 .../hermione/config/ChromaWarmupService.java  | 91 ++++++++++++++-----
 1 file changed, 70 insertions(+), 21 deletions(-)

diff --git a/src/main/java/com/olympus/hermione/config/ChromaWarmupService.java b/src/main/java/com/olympus/hermione/config/ChromaWarmupService.java
index a59bc5d..b9e997f 100644
--- a/src/main/java/com/olympus/hermione/config/ChromaWarmupService.java
+++ b/src/main/java/com/olympus/hermione/config/ChromaWarmupService.java
@@ -5,6 +5,9 @@ import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.context.event.ApplicationReadyEvent;
 import org.springframework.context.event.EventListener;
+import org.springframework.http.HttpEntity;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
 import org.springframework.scheduling.annotation.Async;
 import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Component;
@@ -33,6 +36,9 @@ public class ChromaWarmupService {
     @Value("${spring.ai.vectorstore.chroma.collection-name:olympus_collection}")
     private String collectionName;
     
+    @Value("${spring.ai.vectorstore.chroma.client.key-token:}")
+    private String keyToken;
+    
     private final RestTemplate restTemplate;
     private volatile boolean chromaReady = false;
     
@@ -103,44 +109,87 @@ public class ChromaWarmupService {
     }
     
     /**
-     * Esegue un warmup completo: heartbeat + operazioni sulla collezione
+     * Esegue un warmup completo: operazioni sulla collezione (usa API v2)
      */
     private boolean performCompleteWarmup() {
         long startTime = System.currentTimeMillis();
-        boolean success = true;
+        boolean success = false;
         
         try {
-            // 1. Health check (updated to v2 API)
-            String healthUrl = chromaHost + ":" + chromaPort + "/api/v2/heartbeat";
-            restTemplate.getForObject(healthUrl, String.class);
-            logger.debug("Heartbeat successful");
+            HttpHeaders headers = createAuthHeaders();
+            HttpEntity<String> entity = new HttpEntity<>(headers);
             
-            // 2. Lista collezioni (operazione più pesante, updated to v2 API)
-            String collectionsUrl = chromaHost + ":" + chromaPort + "/api/v2/collections";
-            restTemplate.getForObject(collectionsUrl, String.class);
-            logger.debug("Collections list successful");
+            // Lista collezioni - prova diversi endpoint API v2
+            String[] collectionsEndpoints = {
+                "/api/v2/tenants/default_tenant/databases/default_database/collections",
+                "/api/v1/collections"  // fallback in caso v2 non funzioni
+            };
             
-            // 3. Prova a contare elementi nella collezione principale
-            try {
-                String countUrl = chromaHost + ":" + chromaPort + "/api/v2/collections/" + collectionName + "/count";
-                restTemplate.getForObject(countUrl, String.class);
-                logger.debug("Collection count successful");
-            } catch (Exception e) {
-                // La collezione potrebbe non esistere ancora
-                logger.debug("Collection count failed (collection may not exist): {}", e.getMessage());
+            for (String endpoint : collectionsEndpoints) {
+                try {
+                    String collectionsUrl = chromaHost + ":" + chromaPort + endpoint;
+                    restTemplate.exchange(collectionsUrl, HttpMethod.GET, entity, String.class);
+                    logger.debug("Collections list successful on endpoint: {}", endpoint);
+                    success = true;
+                    
+                    // Se ha funzionato, prova anche il count
+                    if (endpoint.contains("/tenants/")) {
+                        tryCountV2(entity);
+                    } else {
+                        tryCountV1(entity);
+                    }
+                    break;
+                } catch (Exception e) {
+                    logger.debug("Collections endpoint {} failed: {}", endpoint, e.getMessage());
+                }
             }
             
-            long duration = System.currentTimeMillis() - startTime;
-            logger.info("Chroma warmup completed successfully ({}ms)", duration);
+            if (!success) {
+                logger.warn("All collections endpoints failed");
+            } else {
+                long duration = System.currentTimeMillis() - startTime;
+                logger.info("Chroma warmup completed successfully ({}ms)", duration);
+            }
             
         } catch (Exception e) {
             logger.warn("Chroma warmup failed: {}", e.getMessage());
-            success = false;
         }
         
         return success;
     }
     
+    private void tryCountV2(HttpEntity<String> entity) {
+        try {
+            String countUrl = chromaHost + ":" + chromaPort + 
+                "/api/v2/tenants/default_tenant/databases/default_database/collections/" + collectionName + "/count";
+            restTemplate.exchange(countUrl, HttpMethod.GET, entity, String.class);
+            logger.debug("Collection count successful (v2)");
+        } catch (Exception e) {
+            logger.debug("Collection count failed (v2): {}", e.getMessage());
+        }
+    }
+    
+    private void tryCountV1(HttpEntity<String> entity) {
+        try {
+            String countUrl = chromaHost + ":" + chromaPort + "/api/v1/collections/" + collectionName + "/count";
+            restTemplate.exchange(countUrl, HttpMethod.GET, entity, String.class);
+            logger.debug("Collection count successful (v1)");
+        } catch (Exception e) {
+            logger.debug("Collection count failed (v1): {}", e.getMessage());
+        }
+    }
+    
+    /**
+     * Crea gli headers HTTP con autenticazione per Chroma
+     */
+    private HttpHeaders createAuthHeaders() {
+        HttpHeaders headers = new HttpHeaders();
+        if (keyToken != null && !keyToken.isEmpty()) {
+            headers.set("Authorization", "Bearer " + keyToken);
+        }
+        return headers;
+    }
+    
     public boolean isChromaReady() {
         return chromaReady;
     }