OPIK-744 ignore default

comet-ml · Jan 15, 2025 · 000354e · 000354e
1 parent d18f06c
commit 000354e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderFactory.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderFactory.java
@@ -14,6 +14,7 @@
 import lombok.RequiredArgsConstructor;
 import org.apache.commons.lang3.EnumUtils;
 
+import java.util.Set;
 import java.util.function.Function;
 
 @Singleton
@@ -50,13 +51,13 @@ public ChatLanguageModel getLanguageModel(@NonNull String workspaceId,
      * The agreed requirement is to resolve the LLM provider and its API key based on the model.
      */
     private LlmProvider getLlmProvider(String model) {
-        if (isModelBelongToProvider(model, ModelPrice.class, ModelPrice::getName)) {
+        if (isModelBelongToProvider(model, ModelPrice.class, ModelPrice::getName, Set.of(ModelPrice.DEFAULT))) {
             return LlmProvider.OPEN_AI;
         }
-        if (isModelBelongToProvider(model, AnthropicModelName.class, AnthropicModelName::toString)) {
+        if (isModelBelongToProvider(model, AnthropicModelName.class, AnthropicModelName::toString, Set.of())) {
             return LlmProvider.ANTHROPIC;
         }
-        if (isModelBelongToProvider(model, GeminiModelName.class, GeminiModelName::toString)) {
+        if (isModelBelongToProvider(model, GeminiModelName.class, GeminiModelName::toString, Set.of())) {
             return LlmProvider.GEMINI;
         }
 
@@ -77,8 +78,9 @@ private String getEncryptedApiKey(String workspaceId, LlmProvider llmProvider) {
     }
 
     private static <E extends Enum<E>> boolean isModelBelongToProvider(
-            String model, Class<E> enumClass, Function<E, String> valueGetter) {
+            String model, Class<E> enumClass, Function<E, String> valueGetter, Set<Object> exclude) {
         return EnumUtils.getEnumList(enumClass).stream()
+                .filter(value -> !exclude.contains(value))
                 .map(valueGetter)
                 .anyMatch(model::equals);
     }

diff --git a/...opik-backend/src/test/java/com/comet/opik/domain/llmproviders/LlmProviderFactoryTest.java b/...opik-backend/src/test/java/com/comet/opik/domain/llmproviders/LlmProviderFactoryTest.java
@@ -79,6 +79,7 @@ void testGetService(String model, LlmProvider llmProvider, Class<? extends LlmPr
 
     private static Stream<Arguments> testGetService() {
         var openAiModels = EnumUtils.getEnumList(ModelPrice.class).stream()
+                .filter(value -> value != ModelPrice.DEFAULT)
                 .map(model -> arguments(model.getName(), LlmProvider.OPEN_AI, LlmProviderOpenAi.class));
         var anthropicModels = EnumUtils.getEnumList(AnthropicModelName.class).stream()
                 .map(model -> arguments(model.toString(), LlmProvider.ANTHROPIC, LlmProviderAnthropic.class));