From 49e26c1e857f4b31ea46e4b20e54e0036d93a823 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Tue, 13 Aug 2024 20:02:31 +0200
Subject: [PATCH 1/3] chore: update vercel integration

---
 src/instrumentation/vercel-sdk.ts    | 11 +++--
 tests/integration/vercel-sdk.test.ts | 73 ++++++++++++++--------------
 2 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/src/instrumentation/vercel-sdk.ts b/src/instrumentation/vercel-sdk.ts
index c122abb..be6b858 100644
--- a/src/instrumentation/vercel-sdk.ts
+++ b/src/instrumentation/vercel-sdk.ts
@@ -256,13 +256,16 @@ export const makeInstrumentVercelSDK = (
     type TOptions = Options<TFunction>;
     type TResult = Result<TFunction>;
 
-    return async (
-      options: TOptions & { literalAiParent?: Step | Thread }
-    ): Promise<TResult> => {
-      const { literalAiParent: parent, ...originalOptions } = options;
+    return async (options: TOptions): Promise<TResult> => {
+      const { ...originalOptions } = options;
       const startTime = Date.now();
       const result: TResult = await (fn as any)(originalOptions);
 
+      const threadFromStore = client._currentThread();
+      const stepFromStore = client._currentStep();
+
+      const parent = stepFromStore || threadFromStore;
+
       // Fork the stream to compute metrics
       let stream: ReadableStream<OriginalStreamPart>;
       if ('originalStream' in result) {
diff --git a/tests/integration/vercel-sdk.test.ts b/tests/integration/vercel-sdk.test.ts
index 10274bc..e208dd0 100644
--- a/tests/integration/vercel-sdk.test.ts
+++ b/tests/integration/vercel-sdk.test.ts
@@ -204,44 +204,43 @@ describe('Vercel SDK Instrumentation', () => {
     it('should observe on a given thread', async () => {
       const spy = jest.spyOn(client.api, 'sendSteps');
 
-      const thread = await client.thread({ name: 'VercelSDK Test' }).upsert();
-
-      const generateTextWithLiteralAI =
-        client.instrumentation.vercel.instrument(generateText);
-
-      const result = await generateTextWithLiteralAI({
-        model: openai('gpt-3.5-turbo'),
-        prompt: 'Write a vegetarian lasagna recipe for 4 people.',
-        literalAiParent: thread
+      await client.thread({ name: 'VercelSDK Test' }).wrap(async () => {
+        const generateTextWithLiteralAI =
+          client.instrumentation.vercel.instrument(generateText);
+
+        const result = await generateTextWithLiteralAI({
+          model: openai('gpt-3.5-turbo'),
+          prompt: 'Write a vegetarian lasagna recipe for 4 people.'
+        });
+
+        expect(result.text).toBeTruthy();
+
+        // Sending message is done asynchronously
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        expect(spy).toHaveBeenCalledWith([
+          expect.objectContaining({
+            type: 'llm',
+            name: 'gpt-3.5-turbo',
+            threadId: client._currentThread()?.id,
+            generation: expect.any(Object),
+            input: {
+              content: [
+                {
+                  role: 'user',
+                  content: [
+                    {
+                      text: 'Write a vegetarian lasagna recipe for 4 people.',
+                      type: 'text'
+                    }
+                  ]
+                }
+              ]
+            },
+            output: { role: 'assistant', content: result.text }
+          })
+        ]);
       });
-
-      expect(result.text).toBeTruthy();
-
-      // Sending message is done asynchronously
-      await new Promise((resolve) => setTimeout(resolve, 10));
-
-      expect(spy).toHaveBeenCalledWith([
-        expect.objectContaining({
-          type: 'llm',
-          name: 'gpt-3.5-turbo',
-          threadId: thread.id,
-          generation: expect.any(Object),
-          input: {
-            content: [
-              {
-                role: 'user',
-                content: [
-                  {
-                    text: 'Write a vegetarian lasagna recipe for 4 people.',
-                    type: 'text'
-                  }
-                ]
-              }
-            ]
-          },
-          output: { role: 'assistant', content: result.text }
-        })
-      ]);
     });
 
     it('should monitor tools', async () => {

From 2a387b4207285207e77c95df5f0c76fad8daded8 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Tue, 13 Aug 2024 20:02:54 +0200
Subject: [PATCH 2/3] chore: bump version

---
 package.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/package.json b/package.json
index 4f88ce3..be0a7e9 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@literalai/client",
-  "version": "0.0.513",
+  "version": "0.0.514",
   "description": "",
   "exports": {
     ".": {

From 28ee60e11b329bb3ee73d487b75c710c1113e591 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Wed, 14 Aug 2024 10:26:58 +0200
Subject: [PATCH 3/3] fix: comments

---
 src/instrumentation/vercel-sdk.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/instrumentation/vercel-sdk.ts b/src/instrumentation/vercel-sdk.ts
index be6b858..2bdde9e 100644
--- a/src/instrumentation/vercel-sdk.ts
+++ b/src/instrumentation/vercel-sdk.ts
@@ -257,9 +257,8 @@ export const makeInstrumentVercelSDK = (
     type TResult = Result<TFunction>;
 
     return async (options: TOptions): Promise<TResult> => {
-      const { ...originalOptions } = options;
       const startTime = Date.now();
-      const result: TResult = await (fn as any)(originalOptions);
+      const result: TResult = await (fn as any)(options);
 
       const threadFromStore = client._currentThread();
       const stepFromStore = client._currentStep();