feat: add google module

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. |
 

module/pom.xml

@@ -16,6 +16,7 @@
     <properties>
         <jackson-annotations.version>2.22</jackson-annotations.version>
         <fabric8-kubernetes-client.version>8.0.0</fabric8-kubernetes-client.version>
+        <google-auth-library.version>1.37.1</google-auth-library.version>
         <mockito.version>5.24.0</mockito.version>
     </properties>
 
@@ -49,6 +50,13 @@
             <version>${fabric8-kubernetes-client.version}</version>
         </dependency>
 
+        <!-- Google Cloud auth (OAuth2 access tokens for the Cloud Monitoring API, used by llm_google) -->
+        <dependency>
+            <groupId>com.google.auth</groupId>
+            <artifactId>google-auth-library-oauth2-http</artifactId>
+            <version>${google-auth-library.version}</version>
+        </dependency>
+
         <!-- Test deps -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -63,4 +71,4 @@
             <scope>test</scope>
         </dependency>
     </dependencies>
-</project>
+</project>

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<String> requiredSecretKeys()
+    {
+        return this.usesServiceAccountKey() ? Set.of(this.credentialsKey) : Set.of();
+    }
+
+    public boolean usesServiceAccountKey()
+    {
+        return this.credentialsKey != null && !this.credentialsKey.isBlank();
+    }
+}

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<? extends Asset> assetClass()
+    {
+        return LLMGoogleAsset.class;
+    }
+
+    @Override
+    public Optional<WombatSource> createSource(WombatContext context)
+    {
+        return Optional.of(new LLMGoogleSource(GoogleMultiClientConfig.create(context)));
+    }
+}

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;
+        }
+    }
+}

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<String> 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);
+        }
+    }
+}

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
+    );
+}

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<String, GoogleClient> clients;
+
+    public GoogleMultiClient()
+    {
+        this.clients = new HashMap<>();
+    }
+
+    public Optional<GoogleClient> get(String id)
+    {
+        return Optional.ofNullable(this.clients.get(id));
+    }
+
+    public void register(String id, GoogleClient client)
+    {
+        this.clients.put(id, client);
+    }
+}

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());
+    }
+}

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> 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<Point> 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;
+        }
+    }
+}