js feat: Support attachments in traceable (#1169)

Also add warning & handling of invalid attachments. Proposed API: getAttachments(args) => [attachments, otherArgs]; Returning the filtered args means that you're less likely to accidentally double-log attachments in inputs.
langchain-ai · Nov 6, 2024 · f23a832 · f23a832
1 parent bb09ced
commit f23a832
Show file tree

Hide file tree

Showing 6 changed files with 286 additions and 78 deletions.
diff --git a/js/src/client.ts b/js/src/client.ts
@@ -34,6 +34,7 @@ import {
   ValueType,
   AnnotationQueue,
   RunWithAnnotationQueueInfo,
+  Attachments,
 } from "./schemas.js";
 import {
   convertLangChainMessageToExample,
@@ -240,7 +241,7 @@ interface CreateRunParams {
   revision_id?: string;
   trace_id?: string;
   dotted_order?: string;
-  attachments?: Record<string, [string, Uint8Array]>;
+  attachments?: Attachments;
 }
 
 interface UpdateRunParams extends RunUpdate {
@@ -1032,10 +1033,7 @@ export class Client {
       return;
     }
     // transform and convert to dicts
-    const allAttachments: Record<
-      string,
-      Record<string, [string, Uint8Array]>
-    > = {};
+    const allAttachments: Record<string, Attachments> = {};
     let preparedCreateParams = [];
     for (const create of runCreates ?? []) {
       const preparedCreate = this.prepareRunCreateOrUpdateInputs(create);
@@ -1048,7 +1046,6 @@ export class Client {
       delete preparedCreate.attachments;
       preparedCreateParams.push(preparedCreate);
     }
-
     let preparedUpdateParams = [];
     for (const update of runUpdates ?? []) {
       preparedUpdateParams.push(this.prepareRunCreateOrUpdateInputs(update));
@@ -1116,7 +1113,8 @@ export class Client {
     ]) {
       for (const originalPayload of payloads) {
         // collect fields to be sent as separate parts
-        const { inputs, outputs, events, ...payload } = originalPayload;
+        const { inputs, outputs, events, attachments, ...payload } =
+          originalPayload;
         const fields = { inputs, outputs, events };
         // encode the main run payload
         const stringifiedPayload = stringifyForTracing(payload);
@@ -1147,10 +1145,18 @@ export class Client {
             for (const [name, [contentType, content]] of Object.entries(
               attachments
             )) {
+              // 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.length}`,
+                  type: `${contentType}; length=${content.byteLength}`,
                 }),
               });
             }
@@ -1172,6 +1178,7 @@ export class Client {
       for (const part of parts) {
         formData.append(part.name, part.payload);
       }
+      // Log the form data
       await this.batchIngestCaller.call(
         _getFetchImplementation(),
         `${this.apiUrl}/runs/multipart`,

diff --git a/js/src/run_trees.ts b/js/src/run_trees.ts
@@ -1,5 +1,11 @@
 import * as uuid from "uuid";
-import { BaseRun, KVMap, RunCreate, RunUpdate } from "./schemas.js";
+import {
+  Attachments,
+  BaseRun,
+  KVMap,
+  RunCreate,
+  RunUpdate,
+} from "./schemas.js";
 import {
   RuntimeEnvironment,
   getEnvironmentVariable,
@@ -55,6 +61,7 @@ export interface RunTreeConfig {
 
   trace_id?: string;
   dotted_order?: string;
+  attachments?: Attachments;
 }
 
 export interface RunnableConfigLike {
@@ -172,6 +179,11 @@ export class RunTree implements BaseRun {
   tracingEnabled?: boolean;
   execution_order: number;
   child_execution_order: number;
+  /**
+   * Attachments associated with the run.
+   * Each entry is a tuple of [mime_type, bytes]
+   */
+  attachments?: Attachments;
 
   constructor(originalConfig: RunTreeConfig | RunTree) {
     // If you pass in a run tree directly, return a shallow clone
@@ -370,6 +382,7 @@ export class RunTree implements BaseRun {
       trace_id: run.trace_id,
       dotted_order: run.dotted_order,
       tags: run.tags,
+      attachments: run.attachments,
     };
     return persistedRun;
   }
@@ -407,6 +420,7 @@ export class RunTree implements BaseRun {
         dotted_order: this.dotted_order,
         trace_id: this.trace_id,
         tags: this.tags,
+        attachments: this.attachments,
       };
 
       await this.client.updateRun(this.id, runUpdate);

diff --git a/js/src/schemas.ts b/js/src/schemas.ts
@@ -63,6 +63,9 @@ export interface BaseExample {
   source_run_id?: string;
 }
 
+export type AttachmentData = Uint8Array | ArrayBuffer;
+export type Attachments = Record<string, [string, AttachmentData]>;
+
 /**
  * A run can represent either a trace (root run)
  * or a child run (~span).
@@ -131,7 +134,7 @@ export interface BaseRun {
    * Attachments associated with the run.
    * Each entry is a tuple of [mime_type, bytes]
    */
-  attachments?: Record<string, [string, Uint8Array]>;
+  attachments?: Attachments;
 }
 
 type S3URL = {
@@ -231,6 +234,12 @@ export interface RunUpdate {
    *   - 20230915T223155647Z1b64098b-4ab7-43f6-afee-992304f198d8.20230914T223155650Zc8d9f4c5-6c5a-4b2d-9b1c-3d9d7a7c5c7c
    */
   dotted_order?: string;
+
+  /**
+   * Attachments associated with the run.
+   * Each entry is a tuple of [mime_type, bytes]
+   */
+  attachments?: Attachments;
 }
 
 export interface ExampleCreate extends BaseExample {