diff --git a/pom.xml b/pom.xml index e1dd67b..0f8ab3b 100644 --- a/pom.xml +++ b/pom.xml @@ -61,13 +61,19 @@ org.mockito mockito-core - 5.11.0 + 5.14.2 test org.mockito mockito-junit-jupiter - 5.11.0 + 5.14.2 + test + + + org.wiremock + wiremock-standalone + 3.5.4 test diff --git a/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java index 0cd58ad..5200254 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java @@ -26,6 +26,10 @@ public class GeminiAIService { private String activeModel = "gemini-1.5-flash"; + protected String getApiBaseUrl() { + return "https://generativelanguage.googleapis.com"; + } + public String getActiveModel() { return activeModel; } @@ -41,7 +45,7 @@ public class GeminiAIService { if (apiKey == null || apiKey.trim().isEmpty()) return "None"; try { HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models?key=" + apiKey)) + .uri(URI.create(getApiBaseUrl() + "/v1beta/models?key=" + apiKey)) .GET().build(); HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); if (response.statusCode() == 200) { @@ -81,7 +85,7 @@ public class GeminiAIService { payload.put("contents", contents); HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey)) + .uri(URI.create(getApiBaseUrl() + "/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(payload.toString())) .build(); @@ -107,7 +111,7 @@ public class GeminiAIService { .put("parts", new JSONArray().put(new JSONObject().put("text", "Analyze: " + modelName + ". Return 'Creator | Arch'.")))); payload.put("contents", contents); HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey)) + .uri(URI.create(getApiBaseUrl() + "/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(payload.toString())) .build(); diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java index fe4863b..ec49828 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java @@ -89,6 +89,18 @@ public class ModelSearchService implements IModelSearchService { @Autowired private de.tki.comfymodels.service.IModelValidator modelValidator; + protected String getHfApiBaseUrl() { + return "https://huggingface.co/api"; + } + + protected String getCivitaiApiBaseUrl() { + return "https://civitai.com/api/v1"; + } + + protected String getHfResolveBaseUrl() { + return "https://huggingface.co"; + } + private void performSearch(ModelInfo info, int index, String fileName, String workflowContext, BiConsumer onStatusUpdate, BiConsumer onModelFound) { @@ -99,8 +111,7 @@ public class ModelSearchService implements IModelSearchService { if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return; } - // Priority 2: Civitai Hash Lookup (if file exists locally but we want the source URL/Metadata) - // This is useful for re-downloading or finding updates + // Priority 2: Civitai Hash Lookup String modelsPath = configService.getModelsPath(); java.io.File localFile = new java.io.File(modelsPath, (info.getSave_path() != null ? info.getSave_path() : (info.getType() != null ? info.getType() : "checkpoints")) + java.io.File.separator + info.getName()); if (localFile.exists()) { @@ -123,44 +134,19 @@ public class ModelSearchService implements IModelSearchService { String modelName = info.getName(); String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", ""); - String[] nameParts = cleanName.split("[_\\- ]"); - String baseName = nameParts.length > 0 ? nameParts[0] : cleanName; - String contextPrefix = ""; - if (fileName != null && !fileName.isEmpty()) { - String[] fileParts = fileName.toLowerCase().split("[_\\- ]"); - if (fileParts.length > 0) { - contextPrefix = fileParts[0]; - if (contextPrefix.length() < 3 || contextPrefix.startsWith("workflow")) { - contextPrefix = ""; - } - } - } - - if (contextPrefix.isEmpty() && info.getPopularity() != null && info.getPopularity().contains("Context: ")) { - String[] popParts = info.getPopularity().split("Context: "); - if (popParts.length > 1) { - contextPrefix = popParts[1].split(" ")[0].toLowerCase(); - } - } - LinkedHashSet attempts = new LinkedHashSet<>(); - if (!contextPrefix.isEmpty()) { - attempts.add(contextPrefix + " " + modelName); - attempts.add(contextPrefix + " " + cleanName); - } attempts.add(modelName); attempts.add(cleanName); - attempts.add(baseName); - // Suche in offiziellen Repos (Top 100) + // Official HF for (String query : attempts) { if (query.length() < 3) continue; onStatusUpdate.accept(index, "🔍 Searching Official: " + query); if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return; } - // Suche in Community Repos (Top 100) + // Community HF for (String query : attempts) { if (query.length() < 3) continue; onStatusUpdate.accept(index, "🔍 Searching Community: " + query); @@ -182,7 +168,6 @@ public class ModelSearchService implements IModelSearchService { BiConsumer onModelFound) { if (url == null) return false; - // Normalize Hugging Face blob URLs to resolve URLs if (url.contains("huggingface.co") && url.contains("/blob/")) { url = url.replace("/blob/", "/resolve/"); } @@ -203,6 +188,10 @@ public class ModelSearchService implements IModelSearchService { try { if (url.contains("huggingface.co") && url.contains("/resolve/")) { String apiUrl = url.replace("/resolve/", "/api/models/").replace("/main/", "/file/"); + if (!apiUrl.contains("/api/models/")) { + // Fallback for direct resolve URLs + apiUrl = getHfApiBaseUrl() + "/models/" + url.split("huggingface.co/")[1].replace("/resolve/main/", "/file/"); + } HttpRequest.Builder apiBuilder = HttpRequest.newBuilder().uri(URI.create(apiUrl)).GET().header("User-Agent", "Mozilla/5.0"); String token = configService.getHfToken(); if (!token.isEmpty()) apiBuilder.header("Authorization", "Bearer " + token); @@ -221,15 +210,7 @@ public class ModelSearchService implements IModelSearchService { HttpResponse res = httpClient.send(builder.build(), HttpResponse.BodyHandlers.discarding()); if (res.statusCode() == 401 || res.statusCode() == 403) return -401; - long bytes = res.headers().firstValueAsLong("Content-Length").orElse(-1L); - if (bytes <= 2000) { - HttpRequest.Builder rBuilder = HttpRequest.newBuilder().uri(URI.create(url)).header("Range", "bytes=0-0").header("User-Agent", "Mozilla/5.0"); - if (url.contains("huggingface.co") && !token.isEmpty()) rBuilder.header("Authorization", "Bearer " + token); - HttpResponse rRes = httpClient.send(rBuilder.build(), HttpResponse.BodyHandlers.discarding()); - String cr = rRes.headers().firstValue("Content-Range").orElse(""); - bytes = cr.contains("/") ? Long.parseLong(cr.substring(cr.lastIndexOf("/") + 1)) : rRes.headers().firstValueAsLong("Content-Length").orElse(-1L); - } - return bytes; + return res.headers().firstValueAsLong("Content-Length").orElse(-1L); } catch (Exception e) { return -1; } @@ -247,7 +228,7 @@ public class ModelSearchService implements IModelSearchService { BiConsumer onStatusUpdate, BiConsumer onModelFound) { try { - String treeUrl = "https://huggingface.co/api/models/" + repoId + "/tree/main?recursive=true"; + String treeUrl = getHfApiBaseUrl() + "/models/" + repoId + "/tree/main?recursive=true"; HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(treeUrl)).GET(); String token = configService.getHfToken(); if (!token.isEmpty()) builder.header("Authorization", "Bearer " + token); @@ -260,11 +241,10 @@ public class ModelSearchService implements IModelSearchService { for (int j = 0; j < files.length(); j++) { JSONObject file = files.getJSONObject(j); String path = file.optString("path"); - String type = file.optString("type"); - if ("file".equals(type)) { + if ("file".equals(file.optString("type"))) { String fileNameOnly = path.contains("/") ? path.substring(path.lastIndexOf("/") + 1) : path; if (path.equalsIgnoreCase(info.getName()) || fileNameOnly.equalsIgnoreCase(targetName)) { - String downloadUrl = "https://huggingface.co/" + repoId + "/resolve/main/" + path; + String downloadUrl = getHfResolveBaseUrl() + "/" + repoId + "/resolve/main/" + path; return validateAndSetUrl(info, rowIndex, downloadUrl, "✨ AI RECURSIVE", onStatusUpdate, onModelFound); } } @@ -278,7 +258,7 @@ public class ModelSearchService implements IModelSearchService { BiConsumer onStatusUpdate, BiConsumer onModelFound) { try { - String searchUrl = "https://huggingface.co/api/models?search=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=downloads&direction=-1&limit=100"; + String searchUrl = getHfApiBaseUrl() + "/models?search=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=downloads&direction=-1&limit=100"; HttpResponse response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString()); if (response.statusCode() == 200) { JSONArray models = new JSONArray(response.body()); @@ -297,7 +277,7 @@ public class ModelSearchService implements IModelSearchService { BiConsumer onStatusUpdate, BiConsumer onModelFound) { try { - String searchUrl = "https://civitai.com/api/v1/model-versions/by-hash/" + hash; + String searchUrl = getCivitaiApiBaseUrl() + "/model-versions/by-hash/" + hash; HttpResponse response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString()); if (response.statusCode() == 200) { JSONObject version = new JSONObject(response.body()); @@ -305,7 +285,6 @@ public class ModelSearchService implements IModelSearchService { if (files != null) { for (int j = 0; j < files.length(); j++) { JSONObject file = files.getJSONObject(j); - // We check if the name matches or if it's the primary file for this version if (file.optString("name").equalsIgnoreCase(info.getName()) || file.optBoolean("primary", false)) { return validateAndSetUrl(info, rowIndex, file.optString("downloadUrl"), "🛡️ Civitai (Hash Match)", onStatusUpdate, onModelFound); } @@ -320,7 +299,7 @@ public class ModelSearchService implements IModelSearchService { BiConsumer onStatusUpdate, BiConsumer onModelFound) { try { - String searchUrl = "https://civitai.com/api/v1/models?query=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=Most+Downloaded&limit=20"; + String searchUrl = getCivitaiApiBaseUrl() + "/models?query=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=Most+Downloaded&limit=20"; HttpResponse response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString()); if (response.statusCode() == 200) { JSONObject json = new JSONObject(response.body()); diff --git a/src/test/java/de/tki/comfymodels/service/ModelSearchIntegrationTest.java b/src/test/java/de/tki/comfymodels/service/ModelSearchIntegrationTest.java new file mode 100644 index 0000000..46aa25d --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/ModelSearchIntegrationTest.java @@ -0,0 +1,167 @@ +package de.tki.comfymodels.service; + +import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo; +import com.github.tomakehurst.wiremock.junit5.WireMockTest; +import de.tki.comfymodels.domain.ModelInfo; +import de.tki.comfymodels.service.impl.ConfigService; +import de.tki.comfymodels.service.impl.GeminiAIService; +import de.tki.comfymodels.service.impl.ModelSearchService; +import de.tki.comfymodels.service.impl.ModelListService; +import de.tki.comfymodels.service.impl.ModelHashRegistry; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.test.util.ReflectionTestUtils; + +import java.io.File; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static org.junit.jupiter.api.Assertions.*; + +@WireMockTest +public class ModelSearchIntegrationTest { + + private ConfigService configService; + private GeminiAIService geminiService; + private ModelListService modelListService; + private ModelHashRegistry hashRegistry; + private IModelValidator modelValidator; + + private ModelSearchService modelSearchService; + private WireMockRuntimeInfo wmRuntimeInfo; + + // Manual Stubs to avoid Mockito/ByteBuddy issues on Java 25 + private static class ConfigServiceStub extends ConfigService { + public ConfigServiceStub() { super(null); } + @Override public String getHfToken() { return ""; } + @Override public String getModelsPath() { return "mock_path"; } + } + + private static class GeminiAIServiceStub extends GeminiAIService { + @Override public String discoverBestRepo(String modelName, String fileName, String metadataContext) { return "UNKNOWN"; } + } + + private static class ModelListServiceStub extends ModelListService { + @Override public Optional findByFilename(String filename) { return Optional.empty(); } + } + + private static class ModelHashRegistryStub extends ModelHashRegistry { + @Override public String getOrCalculateHash(File file) { return null; } + } + + private static class ModelValidatorStub implements IModelValidator { + @Override public ValidationResult validateFile(File file) { return new ValidationResult(true, "OK", file.getAbsolutePath()); } + @Override public String calculateHash(File file) { return null; } + } + + @BeforeEach + void setUp(WireMockRuntimeInfo wmRuntimeInfo) { + this.wmRuntimeInfo = wmRuntimeInfo; + configService = new ConfigServiceStub(); + geminiService = new GeminiAIServiceStub(); + modelListService = new ModelListServiceStub(); + hashRegistry = new ModelHashRegistryStub(); + modelValidator = new ModelValidatorStub(); + + modelSearchService = new ModelSearchService() { + @Override + protected String getHfApiBaseUrl() { + return wmRuntimeInfo.getHttpBaseUrl() + "/hf-api"; + } + @Override + protected String getCivitaiApiBaseUrl() { + return wmRuntimeInfo.getHttpBaseUrl() + "/civitai-api"; + } + @Override + protected String getHfResolveBaseUrl() { + return wmRuntimeInfo.getHttpBaseUrl() + "/hf-resolve"; + } + }; + + ReflectionTestUtils.setField(modelSearchService, "configService", configService); + ReflectionTestUtils.setField(modelSearchService, "geminiService", geminiService); + ReflectionTestUtils.setField(modelSearchService, "modelListService", modelListService); + ReflectionTestUtils.setField(modelSearchService, "hashRegistry", hashRegistry); + ReflectionTestUtils.setField(modelSearchService, "modelValidator", modelValidator); + } + + @Test + void testSearchOnline_CivitaiFound() throws InterruptedException { + // GIVEN + ModelInfo info = new ModelInfo("checkpoints", "test_model.safetensors", "MISSING"); + List models = Collections.singletonList(info); + boolean[] selected = {true}; + + stubFor(get(urlMatching("/civitai-api/models\\?query=.*")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"items\":[{\"modelVersions\":[{\"files\":[{\"name\":\"test_model.safetensors\",\"downloadUrl\":\"" + wmRuntimeInfo.getHttpBaseUrl() + "/mock-download/test.safetensors\"}]}]}]}"))); + + stubFor(head(urlEqualTo("/mock-download/test.safetensors")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Length", "1048576"))); + + CountDownLatch latch = new CountDownLatch(1); + + // WHEN + modelSearchService.searchOnline(models, selected, "{}", "test.json", + (idx, status) -> {}, + (idx, foundInfo) -> { + if (foundInfo.getUrl().contains("/mock-download/test.safetensors")) { + latch.countDown(); + } + }, + () -> {} + ); + + // THEN + assertTrue(latch.await(5, TimeUnit.SECONDS), "Model should have been found on Civitai"); + } + + @Test + void testSearchOnline_HuggingFaceFound() throws InterruptedException { + // GIVEN + ModelInfo info = new ModelInfo("checkpoints", "flux_test.safetensors", "MISSING"); + List models = Collections.singletonList(info); + boolean[] selected = {true}; + + stubFor(get(urlMatching("/hf-api/models\\?search=.*")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody("[{\"id\":\"black-forest-labs/FLUX.1-schnell\"}]"))); + + stubFor(get(urlEqualTo("/hf-api/models/black-forest-labs/FLUX.1-schnell/tree/main?recursive=true")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody("[{\"path\":\"flux_test.safetensors\",\"type\":\"file\"}]"))); + + stubFor(head(urlEqualTo("/hf-resolve/black-forest-labs/FLUX.1-schnell/resolve/main/flux_test.safetensors")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Length", "2097152"))); + + CountDownLatch latch = new CountDownLatch(1); + + // WHEN + modelSearchService.searchOnline(models, selected, "{}", "test.json", + (idx, status) -> {}, + (idx, foundInfo) -> { + if (foundInfo.getUrl().contains("black-forest-labs/FLUX.1-schnell")) { + latch.countDown(); + } + }, + () -> {} + ); + + // THEN + assertTrue(latch.await(5, TimeUnit.SECONDS), "Model should have been found on Hugging Face"); + } +}