fix: review

src/api.ts

@@ -726,19 +726,8 @@ export class API {
       params.content = Buffer.from(params.content);
     }
 
-    let threadFromStore: Thread | null = null;
-    try {
-      threadFromStore = this.client.getCurrentThread();
-    } catch (error) {
-      // Ignore error thrown if getCurrentThread is called outside of a context
-    }
-
-    let stepFromStore: Step | null = null;
-    try {
-      stepFromStore = this.client.getCurrentStep();
-    } catch (error) {
-      // Ignore error thrown if getCurrentStep is called outside of a context
-    }
+    const threadFromStore = this.client._currentThread();
+    const stepFromStore = this.client._currentStep();
 
     if (threadFromStore) {
       params.threadId = threadFromStore.id;

src/index.ts

@@ -49,6 +49,18 @@ export class LiteralClient {
     return this.step({ ...data, type: 'run' });
   }
 
+  _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.

src/instrumentation/openai.ts

@@ -13,9 +13,7 @@ import {
   IGenerationMessage,
   LiteralClient,
   Maybe,
-  Step,
-  StepConstructor,
-  Thread
+  StepConstructor
 } from '..';
 
 // Define a generic type for the original function to be wrapped
@@ -310,25 +308,13 @@ const processOpenAIOutput = async (
     tags: tags
   };
 
-  let threadFromStore: Thread | null = null;
-  try {
-    threadFromStore = client.getCurrentThread();
-  } catch (error) {
-    // Ignore error thrown if getCurrentThread is called outside of a context
-  }
-
-  let stepFromStore: Step | null = null;
-  try {
-    stepFromStore = client.getCurrentStep();
-  } catch (error) {
-    // Ignore error thrown if getCurrentStep is called outside of a context
-  }
+  const threadFromStore = client._currentThread();
+  const stepFromStore = client._currentStep();
 
   const parent = stepFromStore || threadFromStore;
 
   if ('data' in output) {
     // Image Generation
-
     const stepData: StepConstructor = {
       name: inputs.model || 'openai',
       type: 'llm',

tests/attachments.test.ts

@@ -24,7 +24,7 @@ describe('Attachments', () => {
     const arrayBuffer = buffer.buffer;
     const blob = new Blob([buffer]);
     // We wrap the blob in a blob and simulate the structure of a File
-    const file = new Blob([blob], { type: 'image/jpeg' });
+    const file = new Blob([blob], { type: 'image/png' });
 
     it.each([
       { type: 'Stream', content: stream! },