diff --git a/MODULE.bazel b/MODULE.bazel index d9e1ea494..3b5ac9266 100644 --- a/MODULE.bazel +++ b/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", diff --git a/smart_tests/commands/record/commit.py b/smart_tests/commands/record/commit.py index 9ce36a951..2a09fb042 100644 --- a/smart_tests/commands/record/commit.py +++ b/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)) diff --git a/smart_tests/jar/exe_deploy.jar b/smart_tests/jar/exe_deploy.jar index b0c3e41cd..cfd8a3eec 100755 Binary files a/smart_tests/jar/exe_deploy.jar and b/smart_tests/jar/exe_deploy.jar differ diff --git a/smart_tests/utils/env_keys.py b/smart_tests/utils/env_keys.py index 5d9f8e408..7904ea421 100644 --- a/smart_tests/utils/env_keys.py +++ b/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" diff --git a/src/main/java/com/launchableinc/ingest/commits/BUILD b/src/main/java/com/launchableinc/ingest/commits/BUILD index 91579a240..51258904e 100644 --- a/src/main/java/com/launchableinc/ingest/commits/BUILD +++ b/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", diff --git a/src/main/java/com/launchableinc/ingest/commits/CommitGraphCollector.java b/src/main/java/com/launchableinc/ingest/commits/CommitGraphCollector.java index fcfb1d2ee..525e1e0a8 100644 --- a/src/main/java/com/launchableinc/ingest/commits/CommitGraphCollector.java +++ b/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 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 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 parallel = new ConcurrentConsumer<>((ContentProducer cp) -> {}, transferPool); + FlushableConsumer 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 { + private final EmbeddingStrategy strategy; + private final EmbeddingUploader uploader; + private final URL service; + private final LaunchableHttpClient client; + final List 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 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. */ diff --git a/src/main/java/com/launchableinc/ingest/commits/EmbeddingUploader.java b/src/main/java/com/launchableinc/ingest/commits/EmbeddingUploader.java new file mode 100644 index 000000000..1ddb4ea95 --- /dev/null +++ b/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 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 readResponse(CloseableHttpResponse response, Class 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 files; + } + + static class JSFileEmbedding { + @JsonProperty public String fileName; + @JsonProperty public String blobSha; + @JsonProperty public float[] embedding; + } +} diff --git a/src/main/java/com/launchableinc/ingest/commits/LaunchableHttpClient.java b/src/main/java/com/launchableinc/ingest/commits/LaunchableHttpClient.java index 5274c0c02..e12299d87 100644 --- a/src/main/java/com/launchableinc/ingest/commits/LaunchableHttpClient.java +++ b/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; } diff --git a/src/main/java/com/launchableinc/ingest/commits/Main.java b/src/main/java/com/launchableinc/ingest/commits/Main.java index ef6b88f07..1cfc42747 100644 --- a/src/main/java/com/launchableinc/ingest/commits/Main.java +++ b/src/main/java/com/launchableinc/ingest/commits/Main.java @@ -1,12 +1,16 @@ package com.launchableinc.ingest.commits; import com.google.common.annotations.VisibleForTesting; +import com.launchableinc.ingest.embedding.EmbeddingStrategyFactory; import java.io.File; import java.io.IOException; import java.net.MalformedURLException; import java.net.URL; import java.util.logging.Logger; import java.util.logging.Level; +import org.apache.http.client.config.RequestConfig; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClientBuilder; import org.eclipse.jgit.api.Git; import org.eclipse.jgit.lib.Repository; import org.eclipse.jgit.lib.RepositoryBuilder; @@ -56,6 +60,21 @@ public class Main { @Option(name = "-enable-timeout", usage = "Enable timeout for the HTTP requests") public boolean enableTimeout; + @Option(name = "-embedding-endpoint", usage = "OpenAI-compatible endpoint for client-side embeddings") + public URL embeddingEndpoint; + + @Option(name = "-embedding-model", usage = "Embedding model name (e.g. text-embedding-3-small)") + public String embeddingModel; + + @Option(name = "-embedding-dimensions", usage = "Expected vector dimensions (e.g. 1536)") + public int embeddingDimensions = 0; + + @Option(name = "-embedding-augmentation", usage = "Prepend server-generated summaries to file content before embedding") + public boolean embeddingAugmentation; + + @Option(name = "-embedding-provider", usage = "Embedding provider: openai, azure_openai, or custom") + public String embeddingProvider; + private Authenticator authenticator; @VisibleForTesting String launchableToken = null; @@ -141,6 +160,25 @@ void run() throws CmdLineException, IOException { cgc.setDryRun(dryRun); cgc.collectCommitMessage(commitMessage); cgc.collectFiles(collectFiles); + if (embeddingEndpoint != null && embeddingModel != null && embeddingDimensions > 0) { + String apiKey = System.getenv("SMART_TESTS_EMBEDDING_API_KEY"); + if (apiKey != null && !apiKey.isEmpty()) { + URL summariesUrl = new URL(endpoint, "collect/summaries"); + // Plain client for the customer's LLM endpoint — must NOT carry Launchable auth. + CloseableHttpClient embeddingClient = HttpClientBuilder.create() + .useSystemProperties() + .build(); + // Auth-configured client for the summaries endpoint (Launchable-gated). + CloseableHttpClient summariesClient = HttpClientBuilder.create() + .useSystemProperties() + .setDefaultHeaders(authenticator.getAuthenticationHeaders()) + .build(); + cgc.setEmbeddingStrategy(EmbeddingStrategyFactory.create( + embeddingEndpoint, embeddingModel, embeddingDimensions, apiKey, + embeddingAugmentation, embeddingProvider, summariesUrl, + embeddingClient, summariesClient)); + } + } cgc.transfer(endpoint, authenticator, enableTimeout); int numCommits = cgc.getCommitsSent(); int numFiles = cgc.getFilesSent(); diff --git a/src/main/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategy.java b/src/main/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategy.java new file mode 100644 index 000000000..9e6853db9 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategy.java @@ -0,0 +1,55 @@ +package com.launchableinc.ingest.embedding; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +/** + * Decorator that prepends server-generated repo/directory summaries to each file's content + * before passing to the inner strategy. The prepend format matches the server template exactly + * (FileEmbeddingService augmentation logic). + */ +public class AugmentedEmbeddingStrategy implements EmbeddingStrategy { + private final EmbeddingStrategy inner; + private final RepoContextProvider contextProvider; + + public AugmentedEmbeddingStrategy(EmbeddingStrategy inner, RepoContextProvider contextProvider) { + this.inner = inner; + this.contextProvider = contextProvider; + } + + @Override + public String modelName() { + return inner.modelName(); + } + + @Override + public int dimensions() { + return inner.dimensions(); + } + + @Override + public List embed(List files) throws IOException { + Summaries summaries = contextProvider.getSummaries(files); + List augmented = new ArrayList<>(files.size()); + for (FileToEmbed f : files) { + augmented.add(new FileToEmbed(f.fileName, buildAugmentedContent(summaries, f), f.blobSha)); + } + return inner.embed(augmented); + } + + private String buildAugmentedContent(Summaries s, FileToEmbed f) { + String dir = extractParentDir(f.fileName); + return String.format( + "Repository summary: %s\nDirectory summary: %s\nFile name: %s\n----\n%s", + s.repoSummary, + s.dirSummaries.getOrDefault(dir, ""), + f.fileName, + f.content); + } + + private static String extractParentDir(String fileName) { + int slash = fileName.lastIndexOf('/'); + return slash > 0 ? fileName.substring(0, slash) : ""; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/BUILD b/src/main/java/com/launchableinc/ingest/embedding/BUILD new file mode 100644 index 000000000..39a10cdb4 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/BUILD @@ -0,0 +1,18 @@ +package( + default_visibility = ["//visibility:public"], +) + +java_library( + name = "embedding", + srcs = glob(["*.java"]), + deps = [ + "@maven//:com_fasterxml_jackson_core_jackson_annotations", + "@maven//:com_google_guava_guava", + "@maven//:com_fasterxml_jackson_core_jackson_core", + "@maven//:com_fasterxml_jackson_core_jackson_databind", + "@maven//:com_knuddels_jtokkit", + "@maven//:org_apache_httpcomponents_httpclient", + "@maven//:org_apache_httpcomponents_httpcore", + "@maven//:org_slf4j_slf4j_api", + ], +) diff --git a/src/main/java/com/launchableinc/ingest/embedding/Cl100kTokenizer.java b/src/main/java/com/launchableinc/ingest/embedding/Cl100kTokenizer.java new file mode 100644 index 000000000..1788c2a4b --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/Cl100kTokenizer.java @@ -0,0 +1,23 @@ +package com.launchableinc.ingest.embedding; + +import com.knuddels.jtokkit.Encodings; +import com.knuddels.jtokkit.api.Encoding; +import com.knuddels.jtokkit.api.EncodingType; + +public class Cl100kTokenizer implements Tokenizer { + private final Encoding encoding; + + public Cl100kTokenizer() { + this.encoding = Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.CL100K_BASE); + } + + @Override + public int countTokens(String text) { + return encoding.countTokens(text); + } + + @Override + public boolean isAccurate() { + return true; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategy.java b/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategy.java new file mode 100644 index 000000000..e5e1de757 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategy.java @@ -0,0 +1,12 @@ +package com.launchableinc.ingest.embedding; + +import java.io.IOException; +import java.util.List; + +public interface EmbeddingStrategy { + List embed(List files) throws IOException; + + String modelName(); + + int dimensions(); +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactory.java b/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactory.java new file mode 100644 index 000000000..93d65169a --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactory.java @@ -0,0 +1,45 @@ +package com.launchableinc.ingest.embedding; + +import org.apache.http.impl.client.CloseableHttpClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.MalformedURLException; +import java.net.URL; + +/** + * Builds the right EmbeddingStrategy from workspace options + env vars. + * When embeddingAugmentation is true, wraps RemoteEmbeddingStrategy in AugmentedEmbeddingStrategy. + */ +public class EmbeddingStrategyFactory { + private static final Logger logger = LoggerFactory.getLogger(EmbeddingStrategyFactory.class); + + /** + * @param provider embedding provider from server options ("openai", "azure_openai", "custom", or null) + * @param embeddingHttpClient plain unauthenticated client for the customer's LLM endpoint + * @param summariesHttpClient Launchable-authenticated client for the summaries endpoint + */ + public static EmbeddingStrategy create( + URL embeddingEndpoint, + String model, + int dimensions, + String apiKey, + boolean embeddingAugmentation, + String provider, + URL summariesUrl, + CloseableHttpClient embeddingHttpClient, + CloseableHttpClient summariesHttpClient) throws MalformedURLException { + + Tokenizer tokenizer = TokenizerFactory.create(provider, embeddingEndpoint); + EmbeddingStrategy strategy = new RemoteEmbeddingStrategy( + embeddingEndpoint, model, dimensions, apiKey, embeddingHttpClient, tokenizer); + + if (embeddingAugmentation) { + logger.info("Embedding augmentation enabled; will fetch summaries from server"); + RepoContextProvider contextProvider = new ServerRepoContextProvider(summariesUrl, summariesHttpClient); + strategy = new AugmentedEmbeddingStrategy(strategy, contextProvider); + } + + return strategy; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/FileEmbeddingResult.java b/src/main/java/com/launchableinc/ingest/embedding/FileEmbeddingResult.java new file mode 100644 index 000000000..f7679c460 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/FileEmbeddingResult.java @@ -0,0 +1,13 @@ +package com.launchableinc.ingest.embedding; + +public class FileEmbeddingResult { + public final String fileName; + public final String blobSha; + public final float[] embedding; + + public FileEmbeddingResult(String fileName, String blobSha, float[] embedding) { + this.fileName = fileName; + this.blobSha = blobSha; + this.embedding = embedding; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/FileToEmbed.java b/src/main/java/com/launchableinc/ingest/embedding/FileToEmbed.java new file mode 100644 index 000000000..df0305ab1 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/FileToEmbed.java @@ -0,0 +1,13 @@ +package com.launchableinc.ingest.embedding; + +public class FileToEmbed { + public final String fileName; + public final String content; + public final String blobSha; + + public FileToEmbed(String fileName, String content, String blobSha) { + this.fileName = fileName; + this.content = content; + this.blobSha = blobSha; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/NoopTokenizer.java b/src/main/java/com/launchableinc/ingest/embedding/NoopTokenizer.java new file mode 100644 index 000000000..894be9aad --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/NoopTokenizer.java @@ -0,0 +1,13 @@ +package com.launchableinc.ingest.embedding; + +public class NoopTokenizer implements Tokenizer { + @Override + public int countTokens(String text) { + return 0; + } + + @Override + public boolean isAccurate() { + return false; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategy.java b/src/main/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategy.java new file mode 100644 index 000000000..55a72202e --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategy.java @@ -0,0 +1,232 @@ +package com.launchableinc.ingest.embedding; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +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.google.common.io.CharStreams; +import com.google.common.util.concurrent.RateLimiter; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpPost; +import org.apache.http.entity.StringEntity; +import org.apache.http.impl.client.CloseableHttpClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.io.InputStreamReader; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public class RemoteEmbeddingStrategy implements EmbeddingStrategy { + private static final Logger logger = LoggerFactory.getLogger(RemoteEmbeddingStrategy.class); + + static final int MAX_FILES_PER_BATCH = 1900; + static final int MAX_TOKENS_PER_BATCH = 210_000; + static final int MAX_TOKENS_PER_FILE = 8100; + private static final int TRIM_LINES = 30; + private static final int MAX_RETRIES = 5; + static final long RETRY_BASE_MS = 1000; + static final double DEFAULT_RATE_LIMIT_TOKENS_PER_SEC = 150_000; + + protected long retryBaseMs() { return RETRY_BASE_MS; } + + private static final ObjectMapper objectMapper = new ObjectMapper(); + + private final URL endpoint; + private final String model; + private final int dims; + private final String apiKey; + private final CloseableHttpClient client; + private final Tokenizer tokenizer; + private final RateLimiter rateLimiter; + + public RemoteEmbeddingStrategy(URL endpoint, String model, int dims, String apiKey, + CloseableHttpClient client, Tokenizer tokenizer) { + this(endpoint, model, dims, apiKey, client, tokenizer, DEFAULT_RATE_LIMIT_TOKENS_PER_SEC); + } + + RemoteEmbeddingStrategy(URL endpoint, String model, int dims, String apiKey, + CloseableHttpClient client, Tokenizer tokenizer, double rateLimitTokensPerSec) { + this.endpoint = endpoint; + this.model = model; + this.dims = dims; + this.apiKey = apiKey; + this.client = client; + this.tokenizer = tokenizer; + this.rateLimiter = RateLimiter.create(rateLimitTokensPerSec); + } + + @Override + public String modelName() { + return model; + } + + @Override + public int dimensions() { + return dims; + } + + @Override + public List embed(List files) throws IOException { + List results = new ArrayList<>(); + List batch = new ArrayList<>(); + int batchTokens = 0; + + for (FileToEmbed file : files) { + if (file.content == null || file.content.isBlank()) continue; + FileToEmbed trimmed = trimToTokenLimit(file); + if (trimmed.content == null || trimmed.content.isBlank()) continue; + int fileTokens = tokenizer.isAccurate() ? tokenizer.countTokens(trimmed.content) : 0; + + boolean batchFull = batch.size() >= MAX_FILES_PER_BATCH + || (tokenizer.isAccurate() && batchTokens + fileTokens > MAX_TOKENS_PER_BATCH); + + if (!batch.isEmpty() && batchFull) { + results.addAll(flushBatch(batch, batchTokens)); + batch.clear(); + batchTokens = 0; + } + batch.add(trimmed); + batchTokens += fileTokens; + } + + if (!batch.isEmpty()) { + results.addAll(flushBatch(batch, batchTokens)); + } + return results; + } + + private List flushBatch(List batch, int tokenCount) throws IOException { + if (tokenizer.isAccurate()) { + rateLimiter.acquire(Math.max(1, tokenCount)); + } + return embedBatch(batch); + } + + /** Drops trailing TRIM_LINES lines until content fits within MAX_TOKENS_PER_FILE, matching server behavior. */ + private FileToEmbed trimToTokenLimit(FileToEmbed file) { + if (!tokenizer.isAccurate()) return file; + if (tokenizer.countTokens(file.content) <= MAX_TOKENS_PER_FILE) return file; + + List lines = new ArrayList<>(Arrays.asList(file.content.split("\n", -1))); + while (!lines.isEmpty() && tokenizer.countTokens(String.join("\n", lines)) > MAX_TOKENS_PER_FILE) { + int removeCount = Math.min(TRIM_LINES, lines.size()); + lines = lines.subList(0, lines.size() - removeCount); + } + return new FileToEmbed(file.fileName, String.join("\n", lines), file.blobSha); + } + + private List embedBatch(List batch) throws IOException { + String[] inputs = new String[batch.size()]; + for (int i = 0; i < batch.size(); i++) { + inputs[i] = batch.get(i).content; + } + + JSEmbeddingRequest requestBody = new JSEmbeddingRequest(); + requestBody.model = model; + requestBody.input = inputs; + + String json = objectMapper.writeValueAsString(requestBody); + + JSEmbeddingResponse response; + try { + response = executeWithRetry(json); + } catch (IOException e) { + // NoopTokenizer can't pre-count; split on 400 too_many_tokens as a last resort + if (!tokenizer.isAccurate() && batch.size() > 1 + && e.getMessage() != null && e.getMessage().contains("too_many_tokens")) { + logger.warn("too_many_tokens for batch of {}; splitting in half", batch.size()); + int mid = batch.size() / 2; + List combined = new ArrayList<>(); + combined.addAll(embedBatch(batch.subList(0, mid))); + combined.addAll(embedBatch(batch.subList(mid, batch.size()))); + return combined; + } + throw e; + } + + List results = new ArrayList<>(); + for (JSEmbeddingResponse.EmbeddingData data : response.data) { + FileToEmbed src = batch.get(data.index); + float[] normalized = l2Normalize(data.embedding); + results.add(new FileEmbeddingResult(src.fileName, src.blobSha, normalized)); + } + return results; + } + + private JSEmbeddingResponse executeWithRetry(String json) throws IOException { + IOException lastException = null; + for (int attempt = 0; attempt <= MAX_RETRIES; attempt++) { + if (attempt > 0) { + long delayMs = retryBaseMs() << (attempt - 1); + logger.warn("Retrying embedding request (attempt {}/{}), waiting {}ms", attempt, MAX_RETRIES, delayMs); + try { + Thread.sleep(delayMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted during retry backoff", e); + } + } + + HttpPost request = new HttpPost(endpoint.toExternalForm()); + request.setHeader("Content-Type", "application/json"); + request.setHeader("Authorization", "Bearer " + apiKey); + request.setEntity(new StringEntity(json, StandardCharsets.UTF_8)); + + try (CloseableHttpResponse response = client.execute(request)) { + int code = response.getStatusLine().getStatusCode(); + if (code >= 500 || code == 429) { + String body = CharStreams.toString( + new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8)); + lastException = new IOException(String.format( + "Embedding request to %s failed (attempt %d): %s%n%s", + endpoint, attempt + 1, response.getStatusLine(), body)); + continue; + } + if (code >= 400) { + String body = CharStreams.toString( + new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8)); + throw new IOException(String.format( + "Embedding request to %s failed: %s%n%s", endpoint, response.getStatusLine(), body)); + } + try (JsonParser parser = new JsonFactory().createParser(response.getEntity().getContent())) { + return objectMapper.readValue(parser, JSEmbeddingResponse.class); + } + } + } + throw new IOException("Embedding request to " + endpoint + " failed after " + MAX_RETRIES + " retries", lastException); + } + + private static float[] l2Normalize(float[] v) { + double sumSq = 0; + for (float x : v) sumSq += (double) x * x; + if (sumSq == 0) return v; + float norm = (float) Math.sqrt(sumSq); + float[] out = new float[v.length]; + for (int i = 0; i < v.length; i++) out[i] = v[i] / norm; + return out; + } + + // --- Jackson DTOs --- + + static class JSEmbeddingRequest { + @JsonProperty public String model; + @JsonProperty public String[] input; + } + + @JsonIgnoreProperties(ignoreUnknown = true) + static class JSEmbeddingResponse { + @JsonProperty public List data; + + @JsonIgnoreProperties(ignoreUnknown = true) + static class EmbeddingData { + @JsonProperty public int index; + @JsonProperty public float[] embedding; + } + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/RepoContextProvider.java b/src/main/java/com/launchableinc/ingest/embedding/RepoContextProvider.java new file mode 100644 index 000000000..5301def90 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/RepoContextProvider.java @@ -0,0 +1,8 @@ +package com.launchableinc.ingest.embedding; + +import java.io.IOException; +import java.util.List; + +public interface RepoContextProvider { + Summaries getSummaries(List files) throws IOException; +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/ServerRepoContextProvider.java b/src/main/java/com/launchableinc/ingest/embedding/ServerRepoContextProvider.java new file mode 100644 index 000000000..748a86692 --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/ServerRepoContextProvider.java @@ -0,0 +1,77 @@ +package com.launchableinc.ingest.embedding; + +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.google.common.io.CharStreams; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpPost; +import org.apache.http.entity.StringEntity; +import org.apache.http.impl.client.CloseableHttpClient; + +import java.io.IOException; +import java.io.InputStreamReader; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * Calls POST .../commits/collect/summaries to get server-generated repo and directory summaries. + * The server returns 404 when augmentation is disabled for the workspace. + */ +public class ServerRepoContextProvider implements RepoContextProvider { + private static final ObjectMapper objectMapper = new ObjectMapper(); + + private final URL summariesUrl; + private final CloseableHttpClient client; + + public ServerRepoContextProvider(URL summariesUrl, CloseableHttpClient client) { + this.summariesUrl = summariesUrl; + this.client = client; + } + + @Override + public Summaries getSummaries(List files) throws IOException { + JSRequest req = new JSRequest(); + req.tree = new ArrayList<>(files.size()); + for (FileToEmbed f : files) { + JSTreeEntry e = new JSTreeEntry(); + e.path = f.fileName; + req.tree.add(e); + } + + HttpPost request = new HttpPost(summariesUrl.toExternalForm()); + request.setHeader("Content-Type", "application/json"); + request.setEntity(new StringEntity(objectMapper.writeValueAsString(req), StandardCharsets.UTF_8)); + + try (CloseableHttpResponse response = client.execute(request)) { + int code = response.getStatusLine().getStatusCode(); + if (code >= 400) { + String body = CharStreams.toString( + new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8)); + throw new IOException(String.format( + "Summaries request failed: %s%n%s", response.getStatusLine(), body)); + } + try (JsonParser parser = new JsonFactory().createParser(response.getEntity().getContent())) { + JSResponse resp = objectMapper.readValue(parser, JSResponse.class); + return new Summaries(resp.repoSummary, resp.dirSummaries); + } + } + } + + static class JSRequest { + @JsonProperty public List tree; + } + + static class JSTreeEntry { + @JsonProperty public String path; + } + + static class JSResponse { + @JsonProperty public String repoSummary; + @JsonProperty public Map dirSummaries; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/Summaries.java b/src/main/java/com/launchableinc/ingest/embedding/Summaries.java new file mode 100644 index 000000000..73a015aeb --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/Summaries.java @@ -0,0 +1,13 @@ +package com.launchableinc.ingest.embedding; + +import java.util.Map; + +public class Summaries { + public final String repoSummary; + public final Map dirSummaries; + + public Summaries(String repoSummary, Map dirSummaries) { + this.repoSummary = repoSummary; + this.dirSummaries = dirSummaries; + } +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/Tokenizer.java b/src/main/java/com/launchableinc/ingest/embedding/Tokenizer.java new file mode 100644 index 000000000..54132a0af --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/Tokenizer.java @@ -0,0 +1,8 @@ +package com.launchableinc.ingest.embedding; + +public interface Tokenizer { + int countTokens(String text); + + /** Whether this counter is accurate (jtokkit) or a no-op (unknown provider). */ + boolean isAccurate(); +} diff --git a/src/main/java/com/launchableinc/ingest/embedding/TokenizerFactory.java b/src/main/java/com/launchableinc/ingest/embedding/TokenizerFactory.java new file mode 100644 index 000000000..75ddd2d6f --- /dev/null +++ b/src/main/java/com/launchableinc/ingest/embedding/TokenizerFactory.java @@ -0,0 +1,62 @@ +package com.launchableinc.ingest.embedding; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.URL; + +public class TokenizerFactory { + private static final Logger logger = LoggerFactory.getLogger(TokenizerFactory.class); + + /** + * Creates a tokenizer using the provider name returned by the server options response. + * Falls back to URL-host sniffing when provider is null. + */ + public static Tokenizer create(String provider, URL endpoint) { + String override = System.getenv("SMART_TESTS_EMBEDDING_TOKENIZER"); + if (override != null) { + switch (override) { + case "cl100k_base": + logger.info("Tokenizer: cl100k_base (forced via SMART_TESTS_EMBEDDING_TOKENIZER)"); + return new Cl100kTokenizer(); + case "none": + logger.info("Tokenizer: none (forced via SMART_TESTS_EMBEDDING_TOKENIZER)"); + return new NoopTokenizer(); + default: + throw new IllegalArgumentException( + "Unknown SMART_TESTS_EMBEDDING_TOKENIZER value: " + override + + ". Valid values: cl100k_base, none"); + } + } + + if (provider != null) { + switch (provider) { + case "openai": + case "azure_openai": + logger.info("Tokenizer: cl100k_base (provider={})", provider); + return new Cl100kTokenizer(); + case "custom": + logger.info("Tokenizer: none (provider=custom)"); + return new NoopTokenizer(); + default: + logger.info("Tokenizer: none (unknown provider {}). Set SMART_TESTS_EMBEDDING_TOKENIZER=cl100k_base if this endpoint serves an OpenAI-family model.", provider); + return new NoopTokenizer(); + } + } + + // Fallback: sniff from URL host + String host = endpoint.getHost(); + if ("api.openai.com".equals(host) || host.endsWith(".openai.azure.com")) { + logger.info("Tokenizer: cl100k_base (detected OpenAI/Azure endpoint)"); + return new Cl100kTokenizer(); + } + + logger.info("Tokenizer: none (unrecognized host {}). Set SMART_TESTS_EMBEDDING_TOKENIZER=cl100k_base if this endpoint serves an OpenAI-family model.", host); + return new NoopTokenizer(); + } + + /** Convenience overload when no provider is available (URL-sniffing only). */ + public static Tokenizer create(URL endpoint) { + return create(null, endpoint); + } +} diff --git a/src/maven_install.json b/src/maven_install.json index d701100a2..ab7d209ca 100644 --- a/src/maven_install.json +++ b/src/maven_install.json @@ -1,7 +1,7 @@ { "__AUTOGENERATED_FILE_DO_NOT_MODIFY_THIS_FILE_MANUALLY": "THERE_IS_NO_DATA_ONLY_ZUUL", - "__INPUT_ARTIFACTS_HASH": 1316397453, - "__RESOLVED_ARTIFACTS_HASH": -1407570614, + "__INPUT_ARTIFACTS_HASH": 894315978, + "__RESOLVED_ARTIFACTS_HASH": 1556218802, "conflict_resolution": { "com.google.errorprone:error_prone_annotations:2.3.2": "com.google.errorprone:error_prone_annotations:2.28.0", "com.google.guava:guava:31.1-jre": "com.google.guava:guava:33.3.1-jre", @@ -117,6 +117,12 @@ }, "version": "0.1.23" }, + "com.knuddels:jtokkit": { + "shasums": { + "jar": "1501ce0259ab897c6746ccfafa1d208acd404fb17e1ac62e157172f2678b1183" + }, + "version": "1.1.0" + }, "commons-cli:commons-cli": { "shasums": { "jar": "43f24850b7b7b7d79c5fa652418518fbdf427e602b1edabe6f11b85fb93eb013" @@ -711,6 +717,10 @@ "com.jcraft.jsch.jce", "com.jcraft.jsch.jcraft" ], + "com.knuddels:jtokkit": [ + "com.knuddels.jtokkit", + "com.knuddels.jtokkit.api" + ], "commons-cli:commons-cli": [ "org.apache.commons.cli" ], @@ -2408,6 +2418,7 @@ "com.google.truth:truth", "com.googlecode.javaewah:JavaEWAH", "com.jcraft:jsch", + "com.knuddels:jtokkit", "commons-cli:commons-cli", "commons-codec:commons-codec", "commons-io:commons-io", diff --git a/src/test/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategyTest.java b/src/test/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategyTest.java new file mode 100644 index 000000000..3421f13ec --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/AugmentedEmbeddingStrategyTest.java @@ -0,0 +1,115 @@ +package com.launchableinc.ingest.embedding; + +import static com.google.common.truth.Truth.assertThat; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +@RunWith(JUnit4.class) +public class AugmentedEmbeddingStrategyTest { + + /** Captures the augmented inputs that were passed to the inner strategy. */ + private static class CapturingStrategy implements EmbeddingStrategy { + List capturedFiles; + + @Override + public List embed(List files) { + this.capturedFiles = files; + return Collections.emptyList(); + } + + @Override public String modelName() { return "test"; } + @Override public int dimensions() { return 2; } + } + + @Test + public void prependsRepoAndDirSummaryWithExactFormat() throws IOException { + CapturingStrategy inner = new CapturingStrategy(); + + Map dirSummaries = new HashMap<>(); + dirSummaries.put("src/foo", "Foo utilities"); + Summaries summaries = new Summaries("Repo about testing", dirSummaries); + + RepoContextProvider provider = files -> summaries; + + AugmentedEmbeddingStrategy strategy = new AugmentedEmbeddingStrategy(inner, provider); + + List files = Collections.singletonList( + new FileToEmbed("src/foo/Bar.java", "class Bar {}", "sha1")); + + strategy.embed(files); + + assertThat(inner.capturedFiles).hasSize(1); + String augmented = inner.capturedFiles.get(0).content; + assertThat(augmented).isEqualTo( + "Repository summary: Repo about testing\n" + + "Directory summary: Foo utilities\n" + + "File name: src/foo/Bar.java\n" + + "----\n" + + "class Bar {}"); + } + + @Test + public void usesEmptyDirSummaryWhenNoMatchingDirectory() throws IOException { + CapturingStrategy inner = new CapturingStrategy(); + + Summaries summaries = new Summaries("Repo summary", Collections.emptyMap()); + RepoContextProvider provider = files -> summaries; + + AugmentedEmbeddingStrategy strategy = new AugmentedEmbeddingStrategy(inner, provider); + + List files = Collections.singletonList( + new FileToEmbed("Root.java", "class Root {}", "sha1")); + + strategy.embed(files); + + String augmented = inner.capturedFiles.get(0).content; + assertThat(augmented).isEqualTo( + "Repository summary: Repo summary\n" + + "Directory summary: \n" + + "File name: Root.java\n" + + "----\n" + + "class Root {}"); + } + + @Test + public void preservesBlobShaAndFileName() throws IOException { + CapturingStrategy inner = new CapturingStrategy(); + RepoContextProvider provider = files -> new Summaries("r", Collections.emptyMap()); + + AugmentedEmbeddingStrategy strategy = new AugmentedEmbeddingStrategy(inner, provider); + + List files = Arrays.asList( + new FileToEmbed("a/A.java", "A content", "sha-a"), + new FileToEmbed("b/B.java", "B content", "sha-b")); + + strategy.embed(files); + + assertThat(inner.capturedFiles.get(0).fileName).isEqualTo("a/A.java"); + assertThat(inner.capturedFiles.get(0).blobSha).isEqualTo("sha-a"); + assertThat(inner.capturedFiles.get(1).fileName).isEqualTo("b/B.java"); + assertThat(inner.capturedFiles.get(1).blobSha).isEqualTo("sha-b"); + } + + @Test + public void delegatesModelNameAndDimensions() { + EmbeddingStrategy inner = new CapturingStrategy() { + @Override public String modelName() { return "text-embedding-3-small"; } + @Override public int dimensions() { return 1536; } + }; + RepoContextProvider provider = files -> new Summaries("", Collections.emptyMap()); + + AugmentedEmbeddingStrategy strategy = new AugmentedEmbeddingStrategy(inner, provider); + + assertThat(strategy.modelName()).isEqualTo("text-embedding-3-small"); + assertThat(strategy.dimensions()).isEqualTo(1536); + } +} diff --git a/src/test/java/com/launchableinc/ingest/embedding/BUILD b/src/test/java/com/launchableinc/ingest/embedding/BUILD new file mode 100644 index 000000000..591968bbc --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/BUILD @@ -0,0 +1,12 @@ +java_test( + name = "EmbeddingTests", + srcs = glob(["*.java"]), + test_class = "com.launchableinc.ingest.embedding.EmbeddingAllTests", + deps = [ + "//src/main/java/com/launchableinc/ingest/embedding", + "@maven//:com_google_truth_truth", + "@maven//:junit_junit", + "@maven//:org_apache_httpcomponents_httpclient", + "@maven//:org_mock_server_mockserver_junit_rule_no_dependencies", + ], +) diff --git a/src/test/java/com/launchableinc/ingest/embedding/EmbeddingAllTests.java b/src/test/java/com/launchableinc/ingest/embedding/EmbeddingAllTests.java new file mode 100644 index 000000000..77bc48e34 --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/EmbeddingAllTests.java @@ -0,0 +1,14 @@ +package com.launchableinc.ingest.embedding; + +import org.junit.runner.RunWith; +import org.junit.runners.Suite; +import org.junit.runners.Suite.SuiteClasses; + +@RunWith(Suite.class) +@SuiteClasses({ + AugmentedEmbeddingStrategyTest.class, + EmbeddingStrategyFactoryTest.class, + RemoteEmbeddingStrategyTest.class, + ServerRepoContextProviderTest.class, +}) +public class EmbeddingAllTests {} diff --git a/src/test/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactoryTest.java b/src/test/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactoryTest.java new file mode 100644 index 000000000..557b98e21 --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/EmbeddingStrategyFactoryTest.java @@ -0,0 +1,106 @@ +package com.launchableinc.ingest.embedding; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockserver.model.HttpRequest.request; +import static org.mockserver.model.HttpResponse.response; + +import org.apache.http.impl.client.HttpClients; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockserver.client.MockServerClient; +import org.mockserver.junit.MockServerRule; + +import java.net.InetSocketAddress; +import java.net.URL; + +@RunWith(JUnit4.class) +public class EmbeddingStrategyFactoryTest { + + @Rule public MockServerRule mockServerRule = new MockServerRule(this); + private MockServerClient mockServerClient; + + @Test + public void createReturnsRemoteStrategyWithoutAugmentation() throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + URL summariesUrl = new URL(String.format("http://%s:%d/collect/summaries", addr.getHostString(), addr.getPort())); + + EmbeddingStrategy strategy = EmbeddingStrategyFactory.create( + endpoint, "text-embedding-3-small", 1536, "api-key", + false, null, summariesUrl, HttpClients.createDefault(), HttpClients.createDefault()); + + assertThat(strategy).isInstanceOf(RemoteEmbeddingStrategy.class); + assertThat(strategy.modelName()).isEqualTo("text-embedding-3-small"); + assertThat(strategy.dimensions()).isEqualTo(1536); + } + + @Test + public void createWrapsInAugmentedStrategyWhenEnabled() throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + URL summariesUrl = new URL(String.format("http://%s:%d/collect/summaries", addr.getHostString(), addr.getPort())); + + EmbeddingStrategy strategy = EmbeddingStrategyFactory.create( + endpoint, "text-embedding-3-small", 1536, "api-key", + true, null, summariesUrl, HttpClients.createDefault(), HttpClients.createDefault()); + + assertThat(strategy).isInstanceOf(AugmentedEmbeddingStrategy.class); + assertThat(strategy.modelName()).isEqualTo("text-embedding-3-small"); + assertThat(strategy.dimensions()).isEqualTo(1536); + } + + @Test + public void endToEndWithoutAugmentation() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[0.6,0.0,0.8]}]}")); + + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + URL summariesUrl = new URL(String.format("http://%s:%d/collect/summaries", addr.getHostString(), addr.getPort())); + + EmbeddingStrategy strategy = EmbeddingStrategyFactory.create( + endpoint, "text-embedding-3-small", 3, "api-key", + false, null, summariesUrl, HttpClients.createDefault(), HttpClients.createDefault()); + + java.util.List results = strategy.embed( + java.util.Collections.singletonList(new FileToEmbed("A.java", "class A {}", "sha1"))); + + assertThat(results).hasSize(1); + assertThat(results.get(0).fileName).isEqualTo("A.java"); + // vector [0.6, 0.0, 0.8] has norm 1.0, so already normalized + assertThat(results.get(0).embedding[0]).isWithin(1e-5f).of(0.6f); + } + + @Test + public void providerOpenaiPicksCl100kTokenizer() throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + URL summariesUrl = new URL(String.format("http://%s:%d/collect/summaries", addr.getHostString(), addr.getPort())); + + // localhost would normally pick NoopTokenizer; provider="openai" overrides that + EmbeddingStrategy strategy = EmbeddingStrategyFactory.create( + endpoint, "text-embedding-3-small", 1536, "api-key", + false, "openai", summariesUrl, HttpClients.createDefault(), HttpClients.createDefault()); + + assertThat(strategy).isInstanceOf(RemoteEmbeddingStrategy.class); + } + + @Test + public void providerCustomPicksNoopTokenizer() throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + URL summariesUrl = new URL(String.format("http://%s:%d/collect/summaries", addr.getHostString(), addr.getPort())); + + EmbeddingStrategy strategy = EmbeddingStrategyFactory.create( + endpoint, "my-model", 512, "api-key", + false, "custom", summariesUrl, HttpClients.createDefault(), HttpClients.createDefault()); + + assertThat(strategy).isInstanceOf(RemoteEmbeddingStrategy.class); + } +} diff --git a/src/test/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategyTest.java b/src/test/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategyTest.java new file mode 100644 index 000000000..620df27d7 --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/RemoteEmbeddingStrategyTest.java @@ -0,0 +1,274 @@ +package com.launchableinc.ingest.embedding; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockserver.model.HttpRequest.request; +import static org.mockserver.model.HttpResponse.response; + +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockserver.client.MockServerClient; +import org.mockserver.junit.MockServerRule; + +import java.net.InetSocketAddress; +import java.net.URL; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +@RunWith(JUnit4.class) +public class RemoteEmbeddingStrategyTest { + @Rule public MockServerRule mockServerRule = new MockServerRule(this); + private MockServerClient mockServerClient; + + private static final int DIMS = 2; + + @Test + public void embedsSingleBatch() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[1.0,0.0]},{\"index\":1,\"embedding\":[0.0,1.0]}]}")); + + RemoteEmbeddingStrategy strategy = buildStrategy(false); + + List files = Arrays.asList( + new FileToEmbed("a.java", "class A {}", "abc1"), + new FileToEmbed("b.java", "class B {}", "abc2") + ); + + List results = strategy.embed(files); + + assertThat(results).hasSize(2); + assertThat(results.get(0).fileName).isEqualTo("a.java"); + assertThat(results.get(0).blobSha).isEqualTo("abc1"); + assertThat(results.get(1).fileName).isEqualTo("b.java"); + assertThat(results.get(1).blobSha).isEqualTo("abc2"); + // [1.0, 0.0] is already unit length + assertThat(results.get(0).embedding[0]).isWithin(1e-6f).of(1.0f); + assertThat(results.get(0).embedding[1]).isWithin(1e-6f).of(0.0f); + } + + @Test + public void l2NormalizesVectors() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[3.0,4.0]}]}")); + + RemoteEmbeddingStrategy strategy = buildStrategy(false); + + List results = strategy.embed( + Collections.singletonList(new FileToEmbed("f.java", "content", "sha1"))); + + assertThat(results).hasSize(1); + float[] v = results.get(0).embedding; + // [3,4] normalized = [0.6, 0.8] + assertThat(v[0]).isWithin(1e-5f).of(0.6f); + assertThat(v[1]).isWithin(1e-5f).of(0.8f); + } + + @Test + public void splitsOnTooManyTokensWhenNoopTokenizer() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings"), org.mockserver.matchers.Times.once()) + .respond(response() + .withStatusCode(400) + .withHeader("Content-Type", "application/json") + .withBody("{\"error\":{\"message\":\"too_many_tokens\",\"type\":\"invalid_request_error\"}}")); + + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[0.6,0.8]}]}")); + + RemoteEmbeddingStrategy strategy = buildStrategy(false); + + List files = Arrays.asList( + new FileToEmbed("a.java", "class A {}", "sha1"), + new FileToEmbed("b.java", "class B {}", "sha2") + ); + + List results = strategy.embed(files); + assertThat(results).hasSize(2); + // [0.6, 0.8] is already unit length + assertThat(results.get(0).embedding[0]).isWithin(1e-5f).of(0.6f); + assertThat(results.get(0).embedding[1]).isWithin(1e-5f).of(0.8f); + } + + @Test + public void retriesOn5xxThenSucceeds() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings"), org.mockserver.matchers.Times.once()) + .respond(response().withStatusCode(503).withBody("Service Unavailable")); + + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[1.0,0.0]}]}")); + + // Use a tiny retry base so the test doesn't sleep long + RemoteEmbeddingStrategy strategy = buildStrategyWithRetryBase(false, 1); + + List results = strategy.embed( + Collections.singletonList(new FileToEmbed("f.java", "content", "sha1"))); + + assertThat(results).hasSize(1); + } + + @Test + public void retriesOn429ThenSucceeds() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings"), org.mockserver.matchers.Times.once()) + .respond(response().withStatusCode(429).withBody("Rate limited")); + + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[1.0,0.0]}]}")); + + RemoteEmbeddingStrategy strategy = buildStrategyWithRetryBase(false, 1); + + List results = strategy.embed( + Collections.singletonList(new FileToEmbed("f.java", "content", "sha1"))); + + assertThat(results).hasSize(1); + } + + @Test + public void doesNotRetryOn4xx() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response().withStatusCode(401).withBody("Unauthorized")); + + RemoteEmbeddingStrategy strategy = buildStrategy(false); + + try { + strategy.embed(Collections.singletonList(new FileToEmbed("f.java", "content", "sha1"))); + throw new AssertionError("Expected IOException"); + } catch (java.io.IOException e) { + assertThat(e.getMessage()).contains("401"); + } + } + + @Test + public void trimsFileExceedingTokenLimit() throws Exception { + // Cl100kTokenizer is accurate; build a file that is over 8100 tokens + // We can't easily make a real 8100-token file in a unit test, so we use + // a strategy with a custom tokenizer that reports the file as over-limit + // on first call but under-limit after trimming. + mockServerClient + .when(request().withMethod("POST").withPath("/v1/embeddings")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"data\":[{\"index\":0,\"embedding\":[1.0,0.0]}]}")); + + // Tokenizer that reports content over limit until it is short enough + Tokenizer stubbedTokenizer = new Tokenizer() { + @Override public boolean isAccurate() { return true; } + @Override public int countTokens(String text) { + // Count newlines as a proxy: >32 newlines = over limit (stays non-empty after TRIM_LINES=30 removal) + long newlines = text.chars().filter(c -> c == '\n').count(); + return newlines > 32 ? RemoteEmbeddingStrategy.MAX_TOKENS_PER_FILE + 1 : 1; + } + }; + + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + + try (CloseableHttpClient httpClient = HttpClients.createDefault()) { + RemoteEmbeddingStrategy strategy = new RemoteEmbeddingStrategy( + endpoint, "text-embedding-3-small", DIMS, "test-key", httpClient, stubbedTokenizer); + + // 40 lines so that after one TRIM_LINES=30 removal 10 remain (non-empty) + StringBuilder sb = new StringBuilder(); + for (int i = 1; i <= 40; i++) sb.append("line").append(i).append("\n"); + String longContent = sb.toString().stripTrailing(); + List results = strategy.embed( + Collections.singletonList(new FileToEmbed("f.java", longContent, "sha1"))); + + assertThat(results).hasSize(1); + } + } + + @Test + public void tokenizerFactoryPicksCl100kForOpenAI() throws Exception { + Tokenizer t = TokenizerFactory.create(new URL("https://api.openai.com/v1/embeddings")); + assertThat(t).isInstanceOf(Cl100kTokenizer.class); + assertThat(t.isAccurate()).isTrue(); + } + + @Test + public void tokenizerFactoryPicksCl100kForAzure() throws Exception { + Tokenizer t = TokenizerFactory.create(new URL("https://mydeployment.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings")); + assertThat(t).isInstanceOf(Cl100kTokenizer.class); + } + + @Test + public void tokenizerFactoryPicksNoopForUnknownHost() throws Exception { + Tokenizer t = TokenizerFactory.create(new URL("http://localhost:8080/v1/embeddings")); + assertThat(t).isInstanceOf(NoopTokenizer.class); + assertThat(t.isAccurate()).isFalse(); + } + + @Test + public void tokenizerFactoryProviderOpenaiPicksCl100k() throws Exception { + URL localEndpoint = new URL("http://localhost:8080/v1/embeddings"); + Tokenizer t = TokenizerFactory.create("openai", localEndpoint); + assertThat(t).isInstanceOf(Cl100kTokenizer.class); + assertThat(t.isAccurate()).isTrue(); + } + + @Test + public void tokenizerFactoryProviderAzureOpenaiPicksCl100k() throws Exception { + URL localEndpoint = new URL("http://localhost:8080/v1/embeddings"); + Tokenizer t = TokenizerFactory.create("azure_openai", localEndpoint); + assertThat(t).isInstanceOf(Cl100kTokenizer.class); + assertThat(t.isAccurate()).isTrue(); + } + + @Test + public void tokenizerFactoryProviderCustomPicksNoop() throws Exception { + Tokenizer t = TokenizerFactory.create("custom", new URL("https://api.openai.com/v1/embeddings")); + assertThat(t).isInstanceOf(NoopTokenizer.class); + assertThat(t.isAccurate()).isFalse(); + } + + @Test + public void tokenizerFactoryNullProviderFallsBackToUrlSniffing() throws Exception { + Tokenizer t = TokenizerFactory.create(null, new URL("https://api.openai.com/v1/embeddings")); + assertThat(t).isInstanceOf(Cl100kTokenizer.class); + } + + // --- helpers --- + + private RemoteEmbeddingStrategy buildStrategy(boolean accurate) throws Exception { + return buildStrategyWithRetryBase(accurate, RemoteEmbeddingStrategy.RETRY_BASE_MS); + } + + private RemoteEmbeddingStrategy buildStrategyWithRetryBase(boolean accurate, long retryBaseMs) throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL endpoint = new URL(String.format("http://%s:%d/v1/embeddings", addr.getHostString(), addr.getPort())); + CloseableHttpClient httpClient = HttpClients.createDefault(); + Tokenizer tokenizer = accurate ? new Cl100kTokenizer() : new NoopTokenizer(); + return new RemoteEmbeddingStrategy(endpoint, "text-embedding-3-small", DIMS, "test-key", + httpClient, tokenizer, RemoteEmbeddingStrategy.DEFAULT_RATE_LIMIT_TOKENS_PER_SEC) { + @Override + protected long retryBaseMs() { return retryBaseMs; } + }; + } +} diff --git a/src/test/java/com/launchableinc/ingest/embedding/ServerRepoContextProviderTest.java b/src/test/java/com/launchableinc/ingest/embedding/ServerRepoContextProviderTest.java new file mode 100644 index 000000000..79f045ab9 --- /dev/null +++ b/src/test/java/com/launchableinc/ingest/embedding/ServerRepoContextProviderTest.java @@ -0,0 +1,92 @@ +package com.launchableinc.ingest.embedding; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import static org.mockserver.model.HttpRequest.request; +import static org.mockserver.model.HttpResponse.response; + +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockserver.client.MockServerClient; +import org.mockserver.junit.MockServerRule; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.URL; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +@RunWith(JUnit4.class) +public class ServerRepoContextProviderTest { + + @Rule public MockServerRule mockServerRule = new MockServerRule(this); + private MockServerClient mockServerClient; + + @Test + public void returnsSummariesFromServer() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/collect/summaries")) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"repoSummary\":\"A Java repo\",\"dirSummaries\":{\"src\":\"Source files\",\"src/foo\":\"Foo package\"}}")); + + ServerRepoContextProvider provider = buildProvider("/collect/summaries"); + + List files = Arrays.asList( + new FileToEmbed("src/A.java", "class A {}", "sha1"), + new FileToEmbed("src/foo/B.java", "class B {}", "sha2")); + + Summaries summaries = provider.getSummaries(files); + + assertThat(summaries.repoSummary).isEqualTo("A Java repo"); + assertThat(summaries.dirSummaries).containsEntry("src", "Source files"); + assertThat(summaries.dirSummaries).containsEntry("src/foo", "Foo package"); + } + + @Test + public void sendsFilePathsInRequestBody() throws Exception { + mockServerClient + .when(request() + .withMethod("POST") + .withPath("/collect/summaries") + .withBody(org.mockserver.model.JsonBody.json( + "{\"tree\":[{\"path\":\"src/Foo.java\"},{\"path\":\"Bar.java\"}]}"))) + .respond(response() + .withStatusCode(200) + .withHeader("Content-Type", "application/json") + .withBody("{\"repoSummary\":\"\",\"dirSummaries\":{}}")); + + ServerRepoContextProvider provider = buildProvider("/collect/summaries"); + + provider.getSummaries(Arrays.asList( + new FileToEmbed("src/Foo.java", "content", "sha1"), + new FileToEmbed("Bar.java", "content", "sha2"))); + // If the body didn't match, mockserver would return 404 and the above would throw + } + + @Test + public void throwsOnNon2xx() throws Exception { + mockServerClient + .when(request().withMethod("POST").withPath("/collect/summaries")) + .respond(response().withStatusCode(404).withBody("Not found")); + + ServerRepoContextProvider provider = buildProvider("/collect/summaries"); + + assertThrows(IOException.class, () -> + provider.getSummaries(Collections.singletonList( + new FileToEmbed("A.java", "content", "sha1")))); + } + + private ServerRepoContextProvider buildProvider(String path) throws Exception { + InetSocketAddress addr = mockServerClient.remoteAddress(); + URL url = new URL(String.format("http://%s:%d%s", addr.getHostString(), addr.getPort(), path)); + CloseableHttpClient client = HttpClients.createDefault(); + return new ServerRepoContextProvider(url, client); + } +} diff --git a/tests/commands/record/test_commit_byollm.py b/tests/commands/record/test_commit_byollm.py new file mode 100644 index 000000000..0e118c14c --- /dev/null +++ b/tests/commands/record/test_commit_byollm.py @@ -0,0 +1,208 @@ +""" +Smoke tests for the BYOLLM (bring-your-own-LLM) embedding path in `record commit`. + +These tests mock: +- The options endpoint (returns embeddingMode=client + model/dimensions/augmentation) +- exec_jar (captures the arguments without actually running Java) + +They verify that the Python CLI: +1. Passes embedding_endpoint / embedding_model / embedding_dimensions to exec_jar +2. Passes embedding_augmentation=True only when the workspace has it enabled +3. Skips embeddings (passes None) when SMART_TESTS_EMBEDDING_ENDPOINT is unset +4. Skips embeddings when SMART_TESTS_EMBEDDING_API_KEY is unset +5. SMART_TESTS_EMBEDDING_MODEL env var overrides the model name from the options endpoint +6. embeddingMode=server → no embedding args passed +""" + +import os +from unittest import mock + +import responses + +# Ensure the record package is imported so that sys.modules contains +# smart_tests.commands.record.commit as a module (not the Command object). +import smart_tests.commands.record # noqa: F401 +from smart_tests.utils.env_keys import EMBEDDING_API_KEY_KEY, EMBEDDING_ENDPOINT_KEY, EMBEDDING_MODEL_KEY +from smart_tests.utils.http_client import get_base_url +from tests.cli_test_case import CliTestCase + +_EXEC_JAR = "smart_tests.commands.record.commit.exec_jar" + + +def _options_response(embedding_mode=None, embedding_model=None, + embedding_dimensions=None, embedding_augmentation=False): + body = {"commitMessage": False, "files": False} + if embedding_mode is not None: + body["embeddingMode"] = embedding_mode + if embedding_model is not None: + body["embeddingModel"] = embedding_model + if embedding_dimensions is not None: + body["embeddingDimensions"] = embedding_dimensions + if embedding_augmentation: + body["embeddingAugmentation"] = True + return body + + +class CommitByollmTest(CliTestCase): + + def _replace_options(self, body): + options_url = ( + f"{get_base_url()}/intake/organizations/{self.organization}" + f"/workspaces/{self.workspace}/commits/collect/options" + ) + responses.replace(responses.GET, options_url, json=body, status=200) + + @responses.activate + def test_embedding_flags_passed_to_jar(self): + """When embeddingMode=client and env vars are set, all embedding flags reach exec_jar.""" + self._replace_options(_options_response( + embedding_mode="client", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + embedding_augmentation=False, + )) + + env = { + "SMART_TESTS_TOKEN": self.smart_tests_token, + EMBEDDING_ENDPOINT_KEY: "https://api.openai.com/v1/embeddings", + EMBEDDING_API_KEY_KEY: "sk-test", + } + + with mock.patch.dict(os.environ, env): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + _, kwargs = mock_exec_jar.call_args + args = mock_exec_jar.call_args[0] + # exec_jar(name, source, max_days, app, is_collect_message, is_collect_files, + # embedding_endpoint, embedding_model, embedding_dimensions, embedding_augmentation) + self.assertEqual(args[6], "https://api.openai.com/v1/embeddings") + self.assertEqual(args[7], "text-embedding-3-small") + self.assertEqual(args[8], 1536) + self.assertFalse(args[9]) + + @responses.activate + def test_embedding_augmentation_flag_passed_when_enabled(self): + """When embeddingAugmentation=true in workspace options, augmentation=True is passed.""" + self._replace_options(_options_response( + embedding_mode="client", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + embedding_augmentation=True, + )) + + env = { + "SMART_TESTS_TOKEN": self.smart_tests_token, + EMBEDDING_ENDPOINT_KEY: "https://api.openai.com/v1/embeddings", + EMBEDDING_API_KEY_KEY: "sk-test", + } + + with mock.patch.dict(os.environ, env): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + args = mock_exec_jar.call_args[0] + self.assertTrue(args[9]) # embedding_augmentation + + @responses.activate + def test_skips_embeddings_when_endpoint_env_var_missing(self): + """When SMART_TESTS_EMBEDDING_ENDPOINT is not set, embedding args are None.""" + self._replace_options(_options_response( + embedding_mode="client", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + )) + + env_without_endpoint = {k: v for k, v in os.environ.items() + if k != EMBEDDING_ENDPOINT_KEY} + env_without_endpoint["SMART_TESTS_TOKEN"] = self.smart_tests_token + env_without_endpoint[EMBEDDING_API_KEY_KEY] = "sk-test" + + with mock.patch.dict(os.environ, env_without_endpoint, clear=True): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + args = mock_exec_jar.call_args[0] + self.assertIsNone(args[6]) # embedding_endpoint + self.assertIsNone(args[7]) # embedding_model + self.assertIsNone(args[8]) # embedding_dimensions + + @responses.activate + def test_skips_embeddings_when_api_key_env_var_missing(self): + """When SMART_TESTS_EMBEDDING_API_KEY is not set, embedding args are None.""" + self._replace_options(_options_response( + embedding_mode="client", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + )) + + env_without_key = {k: v for k, v in os.environ.items() + if k != EMBEDDING_API_KEY_KEY} + env_without_key["SMART_TESTS_TOKEN"] = self.smart_tests_token + env_without_key[EMBEDDING_ENDPOINT_KEY] = "https://api.openai.com/v1/embeddings" + + with mock.patch.dict(os.environ, env_without_key, clear=True): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + args = mock_exec_jar.call_args[0] + self.assertIsNone(args[6]) # embedding_endpoint + + @responses.activate + def test_env_var_overrides_server_model(self): + """SMART_TESTS_EMBEDDING_MODEL env var takes precedence over the model from options.""" + self._replace_options(_options_response( + embedding_mode="client", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + )) + + env = { + "SMART_TESTS_TOKEN": self.smart_tests_token, + EMBEDDING_ENDPOINT_KEY: "https://api.openai.com/v1/embeddings", + EMBEDDING_API_KEY_KEY: "sk-test", + EMBEDDING_MODEL_KEY: "text-embedding-ada-002", + } + + with mock.patch.dict(os.environ, env): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + args = mock_exec_jar.call_args[0] + self.assertEqual(args[7], "text-embedding-ada-002") + + @responses.activate + def test_no_embedding_flags_when_mode_is_server(self): + """When embeddingMode=server, embedding args are None.""" + self._replace_options(_options_response( + embedding_mode="server", + embedding_model="text-embedding-3-small", + embedding_dimensions=1536, + )) + + env = { + "SMART_TESTS_TOKEN": self.smart_tests_token, + EMBEDDING_ENDPOINT_KEY: "https://api.openai.com/v1/embeddings", + EMBEDDING_API_KEY_KEY: "sk-test", + } + + with mock.patch.dict(os.environ, env): + with mock.patch(_EXEC_JAR) as mock_exec_jar: + mock_exec_jar.return_value = None + result = self.cli("record", "commit", "--name", "test-repo") + + self.assert_success(result) + args = mock_exec_jar.call_args[0] + self.assertIsNone(args[6]) # embedding_endpoint + self.assertIsNone(args[7]) # embedding_model + self.assertIsNone(args[8]) # embedding_dimensions