feat(search): add Google Gemini embedding provider

docs/superpowers/specs/2026-05-07-google-gemini-embedding-client-design.md

@@ -0,0 +1,169 @@
+# Google Gemini Embedding Client — Design
+
+**Date:** 2026-05-07
+**Status:** Proposed
+**Scope:** Add a fourth embedding provider — Google's Generative Language API (Gemini) — alongside the existing `openai`, `bedrock`, and `djl` providers used by OpenMetadata's vector search.
+
+## Goal
+
+Allow OpenMetadata operators to point natural-language / semantic search at Google's Gemini embedding models using a single API key from Google AI Studio, with no GCP project, service account, or OAuth setup required.
+
+## Non-goals
+
+- Vertex AI on GCP (project + service-account auth) — separate provider, future work.
+- Gemini chat/completions for NLQ query transformation — only the embedding side is covered here.
+- Dynamic dimension reduction via `outputDimensionality` — operators set `embeddingDimension` to match the model's native output, same contract as the existing providers.
+
+## Context
+
+The codebase already has a clean extension point for embedding providers:
+
+- **Base class:** `EmbeddingClient` (`openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/EmbeddingClient.java`) — template-method pattern with semaphore-based concurrency limiting.
+- **Three siblings:** `OpenAIEmbeddingClient`, `BedrockEmbeddingClient`, `DjlEmbeddingClient` in the same package.
+- **Schema:** `naturalLanguageSearch` block in `openmetadata-spec/src/main/resources/json/schema/configuration/elasticSearchConfiguration.json` defines a per-provider sub-block (`openai`, `bedrock`, `djl`).
+- **Wiring:** `SearchRepository.createEmbeddingClient(ElasticSearchConfiguration)` is a single switch on `embeddingProvider`.
+
+The Google provider slots into all four touch points without changing the contract.
+
+## Design
+
+### Provider name
+
+`google` — used as both the JSON schema field name (`naturalLanguageSearch.google`) and the `embeddingProvider` value. Matches the flatness of `openai` / `bedrock`. A future Vertex AI provider would land as a separate `googleVertexAi` block; auth model is different enough that fusing them would muddy the config.
+
+### Schema (JSON)
+
+Added under `naturalLanguageSearch.properties` in `elasticSearchConfiguration.json`, alongside `openai`/`bedrock`/`djl`:
+
+```json
+"google": {
+  "description": "Google Gemini configuration for embedding generation via the Generative Language API.",
+  "type": "object",
+  "javaType": "org.openmetadata.schema.service.configuration.elasticsearch.Google",
+  "properties": {
+    "apiKey": {
+      "description": "API key from Google AI Studio for authenticating with the Generative Language API.",
+      "type": "string"
+    },
+    "embeddingModelId": {
+      "description": "Gemini embedding model identifier (e.g., text-embedding-004, gemini-embedding-001).",
+      "type": "string",
+      "default": "text-embedding-004"
+    },
+    "embeddingDimension": {
+      "description": "Dimension of the embedding vector. Must match the model's native output dimension (e.g., 768 for text-embedding-004).",
+      "type": "integer",
+      "default": 768
+    },
+    "endpoint": {
+      "description": "Custom endpoint URL. Leave empty for the default Generative Language API.",
+      "type": "string"
+    }
+  },
+  "additionalProperties": false
+}
+```
+
+The `embeddingProvider` description is updated to mention `google` as a valid value.
+
+### Client implementation
+
+New class `GoogleEmbeddingClient` in `openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/GoogleEmbeddingClient.java`. Mirrors `OpenAIEmbeddingClient` structure:
+
+- `final` class extending `EmbeddingClient`.
+- Public constructor `GoogleEmbeddingClient(ElasticSearchConfiguration config)` — pulls the `Google` sub-config, validates required fields, builds an `HttpClient`.
+- Package-private constructor `GoogleEmbeddingClient(HttpClient, String apiKey, String modelId, int dimension, String endpoint)` for unit-test injection (also mirrors OpenAI).
+- `protected float[] doEmbed(String text)` — issues the HTTP call.
+- `getDimension()` / `getModelId()` return the configured values.
+
+#### Validation
+
+In the public constructor — same shape as `OpenAIEmbeddingClient`:
+
+- `googleCfg == null` → `IllegalArgumentException("Google configuration is required")`
+- `apiKey` null/blank → `IllegalArgumentException("Google API key is required")`
+- `embeddingModelId` null/blank → `IllegalArgumentException("Google embedding model ID is required")`
+- `embeddingDimension == null || <= 0` → `IllegalArgumentException("Google embedding dimension must be positive")`
+
+#### HTTP call
+
+- **Method:** `POST`
+- **URL:** `{endpoint || "https://generativelanguage.googleapis.com/v1beta/models/" + modelId + ":embedContent"}?key={apiKey}` — endpoint resolution mirrors the OpenAI pattern (custom override wins, otherwise default). When the user supplies `endpoint`, we treat it as the full URL up to (but not including) the query string and append `?key={apiKey}`.
+- **Headers:** `Content-Type: application/json` only. Auth travels in the query string per Google's standard for the Generative Language API. (We deliberately do not use `Authorization: Bearer` — that's the OAuth/Vertex pattern.)
+- **Body:**
+  ```json
+  {
+    "model": "models/{modelId}",
+    "content": { "parts": [{ "text": "<input>" }] }
+  }
+  ```
+  No `outputDimensionality` field — keeps the contract aligned with `OpenAIEmbeddingClient`, which also does not pass `dimensions` despite OpenAI supporting it. Operators must set `embeddingDimension` to match the model's native output.
+- **Timeout:** `Duration.ofSeconds(30)` — matches OpenAI client.
+
+#### Response handling
+
+Success body shape:
+```json
+{ "embedding": { "values": [0.013, -0.008, ...] } }
+```
+
+Parser walks `embedding.values`, builds `float[]`. Same defensive checks as the OpenAI parser:
+
+- `embedding` node missing or non-object → `RuntimeException("Invalid Google response: no embedding object found")`
+- `values` array missing or empty → `RuntimeException("Invalid Google response: no values array found")`
+
+Non-200: extract `error.message` from the body if present (Google's standard error envelope is `{ "error": { "code": ..., "message": ..., "status": ... } }`), otherwise echo the raw body. Throw `RuntimeException("Google API returned status " + status + ": " + msg)`.
+
+`IOException` and `InterruptedException` are caught and re-thrown as `RuntimeException`, exactly like the OpenAI client. `InterruptedException` re-sets the interrupt flag.
+
+### Registration
+
+In `SearchRepository.createEmbeddingClient(ElasticSearchConfiguration)` — one new case:
+
+```java
+case "google" -> new GoogleEmbeddingClient(esConfig);
+```
+
+No other call sites change. `VectorDocBuilder`, `OpenSearchVectorService`, `VectorEmbeddingHandler` work against the `EmbeddingClient` base class and need no edits.
+
+### Tests
+
+New file `openmetadata-service/src/test/java/org/openmetadata/service/search/vector/client/GoogleEmbeddingClientTest.java`, mirroring `OpenAIEmbeddingClientTest`. Cases:
+
+- **Construction**
+  - Valid config → succeeds, `getDimension()` and `getModelId()` return configured values.
+  - Missing `Google` sub-config → `IllegalArgumentException`.
+  - Missing/blank `apiKey` → `IllegalArgumentException`.
+  - Missing/blank `embeddingModelId` → `IllegalArgumentException`.
+  - Null or non-positive `embeddingDimension` → `IllegalArgumentException`.
+- **Embedding generation** (test-only constructor with a mocked `HttpClient`)
+  - 200 with valid body → returns the parsed `float[]` with the expected length and values.
+  - 200 with malformed body (missing `embedding` / `values`) → `RuntimeException`.
+  - Non-200 with Google error envelope → `RuntimeException` whose message contains the extracted `error.message`.
+  - Non-200 with unparseable body → `RuntimeException` whose message contains the raw body.
+- **Request shape**
+  - The captured request URL ends with `:embedContent?key={apiKey}`.
+  - The captured request body contains `"model":"models/<modelId>"` and the input text under `content.parts[0].text`.
+  - The `Authorization` header is **not** set; only `Content-Type` and the API key in the URL.
+- **Blank input**
+  - Calling `embed("")` → `IllegalArgumentException` (mirrors OpenAI client's null/blank guard).
+
+Tests use the hand-rolled `StubHttpResponse` + custom `HttpClient` subclass pattern from `OpenAIEmbeddingClientTest` (no Mockito for the HTTP layer). Captured `HttpRequest` is inspected directly for URL, headers, and body content.
+
+## Out of scope (explicit)
+
+- **Batch endpoint.** Google offers `:batchEmbedContents` but the base class's default `embedBatch` (loop + serial calls) is what the other providers use today. Adding a true batch path is a separate optimization.
+- **Streaming, retries, backoff.** Not present in the existing providers; this client matches them.
+- **NLQ chat-completion provider.** This design covers only the `EmbeddingClient` slot. The `NLQService` (`org.openmetadata.service.search.nlq.NLQService`) remains untouched.
+- **Vertex AI / service-account auth.** A separate `googleVertexAi` provider can be added later without disturbing this one.
+
+## Touched files
+
+```
+openmetadata-spec/src/main/resources/json/schema/configuration/elasticSearchConfiguration.json   [MODIFY]
+openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/GoogleEmbeddingClient.java   [CREATE]
+openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java   [MODIFY — one case in switch]
+openmetadata-service/src/test/java/org/openmetadata/service/search/vector/client/GoogleEmbeddingClientTest.java   [CREATE]
+```
+
+The generated `Google.java` schema class lands under `openmetadata-spec` automatically when `mvn clean install` runs against `openmetadata-spec` after the schema change — same flow as `Openai.java`.

openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java

@@ -142,6 +142,7 @@
 import org.openmetadata.service.search.vector.client.BedrockEmbeddingClient;
 import org.openmetadata.service.search.vector.client.DjlEmbeddingClient;
 import org.openmetadata.service.search.vector.client.EmbeddingClient;
+import org.openmetadata.service.search.vector.client.GoogleEmbeddingClient;
 import org.openmetadata.service.search.vector.client.OpenAIEmbeddingClient;
 import org.openmetadata.service.security.policyevaluator.SubjectContext;
 import org.openmetadata.service.util.EntityUtil;
@@ -3227,6 +3228,13 @@ protected EmbeddingClient createEmbeddingClient(ElasticSearchConfiguration esCon
         }
         yield new OpenAIEmbeddingClient(esConfig);
       }
+      case "google" -> {
+        if (config.getGoogle() == null) {
+          throw new IllegalStateException(
+              "Google configuration is required when using google provider");
+        }
+        yield new GoogleEmbeddingClient(esConfig);
+      }
       case "djl" -> {
         if (config.getDjl() == null) {
           throw new IllegalStateException("DJL configuration is required when using djl provider");

openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/GoogleEmbeddingClient.java

@@ -0,0 +1,192 @@
+/*
+ *  Copyright 2024 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ *  except in compliance with the License. You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software distributed under the License
+ *  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ *  or implied. See the License for the specific language governing permissions and limitations under
+ *  the License.
+ */
+package org.openmetadata.service.search.vector.client;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URLEncoder;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
+import org.openmetadata.schema.service.configuration.elasticsearch.Google;
+import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
+
+@Slf4j
+public final class GoogleEmbeddingClient extends EmbeddingClient {
+  private static final ObjectMapper MAPPER = new ObjectMapper();
+  private static final String MODELS_PREFIX = "models/";
+  private static final String DEFAULT_BASE_URL =
+      "https://generativelanguage.googleapis.com/v1beta/" + MODELS_PREFIX;
+
+  private final HttpClient httpClient;
+  private final String apiKey;
+  private final String modelId;
+  private final int dimension;
+  private final String endpoint;
+
+  public GoogleEmbeddingClient(ElasticSearchConfiguration config) {
+    super(resolveMaxConcurrent(config));
+    NaturalLanguageSearchConfiguration nlsCfg = config.getNaturalLanguageSearch();
+    Google googleCfg = nlsCfg.getGoogle();
+    if (googleCfg == null) {
+      throw new IllegalArgumentException("Google configuration is required");
+    }
+    if (googleCfg.getApiKey() == null || googleCfg.getApiKey().isBlank()) {
+      throw new IllegalArgumentException("Google API key is required");
+    }
+    if (googleCfg.getEmbeddingModelId() == null || googleCfg.getEmbeddingModelId().isBlank()) {
+      throw new IllegalArgumentException("Google embedding model ID is required");
+    }
+    if (googleCfg.getEmbeddingDimension() == null || googleCfg.getEmbeddingDimension() <= 0) {
+      throw new IllegalArgumentException("Google embedding dimension must be positive");
+    }
+
+    this.apiKey = googleCfg.getApiKey();
+    this.modelId = googleCfg.getEmbeddingModelId();
+    this.dimension = googleCfg.getEmbeddingDimension();
+    this.endpoint = resolveEndpoint(googleCfg);
+    this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build();
+
+    LOG.info(
+        "Initialized GoogleEmbeddingClient with model={}, dimension={}, endpoint={}",
+        modelId,
+        dimension,
+        endpoint);
+  }
+
+  GoogleEmbeddingClient(
+      HttpClient httpClient, String apiKey, String modelId, int dimension, String endpoint) {
+    this(httpClient, apiKey, modelId, dimension, endpoint, DEFAULT_MAX_CONCURRENT_REQUESTS);
+  }
+
+  GoogleEmbeddingClient(
+      HttpClient httpClient,
+      String apiKey,
+      String modelId,
+      int dimension,
+      String endpoint,
+      int maxConcurrentRequests) {
+    super(maxConcurrentRequests);
+    this.httpClient = httpClient;
+    this.apiKey = apiKey;
+    this.modelId = modelId;
+    this.dimension = dimension;
+    this.endpoint = endpoint;
+  }
+
+  private String resolveEndpoint(Google config) {
+    String configured = config.getEndpoint();
+    if (configured != null && !configured.isBlank()) {
+      return configured.replaceAll("/+$", "");
+    }
+    return DEFAULT_BASE_URL + config.getEmbeddingModelId() + ":embedContent";
+  }
+
+  @Override
+  protected float[] doEmbed(String text) {
+    if (text == null || text.isBlank()) {
+      throw new IllegalArgumentException("Input text must not be null or blank");
+    }
+
+    try {
+      String body = buildRequestBody(text);
+      HttpRequest request = buildRequest(body);
+      HttpResponse<String> response =
+          httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+
+      if (response.statusCode() != 200) {
+        String errorMsg = extractErrorMessage(response.body());
+        throw new RuntimeException(
+            "Google API returned status " + response.statusCode() + ": " + errorMsg);
+      }
+
+      return parseEmbeddingResponse(response.body());
+    } catch (IOException e) {
+      LOG.error("IO error calling Google API: {}", e.getMessage(), e);
+      throw new RuntimeException("Google embedding generation failed due to IO error", e);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException("Google embedding generation was interrupted", e);
+    }
+  }
+
+  private String buildRequestBody(String text) throws IOException {
+    ObjectNode payload = MAPPER.createObjectNode();
+    payload.put("model", MODELS_PREFIX + modelId);
+    ObjectNode content = payload.putObject("content");
+    ArrayNode parts = content.putArray("parts");
+    ObjectNode part = parts.addObject();
+    part.put("text", text);
+    return MAPPER.writeValueAsString(payload);
+  }
+
+  private HttpRequest buildRequest(String body) {
+    String encodedKey = URLEncoder.encode(apiKey, StandardCharsets.UTF_8);
+    String url = endpoint + "?key=" + encodedKey;
+    return HttpRequest.newBuilder()
+        .uri(URI.create(url))
+        .header("Content-Type", "application/json")
+        .timeout(Duration.ofSeconds(30))
+        .POST(HttpRequest.BodyPublishers.ofString(body))
+        .build();
+  }
+
+  @Override
+  public int getDimension() {
+    return dimension;
+  }
+
+  @Override
+  public String getModelId() {
+    return modelId;
+  }
+
+  private float[] parseEmbeddingResponse(String responseBody) {
+    try {
+      JsonNode root = MAPPER.readTree(responseBody);
+      JsonNode embedding = root.get("embedding");
+      if (embedding == null || !embedding.isObject()) {
+        throw new RuntimeException("Invalid Google response: no embedding object found");
+      }
+      JsonNode values = embedding.get("values");
+      if (values == null || !values.isArray() || values.isEmpty()) {
+        throw new RuntimeException("Invalid Google response: no values array found");
+      }
+      float[] result = new float[values.size()];
+      for (int i = 0; i < values.size(); i++) {
+        result[i] = (float) values.get(i).asDouble();
+      }
+      return result;
+    } catch (IOException e) {
+      throw new RuntimeException("Failed to parse Google embedding response", e);
+    }
+  }
+
+  private String extractErrorMessage(String responseBody) {
+    try {
+      JsonNode root = MAPPER.readTree(responseBody);
+      JsonNode error = root.get("error");
+      if (error != null && error.has("message")) {
+        return error.get("message").asText();
+      }
+    } catch (Exception ignored) {
+    }
+    return responseBody;
+  }
+}