This commit is contained in:
2025-12-17 10:25:07 +01:00
parent 09113ee3f3
commit 209341e546

View File

@@ -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;
}