feat(js): Initial implementation of compression

js/src/client.ts

@@ -53,7 +53,11 @@
   getLangSmithEnvironmentVariable,
   getRuntimeEnvironment,
 } from "./utils/env.js";
-
+import {
+  convertToMultipartBlobChunks,
+  serializeRunOperationAsMultipart,
+  type MultipartPart,
+} from "./utils/multipart.js";
 import {
   EvaluationResult,
   EvaluationResults,
@@ -87,6 +91,10 @@
    * Useful if encountering network rate limits at trace high volumes.
    */
   manualFlushMode?: boolean;
+  /**
+   * Whether to enable compression of batched runs before sending to LangSmith.
+   */
+  useRunCompression?: boolean;
 }
 
 /**
@@ -284,16 +292,20 @@
   sourceRunId?: string;
 };
 
-type AutoBatchQueueItem = {
+type AutoBatchQueueItemInput = {
   action: "create" | "update";
   item: RunCreate | RunUpdate;
 };
 
-type MultipartPart = {
-  name: string;
-  payload: Blob;
+type AutoBatchQueueItem = {
+  action: "create" | "update";
+  item: RunCreate | RunUpdate | Blob;
 };
 
+function isBlob(x: unknown): x is Blob {
+  return x != null && typeof x === "object" && "size" in x;
+}
+
 export function mergeRuntimeEnvIntoRunCreate(run: RunCreate) {
   const runtimeEnv = getRuntimeEnvironment();
   const envVars = getLangChainEnvVarsMetadata();
@@ -373,10 +385,35 @@
   return false;
 };
 
+const _compressPayload = async (stream: CompressionStream, payload: Blob[]) => {
+  const compressedPayloadStream = new Blob(payload)
+    .stream()
+    .pipeThrough(stream);
+  const reader = compressedPayloadStream.getReader();
+  const chunks = [];
+  let totalLength = 0;
+  try {
+    // eslint-disable-next-line no-constant-condition
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      chunks.push(value);
+      totalLength += value.length;
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  return new Blob(chunks);
+};
+
+function _createNewCompressionStream() {
+  return new CompressionStream("gzip");
+}
+
 export class AutoBatchQueue {
   items: {
     action: "create" | "update";
-    payload: RunCreate | RunUpdate;
+    payload: RunCreate | RunUpdate | Blob;
     itemPromiseResolve: () => void;
     itemPromise: Promise<void>;
     size: number;
@@ -388,17 +425,20 @@
     return this.items[0];
   }
 
-  push(item: AutoBatchQueueItem): Promise<void> {
+  async push(item: AutoBatchQueueItemInput): Promise<void> {
     let itemPromiseResolve;
     const itemPromise = new Promise<void>((resolve) => {
       // Setting itemPromiseResolve is synchronous with promise creation:
       // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Promise/Promise
       itemPromiseResolve = resolve;
     });
-    const size = stringifyForTracing(item.item).length;
+    let payload: RunCreate | RunUpdate | Blob = item.item;
+    const size = isBlob(payload)
+      ? payload.size
+      : stringifyForTracing(payload).length;
     this.items.push({
       action: item.action,
-      payload: item.item,
+      payload,
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       itemPromiseResolve: itemPromiseResolve!,
       itemPromise,
@@ -429,7 +469,7 @@
     // If there is an item on the queue we were unable to pop,
     // just return it as a single batch.
     if (popped.length === 0 && this.items.length > 0) {
       const item = this.items.shift()!;
       popped.push(item);
       poppedSizeBytes += item.size;
       this.sizeBytes -= item.size;
@@ -495,6 +535,11 @@
 
   private manualFlushMode = false;
 
+  private useRunCompression =
+    getEnvironmentVariable("LANGSMITH_USE_RUN_COMPRESSION") === "true";
+
+  private compressionStream = _createNewCompressionStream();
+
   constructor(config: ClientConfig = {}) {
     const defaultConfig = Client.getDefaultClientConfig();
 
@@ -533,6 +578,7 @@
     this.batchSizeBytesLimit = config.batchSizeBytesLimit;
     this.fetchOptions = config.fetchOptions || {};
     this.manualFlushMode = config.manualFlushMode ?? this.manualFlushMode;
+    this.useRunCompression = config.useRunCompression ?? this.useRunCompression;
   }
 
   public static getDefaultClientConfig(): {
@@ -821,18 +867,48 @@
     }
   }
 
-  private async processRunOperation(item: AutoBatchQueueItem) {
+  private async processRunOperation(
+    item: AutoBatchQueueItemInput
+  ): Promise<void> {
     clearTimeout(this.autoBatchTimeout);
     this.autoBatchTimeout = undefined;
+    let queueItem: AutoBatchQueueItem = item;
     if (item.action === "create") {
       item.item = mergeRuntimeEnvIntoRunCreate(item.item as RunCreate);
     }
+    const sizeLimitBytes = await this._getBatchSizeLimitBytes();
+    if (this.useRunCompression) {
+      const { parts } = serializeRunOperationAsMultipart(
+        item.action,
+        item.item
+      );
+      const blobChunks = convertToMultipartBlobChunks(parts);
+      // It would be great if you could get the compressed size from a ReadableStream
+      // so that we could pass the CompressionStream directly into a "fetch" call,
+      // but this doesn't seem like it's possible.
+      const compressedItem = await _compressPayload(
+        this.compressionStream,
+        blobChunks
+      );
+      if (
+        this.autoBatchQueue.sizeBytes + compressedItem.size >
+        sizeLimitBytes
+      ) {
+        this.compressionStream = _createNewCompressionStream();
+        // Drain the entire queue since items compressed with one stream are not compatible with others
+        await this.drainAutoBatchQueue(Number.MAX_SAFE_INTEGER);
+        // Throw out the compressed value and retry
+        // TODO: Instead, should we estimate compressed size so as not to throw away compression work?
+        return this.processRunOperation(item);
+      } else {
+        queueItem.item = compressedItem;
+      }
+    }
     const itemPromise = this.autoBatchQueue.push(item);
     if (this.manualFlushMode) {
       // Rely on manual flushing in serverless environments
       return itemPromise;
     }
-    const sizeLimitBytes = await this._getBatchSizeLimitBytes();
     if (this.autoBatchQueue.sizeBytes > sizeLimitBytes) {
       void this.drainAutoBatchQueue(sizeLimitBytes);
     }
@@ -862,7 +938,7 @@
         if (this._serverInfo === undefined) {
           try {
             this._serverInfo = await this._getServerInfo();
           } catch (e) {
             console.warn(
               `[WARNING]: LangSmith failed to fetch info on supported operations. Falling back to batch operations and default limits.`
             );
@@ -1597,7 +1673,7 @@
     treeFilter?: string;
     isRoot?: boolean;
     dataSourceType?: string;
   }): Promise<any> {
     let projectIds_ = projectIds || [];
     if (projectNames) {
       projectIds_ = [
@@ -1885,7 +1961,7 @@
         `Failed to list shared examples: ${response.status} ${response.statusText}`
       );
     }
     return result.map((example: any) => ({
       ...example,
       _hostUrl: this.getHostUrl(),
     }));
@@ -2022,7 +2098,7 @@
       }
       // projectId querying
       return true;
     } catch (e) {
       return false;
     }
   }
@@ -3399,7 +3475,7 @@
   async _logEvaluationFeedback(
     evaluatorResponse: EvaluationResult | EvaluationResults,
     run?: Run,
     sourceInfo?: { [key: string]: any }
   ): Promise<[results: EvaluationResult[], feedbacks: Feedback[]]> {
     const evalResults: Array<EvaluationResult> =
       this._selectEvalResults(evaluatorResponse);
@@ -3438,7 +3514,7 @@
   public async logEvaluationFeedback(
     evaluatorResponse: EvaluationResult | EvaluationResults,
     run?: Run,
     sourceInfo?: { [key: string]: any }
   ): Promise<EvaluationResult[]> {
     const [results] = await this._logEvaluationFeedback(
       evaluatorResponse,
@@ -3934,7 +4010,7 @@
 
   public async createCommit(
     promptIdentifier: string,
     object: any,
     options?: {
       parentCommitHash?: string;
     }
@@ -4166,7 +4242,7 @@
       isPublic?: boolean;
       isArchived?: boolean;
     }
   ): Promise<Record<string, any>> {
     if (!(await this.promptExists(promptIdentifier))) {
       throw new Error("Prompt does not exist, you must create it first.");
     }

js/src/utils/multipart.ts

@@ -0,0 +1,103 @@
+import { AttachmentData, RunCreate, RunUpdate } from "../schemas.js";
+import { stringify as stringifyForTracing } from "./fast-safe-stringify/index.js";
+
+export const MULTIPART_BOUNDARY =
+  "----LangSmithFormBoundary" + Math.random().toString(36).slice(2);
+
+export type MultipartPart = {
+  name: string;
+  payload: Blob;
+};
+
+export function serializeRunOperationAsMultipart(
+  method: string,
+  originalPayload: RunCreate | RunUpdate
+) {
+  const accumulatedParts: MultipartPart[] = [];
+  const accumulatedContext: string[] = [];
+  // collect fields to be sent as separate parts
+  const { inputs, outputs, events, attachments, ...payload } = originalPayload;
+  const fields = { inputs, outputs, events };
+  // encode the main run payload
+  const stringifiedPayload = stringifyForTracing(payload);
+  accumulatedParts.push({
+    name: `${method}.${payload.id}`,
+    payload: new Blob([stringifiedPayload], {
+      type: `application/json; length=${stringifiedPayload.length}`,
+    }),
+  });
+  // encode the fields we collected
+  for (const [key, value] of Object.entries(fields)) {
+    if (value === undefined) {
+      continue;
+    }
+    const stringifiedValue = stringifyForTracing(value);
+    accumulatedParts.push({
+      name: `${method}.${payload.id}.${key}`,
+      payload: new Blob([stringifiedValue], {
+        type: `application/json; length=${stringifiedValue.length}`,
+      }),
+    });
+  }
+  // encode the attachments
+  if (payload.id !== undefined) {
+    if (attachments) {
+      for (const [name, attachment] of Object.entries(attachments)) {
+        let contentType: string;
+        let content: AttachmentData;
+
+        if (Array.isArray(attachment)) {
+          [contentType, content] = attachment;
+        } else {
+          contentType = attachment.mimeType;
+          content = attachment.data;
+        }
+
+        // Validate that the attachment name doesn't contain a '.'
+        if (name.includes(".")) {
+          console.warn(
+            `Skipping attachment '${name}' for run ${payload.id}: Invalid attachment name. ` +
+              `Attachment names must not contain periods ('.'). Please rename the attachment and try again.`
+          );
+          continue;
+        }
+        accumulatedParts.push({
+          name: `attachment.${payload.id}.${name}`,
+          payload: new Blob([content], {
+            type: `${contentType}; length=${content.byteLength}`,
+          }),
+        });
+      }
+    }
+  }
+  // compute context
+  accumulatedContext.push(`trace=${payload.trace_id},id=${payload.id}`);
+  return {
+    parts: accumulatedParts,
+    context: accumulatedContext,
+  };
+}
+
+export function convertToMultipartBlobChunks(parts: MultipartPart[]) {
+  // Create multipart form data manually using Blobs
+  const chunks: Blob[] = [];
+
+  for (const part of parts) {
+    // Add field boundary
+    chunks.push(new Blob([`--${MULTIPART_BOUNDARY}\r\n`]));
+    chunks.push(
+      new Blob([
+        `Content-Disposition: form-data; name="${part.name}"\r\n`,
+        `Content-Type: ${part.payload.type}\r\n\r\n`,
+      ])
+    );
+    chunks.push(part.payload);
+    chunks.push(new Blob(["\r\n"]));
+  }
+
+  // Do once at the end?
+  // // Add final boundary
+  // chunks.push(new Blob([`--${MULTIPART_BOUNDARY}--\r\n`]));
+
+  return chunks;
+}