Update
This commit is contained in:
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user