Merge branch 'main' of github.com:Chainlit/typescript-client into wil…

…ly/eng-1669-experiments-features
Chainlit · Jul 26, 2024 · 2d877a5 · 2d877a5
2 parents 1d79bae + 2ac8082
commit 2d877a5
Show file tree

Hide file tree

Showing 9 changed files with 398 additions and 191 deletions.
diff --git a/src/api.ts b/src/api.ts
@@ -1,6 +1,7 @@
 import axios, { AxiosError } from 'axios';
 import FormData from 'form-data';
 import { createReadStream } from 'fs';
+import { ReadStream } from 'fs';
 import { v4 as uuidv4 } from 'uuid';
 
 import { LiteralClient } from '.';
@@ -21,6 +22,7 @@ import {
   PersistedGeneration
 } from './generation';
 import {
+  Attachment,
   CleanThreadFields,
   Dataset,
   DatasetExperiment,
@@ -330,6 +332,28 @@ function addGenerationsToDatasetQueryBuilder(generationIds: string[]) {
     `;
 }
 
+type UploadFileBaseParams = {
+  id?: Maybe<string>;
+  threadId?: string;
+  mime?: Maybe<string>;
+};
+type UploadFileParamsWithPath = UploadFileBaseParams & {
+  path: string;
+};
+type UploadFileParamsWithContent = UploadFileBaseParams & {
+  content:
+    | ReadableStream<any>
+    | ReadStream
+    | Buffer
+    | File
+    | Blob
+    | ArrayBuffer;
+};
+type CreateAttachmentParams = {
+  name?: string;
+  metadata?: Maybe<Record<string, any>>;
+};
+
 export class API {
   /** @ignore */
   public client: LiteralClient;
@@ -408,7 +432,6 @@ export class API {
 
       return response.data;
     } catch (e) {
-      console.error(e);
       if (e instanceof AxiosError) {
         throw new Error(JSON.stringify(e.response?.data.errors));
       } else {
@@ -436,7 +459,6 @@ export class API {
 
       return response.data;
     } catch (e) {
-      console.error(e);
       if (e instanceof AxiosError) {
         throw new Error(JSON.stringify(e.response?.data));
       } else {
@@ -605,19 +627,25 @@ export class API {
    * @returns An object containing the `objectKey` of the uploaded file and the signed `url`, or `null` values if the upload fails.
    * @throws {Error} Throws an error if neither `content` nor `path` is provided, or if the server response is invalid.
    */
+
+  async uploadFile(params: UploadFileParamsWithContent): Promise<{
+    objectKey: Maybe<string>;
+    url: Maybe<string>;
+  }>;
+  async uploadFile(params: UploadFileParamsWithPath): Promise<{
+    objectKey: Maybe<string>;
+    url: Maybe<string>;
+  }>;
   async uploadFile({
     content,
     path,
     id,
     threadId,
     mime
-  }: {
-    content?: Maybe<any>;
-    path?: Maybe<string>;
-    id?: Maybe<string>;
-    threadId: string;
-    mime?: Maybe<string>;
-  }) {
+  }: UploadFileParamsWithContent & UploadFileParamsWithPath): Promise<{
+    objectKey: Maybe<string>;
+    url: Maybe<string>;
+  }> {
     if (!content && !path) {
       throw new Error('Either content or path must be provided');
     }
@@ -687,6 +715,52 @@ export class API {
     }
   }
 
+  async createAttachment(
+    params: UploadFileParamsWithContent & CreateAttachmentParams
+  ): Promise<Attachment>;
+  async createAttachment(
+    params: UploadFileParamsWithPath & CreateAttachmentParams
+  ): Promise<Attachment>;
+  async createAttachment(
+    params: UploadFileParamsWithContent &
+      UploadFileParamsWithPath &
+      CreateAttachmentParams
+  ): Promise<Attachment> {
+    if (params.content instanceof Blob) {
+      params.content = Buffer.from(await params.content.arrayBuffer());
+    }
+    if (params.content instanceof ArrayBuffer) {
+      params.content = Buffer.from(params.content);
+    }
+
+    const threadFromStore = this.client._currentThread();
+    const stepFromStore = this.client._currentStep();
+
+    if (threadFromStore) {
+      params.threadId = threadFromStore.id;
+    }
+
+    const { objectKey, url } = await this.uploadFile(params);
+
+    const attachment = new Attachment({
+      name: params.name,
+      objectKey,
+      mime: params.mime,
+      metadata: params.metadata,
+      url
+    });
+
+    if (stepFromStore) {
+      if (!stepFromStore.attachments) {
+        stepFromStore.attachments = [];
+      }
+
+      stepFromStore.attachments.push(attachment);
+    }
+
+    return attachment;
+  }
+
   // Generation
   /**
    * Retrieves a paginated list of Generations based on the provided filters and sorting order.

diff --git a/src/index.ts b/src/index.ts
@@ -75,6 +75,18 @@ export class LiteralClient {
     });
   }
 
+  _currentThread(): Thread | null {
+    const store = storage.getStore();
+
+    return store?.currentThread || null;
+  }
+
+  _currentStep(): Step | null {
+    const store = storage.getStore();
+
+    return store?.currentStep || null;
+  }
+
   /**
    * Gets the current thread from the context.
    * WARNING : this will throw if run outside of a thread context.

diff --git a/src/instrumentation/index.ts b/src/instrumentation/index.ts
@@ -1,15 +1,16 @@
-import { LiteralClient } from '..';
+import { LiteralClient, Maybe } from '..';
 import { LiteralCallbackHandler } from './langchain';
 import { instrumentLlamaIndex, withThread } from './llamaindex';
 import instrumentOpenAI from './openai';
-import { InstrumentOpenAIOptions } from './openai';
 import { makeInstrumentVercelSDK } from './vercel-sdk';
 
-export type { InstrumentOpenAIOptions } from './openai';
+export type OpenAIGlobalOptions = {
+  tags?: Maybe<string[]>;
+  metadata?: Maybe<Record<string, any>>;
+};
 
 export default (client: LiteralClient) => ({
-  openai: (options?: InstrumentOpenAIOptions) =>
-    instrumentOpenAI(client, options),
+  openai: (options?: OpenAIGlobalOptions) => instrumentOpenAI(client, options),
   langchain: {
     literalCallback: (threadId?: string) => {
       try {