From bc5df84aa5873a2b6058b221db08acb4f3dded28 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Sun, 15 Sep 2024 16:04:56 +0200
Subject: [PATCH] fix: require instrumentation dependencies at runtime

---
 src/instrumentation/index.ts         | 72 +++++++++++++++++++---------
 src/instrumentation/vercel-sdk.ts    |  4 +-
 tests/integration/vercel-sdk.test.ts |  4 +-
 3 files changed, 53 insertions(+), 27 deletions(-)

diff --git a/src/instrumentation/index.ts b/src/instrumentation/index.ts
index 1a66344..c50da06 100644
--- a/src/instrumentation/index.ts
+++ b/src/instrumentation/index.ts
@@ -1,15 +1,11 @@
-import OpenAI from 'openai';
-
 import { LiteralClient, Maybe } from '..';
-import { LiteralCallbackHandler } from './langchain';
-import { instrumentLlamaIndex, withThread } from './llamaindex';
-import instrumentOpenAI from './openai';
-import { makeInstrumentVercelSDK } from './vercel-sdk';
+import { Thread } from '../observability/thread';
+import type { AllVercelFn } from './vercel-sdk';
 
 export type OpenAIGlobalOptions = {
   tags?: Maybe<string[]>;
   metadata?: Maybe<Record<string, any>>;
-  client?: OpenAI;
+  client?: any;
 };
 
 export default (client: LiteralClient) => ({
@@ -23,7 +19,17 @@ export default (client: LiteralClient) => ({
    * @param options.metadata Metadata to attach to all generations.
    * @returns
    */
-  openai: (options?: OpenAIGlobalOptions) => instrumentOpenAI(client, options),
+  openai: (options?: OpenAIGlobalOptions) => {
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-var-requires
+      const { default: instrumentOpenAI } = require('./openai');
+      return instrumentOpenAI(client, options);
+    } catch (error) {
+      throw new Error(
+        'Failed to load OpenAI. Please ensure openai is installed.'
+      );
+    }
+  },
 
   langchain: {
     literalCallback: (params?: {
@@ -31,6 +37,8 @@ export default (client: LiteralClient) => ({
       chainTypesToIgnore?: string[];
     }) => {
       try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const { LiteralCallbackHandler } = require('./langchain');
         return new LiteralCallbackHandler(
           client,
           params?.threadId,
@@ -45,23 +53,41 @@ export default (client: LiteralClient) => ({
   },
 
   vercel: {
-    /**
-     * Instrument a Vercel SDK function to log all invocations to the Literal AI API.
-     * It will be augmented with a `literalAiParent` option that allows you to specify a parent step or thread.
-     * @param fn The function to instrument. This can be Vercel SDK's `generateText` or `streamText` functions.
-     * @returns A new function that logs all invocations to the Literal AI API.
-     */
-    instrument: makeInstrumentVercelSDK(client)
+    instrument: <TFunction extends AllVercelFn>(fn: TFunction) => {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const { makeInstrumentVercelSDK } = require('./vercel-sdk');
+        return makeInstrumentVercelSDK(client)(fn);
+      } catch (error) {
+        throw new Error(
+          'Failed to load Vercel SDK. Please ensure @vercel/ai is installed.'
+        );
+      }
+    }
   },
 
   llamaIndex: {
-    /**
-     * Instrument the LlamaIndex client to log all generations to the Literal AI API.
-     */
-    instrument: () => instrumentLlamaIndex(client),
-    /**
-     * Run a callback with a thread context.
-     */
-    withThread
+    instrument: () => {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const { instrumentLlamaIndex } = require('./llamaindex');
+        return instrumentLlamaIndex(client);
+      } catch (error) {
+        throw new Error(
+          'Failed to load LlamaIndex. Please ensure llamaindex is installed.'
+        );
+      }
+    },
+    withThread: <R>(thread: Thread, callback: () => R) => {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-var-requires
+        const { withThread } = require('./llamaindex');
+        return withThread(thread, callback);
+      } catch (error) {
+        throw new Error(
+          'Failed to load LlamaIndex. Please ensure llamaindex is installed.'
+        );
+      }
+    }
   }
 });
diff --git a/src/instrumentation/vercel-sdk.ts b/src/instrumentation/vercel-sdk.ts
index 01409ca..6693205 100644
--- a/src/instrumentation/vercel-sdk.ts
+++ b/src/instrumentation/vercel-sdk.ts
@@ -26,7 +26,7 @@ type Result<T extends (...args: any[]) => any> = Awaited<ReturnType<T>>;
 
 type GenerateFn = typeof generateObject | typeof generateText;
 type StreamFn = typeof streamObject | typeof streamText;
-type AllVercelFn = GenerateFn | StreamFn;
+export type AllVercelFn = GenerateFn | StreamFn;
 
 type OriginalStreamPart = string | ObjectStreamPart<any> | TextStreamPart<any>;
 
@@ -289,7 +289,7 @@ export type InstrumentationVercelMethod = {
     options: Parameters<typeof streamText<TOOLS>>[0] & VercelExtraOptions
   ) => ReturnType<typeof streamText<TOOLS>>;
 
-  (fn: typeof generateObject): <T>(
+  (fn: typeof generateObject): <T extends string>(
     options: Parameters<typeof generateObject<T>>[0] & VercelExtraOptions
   ) => ReturnType<typeof generateObject<T>>;
 
diff --git a/tests/integration/vercel-sdk.test.ts b/tests/integration/vercel-sdk.test.ts
index 76e88b6..ffd110a 100644
--- a/tests/integration/vercel-sdk.test.ts
+++ b/tests/integration/vercel-sdk.test.ts
@@ -257,7 +257,7 @@ describe.skip('Vercel SDK Instrumentation', () => {
             parameters: z.object({
               value: z.number().describe('The value in celsius')
             }),
-            execute: async ({ value }) => {
+            execute: async ({ value }: any) => {
               const celsius = parseFloat(value);
               const fahrenheit = celsius * (9 / 5) + 32;
               return fahrenheit;
@@ -345,7 +345,7 @@ describe.skip('Vercel SDK Instrumentation', () => {
             parameters: z.object({
               value: z.number().describe('The value in celsius')
             }),
-            execute: async ({ value }) => {
+            execute: async ({ value }: any) => {
               const celsius = parseFloat(value);
               const fahrenheit = celsius * (9 / 5) + 32;
               return fahrenheit;