Skip to content

feat(module): load google module #28

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mohamedalibarkailluin wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat(module): load google module #28

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@

<properties>
<revision>0.0.0-SNAPSHOT</revision>
<!-- TODO : Update once google module is integrated in wombat-core-->
<wombat-core.version>0.12.0</wombat-core.version>
<maven.deploy.skip>true</maven.deploy.skip>
<compiler-plugin.version>3.16.0</compiler-plugin.version>
Expand Down
3 changes: 3 additions & 0 deletions src/main/java/tech/illuin/wombat/engine/module/extension/CoreExtensionLoader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import tech.illuin.wombat.core.module.WombatModule;
import tech.illuin.wombat.module.kubernetes_api.KubernetesAPIModule;
import tech.illuin.wombat.module.kubernetes_simulated.KubernetesSimulatedModule;
import tech.illuin.wombat.module.llm_google.LLMGoogleModule;
import tech.illuin.wombat.module.llm_prometheus.LLMPrometheusModule;
import tech.illuin.wombat.module.llm_simulated.LLMSimulatedModule;
import tech.illuin.wombat.module.llm_static.LLMStaticModule;
Expand Down Expand Up @@ -33,6 +34,8 @@ public List<WombatModule> load()
modules.add(new LLMStaticModule());
if (this.properties.kubernetesApi().enabled())
modules.add(new KubernetesAPIModule());
if (this.properties.llmGoogle().enabled())
modules.add(new LLMGoogleModule());

return modules;
}
Expand Down
9 changes: 9 additions & 0 deletions src/main/java/tech/illuin/wombat/engine/module/extension/WombatModuleProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ public interface WombatModuleProperties
@WithName("llm-static")
LlmStaticProperties llmStatic();

@WithName("llm-google")
LlmGoogleProperties llmGoogle();

@WithName("extension")
ExtensionProperties extension();

Expand All @@ -39,6 +42,12 @@ interface LlmStaticProperties
boolean enabled();
}

interface LlmGoogleProperties
{
@WithDefault("true")
boolean enabled();
}

