WIP: feat: create very simple wrappers

.gitignore

@@ -1,3 +1,5 @@
+literalai-client-*.tgz
+
 # Logs
 logs
 *.log

package.json

@@ -14,7 +14,7 @@
   "types": "./dist/index.d.ts",
   "scripts": {
     "build": "tsup ./src",
-    "test": "jest --runInBand --watchAll=false",
+    "test": "jest --detectOpenHandles --runInBand --watchAll=false",
     "prepare": "husky install"
   },
   "author": "Literal AI",

src/api.ts

@@ -3,6 +3,7 @@ import FormData from 'form-data';
 import { createReadStream } from 'fs';
 import { v4 as uuidv4 } from 'uuid';
 
+import { LiteralClient } from '.';
 import {
   GenerationsFilter,
   GenerationsOrderBy,
@@ -327,6 +328,8 @@ function addGenerationsToDatasetQueryBuilder(generationIds: string[]) {
 }
 
 export class API {
+  /** @ignore */
+  private client: LiteralClient;
   /** @ignore */
   private apiKey: string;
   /** @ignore */
@@ -339,7 +342,13 @@ export class API {
   public disabled: boolean;
 
   /** @ignore */
-  constructor(apiKey: string, url: string, disabled?: boolean) {
+  constructor(
+    client: LiteralClient,
+    apiKey: string,
+    url: string,
+    disabled?: boolean
+  ) {
+    this.client = client;
     this.apiKey = apiKey;
     this.url = url;
     this.graphqlEndpoint = `${url}/api/graphql`;
@@ -509,7 +518,9 @@ export class API {
 
     const response = result.data.steps;
 
-    response.data = response.edges.map((x: any) => new Step(this, x.node));
+    response.data = response.edges.map(
+      (x: any) => new Step(this.client, x.node, true)
+    );
     delete response.edges;
 
     return response;
@@ -541,7 +552,7 @@ export class API {
       return null;
     }
 
-    return new Step(this, result.data.step);
+    return new Step(this.client, result.data.step, true);
   }
 
   /**
@@ -878,7 +889,7 @@ export class API {
     };
 
     const response = await this.makeGqlCall(query, variables);
-    return new Thread(this, response.data.upsertThread);
+    return new Thread(this.client, response.data.upsertThread);
   }
 
   /**
@@ -943,7 +954,9 @@ export class API {
 
     const response = result.data.threads;
 
-    response.data = response.edges.map((x: any) => new Thread(this, x.node));
+    response.data = response.edges.map(
+      (x: any) => new Thread(this.client, x.node)
+    );
     delete response.edges;
 
     return response;
@@ -967,12 +980,11 @@ export class API {
     const variables = { id };
 
     const response = await this.makeGqlCall(query, variables);
-
     if (!response.data.threadDetail) {
       return null;
     }
 
-    return new Thread(this, response.data.threadDetail);
+    return new Thread(this.client, response.data.threadDetail);
   }
 
   /**

src/index.ts

@@ -1,17 +1,33 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
 import { API } from './api';
 import instrumentation from './instrumentation';
 import openai from './openai';
-import { Step, StepConstructor, Thread, ThreadConstructor } from './types';
+import {
+  Maybe,
+  Step,
+  StepConstructor,
+  Thread,
+  ThreadConstructor
+} from './types';
 
 export * from './types';
 export * from './generation';
 
 export type * from './instrumentation';
 
+type StoredContext = {
+  currentThread: Thread | null;
+  currentStep: Step | null;
+};
+
+const storage = new AsyncLocalStorage<StoredContext>();
+
 export class LiteralClient {
   api: API;
   openai: ReturnType<typeof openai>;
   instrumentation: ReturnType<typeof instrumentation>;
+  store: AsyncLocalStorage<StoredContext> = storage;
 
   constructor(apiKey?: string, apiUrl?: string, disabled?: boolean) {
     if (!apiKey) {
@@ -22,21 +38,28 @@ export class LiteralClient {
       apiUrl = process.env.LITERAL_API_URL || 'https://cloud.getliteral.ai';
     }
 
-    this.api = new API(apiKey!, apiUrl!, disabled);
+    this.api = new API(this, apiKey!, apiUrl!, disabled);
     this.openai = openai(this);
     this.instrumentation = instrumentation(this);
   }
 
   thread(data?: ThreadConstructor) {
-    return new Thread(this.api, data);
+    return new Thread(this, data);
   }
 
   step(data: StepConstructor) {
-    return new Step(this.api, data);
+    return new Step(this, data);
   }
 
   run(data: Omit<StepConstructor, 'type'>) {
-    const runData = { ...data, type: 'run' as const };
-    return new Step(this.api, runData);
+    return this.step({ ...data, type: 'run' });
+  }
+
+  getCurrentThread(): Maybe<Thread> {
+    return storage.getStore()?.currentThread ?? null;
+  }
+
+  getCurrentStep(): Maybe<Step> {
+    return storage.getStore()?.currentStep ?? null;
   }
 }