[AIENG-601] Client-Side Embedding generation poc implementation

MODULE.bazel

@@ -11,6 +11,7 @@ maven = use_extension("@rules_jvm_external//:extensions.bzl", "maven")
 maven.install(
     artifacts = [
         "args4j:args4j:2.33",  # can't go higher so long as we stay on Java 8
+        "com.knuddels:jtokkit:1.1.0",
         "com.fasterxml.jackson.core:jackson-annotations:2.18.2",
         "com.fasterxml.jackson.core:jackson-core:2.18.2",
         "com.fasterxml.jackson.core:jackson-databind:2.18.2",

smart_tests/commands/record/commit.py

@@ -14,7 +14,7 @@
 from ...app import Application
 from ...utils.commands import Command
 from ...utils.commit_ingester import upload_commits
-from ...utils.env_keys import COMMIT_TIMEOUT, REPORT_ERROR_KEY
+from ...utils.env_keys import COMMIT_TIMEOUT, EMBEDDING_API_KEY_KEY, EMBEDDING_ENDPOINT_KEY, EMBEDDING_MODEL_KEY, REPORT_ERROR_KEY
 from ...utils.fail_fast_mode import set_fail_fast_mode, warn_and_exit_if_fail_fast_mode
 from ...utils.git_log_parser import parse_git_log
 from ...utils.http_client import get_base_url
@@ -63,11 +63,25 @@ def commit(
     # Commit messages are not collected in the default.
     is_collect_message = False
     is_collect_files = False
+    embedding_mode = None
+    embedding_model = None
+    embedding_dimensions = None
+    embedding_augmentation = False
+    embedding_provider = None
+    embedding_endpoint = None
     try:
         res = client.request("get", "commits/collect/options")
         res.raise_for_status()
-        is_collect_message = res.json().get("commitMessage", False)
-        is_collect_files = res.json().get("files", False)
+        opts = res.json()
+        is_collect_message = opts.get("commitMessage", False)
+        is_collect_files = opts.get("files", False)
+        embedding_mode = opts.get("embeddingMode")
+        # env var overrides take precedence over server-provided values
+        embedding_model = os.getenv(EMBEDDING_MODEL_KEY) or opts.get("embeddingModel")
+        embedding_dimensions = opts.get("embeddingDimensions")
+        embedding_augmentation = opts.get("embeddingAugmentation", False)
+        embedding_provider = opts.get("embeddingProvider")
+        embedding_endpoint = os.getenv(EMBEDDING_ENDPOINT_KEY) or opts.get("embeddingEndpoint")
     except Exception as e:
         tracking_client.send_error_event(
             event_name=Tracking.ErrorEvent.INTERNAL_CLI_ERROR,
@@ -79,8 +93,36 @@ def commit(
     cwd = os.path.abspath(source)
     if not name:
         name = os.path.basename(cwd)
+
+    embedding_api_key = os.getenv(EMBEDDING_API_KEY_KEY)
+
+    if embedding_mode == "client":
+        if not embedding_endpoint:
+            warn_and_exit_if_fail_fast_mode(
+                f"Workspace requires client-side embeddings but no endpoint is configured. "
+                f"Set {EMBEDDING_ENDPOINT_KEY} to override.")
+            click.secho(
+                f"Warning: workspace requires client-side embeddings but no endpoint is configured. "
+                f"Set {EMBEDDING_ENDPOINT_KEY} to override. Skipping embeddings.",
+                fg="yellow", err=True)
+            embedding_mode = None
+        elif not embedding_api_key:
+            warn_and_exit_if_fail_fast_mode(
+                f"Workspace requires client-side embeddings but {EMBEDDING_API_KEY_KEY} is not set.")
+            click.secho(
+                f"Warning: workspace requires client-side embeddings but "
+                f"{EMBEDDING_API_KEY_KEY} is not set. Skipping embeddings.",
+                fg="yellow",
+                err=True)
+            embedding_mode = None
+
     try:
-        exec_jar(name, cwd, max_days, app, is_collect_message, is_collect_files)
+        exec_jar(name, cwd, max_days, app, is_collect_message, is_collect_files,
+                 embedding_endpoint if embedding_mode == "client" else None,
+                 embedding_model if embedding_mode == "client" else None,
+                 embedding_dimensions if embedding_mode == "client" else None,
+                 embedding_augmentation if embedding_mode == "client" else False,
+                 embedding_provider if embedding_mode == "client" else None)
     except Exception as e:
         if os.getenv(REPORT_ERROR_KEY):
             raise e
@@ -90,7 +132,10 @@ def commit(
                 "If not, please set a directory use by --source option.\nerror: {}".format(cwd, e))
 
 
-def exec_jar(name: str, source: str, max_days: int, app: Application, is_collect_message: bool, is_collect_files: bool):
+def exec_jar(name: str, source: str, max_days: int, app: Application, is_collect_message: bool, is_collect_files: bool,
+             embedding_endpoint: str | None = None, embedding_model: str | None = None,
+             embedding_dimensions: int | None = None, embedding_augmentation: bool = False,
+             embedding_provider: str | None = None):
     java = get_java_command()
 
     if not java:
@@ -126,6 +171,16 @@ def exec_jar(name: str, source: str, max_days: int, app: Application, is_collect
         command.append("-files")
     if os.getenv(COMMIT_TIMEOUT):
         command.append("-enable-timeout")
+    if embedding_endpoint:
+        command.extend(["-embedding-endpoint", embedding_endpoint])
+    if embedding_model:
+        command.extend(["-embedding-model", embedding_model])
+    if embedding_dimensions is not None:
+        command.extend(["-embedding-dimensions", str(embedding_dimensions)])
+    if embedding_augmentation:
+        command.append("-embedding-augmentation")
+    if embedding_provider:
+        command.extend(["-embedding-provider", embedding_provider])
     command.append(name)
     command.append(cygpath(source))
 

smart_tests/utils/env_keys.py

@@ -15,6 +15,11 @@
 SESSION_DIR_KEY = "SMART_TESTS_SESSION_DIR"
 CALLER_KEY = "SMART_TESTS_CALLER"
 
+EMBEDDING_ENDPOINT_KEY = "SMART_TESTS_EMBEDDING_ENDPOINT"
+EMBEDDING_API_KEY_KEY = "SMART_TESTS_EMBEDDING_API_KEY"
+# Optional: overrides the model name returned by the options endpoint
+EMBEDDING_MODEL_KEY = "SMART_TESTS_EMBEDDING_MODEL"
+
 # Legacy token key for backward compatibility
 LEGACY_TOKEN_KEY = "LAUNCHABLE_TOKEN"
 

src/main/java/com/launchableinc/ingest/commits/BUILD

@@ -6,6 +6,7 @@ java_library(
     name = "commits",
     srcs = glob(["*.java"]),
     deps = [
+        "//src/main/java/com/launchableinc/ingest/embedding",
         "@maven//:args4j_args4j",
         "@maven//:com_fasterxml_jackson_core_jackson_annotations",
         "@maven//:com_fasterxml_jackson_core_jackson_core",

src/main/java/com/launchableinc/ingest/commits/CommitGraphCollector.java

@@ -1,5 +1,8 @@
 package com.launchableinc.ingest.commits;
 
+import com.launchableinc.ingest.embedding.EmbeddingStrategy;
+import com.launchableinc.ingest.embedding.FileToEmbed;
+import com.launchableinc.ingest.embedding.FileEmbeddingResult;
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParser;
@@ -98,6 +101,8 @@ public class CommitGraphCollector {
 
   private int maxDays;
 
+  private EmbeddingStrategy embeddingStrategy;
+
   private boolean reportAllFiles;
 
   private boolean audit;
@@ -169,14 +174,18 @@ public void transfer(URL service, Authenticator authenticator, boolean enableTim
       ImmutableList<ObjectId> advertised = getAdvertisedRefs(latestResponse);
       honorControlHeaders(latestResponse);
 
-      // every time a new stream is needed, supply ByteArrayOutputStream, and when the data is all
-      // written, turn around and ship that over
-      transfer(
-        advertised,
-        (ContentProducer commits) -> sendCommits(service, client, commits),
-        new TreeReceiverImpl(service, client),
-        (ContentProducer files) -> sendFiles(service, client, files),
-        1024);
+      if (embeddingStrategy != null) {
+        transferWithEmbeddings(advertised, service, client);
+      } else {
+        // every time a new stream is needed, supply ByteArrayOutputStream, and when the data is all
+        // written, turn around and ship that over
+        transfer(
+          advertised,
+          (ContentProducer commits) -> sendCommits(service, client, commits),
+          new TreeReceiverImpl(service, client),
+          (ContentProducer files) -> sendFiles(service, client, files),
+          1024);
+      }
     }
   }
 
@@ -379,6 +388,83 @@ - then record commits
     }
   }
 
+  /**
+   * Embedding mode: collects files via the tree handshake, embeds them, uploads vectors.
+   * Commits are sent normally. TAR file upload is skipped.
+   */
+  private void transferWithEmbeddings(ImmutableList<ObjectId> advertised, URL service, LaunchableHttpClient client) throws IOException {
+    EmbeddingUploader uploader = new EmbeddingUploader();
+    EmbeddingFileConsumer embeddingConsumer = new EmbeddingFileConsumer(
+        embeddingStrategy, uploader, service, client);
+
+    ByRepository r = new ByRepository(root, rootName);
+    ExecutorService scanPool = new BoundedExecutorService(4);
+    ExecutorService transferPool = new BoundedExecutorService(4);
+
+    try {
+      r.forEachSubModule(scanPool, br -> {
+        try (ConcurrentConsumer<ContentProducer> parallel = new ConcurrentConsumer<>((ContentProducer cp) -> {}, transferPool);
+             FlushableConsumer<VirtualFile> fsr = fileTransferProgressReporter.newProducer(embeddingConsumer)) {
+          br.collectFiles(advertised, new TreeReceiverImpl(service, client), fsr);
+        }
+
+        try (CommitChunkStreamer cs = new CommitChunkStreamer((ContentProducer commits) -> sendCommits(service, client, commits), 1024)) {
+          br.collectCommits(advertised, cs);
+        }
+      });
+    } finally {
+      scanPool.shutdown();
+      transferPool.shutdown();
+    }
+
+    if (!embeddingConsumer.pending.isEmpty()) {
+      embeddingConsumer.flush();
+    }
+  }
+
+  /** Accumulates VirtualFile objects, converts to FileToEmbed, and on flush calls embed + upload. */
+  private class EmbeddingFileConsumer implements FlushableConsumer<VirtualFile> {
+    private final EmbeddingStrategy strategy;
+    private final EmbeddingUploader uploader;
+    private final URL service;
+    private final LaunchableHttpClient client;
+    final List<FileToEmbed> pending = new ArrayList<>();
+
+    EmbeddingFileConsumer(EmbeddingStrategy strategy, EmbeddingUploader uploader,
+                          URL service, LaunchableHttpClient client) {
+      this.strategy = strategy;
+      this.uploader = uploader;
+      this.service = service;
+      this.client = client;
+    }
+
+    @Override
+    public void accept(VirtualFile f) {
+      try {
+        ByteArrayOutputStream baos = new ByteArrayOutputStream((int) Math.min(f.size(), 1024 * 1024));
+        f.writeTo(baos);
+        String content = baos.toString("UTF-8");
+        pending.add(new FileToEmbed(f.path(), content, f.blob().name()));
+      } catch (IOException e) {
+        throw new UncheckedIOException(e);
+      }
+    }
+
+    @Override
+    public void flush() throws IOException {
+      if (pending.isEmpty()) return;
+      List<FileEmbeddingResult> results = strategy.embed(pending);
+      uploader.upload(service, client, results, strategy.modelName(), strategy.dimensions());
+      pending.clear();
+      filesSent.addAndGet(results.size());
+    }
+
+    @Override
+    public void close() throws IOException {
+      flush();
+    }
+  }
+
   public void collectCommitMessage(boolean commitMessage) {
     this.collectCommitMessage = commitMessage;
   }
@@ -399,6 +485,10 @@ public void collectFiles(boolean collectFiles) {
     this.collectFiles = collectFiles;
   }
 
+  public void setEmbeddingStrategy(EmbeddingStrategy embeddingStrategy) {
+    this.embeddingStrategy = embeddingStrategy;
+  }
+
   /** Process commits per repository. */
   final class ByRepository implements AutoCloseable {
     /** Names that uniquely identifies this Git repository among other Git repositories collected for the workspace. */

src/main/java/com/launchableinc/ingest/commits/EmbeddingUploader.java

@@ -0,0 +1,92 @@
+package com.launchableinc.ingest.commits;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.launchableinc.ingest.embedding.FileEmbeddingResult;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.entity.StringEntity;
+
+import java.io.IOException;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Uploads a batch of file embeddings to the server and polls until the async work completes.
+ *
+ * POST  .../collect/embeddings  → { "workId": N }
+ * GET   .../collect/files/work/{workId}  (existing poll endpoint, shared with TAR upload)
+ */
+class EmbeddingUploader {
+    private static final int POLL_INTERVAL_MS = 3000;
+    private static final ObjectMapper objectMapper = new ObjectMapper();
+
+    void upload(URL service, LaunchableHttpClient client,
+                List<FileEmbeddingResult> results,
+                String model, int dimensions) throws IOException {
+        URL url = new URL(service, "collect/embeddings");
+
+        JSEmbeddingCollectionRequest req = new JSEmbeddingCollectionRequest();
+        req.model = model;
+        req.dimensions = dimensions;
+        req.files = new ArrayList<>(results.size());
+        for (FileEmbeddingResult r : results) {
+            JSFileEmbedding fe = new JSFileEmbedding();
+            fe.fileName = r.fileName;
+            fe.blobSha = r.blobSha;
+            fe.embedding = r.embedding;
+            req.files.add(fe);
+        }
+
+        HttpPost request = new HttpPost(url.toExternalForm());
+        request.setHeader("Content-Type", "application/json");
+        request.setHeader("Accept", "application/json; mode=async");
+        request.setEntity(new StringEntity(objectMapper.writeValueAsString(req), StandardCharsets.UTF_8));
+
+        int workId = readResponse(client.httpPost(request), JSAsyncFileCollectionResponse.class).workId;
+        URL workUrl = new URL(service, "collect/files/work/" + workId);
+
+        while (true) {
+            try {
+                Thread.sleep(POLL_INTERVAL_MS);
+            } catch (InterruptedException e) {
+                throw new IOException("Interrupted while waiting for embedding upload", e);
+            }
+            JSAsyncFileCollectionProgress status =
+                    readResponse(client.httpGet(workUrl), JSAsyncFileCollectionProgress.class);
+            switch (status.status) {
+                case IN_PROGRESS:
+                    break;
+                case SUCCEEDED:
+                    return;
+                case FAILED:
+                case ABANDONED:
+                    throw new IOException("Embedding upload (workId=" + workId + ") failed: " + status.status);
+            }
+        }
+    }
+
+    private <T> T readResponse(CloseableHttpResponse response, Class<T> type) throws IOException {
+        try (JsonParser parser = new JsonFactory().createParser(response.getEntity().getContent())) {
+            return objectMapper.readValue(parser, type);
+        } finally {
+            response.close();
+        }
+    }
+
+    static class JSEmbeddingCollectionRequest {
+        @JsonProperty public String model;
+        @JsonProperty public int dimensions;
+        @JsonProperty public List<JSFileEmbedding> files;
+    }
+
+    static class JSFileEmbedding {
+        @JsonProperty public String fileName;
+        @JsonProperty public String blobSha;
+        @JsonProperty public float[] embedding;
+    }
+}

src/main/java/com/launchableinc/ingest/commits/LaunchableHttpClient.java

@@ -22,10 +22,10 @@
  * to raise an exception for every 4xx/5xx response. We'll wrap those idioms in this class
  * to keep {@link CommitGraphCollector} DRY and apply consistent behavior.
  */
-final class LaunchableHttpClient implements Closeable {
+public final class LaunchableHttpClient implements Closeable {
   final CloseableHttpClient core;
 
-  LaunchableHttpClient(CloseableHttpClient core) {
+  public LaunchableHttpClient(CloseableHttpClient core) {
     this.core = core;
   }