Merged PR 221: Update
Update
This commit is contained in:
@@ -5,6 +5,9 @@ import org.slf4j.LoggerFactory;
|
|||||||
import org.springframework.beans.factory.annotation.Value;
|
import org.springframework.beans.factory.annotation.Value;
|
||||||
import org.springframework.boot.context.event.ApplicationReadyEvent;
|
import org.springframework.boot.context.event.ApplicationReadyEvent;
|
||||||
import org.springframework.context.event.EventListener;
|
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.Async;
|
||||||
import org.springframework.scheduling.annotation.Scheduled;
|
import org.springframework.scheduling.annotation.Scheduled;
|
||||||
import org.springframework.stereotype.Component;
|
import org.springframework.stereotype.Component;
|
||||||
@@ -33,6 +36,9 @@ public class ChromaWarmupService {
|
|||||||
@Value("${spring.ai.vectorstore.chroma.collection-name:olympus_collection}")
|
@Value("${spring.ai.vectorstore.chroma.collection-name:olympus_collection}")
|
||||||
private String collectionName;
|
private String collectionName;
|
||||||
|
|
||||||
|
@Value("${spring.ai.vectorstore.chroma.client.key-token:}")
|
||||||
|
private String keyToken;
|
||||||
|
|
||||||
private final RestTemplate restTemplate;
|
private final RestTemplate restTemplate;
|
||||||
private volatile boolean chromaReady = false;
|
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() {
|
private boolean performCompleteWarmup() {
|
||||||
long startTime = System.currentTimeMillis();
|
long startTime = System.currentTimeMillis();
|
||||||
boolean success = true;
|
boolean success = false;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
// 1. Health check (updated to v2 API)
|
HttpHeaders headers = createAuthHeaders();
|
||||||
String healthUrl = chromaHost + ":" + chromaPort + "/api/v2/heartbeat";
|
HttpEntity<String> entity = new HttpEntity<>(headers);
|
||||||
restTemplate.getForObject(healthUrl, String.class);
|
|
||||||
logger.debug("Heartbeat successful");
|
|
||||||
|
|
||||||
// 2. Lista collezioni (operazione più pesante, updated to v2 API)
|
// Lista collezioni - prova diversi endpoint API v2
|
||||||
String collectionsUrl = chromaHost + ":" + chromaPort + "/api/v2/collections";
|
String[] collectionsEndpoints = {
|
||||||
restTemplate.getForObject(collectionsUrl, String.class);
|
"/api/v2/tenants/default_tenant/databases/default_database/collections",
|
||||||
logger.debug("Collections list successful");
|
"/api/v1/collections" // fallback in caso v2 non funzioni
|
||||||
|
};
|
||||||
|
|
||||||
// 3. Prova a contare elementi nella collezione principale
|
for (String endpoint : collectionsEndpoints) {
|
||||||
try {
|
try {
|
||||||
String countUrl = chromaHost + ":" + chromaPort + "/api/v2/collections/" + collectionName + "/count";
|
String collectionsUrl = chromaHost + ":" + chromaPort + endpoint;
|
||||||
restTemplate.getForObject(countUrl, String.class);
|
restTemplate.exchange(collectionsUrl, HttpMethod.GET, entity, String.class);
|
||||||
logger.debug("Collection count successful");
|
logger.debug("Collections list successful on endpoint: {}", endpoint);
|
||||||
} catch (Exception e) {
|
success = true;
|
||||||
// La collezione potrebbe non esistere ancora
|
|
||||||
logger.debug("Collection count failed (collection may not exist): {}", e.getMessage());
|
// 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;
|
if (!success) {
|
||||||
logger.info("Chroma warmup completed successfully ({}ms)", duration);
|
logger.warn("All collections endpoints failed");
|
||||||
|
} else {
|
||||||
|
long duration = System.currentTimeMillis() - startTime;
|
||||||
|
logger.info("Chroma warmup completed successfully ({}ms)", duration);
|
||||||
|
}
|
||||||
|
|
||||||
} catch (Exception e) {
|
} catch (Exception e) {
|
||||||
logger.warn("Chroma warmup failed: {}", e.getMessage());
|
logger.warn("Chroma warmup failed: {}", e.getMessage());
|
||||||
success = false;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return success;
|
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() {
|
public boolean isChromaReady() {
|
||||||
return chromaReady;
|
return chromaReady;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user