Merge pull request #54 from Chainlit/clement/eng-1746-update-sdks-to-…

…send-the-root-run feat: support root run
Chainlit · Aug 8, 2024 · adb4d3c · adb4d3c
2 parents 445826f + f73840a
commit adb4d3c
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 4 deletions.
diff --git a/src/api.ts b/src/api.ts
@@ -49,6 +49,7 @@ const version = packageJson.version;
 const stepFields = `
     id
     threadId
+    rootRunId
     parentId
     startTime
     endTime
@@ -152,6 +153,7 @@ function ingestStepsFieldsBuilder(steps: Step[]) {
   for (let id = 0; id < steps.length; id++) {
     generated += `$id_${id}: String!
         $threadId_${id}: String
+        $rootRunId_${id}: String
         $type_${id}: StepType
         $startTime_${id}: DateTime
         $endTime_${id}: DateTime
@@ -177,6 +179,7 @@ function ingestStepsArgsBuilder(steps: Step[]) {
       step${id}: ingestStep(
         id: $id_${id}
         threadId: $threadId_${id}
+        rootRunId: $rootRunId_${id}
         startTime: $startTime_${id}
         endTime: $endTime_${id}
         type: $type_${id}
@@ -1651,6 +1654,7 @@ export class API {
           input
           expectedOutput
           intermediarySteps
+          stepId
         }
       }
     `;
@@ -1676,6 +1680,7 @@ export class API {
           input
           expectedOutput
           intermediarySteps
+          stepId
         }
       }
     `;
@@ -1701,6 +1706,7 @@ export class API {
           input
           expectedOutput
           intermediarySteps
+          stepId
         }
       }
     `;
@@ -1732,6 +1738,7 @@ export class API {
           input
           expectedOutput
           intermediarySteps
+          stepId
         }
       }
     `;
@@ -1767,6 +1774,7 @@ export class API {
           input
           expectedOutput
           intermediarySteps
+          stepId
         }
       }
     `;

diff --git a/src/evaluation/experiment-item-run.ts b/src/evaluation/experiment-item-run.ts
@@ -39,7 +39,12 @@ export class ExperimentItemRun extends Step {
       {
         currentThread: currentStore?.currentThread ?? null,
         currentStep: this,
-        currentExperimentItemRunId: this.id ?? null
+        currentExperimentItemRunId: this.id ?? null,
+        rootRun: currentStore?.rootRun
+          ? currentStore?.rootRun
+          : this.type === 'run'
+          ? this
+          : null
       },
       async () => {
         try {

diff --git a/src/index.ts b/src/index.ts
@@ -17,6 +17,7 @@ type StoredContext = {
   currentThread: Thread | null;
   currentStep: Step | null;
   currentExperimentItemRunId?: string | null;
+  rootRun: Step | null;
 };
 
 const storage = new AsyncLocalStorage<StoredContext>();
@@ -125,6 +126,16 @@ export class LiteralClient {
     return store?.currentExperimentItemRunId || null;
   }
 
+  /**
+   * Returns the root run from the context or null if none.
+   * @returns The root run, if any.
+   */
+  _rootRun(): Step | null {
+    const store = storage.getStore();
+
+    return store?.rootRun || null;
+  }
+
   /**
    * Gets the current thread from the context.
    * WARNING : this will throw if run outside of a thread context.
@@ -175,4 +186,21 @@ export class LiteralClient {
 
     return store?.currentExperimentItemRunId;
   }
+
+  /**
+   * Gets the root run from the context.
+   * WARNING : this will throw if run outside of a step context.
+   * @returns The current step, if any.
+   */
+  getRootRun(): Step {
+    const store = storage.getStore();
+
+    if (!store?.rootRun) {
+      throw new Error(
+        'Literal AI SDK : tried to access root run outside of a context.'
+      );
+    }
+
+    return store.rootRun;
+  }
 }
diff --git a/src/observability/step.ts b/src/observability/step.ts
@@ -23,6 +23,7 @@ class StepFields extends Utils {
   name!: string;
   type!: StepType;
   threadId?: string;
+  rootRunId?: Maybe<string>;
   createdAt?: Maybe<string>;
   startTime?: Maybe<string>;
   id?: Maybe<string>;
@@ -73,9 +74,10 @@ export class Step extends StepFields {
       return;
     }
 
-    // Automatically assign parent thread & step if there are any in the store.
+    // Automatically assign parent thread & step & rootRun if there are any in the store.
     this.threadId = this.threadId ?? this.client._currentThread()?.id;
     this.parentId = this.parentId ?? this.client._currentStep()?.id;
+    this.rootRunId = this.rootRunId ?? this.client._rootRun()?.id;
 
     // Set the creation and start time to the current time if not provided.
     if (!this.createdAt) {
@@ -167,7 +169,12 @@ export class Step extends StepFields {
         currentThread: currentStore?.currentThread ?? null,
         currentExperimentItemRunId:
           currentStore?.currentExperimentItemRunId ?? null,
-        currentStep: this
+        currentStep: this,
+        rootRun: currentStore?.rootRun
+          ? currentStore?.rootRun
+          : this.type === 'run'
+          ? this
+          : null
       },
       () => cb(this)
     );
@@ -197,6 +204,7 @@ export class Step extends StepFields {
       this.scores = updatedStep.scores ?? this.scores;
       this.attachments = updatedStep.attachments ?? this.attachments;
       this.environment = updatedStep.environment ?? this.environment;
+      this.rootRunId = updatedStep.rootRunId ?? this.rootRunId;
     }
 
     this.send().catch(console.error);

diff --git a/src/observability/thread.ts b/src/observability/thread.ts
@@ -109,7 +109,8 @@ export class Thread extends ThreadFields {
         currentThread: this,
         currentExperimentItemRunId:
           currentStore?.currentExperimentItemRunId ?? null,
-        currentStep: null
+        currentStep: null,
+        rootRun: null
       },
       () => cb(this)
     );

diff --git a/tests/wrappers.test.ts b/tests/wrappers.test.ts
@@ -143,6 +143,32 @@ describe('Wrapper', () => {
       });
     });
 
+    it('handles nested runs', async () => {
+      let runId: Maybe<string>;
+      let stepId: Maybe<string>;
+
+      const step = async (_query: string) =>
+        client.step({ name: 'foo', type: 'undefined' }).wrap(async () => {
+          stepId = client.getCurrentStep()!.id;
+        });
+
+      await client.thread({ name: 'Test Wrappers Thread' }).wrap(async () => {
+        return client.run({ name: 'Test Wrappers Run' }).wrap(async () => {
+          runId = client.getCurrentStep()!.id;
+
+          return client.run({ name: 'Test Nested Run' }).wrap(async () => {
+            await step('foo');
+          });
+        });
+      });
+
+      await sleep(1000);
+      const run = await client.api.getStep(runId!);
+      const createdStep = await client.api.getStep(stepId!);
+
+      expect(createdStep!.rootRunId).toEqual(run!.id);
+    });
+
     it('handles steps outside of a thread', async () => {
       let runId: Maybe<string>;
       let stepId: Maybe<string>;
@@ -172,6 +198,7 @@ describe('Wrapper', () => {
       expect(step!.name).toEqual('Test Wrappers Step');
       expect(step!.threadId).toBeNull();
       expect(step!.parentId).toEqual(run!.id);
+      expect(step!.rootRunId).toEqual(run!.id);
     });
 
     it("doesn't leak the current store when getting entities from the API", async () => {