diff --git a/README.md b/README.md index c5fc921..e677acf 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,7 @@ Core modules provide ready-to-use implementations of asset types and their assoc | `tech.illuin.wombat-module.kubernetes-api` | `KUBERNETES_CONTAINER` | `KubernetesAPIModule` | `source` | Connects to Kubernetes cluster APIs to poll CPU and memory resource metrics. | | `tech.illuin.wombat-module.llm-static` | `LLM` | `LLMStaticModule` | `activity-resolver` | Models static LLM workloads based on annual request profiles. | | `tech.illuin.wombat-module.llm-prometheus` | `LLM` | `LLMPrometheusModule` | `source` | Queries Prometheus via PromQL for LLM inference telemetry. | +| `tech.illuin.wombat-module.llm-google` | `LLM` | `LLMGoogleModule` | `source` | Queries the Google Cloud Monitoring API for LLM inference telemetry. | | `tech.illuin.wombat-module.kubernetes-simulated` | `KUBERNETES_CONTAINER` | `KubernetesSimulatedModule` | | Special asset type used by the [Wombat app](https://github.com/illuin-tech/wombat#http-api) for simulation requests. | | `tech.illuin.wombat-module.lmm-simulated` | `LLM` | `KubernetesSimulatedModule` | | Special asset type used by the [Wombat app](https://github.com/illuin-tech/wombat#http-api) for simulation requests. | diff --git a/module/pom.xml b/module/pom.xml index ad528b0..120c90f 100644 --- a/module/pom.xml +++ b/module/pom.xml @@ -16,6 +16,7 @@ 2.22 8.0.0 + 1.37.1 5.24.0 @@ -49,6 +50,13 @@ ${fabric8-kubernetes-client.version} + + + com.google.auth + google-auth-library-oauth2-http + ${google-auth-library.version} + + org.junit.jupiter @@ -63,4 +71,4 @@ test - + \ No newline at end of file diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleAsset.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleAsset.java new file mode 100644 index 0000000..28cc036 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleAsset.java @@ -0,0 +1,40 @@ +package tech.illuin.wombat.module.llm_google; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonUnwrapped; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; +import tech.illuin.wombat.core.asset.Asset; +import tech.illuin.wombat.core.asset.AssetIdentity; +import tech.illuin.wombat.core.asset.type.AssetType; +import tech.illuin.wombat.core.secret.SecretAware; +import tech.illuin.wombat.core.source.Monitorable; + +import java.util.Set; + +public record LLMGoogleAsset( + @JsonUnwrapped AssetIdentity identity, + @NotBlank @JsonProperty("google-url") String googleUrl, + @NotBlank @JsonProperty("project-id") String projectId, + @JsonProperty("credentials-key") String credentialsKey, + @JsonProperty("heartbeat-skip") int heartbeatSkip, + @NotNull @JsonProperty("profile") LLMGoogleProfile profile +) implements Asset, SecretAware, Monitorable +{ + @Override + public AssetType type() + { + return LLMGoogleModule.TYPE; + } + + @Override + public Set requiredSecretKeys() + { + return this.usesServiceAccountKey() ? Set.of(this.credentialsKey) : Set.of(); + } + + public boolean usesServiceAccountKey() + { + return this.credentialsKey != null && !this.credentialsKey.isBlank(); + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleModule.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleModule.java new file mode 100644 index 0000000..1ee51ba --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleModule.java @@ -0,0 +1,36 @@ +package tech.illuin.wombat.module.llm_google; + +import tech.illuin.wombat.core.asset.Asset; +import tech.illuin.wombat.core.asset.type.ActivityRegime; +import tech.illuin.wombat.core.asset.type.AssetType; +import tech.illuin.wombat.core.asset.type.ServiceFamily; +import tech.illuin.wombat.core.context.WombatContext; +import tech.illuin.wombat.core.module.WombatModule; +import tech.illuin.wombat.core.source.WombatSource; +import tech.illuin.wombat.module.llm_google.connector.GoogleMultiClientConfig; +import tech.illuin.wombat.module.llm_google.source.LLMGoogleSource; + +import java.util.Optional; + +public class LLMGoogleModule implements WombatModule +{ + public static final AssetType TYPE = AssetType.of("tech.illuin", "wombat-module", "llm-google", ActivityRegime.MEASURED, ServiceFamily.LLM); + + @Override + public AssetType type() + { + return TYPE; + } + + @Override + public Class assetClass() + { + return LLMGoogleAsset.class; + } + + @Override + public Optional createSource(WombatContext context) + { + return Optional.of(new LLMGoogleSource(GoogleMultiClientConfig.create(context))); + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfile.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfile.java new file mode 100644 index 0000000..304f233 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfile.java @@ -0,0 +1,29 @@ +package tech.illuin.wombat.module.llm_google; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; +import tech.illuin.wombat.core.activity.commons.Country; +import tech.illuin.wombat.core.asset.profile.LLMProfile; +import tech.illuin.wombat.core.asset.profile.LLMProvider; + +import java.time.Duration; + +public record LLMGoogleProfile( + @NotNull @JsonProperty("provider") LLMProvider provider, + @NotBlank @JsonProperty("model") String model, + @NotNull @JsonProperty("location") Country location, + @NotNull @Valid @JsonProperty("dynamic-profile") DynamicProfile dynamicProfile +) implements LLMProfile +{ + public record DynamicProfile( + @NotBlank @JsonProperty("filter") String filter, + @JsonProperty("window") Duration window + ) { + public Duration windowOrDefault() + { + return this.window == null ? Duration.ofMinutes(1) : this.window; + } + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptor.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptor.java new file mode 100644 index 0000000..0df1a23 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptor.java @@ -0,0 +1,67 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import com.google.auth.oauth2.GoogleCredentials; +import feign.RequestInterceptor; +import feign.RequestTemplate; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; + +public class GoogleAuthInterceptor implements RequestInterceptor +{ + private static final String AUTHORIZATION = "Authorization"; + private static final List SCOPES = List.of("https://www.googleapis.com/auth/monitoring.read"); + + private volatile GoogleCredentials credentials; + + public GoogleAuthInterceptor(String serviceAccountJson) + { + if (serviceAccountJson != null && !serviceAccountJson.isBlank()) + this.credentials = fromJson(serviceAccountJson); + } + + GoogleAuthInterceptor(GoogleCredentials credentials) + { + this.credentials = credentials; + } + + @Override + public void apply(RequestTemplate template) + { + try { + GoogleCredentials current = this.credentials(); + current.refreshIfExpired(); + template.header(AUTHORIZATION, "Bearer " + current.getAccessToken().getTokenValue()); + } + catch (IOException e) { + throw new IllegalStateException("Unable to obtain a Google access token: " + e.getMessage(), e); + } + } + + private GoogleCredentials credentials() throws IOException + { + GoogleCredentials current = this.credentials; + if (current == null) + { + synchronized (this) + { + current = this.credentials; + if (current == null) + this.credentials = current = GoogleCredentials.getApplicationDefault().createScoped(SCOPES); + } + } + return current; + } + + private static GoogleCredentials fromJson(String json) + { + try { + return GoogleCredentials.fromStream(new ByteArrayInputStream(json.getBytes(StandardCharsets.UTF_8))).createScoped(SCOPES); + } + catch (IOException e) { + throw new IllegalStateException("Unable to load the Google service-account key: " + e.getMessage(), e); + } + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleClient.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleClient.java new file mode 100644 index 0000000..c15ddd3 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleClient.java @@ -0,0 +1,25 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import feign.Param; +import feign.RequestLine; +import tech.illuin.wombat.module.llm_google.connector.model.GoogleQueryResponse; + +public interface GoogleClient +{ + @RequestLine( + "GET /v3/projects/{projectId}/timeSeries" + + "?filter={filter}" + + "&interval.startTime={start}" + + "&interval.endTime={end}" + + "&aggregation.alignmentPeriod={period}" + + "&aggregation.perSeriesAligner=ALIGN_SUM" + + "&aggregation.crossSeriesReducer=REDUCE_SUM" + ) + GoogleQueryResponse query( + @Param("projectId") String projectId, + @Param("filter") String filter, + @Param("start") String start, + @Param("end") String end, + @Param("period") String period + ); +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClient.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClient.java new file mode 100644 index 0000000..8b05de7 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClient.java @@ -0,0 +1,25 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +public class GoogleMultiClient +{ + private final Map clients; + + public GoogleMultiClient() + { + this.clients = new HashMap<>(); + } + + public Optional get(String id) + { + return Optional.ofNullable(this.clients.get(id)); + } + + public void register(String id, GoogleClient client) + { + this.clients.put(id, client); + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfig.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfig.java new file mode 100644 index 0000000..993fd83 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfig.java @@ -0,0 +1,33 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import feign.Feign; +import feign.jackson.JacksonDecoder; +import tech.illuin.wombat.core.context.WombatContext; +import tech.illuin.wombat.core.secret.SecretResolver; +import tech.illuin.wombat.module.llm_google.LLMGoogleAsset; + +public class GoogleMultiClientConfig +{ + public static GoogleMultiClient create(WombatContext context) + { + SecretResolver secrets = context.secrets(); + GoogleMultiClient multi = new GoogleMultiClient(); + context.assets().stream() + .filter(LLMGoogleAsset.class::isInstance) + .map(LLMGoogleAsset.class::cast) + .forEach(asset -> multi.register(asset.identity().id(), createClient(asset, secrets))); + return multi; + } + + private static GoogleClient createClient(LLMGoogleAsset properties, SecretResolver secrets) + { + String serviceAccountJson = properties.usesServiceAccountKey() + ? secrets.require(properties.credentialsKey()) + : null; + + return Feign.builder() + .decoder(new JacksonDecoder()) + .requestInterceptor(new GoogleAuthInterceptor(serviceAccountJson)) + .target(GoogleClient.class, properties.googleUrl()); + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponse.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponse.java new file mode 100644 index 0000000..6f17333 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponse.java @@ -0,0 +1,51 @@ +package tech.illuin.wombat.module.llm_google.connector.model; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; + +import java.util.List; + +@JsonIgnoreProperties(ignoreUnknown = true) +public record GoogleQueryResponse( + @JsonProperty("timeSeries") List timeSeries, + @JsonProperty("nextPageToken") String nextPageToken +) { + public double sum() + { + if (this.timeSeries == null) + return 0.0d; + + return this.timeSeries.stream() + .filter(ts -> ts.points() != null) + .flatMap(ts -> ts.points().stream()) + .mapToDouble(Point::numericValue) + .sum(); + } + + @JsonIgnoreProperties(ignoreUnknown = true) + public record TimeSeries( + @JsonProperty("metric") JsonNode metric, + @JsonProperty("resource") JsonNode resource, + @JsonProperty("metricKind") String metricKind, + @JsonProperty("valueType") String valueType, + @JsonProperty("points") List points + ) {} + + @JsonIgnoreProperties(ignoreUnknown = true) + public record Point( + @JsonProperty("interval") JsonNode interval, + @JsonProperty("value") JsonNode value + ) { + public double numericValue() + { + if (this.value == null) + return 0.0d; + if (this.value.hasNonNull("int64Value")) + return Double.parseDouble(this.value.get("int64Value").asText()); + if (this.value.hasNonNull("doubleValue")) + return this.value.get("doubleValue").asDouble(); + return 0.0d; + } + } +} diff --git a/module/src/main/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSource.java b/module/src/main/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSource.java new file mode 100644 index 0000000..9720b24 --- /dev/null +++ b/module/src/main/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSource.java @@ -0,0 +1,79 @@ +package tech.illuin.wombat.module.llm_google.source; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import tech.illuin.wombat.core.asset.Asset; +import tech.illuin.wombat.core.source.WombatSource; +import tech.illuin.wombat.core.source.data.LLMData; +import tech.illuin.wombat.core.source.data.MetricData; +import tech.illuin.wombat.module.llm_google.LLMGoogleAsset; +import tech.illuin.wombat.module.llm_google.connector.GoogleClient; +import tech.illuin.wombat.module.llm_google.connector.GoogleMultiClient; +import tech.illuin.wombat.module.llm_google.connector.model.GoogleQueryResponse; + +import java.time.Duration; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.List; + +import static java.util.Collections.emptyList; + +public class LLMGoogleSource implements WombatSource +{ + private final GoogleMultiClient client; + + private static final Logger logger = LoggerFactory.getLogger(LLMGoogleSource.class); + + public LLMGoogleSource(GoogleMultiClient client) + { + this.client = client; + } + + @Override + public boolean accept(Asset asset) + { + return asset instanceof LLMGoogleAsset; + } + + @Override + public List source(Instant heartbeat, Asset asset) + { + LLMGoogleAsset googleAsset = (LLMGoogleAsset) asset; + + String assetId = asset.identity().id(); + GoogleClient googleClient = this.client.get(assetId) + .orElseThrow(() -> new IllegalStateException("No Google client registered for asset " + assetId)); + + String filter = googleAsset.profile().dynamicProfile().filter(); + Duration window = googleAsset.profile().dynamicProfile().windowOrDefault(); + Instant end = heartbeat.truncatedTo(ChronoUnit.SECONDS); + Instant start = end.minus(window); + + logger.debug("Sampling output tokens for LLM asset {} over [{}, {}] with filter '{}'", assetId, start, end, filter); + + GoogleQueryResponse response = googleClient.query( + googleAsset.projectId(), + filter, + start.toString(), + end.toString(), + window.toSeconds() + "s" + ); + + if (response == null) + { + logger.warn("Google query for LLM asset {} returned no response", assetId); + return emptyList(); + } + + long outputTokens = Math.max(0L, Math.round(response.sum())); + + return List.of(new LLMData( + googleAsset.profile().model(), + googleAsset.profile().provider(), + googleAsset.profile().model(), + googleAsset.profile().location(), + outputTokens, + 1 + )); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleAssetTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleAssetTest.java new file mode 100644 index 0000000..19ad29b --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleAssetTest.java @@ -0,0 +1,54 @@ +package tech.illuin.wombat.module.llm_google; + +import org.junit.jupiter.api.Test; +import tech.illuin.wombat.core.asset.AssetIdentity; +import tech.illuin.wombat.core.activity.commons.Country; +import tech.illuin.wombat.core.asset.profile.LLMProvider; + +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class LLMGoogleAssetTest +{ + @Test + void anAssetRelyingOnApplicationDefaultCredentialsNeedsNoSecret() + { + LLMGoogleAsset asset = asset(null); + + assertFalse(asset.usesServiceAccountKey()); + assertEquals(Set.of(), asset.requiredSecretKeys()); + } + + @Test + void aBlankCredentialsKeyIsTreatedAsAbsent() + { + LLMGoogleAsset asset = asset(" "); + + assertFalse(asset.usesServiceAccountKey()); + assertEquals(Set.of(), asset.requiredSecretKeys()); + } + + @Test + void aServiceAccountAssetDeclaresTheVariableItNeeds() + { + LLMGoogleAsset asset = asset("GOOGLE_SA_KEY"); + + assertTrue(asset.usesServiceAccountKey()); + assertEquals(Set.of("GOOGLE_SA_KEY"), asset.requiredSecretKeys()); + } + + @Test + void declaresTheLLMGoogleAssetType() + { + assertEquals(LLMGoogleModule.TYPE, asset(null).type()); + } + + private static LLMGoogleAsset asset(String credentialsKey) + { + return new LLMGoogleAsset(AssetIdentity.of("g", "env", "G"), "https://monitoring.googleapis.com", "my-project", credentialsKey, 0, + new LLMGoogleProfile(LLMProvider.google_genai, "gemini-2.5-flash", Country.US, new LLMGoogleProfile.DynamicProfile("f", null))); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleModuleTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleModuleTest.java new file mode 100644 index 0000000..195b50c --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleModuleTest.java @@ -0,0 +1,77 @@ +package tech.illuin.wombat.module.llm_google; + +import org.junit.jupiter.api.Test; +import tech.illuin.wombat.core.asset.AssetIdentity; +import tech.illuin.wombat.core.asset.Environment; +import tech.illuin.wombat.core.activity.commons.Country; +import tech.illuin.wombat.core.asset.profile.LLMProvider; +import tech.illuin.wombat.core.context.ResolvedContext; +import tech.illuin.wombat.core.secret.EnvironmentSecretResolver; +import tech.illuin.wombat.core.secret.MissingSecretException; +import tech.illuin.wombat.core.source.WombatSource; +import tech.illuin.wombat.module.llm_google.source.LLMGoogleSource; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class LLMGoogleModuleTest +{ + private final LLMGoogleModule module = new LLMGoogleModule(); + + @Test + void declaresTheGoogleLLMAssetType() + { + assertEquals(LLMGoogleModule.TYPE, this.module.type()); + assertEquals(LLMGoogleAsset.class, this.module.assetClass()); + } + + @Test + void providesAGoogleSource() + { + // An empty context registers no Google client, so nothing is dialled. + Optional source = this.module.createSource(new ResolvedContext()); + assertTrue(source.isPresent()); + assertInstanceOf(LLMGoogleSource.class, source.get()); + } + + @Test + void sourcingAnAssetWithoutCredentialsKeyDoesNotRequireGoogleCredentialsAtStartup() + { + // Application Default Credentials are resolved lazily on the first request, so wiring must succeed even on + // a machine (CI, a laptop) that has none. + ResolvedContext context = context(Map.of(), asset(null)); + + assertTrue(this.module.createSource(context).isPresent()); + } + + @Test + void sourcingAServiceAccountAssetFailsWhenItsVariableIsUnset() + { + // Proves the module builds its clients from the context's resolver: the credential is read from the + // environment at source creation, never from the asset the reconciler persisted. + ResolvedContext context = context(Map.of(), asset("GOOGLE_SA_KEY")); + + MissingSecretException e = assertThrows(MissingSecretException.class, () -> this.module.createSource(context)); + + assertEquals("GOOGLE_SA_KEY", e.key()); + } + + private static ResolvedContext context(Map environment, LLMGoogleAsset asset) + { + return new ResolvedContext( + List.of(new Environment("env", List.of(asset))), + new EnvironmentSecretResolver(environment::get)); + } + + private static LLMGoogleAsset asset(String credentialsKey) + { + return new LLMGoogleAsset(AssetIdentity.of("g", "env", "G"), "https://monitoring.googleapis.com", "my-project", credentialsKey, 0, + new LLMGoogleProfile(LLMProvider.google_genai, "gemini-2.5-flash", Country.US, new LLMGoogleProfile.DynamicProfile("f", null))); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfileTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfileTest.java new file mode 100644 index 0000000..0bfb713 --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/LLMGoogleProfileTest.java @@ -0,0 +1,22 @@ +package tech.illuin.wombat.module.llm_google; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class LLMGoogleProfileTest +{ + @Test + void theWindowDefaultsToOneMinute() + { + assertEquals(Duration.ofMinutes(1), new LLMGoogleProfile.DynamicProfile("f", null).windowOrDefault()); + } + + @Test + void anExplicitWindowIsKept() + { + assertEquals(Duration.ofMinutes(5), new LLMGoogleProfile.DynamicProfile("f", Duration.ofMinutes(5)).windowOrDefault()); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptorTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptorTest.java new file mode 100644 index 0000000..715410b --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleAuthInterceptorTest.java @@ -0,0 +1,57 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import com.google.auth.oauth2.AccessToken; +import com.google.auth.oauth2.GoogleCredentials; +import feign.RequestTemplate; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Date; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class GoogleAuthInterceptorTest +{ + @Test + void sendsTheAccessTokenAsABearerHeader() throws IOException + { + GoogleCredentials credentials = mock(GoogleCredentials.class); + when(credentials.getAccessToken()).thenReturn(new AccessToken("ya29.token", new Date(System.currentTimeMillis() + 3_600_000))); + RequestTemplate template = new RequestTemplate(); + + new GoogleAuthInterceptor(credentials).apply(template); + + assertEquals("Bearer ya29.token", template.headers().get("Authorization").iterator().next()); + verify(credentials).refreshIfExpired(); + } + + @Test + void wrapsTokenRefreshFailures() throws IOException + { + GoogleCredentials credentials = mock(GoogleCredentials.class); + doThrow(new IOException("boom")).when(credentials).refreshIfExpired(); + + IllegalStateException e = assertThrows(IllegalStateException.class, + () -> new GoogleAuthInterceptor(credentials).apply(new RequestTemplate())); + + assertEquals(IOException.class, e.getCause().getClass()); + } + + @Test + void rejectsAMalformedServiceAccountKeyAtConstruction() + { + assertThrows(IllegalStateException.class, () -> new GoogleAuthInterceptor("{}")); + } + + @Test + void doesNotTouchApplicationDefaultCredentialsAtConstruction() + { + new GoogleAuthInterceptor((String) null); + new GoogleAuthInterceptor(" "); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfigTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfigTest.java new file mode 100644 index 0000000..f5b6c7f --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/GoogleMultiClientConfigTest.java @@ -0,0 +1,94 @@ +package tech.illuin.wombat.module.llm_google.connector; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; +import tech.illuin.wombat.core.asset.AssetIdentity; +import tech.illuin.wombat.core.asset.Environment; +import tech.illuin.wombat.core.activity.commons.Country; +import tech.illuin.wombat.core.asset.profile.LLMProvider; +import tech.illuin.wombat.core.context.ResolvedContext; +import tech.illuin.wombat.core.context.WombatContext; +import tech.illuin.wombat.core.secret.EnvironmentSecretResolver; +import tech.illuin.wombat.core.secret.MissingSecretException; +import tech.illuin.wombat.module.llm_google.LLMGoogleAsset; +import tech.illuin.wombat.module.llm_google.LLMGoogleProfile; + +import java.security.KeyPairGenerator; +import java.util.Base64; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class GoogleMultiClientConfigTest +{ + @Test + void buildsAClientPerAssetRelyingOnApplicationDefaultCredentials() + { + GoogleMultiClient multi = GoogleMultiClientConfig.create( + context(Map.of(), asset("g", null))); + + assertTrue(multi.get("g").isPresent()); + } + + @Test + void buildsAClientFromAServiceAccountKeyReadFromTheDeclaredVariable() throws Exception + { + GoogleMultiClient multi = GoogleMultiClientConfig.create( + context(Map.of("GOOGLE_SA_KEY", serviceAccountJson()), asset("g", "GOOGLE_SA_KEY"))); + + assertTrue(multi.get("g").isPresent()); + } + + @Test + void failsWhenTheServiceAccountVariableIsUnset() + { + WombatContext context = context(Map.of(), asset("g", "GOOGLE_SA_KEY")); + + MissingSecretException e = assertThrows(MissingSecretException.class, + () -> GoogleMultiClientConfig.create(context)); + + assertEquals("GOOGLE_SA_KEY", e.key()); + } + + @Test + void failsOnAMalformedServiceAccountKey() + { + WombatContext context = context(Map.of("GOOGLE_SA_KEY", "not json"), asset("g", "GOOGLE_SA_KEY")); + + assertThrows(IllegalStateException.class, () -> GoogleMultiClientConfig.create(context)); + } + + private static WombatContext context(Map environment, LLMGoogleAsset asset) + { + return new ResolvedContext( + List.of(new Environment("env", List.of(asset))), + new EnvironmentSecretResolver(environment::get)); + } + + private static LLMGoogleAsset asset(String id, String credentialsKey) + { + return new LLMGoogleAsset(AssetIdentity.of(id, "env", id), "https://monitoring.googleapis.com", "my-project", credentialsKey, 0, + new LLMGoogleProfile(LLMProvider.google_genai, "gemini-2.5-flash", Country.US, new LLMGoogleProfile.DynamicProfile("f", null))); + } + + private static String serviceAccountJson() throws Exception + { + KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA"); + generator.initialize(2048); + String body = Base64.getMimeEncoder(64, "\n".getBytes()).encodeToString(generator.generateKeyPair().getPrivate().getEncoded()); + String pem = "-----BEGIN PRIVATE KEY-----\n" + body + "\n-----END PRIVATE KEY-----\n"; + + return new ObjectMapper().writeValueAsString(Map.of( + "type", "service_account", + "project_id", "my-project", + "private_key_id", "key-id", + "private_key", pem, + "client_email", "wombat@my-project.iam.gserviceaccount.com", + "client_id", "1234567890", + "token_uri", "https://oauth2.googleapis.com/token" + )); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponseTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponseTest.java new file mode 100644 index 0000000..f125207 --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/connector/model/GoogleQueryResponseTest.java @@ -0,0 +1,108 @@ +package tech.illuin.wombat.module.llm_google.connector.model; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +class GoogleQueryResponseTest +{ + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static final String SAMPLE = """ + { + "timeSeries": [ + { + "metric": { + "labels": { + "request_type": "shared", + "source": "us-central1", + "type": "output", + "shared_request_type": "standard" + }, + "type": "aiplatform.googleapis.com/publisher/online_serving/token_count" + }, + "resource": { + "type": "aiplatform.googleapis.com/PublisherModel", + "labels": { + "location": "us-central1", + "publisher": "google", + "model_version_id": "", + "model_user_id": "gemini-2.5-flash", + "project_id": "macro-theater-402700" + } + }, + "metricKind": "DELTA", + "valueType": "INT64", + "points": [ + { + "interval": { + "startTime": "2026-10-06T20:58:03.001Z", + "endTime": "2026-10-06T20:59:03Z" + }, + "value": { "int64Value": "1576" } + } + ] + } + ], + "unit": "1" + } + """; + + @Test + void parsesARealCloudMonitoringPayload() throws Exception + { + GoogleQueryResponse response = MAPPER.readValue(SAMPLE, GoogleQueryResponse.class); + + assertEquals(1, response.timeSeries().size()); + GoogleQueryResponse.TimeSeries series = response.timeSeries().getFirst(); + assertEquals("DELTA", series.metricKind()); + assertEquals("INT64", series.valueType()); + assertEquals("gemini-2.5-flash", series.resource().get("labels").get("model_user_id").asText()); + assertEquals(1576.0d, response.sum()); + } + + @Test + void sumsEveryPointOfEverySeries() throws Exception + { + GoogleQueryResponse response = MAPPER.readValue(""" + {"timeSeries":[ + {"points":[{"value":{"int64Value":"100"}},{"value":{"int64Value":"50"}}]}, + {"points":[{"value":{"doubleValue":0.5}},{"value":{"int64Value":"25"}}]} + ]} + """, GoogleQueryResponse.class); + + assertEquals(175.5d, response.sum()); + } + + @Test + void anEmptyResponseSumsToZero() throws Exception + { + GoogleQueryResponse response = MAPPER.readValue("{}", GoogleQueryResponse.class); + + assertNull(response.timeSeries()); + assertEquals(0.0d, response.sum()); + assertEquals(0.0d, new GoogleQueryResponse(List.of(), null).sum()); + } + + @Test + void aSeriesWithoutPointsOrValuesContributesNothing() throws Exception + { + GoogleQueryResponse response = MAPPER.readValue( + "{\"timeSeries\":[{},{\"points\":[{},{\"value\":{}}]}]}", GoogleQueryResponse.class); + + assertEquals(0.0d, response.sum()); + } + + @Test + void unknownFieldsAreIgnored() throws Exception + { + GoogleQueryResponse response = MAPPER.readValue( + "{\"timeSeries\":[],\"unit\":\"1\",\"nextPageToken\":\"abc\",\"executionErrors\":[]}", GoogleQueryResponse.class); + + assertEquals("abc", response.nextPageToken()); + } +} diff --git a/module/src/test/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSourceTest.java b/module/src/test/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSourceTest.java new file mode 100644 index 0000000..7f0a633 --- /dev/null +++ b/module/src/test/java/tech/illuin/wombat/module/llm_google/source/LLMGoogleSourceTest.java @@ -0,0 +1,157 @@ +package tech.illuin.wombat.module.llm_google.source; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; +import tech.illuin.wombat.core.asset.Asset; +import tech.illuin.wombat.core.asset.AssetIdentity; +import tech.illuin.wombat.core.activity.commons.Country; +import tech.illuin.wombat.core.asset.profile.LLMProvider; +import tech.illuin.wombat.core.source.data.LLMData; +import tech.illuin.wombat.core.source.data.MetricData; +import tech.illuin.wombat.module.llm_google.LLMGoogleAsset; +import tech.illuin.wombat.module.llm_google.LLMGoogleProfile; +import tech.illuin.wombat.module.llm_google.connector.GoogleClient; +import tech.illuin.wombat.module.llm_google.connector.GoogleMultiClient; +import tech.illuin.wombat.module.llm_google.connector.model.GoogleQueryResponse; + +import java.time.Duration; +import java.time.Instant; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class LLMGoogleSourceTest +{ + private static final String PROJECT = "macro-theater-402700"; + private static final String FILTER = "metric.type=\"aiplatform.googleapis.com/publisher/online_serving/token_count\" AND metric.labels.type=\"output\""; + private static final Instant HEARTBEAT = Instant.parse("2026-10-06T21:00:00Z"); + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + void sample_emitsLLMDataWithProviderModelAndLocation() throws Exception + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(PROJECT, FILTER, "2026-10-06T20:59:00Z", "2026-10-06T21:00:00Z", "60s")) + .thenReturn(response("{\"timeSeries\":[{\"points\":[{\"value\":{\"int64Value\":\"1576\"}}]}]}")); + + List dataList = source(client).source(HEARTBEAT, (Asset) asset(null)); + + assertEquals(1, dataList.size()); + assertInstanceOf(LLMData.class, dataList.getFirst()); + + LLMData llmData = (LLMData) dataList.getFirst(); + assertEquals("gemini-2.5-flash", llmData.serviceId()); + assertEquals(LLMProvider.google_genai, llmData.provider()); + assertEquals("gemini-2.5-flash", llmData.model()); + assertEquals(Country.US, llmData.location()); + assertEquals(1576L, llmData.outputTokens()); + assertEquals(1L, llmData.requests()); + } + + @Test + void sample_sumsTheTokensOfEverySeriesAndPoint() throws Exception + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(any(), any(), any(), any(), any())).thenReturn(response(""" + {"timeSeries":[ + {"points":[{"value":{"int64Value":"1000"}},{"value":{"int64Value":"500"}}]}, + {"points":[{"value":{"int64Value":"76"}}]} + ]} + """)); + + LLMData llmData = (LLMData) source(client).source(HEARTBEAT, asset(null)).getFirst(); + + assertEquals(1576L, llmData.outputTokens()); + } + + @Test + void sample_recordsZeroTokensWhenNothingMatchedInTheWindow() throws Exception + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(any(), any(), any(), any(), any())).thenReturn(response("{}")); + + List dataList = source(client).source(HEARTBEAT, asset(null)); + + assertEquals(1, dataList.size()); + assertEquals(0L, ((LLMData) dataList.getFirst()).outputTokens()); + } + + @Test + void sample_emitsNothingWhenTheClientReturnsNoResponse() + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(any(), any(), any(), any(), any())).thenReturn(null); + + assertTrue(source(client).source(HEARTBEAT, asset(null)).isEmpty()); + } + + @Test + void sample_queriesAWindowEndingAtTheHeartbeat() throws Exception + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(PROJECT, FILTER, "2026-10-06T20:55:00Z", "2026-10-06T21:00:00Z", "300s")) + .thenReturn(response("{\"timeSeries\":[{\"points\":[{\"value\":{\"int64Value\":\"42\"}}]}]}")); + + List dataList = source(client).source(HEARTBEAT, asset(Duration.ofMinutes(5))); + + assertEquals(42L, ((LLMData) dataList.getFirst()).outputTokens()); + } + + @Test + void sample_truncatesTheHeartbeatToWholeSeconds() throws Exception + { + GoogleClient client = mock(GoogleClient.class); + when(client.query(PROJECT, FILTER, "2026-10-06T20:59:00Z", "2026-10-06T21:00:00Z", "60s")) + .thenReturn(response("{\"timeSeries\":[{\"points\":[{\"value\":{\"int64Value\":\"7\"}}]}]}")); + + List dataList = source(client).source(Instant.parse("2026-10-06T21:00:00.789Z"), asset(null)); + + assertEquals(7L, ((LLMData) dataList.getFirst()).outputTokens()); + } + + @Test + void sample_failsWhenNoClientIsRegisteredForTheAsset() + { + LLMGoogleSource source = new LLMGoogleSource(new GoogleMultiClient()); + + IllegalStateException e = assertThrows(IllegalStateException.class, () -> source.source(HEARTBEAT, asset(null))); + + assertTrue(e.getMessage().contains("llm-google"), e.getMessage()); + } + + @Test + void acceptsOnlyGoogleAssets() + { + LLMGoogleSource source = source(mock(GoogleClient.class)); + + assertTrue(source.accept(asset(null))); + assertFalse(source.accept(mock(Asset.class))); + } + + private static GoogleQueryResponse response(String json) throws Exception + { + return MAPPER.readValue(json, GoogleQueryResponse.class); + } + + private static LLMGoogleSource source(GoogleClient client) + { + GoogleMultiClient multiClient = new GoogleMultiClient(); + multiClient.register("llm-google", client); + return new LLMGoogleSource(multiClient); + } + + private static LLMGoogleAsset asset(Duration window) + { + return new LLMGoogleAsset( + AssetIdentity.of("llm-google", "prod", "Google LLM"), "https://monitoring.googleapis.com", PROJECT, null, 0, + new LLMGoogleProfile(LLMProvider.google_genai, "gemini-2.5-flash", Country.US, new LLMGoogleProfile.DynamicProfile(FILTER, window)) + ); + } +}