interface ExtensionProperties
{
@WithName("path")
Expand Down
22 changes: 22 additions & 0 deletions src/main/resources/monitored/monitored-environments-dev.yaml.dist
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,25 @@ environments:
# location: FRA
# dynamic-profile:
# query: 'sum(increase(generation_tokens_total{model_name="mistral-large-latest"}[10s]))'
# - type: tech.illuin.wombat-module.llm-google
# id: my-gemini-llm
# environment-id: preprod
# name: My Gemini LLM (preprod)
# google-url: https://monitoring.googleapis.com
# project-id: my-gcp-project-id
# heartbeat-skip: 1
# # Authentication is optional. Without it, Application Default Credentials are used. `credentials-key` names
# # the environment variable holding the service-account key JSON. The service account needs the
# # `roles/monitoring.viewer` role.
# credentials-key: WOMBAT_GOOGLE_CREDENTIALS
# profile:
# provider: google_genai
# model: gemini-2.5-flash
# location: USA
# dynamic-profile:
# filter: >-
# metric.type="aiplatform.googleapis.com/publisher/online_serving/token_count"
# AND metric.labels.type="output"
# AND resource.labels.model_user_id="gemini-2.5-flash"
# # Optional, ISO-8601, defaults to PT1M. Keep it equal to the sampling period.
# window: PT1M
16 changes: 12 additions & 4 deletions src/test/java/tech/illuin/wombat/engine/module/WombatModuleConfigTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import tech.illuin.wombat.engine.module.extension.ExtensionLoader;
import tech.illuin.wombat.engine.module.extension.WombatModuleProperties;
import tech.illuin.wombat.module.kubernetes_api.KubernetesAPIModule;
import tech.illuin.wombat.module.llm_google.LLMGoogleModule;
import tech.illuin.wombat.module.llm_prometheus.LLMPrometheusModule;
import tech.illuin.wombat.module.llm_static.LLMStaticModule;

Expand All @@ -39,7 +40,7 @@ class WombatModuleConfigTest
@Test
void provideWombatExtensionLoader_returnsCompositeExtensionLoader()
{
WombatModuleProperties properties = properties(true, true, true);
WombatModuleProperties properties = properties(true, true, true, true);
ExtensionLoader loader = this.config.provideWombatExtensionLoader(properties);

assertInstanceOf(CompositeExtensionLoader.class, loader);
Expand All @@ -49,7 +50,7 @@ void provideWombatExtensionLoader_returnsCompositeExtensionLoader()
void provideWombatModules_loadsModulesFromExtensionLoader()
{
ExtensionLoader loader = mock(ExtensionLoader.class);
List<WombatModule> expected = List.of(new KubernetesAPIModule(), new LLMStaticModule());
List<WombatModule> expected = List.of(new KubernetesAPIModule(), new LLMStaticModule(), new LLMGoogleModule());
when(loader.load()).thenReturn(expected);

List<WombatModule> result = this.config.provideWombatModules(loader);
Expand All @@ -64,7 +65,8 @@ void provideMappers_registerSubtypesForGivenModules()
List<WombatModule> modules = List.of(
new KubernetesAPIModule(),
new LLMPrometheusModule(),
new LLMStaticModule()
new LLMStaticModule(),
new LLMGoogleModule()
);

JsonMapper jsonMapper = this.config.provideJsonMapper(modules);
Expand Down Expand Up @@ -96,7 +98,7 @@ void provideDefaults_createsHandlersForFamilies()
assertNotNull(llmHandlers.costResolver());
}

private static WombatModuleProperties properties(boolean kubernetesApi, boolean llmPrometheus, boolean llmStatic)
private static WombatModuleProperties properties(boolean kubernetesApi, boolean llmPrometheus, boolean llmStatic, boolean llmGoogle)
{
return new WombatModuleProperties()
{
Expand All @@ -118,6 +120,12 @@ public LlmStaticProperties llmStatic()
return () -> llmStatic;
}

@Override
public LlmGoogleProperties llmGoogle()
{
return () -> llmGoogle;
}

@Override
public ExtensionProperties extension()
{
Expand Down
41 changes: 36 additions & 5 deletions src/test/java/tech/illuin/wombat/engine/module/extension/CoreExtensionLoaderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import tech.illuin.wombat.engine.module.extension.WombatModuleProperties;
import tech.illuin.wombat.module.kubernetes_api.KubernetesAPIModule;
import tech.illuin.wombat.module.kubernetes_simulated.KubernetesSimulatedModule;
import tech.illuin.wombat.module.llm_google.LLMGoogleModule;
import tech.illuin.wombat.module.llm_prometheus.LLMPrometheusModule;
import tech.illuin.wombat.module.llm_simulated.LLMSimulatedModule;
import tech.illuin.wombat.module.llm_static.LLMStaticModule;
Expand All @@ -21,7 +22,7 @@ class CoreExtensionLoaderTest
@Test
void alwaysLoadsSimulatedModulesEvenWhenAllCoreModulesDisabled()
{
WombatModuleProperties properties = properties(false, false, false);
WombatModuleProperties properties = properties(false, false, false, false);
try (CoreExtensionLoader loader = new CoreExtensionLoader(properties))
{
List<WombatModule> modules = loader.load();
Expand All @@ -32,6 +33,7 @@ void alwaysLoadsSimulatedModulesEvenWhenAllCoreModulesDisabled()
assertFalse(modules.stream().anyMatch(KubernetesAPIModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMPrometheusModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMStaticModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMGoogleModule.class::isInstance));
}
catch (IOException e) {
fail(e);
Expand All @@ -41,17 +43,18 @@ void alwaysLoadsSimulatedModulesEvenWhenAllCoreModulesDisabled()
@Test
void loadsAllModulesWhenAllTogglesEnabled()
{
WombatModuleProperties properties = properties(true, true, true);
WombatModuleProperties properties = properties(true, true, true, true);
try (CoreExtensionLoader loader = new CoreExtensionLoader(properties))
{
List<WombatModule> modules = loader.load();

assertEquals(5, modules.size());
assertEquals(6, modules.size());
assertTrue(modules.stream().anyMatch(KubernetesSimulatedModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMSimulatedModule.class::isInstance));
assertTrue(modules.stream().anyMatch(KubernetesAPIModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMPrometheusModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMStaticModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMGoogleModule.class::isInstance));
}
catch (IOException e) {
fail(e);
Expand All @@ -61,7 +64,7 @@ void loadsAllModulesWhenAllTogglesEnabled()
@Test
void loadsOnlySpecificallyEnabledModules()
{
WombatModuleProperties properties = properties(true, false, false);
WombatModuleProperties properties = properties(true, false, false, false);
try (CoreExtensionLoader loader = new CoreExtensionLoader(properties))
{
List<WombatModule> modules = loader.load();
Expand All @@ -72,13 +75,35 @@ void loadsOnlySpecificallyEnabledModules()
assertTrue(modules.stream().anyMatch(KubernetesAPIModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMPrometheusModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMStaticModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMGoogleModule.class::isInstance));
}
catch (IOException e) {
fail(e);
}
}

private static WombatModuleProperties properties(boolean kubernetesApi, boolean llmPrometheus, boolean llmStatic)
@Test
void loadsGoogleModuleWhenOnlyGoogleToggleEnabled()
{
WombatModuleProperties properties = properties(false, false, false, true);
try (CoreExtensionLoader loader = new CoreExtensionLoader(properties))
{
List<WombatModule> modules = loader.load();

assertEquals(3, modules.size());
assertTrue(modules.stream().anyMatch(KubernetesSimulatedModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMSimulatedModule.class::isInstance));
assertTrue(modules.stream().anyMatch(LLMGoogleModule.class::isInstance));
assertFalse(modules.stream().anyMatch(KubernetesAPIModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMPrometheusModule.class::isInstance));
assertFalse(modules.stream().anyMatch(LLMStaticModule.class::isInstance));
}
catch (IOException e) {
fail(e);
}
}

private static WombatModuleProperties properties(boolean kubernetesApi, boolean llmPrometheus, boolean llmStatic, boolean llmGoogle)
{
return new WombatModuleProperties()
{
Expand All @@ -100,6 +125,12 @@ public LlmStaticProperties llmStatic()
return () -> llmStatic;
}

@Override
public LlmGoogleProperties llmGoogle()
{
return () -> llmGoogle;
}

@Override
public ExtensionProperties extension()
{
Expand Down