feat(search): new embedding entity + body text extension hook (#27233)

pmbrull · claude · web-flow · commit de950c63605b · 2026-04-14T08:16:10.000+02:00
* feat(search): enroll ContextMemory in vector embedding + body text extension hook

Add ContextMemory to AvailableEntityTypes.LIST so the existing vector
embedding pipeline (VectorEmbeddingHandler live path and the admin
reembed/reindex CLIs) iterates Collate memories alongside the built-in
data assets.

Memories store their semantic payload in title/question/answer/summary,
not description, so the default buildBodyText would feed an empty string
to the embedder. To fix that without pulling Collate schema classes into
OSS, expose a BodyTextExtractor functional interface plus a
registerBodyTextExtractor(entityType, extractor) hook on VectorDocBuilder.
buildBodyText consults the registry first and falls back to the existing
description-based logic when no extractor is registered. Collate
registers a typed extractor from its repository static initializer, so
both server and CLI paths go through the same code.

No behavior change for the default entity types: LIST gains one entry at
the end, and the registry starts empty so every existing call falls
through to the unchanged default branch.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* feat(search): add VectorBodyTextContributor SPI marker interface

Document the contract for plugging in entity-specific body text extractors
into the vector embedding pipeline. Implementations declare their entity
type, return a typed BodyTextExtractor, and call the default register()
hook from a stable initialization site (typically the owning
EntityRepository static initializer).

Pure documentation of the shape — the backing mechanism is still
VectorDocBuilder.registerBodyTextExtractor(), so callers that use the raw
registration hook keep working unchanged. The interface exists so new
contributors get an IDE nudge ("implement this") instead of having to
grep an existing extension class for the pattern.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* feat(search): enroll ContextMemory in vector embedding + body text extension hook

Add ContextMemory to AvailableEntityTypes.LIST so the existing vector
embedding pipeline (VectorEmbeddingHandler live path and the admin
reembed/reindex CLIs) iterates Collate memories alongside the built-in
data assets.

Memories store their semantic payload in title/question/answer/summary,
not description, so the default buildBodyText would feed an empty string
to the embedder. To fix that without pulling Collate schema classes into
OSS, expose a BodyTextExtractor functional interface plus a
registerBodyTextExtractor(entityType, extractor) hook on VectorDocBuilder.
buildBodyText consults the registry first and falls back to the existing
description-based logic when no extractor is registered. Collate
registers a typed extractor from its repository static initializer, so
both server and CLI paths go through the same code.

No behavior change for the default entity types: LIST gains one entry at
the end, and the registry starts empty so every existing call falls
through to the unchanged default branch.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* feat(search): add VectorBodyTextContributor SPI marker interface

Document the contract for plugging in entity-specific body text extractors
into the vector embedding pipeline. Implementations declare their entity
type, return a typed BodyTextExtractor, and call the default register()
hook from a stable initialization site (typically the owning
EntityRepository static initializer).

Pure documentation of the shape — the backing mechanism is still
VectorDocBuilder.registerBodyTextExtractor(), so callers that use the raw
registration hook keep working unchanged. The interface exists so new
contributors get an IDE nudge ("implement this") instead of having to
grep an existing extension class for the pattern.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* Fix NPE in buildBodyText when entityType is null

ConcurrentHashMap.get(null) throws NullPointerException. GlossaryTerm
tests pass null entityType through buildEmbeddingFields. Guard with a
null check and wrap custom extractor calls in try/catch so a faulty
downstream extractor degrades gracefully instead of crashing the
embedding pipeline.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* fix

* add tests

---------

Co-authored-by: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorBodyTextContributor.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorBodyTextContributor.java
@@ -0,0 +1,55 @@
+/*
+ *  Copyright 2025 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;
+
+import org.openmetadata.service.search.vector.VectorDocBuilder.BodyTextExtractor;
+
+/**
+ * Explicit contract for an entity type that wants to contribute a typed body text extractor to
+ * the vector embedding pipeline. Implementations are the single place where "this entity type
+ * has a custom body text extractor" is declared; downstream distributions implement this
+ * interface once per entity type and invoke {@link #register()} from a stable initialization
+ * site (typically the owning {@code EntityRepository} static initializer so both server and CLI
+ * lifecycles register the extractor exactly once).
+ *
+ * <p>This interface is intentionally tiny — it exists to document the shape of the contract and
+ * let IDEs guide contributors, not to impose an inheritance hierarchy. The backing mechanism is
+ * still {@link VectorDocBuilder#registerBodyTextExtractor(String, BodyTextExtractor)}; callers
+ * that already use the raw registration hook keep working unchanged.
+ */
+public interface VectorBodyTextContributor {
+
+  /**
+   * Entity type key as used in the entity reference (for example {@code "table"},
+   * {@code "contextMemory"}). Must match the value returned by {@code getEntityReference().getType()}
+   * on instances of the contributed entity class and the key used in
+   * {@code AvailableEntityTypes.LIST} — mismatches silently disable the custom extractor.
+   */
+  String entityType();
+
+  /**
+   * Typed body text extractor for this entity type. Runs on the hot path of every create / update
+   * and every reembed iteration; implementations should be fast and side-effect free.
+   */
+  BodyTextExtractor extractor();
+
+  /**
+   * Register this contributor's extractor with the shared {@link VectorDocBuilder} registry.
+   * Registration is idempotent, so calling it from a static initializer is safe even if the
+   * owning class is loaded multiple times (for example, once in the server path and once in a
+   * CLI subcommand).
+   */
+  default void register() {
+    VectorDocBuilder.registerBodyTextExtractor(entityType(), extractor());
+  }
+}
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorDocBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorDocBuilder.java
@@ -7,6 +7,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
 import lombok.extern.slf4j.Slf4j;
@@ -28,6 +29,41 @@
 @UtilityClass
 public class VectorDocBuilder {
 
+  /**
+   * Strategy for producing the semantic "body text" of an entity that will be chunked and fed to
+   * the embedding model. The default implementation concatenates {@code description} and, for
+   * tables, the column names — which works for every entity type whose semantic payload lives in
+   * {@code description}. Entity types whose payload is spread across other fields (for example
+   * Collate's {@code ContextMemory}, with title/question/answer/summary) can provide a typed
+   * extractor via {@link #registerBodyTextExtractor(String, BodyTextExtractor)} so the embedding
+   * pipeline uses their fields instead of an empty description.
+   */
+  @FunctionalInterface
+  public interface BodyTextExtractor {
+    /**
+     * Returns the body text for the given entity, or {@code null} to fall back to the default
+     * behavior. Implementations should be fast and side-effect free; they run on the hot path of
+     * every create/update and every reembed iteration.
+     */
+    String extract(EntityInterface entity);
+  }
+
+  private static final Map<String, BodyTextExtractor> BODY_TEXT_EXTRACTORS =
+      new ConcurrentHashMap<>();
+
+  /**
+   * Register a custom {@link BodyTextExtractor} for an entity type. The registry is consulted by
+   * {@link #buildBodyText(EntityInterface, String)} before the default description-based logic,
+   * so callers can cleanly override body text for their own entity types without patching this
+   * class. Registration is idempotent (last writer wins) and thread-safe.
+   */
+  public static void registerBodyTextExtractor(String entityType, BodyTextExtractor extractor) {
+    if (entityType == null || entityType.isBlank() || extractor == null) {
+      return;
+    }
+    BODY_TEXT_EXTRACTORS.put(entityType, extractor);
+  }
+
   public static List<Map<String, Object>> fromEntity(
       EntityInterface entity, EmbeddingClient embeddingClient) {
     Map<String, Object> doc = new HashMap<>(buildEmbeddingFields(entity, embeddingClient));
@@ -219,6 +255,21 @@ static String buildMetaLightText(EntityInterface entity, String entityType) {
   }
 
   static String buildBodyText(EntityInterface entity, String entityType) {
+    if (entityType != null) {
+      BodyTextExtractor customExtractor = BODY_TEXT_EXTRACTORS.get(entityType);
+      if (customExtractor != null) {
+        try {
+          String custom = customExtractor.extract(entity);
+          if (custom != null) {
+            return custom;
+          }
+        } catch (Exception e) {
+          LOG.warn(
+              "Custom BodyTextExtractor failed for [{}], falling back to default", entityType, e);
+        }
+      }
+    }
+
     List<String> bodyParts = new ArrayList<>();
     bodyParts.add("description: " + removeHtml(orEmpty(entity.getDescription())));
 
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/utils/AvailableEntityTypes.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/utils/AvailableEntityTypes.java
@@ -27,7 +27,8 @@ private AvailableEntityTypes() {}
           "page",
           "storedProcedure",
           "searchIndex",
-          "topic");
+          "topic",
+          "contextMemory");
 
   public static final Set<String> SET =
       LIST.stream().map(s -> s.toLowerCase(Locale.ROOT)).collect(Collectors.toUnmodifiableSet());
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorDocBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorDocBuilderTest.java
@@ -312,6 +312,92 @@ private GlossaryTerm createTestGlossaryTerm(String name, String displayName, Str
     return term;
   }
 
+  @Test
+  void testRegisterCustomExtractorIsUsed() {
+    String type = "customExtractorTest_" + UUID.randomUUID();
+    VectorDocBuilder.registerBodyTextExtractor(
+        type, entity -> "custom body for " + entity.getName());
+
+    Table table = createTestTable("ext_table", null, "original desc");
+    String body = VectorDocBuilder.buildBodyText(table, type);
+
+    assertEquals("custom body for ext_table", body);
+  }
+
+  @Test
+  void testCustomExtractorReturningNullFallsBackToDefault() {
+    String type = "nullExtractorTest_" + UUID.randomUUID();
+    VectorDocBuilder.registerBodyTextExtractor(type, entity -> null);
+
+    Table table = createTestTable("fallback_table", null, "fallback desc");
+    String body = VectorDocBuilder.buildBodyText(table, type);
+
+    assertTrue(body.contains("fallback desc"));
+  }
+
+  @Test
+  void testCustomExtractorThrowingFallsBackToDefault() {
+    String type = "throwExtractorTest_" + UUID.randomUUID();
+    VectorDocBuilder.registerBodyTextExtractor(
+        type,
+        entity -> {
+          throw new RuntimeException("boom");
+        });
+
+    Table table = createTestTable("err_table", null, "safe desc");
+    String body = VectorDocBuilder.buildBodyText(table, type);
+
+    assertTrue(body.contains("safe desc"));
+  }
+
+  @Test
+  void testRegisterExtractorIgnoresNullAndBlank() {
+    String type = "ignoreTest_" + UUID.randomUUID();
+
+    VectorDocBuilder.registerBodyTextExtractor(null, entity -> "nope");
+    VectorDocBuilder.registerBodyTextExtractor("", entity -> "nope");
+    VectorDocBuilder.registerBodyTextExtractor("  ", entity -> "nope");
+    VectorDocBuilder.registerBodyTextExtractor(type, null);
+
+    Table table = createTestTable("guard_table", null, "default desc");
+    String body = VectorDocBuilder.buildBodyText(table, type);
+
+    assertTrue(body.contains("default desc"));
+  }
+
+  @Test
+  void testVectorBodyTextContributorRegister() {
+    String type = "contributorTest_" + UUID.randomUUID();
+
+    VectorBodyTextContributor contributor =
+        new VectorBodyTextContributor() {
+          @Override
+          public String entityType() {
+            return type;
+          }
+
+          @Override
+          public VectorDocBuilder.BodyTextExtractor extractor() {
+            return entity -> "contributed: " + entity.getName();
+          }
+        };
+
+    contributor.register();
+
+    Table table = createTestTable("contrib_table", null, "ignored");
+    String body = VectorDocBuilder.buildBodyText(table, type);
+
+    assertEquals("contributed: contrib_table", body);
+  }
+
+  @Test
+  void testBuildBodyTextWithNullEntityType() {
+    Table table = createTestTable("null_type_table", null, "some desc");
+    String body = VectorDocBuilder.buildBodyText(table, null);
+
+    assertTrue(body.contains("some desc"));
+  }
+
   private Table createTestTable(String name, String displayName, String description) {
     Table table = new Table();
     table.setId(UUID.randomUUID());
