fix: log instrumentation patching (#148)

traceloop · Mar 15, 2024 · 82d82bf · 82d82bf
1 parent 8287a02
commit 82d82bf
Show file tree

Hide file tree

Showing 14 changed files with 198 additions and 62 deletions.
diff --git a/packages/instrumentation-azure/src/instrumentation.ts b/packages/instrumentation-azure/src/instrumentation.ts
@@ -51,6 +51,8 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
   }
 
   public manuallyInstrument(module: typeof azure) {
+    this._diag.debug(`Patching @azure/openai manually`);
+
     this._wrap(
       module.OpenAIClient.prototype,
       "getChatCompletions",
@@ -74,7 +76,9 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
     return module;
   }
 
-  private patch(moduleExports: typeof azure) {
+  private patch(moduleExports: typeof azure, moduleVersion?: string) {
+    this._diag.debug(`Patching @azure/openai@${moduleVersion}`);
+
     this._wrap(
       moduleExports.OpenAIClient.prototype,
       "getChatCompletions",
@@ -88,7 +92,9 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private unpatch(moduleExports: typeof azure): void {
+  private unpatch(moduleExports: typeof azure, moduleVersion?: string): void {
+    this._diag.debug(`Unpatching @azure/openai@${moduleVersion}`);
+
     this._unwrap(moduleExports.OpenAIClient.prototype, "getChatCompletions");
     this._unwrap(moduleExports.OpenAIClient.prototype, "getCompletions");
   }
@@ -127,8 +133,11 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            if (e) {
+              plugin._diag.error("Error in Azure OpenAI instrumentation", e);
+            }
+          },
         );
 
         const wrappedPromise = plugin._wrapPromise(

diff --git a/packages/instrumentation-bedrock/src/instrumentation.ts b/packages/instrumentation-bedrock/src/instrumentation.ts
@@ -58,14 +58,20 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
   }
 
   public manuallyInstrument(module: typeof bedrock) {
+    this._diag.debug(`Patching @aws-sdk/client-bedrock-runtime manually`);
+
     this._wrap(
       module.BedrockRuntimeClient.prototype,
       "send",
       this.wrapperMethod(),
     );
   }
 
-  private wrap(module: typeof bedrock) {
+  private wrap(module: typeof bedrock, moduleVersion?: string) {
+    this._diag.debug(
+      `Patching @aws-sdk/client-bedrock-runtime@${moduleVersion}`,
+    );
+
     this._wrap(
       module.BedrockRuntimeClient.prototype,
       "send",
@@ -75,7 +81,11 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
     return module;
   }
 
-  private unwrap(module: typeof bedrock) {
+  private unwrap(module: typeof bedrock, moduleVersion?: string) {
+    this._diag.debug(
+      `Unpatching @aws-sdk/client-bedrock-runtime@${moduleVersion}`,
+    );
+
     this._unwrap(module.BedrockRuntimeClient.prototype, "send");
   }
 
@@ -95,8 +105,11 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            if (e) {
+              plugin._diag.error(`Error in bedrock instrumentation`, e);
+            }
+          },
         );
         const wrappedPromise = plugin._wrapPromise(span, execPromise);
         return context.bind(execContext, wrappedPromise);

diff --git a/packages/instrumentation-cohere/src/instrumentation.ts b/packages/instrumentation-cohere/src/instrumentation.ts
@@ -60,10 +60,13 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
   }
 
   public manuallyInstrument(module: typeof cohere) {
+    this._diag.debug(`Manually patching cohere-ai`);
     this.wrap(module);
   }
 
-  private wrap(module: typeof cohere) {
+  private wrap(module: typeof cohere, moduleVersion?: string) {
+    this._diag.debug(`Patching cohere-ai@${moduleVersion}`);
+
     this._wrap(
       module.CohereClient.prototype,
       "generate",
@@ -93,7 +96,9 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
     return module;
   }
 
-  private unwrap(module: typeof cohere) {
+  private unwrap(module: typeof cohere, moduleVersion?: string) {
+    this._diag.debug(`Unpatching @cohere-ai@${moduleVersion}`);
+
     this._unwrap(module.CohereClient.prototype, "generateStream");
     this._unwrap(module.CohereClient.prototype, "chat");
     this._unwrap(module.CohereClient.prototype, "chatStream");
@@ -117,8 +122,11 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            if (e) {
+              plugin._diag.error("Error in cohere instrumentation", e);
+            }
+          },
         );
         const wrappedPromise = plugin._wrapPromise(type, span, execPromise);
         return context.bind(execContext, wrappedPromise as any);

diff --git a/packages/instrumentation-langchain/src/instrumentation.ts b/packages/instrumentation-langchain/src/instrumentation.ts
@@ -42,12 +42,15 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     toolsModule?: any;
   }) {
     if (chainsModule) {
+      this._diag.debug("Manually instrumenting langchain chains");
       this.patchChainModule(chainsModule);
     }
     if (agentsModule) {
+      this._diag.debug("Manually instrumenting langchain agents");
       this.patchAgentModule(agentsModule);
     }
     if (toolsModule) {
+      this._diag.debug("Manually instrumenting langchain tools");
       this.patchToolsModule(toolsModule);
     }
   }
@@ -74,7 +77,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     return [chainModule, agentModule, toolsModule];
   }
 
-  private patchChainModule(moduleExports: typeof ChainsModule) {
+  private patchChainModule(
+    moduleExports: typeof ChainsModule,
+    moduleVersion?: string,
+  ) {
+    this._diag.debug(`Patching langchain/chains.cjs@${moduleVersion}`);
+
     this._wrap(
       moduleExports.RetrievalQAChain.prototype,
       "_call",
@@ -92,7 +100,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private patchAgentModule(moduleExports: typeof AgentsModule) {
+  private patchAgentModule(
+    moduleExports: typeof AgentsModule,
+    moduleVersion?: string,
+  ) {
+    this._diag.debug(`Patching langchain/agents.cjs@${moduleVersion}`);
+
     this._wrap(
       moduleExports.AgentExecutor.prototype,
       "_call",
@@ -105,7 +118,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private patchToolsModule(moduleExports: typeof ToolsModule) {
+  private patchToolsModule(
+    moduleExports: typeof ToolsModule,
+    moduleVersion?: string,
+  ) {
+    this._diag.debug(`Patching langchain/tools.cjs@${moduleVersion}`);
+
     this._wrap(
       moduleExports.Tool.prototype,
       "call",
@@ -114,19 +132,31 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private unpatchChainModule(moduleExports: any) {
+  private unpatchChainModule(
+    moduleExports: typeof ChainsModule,
+    moduleVersion?: string,
+  ) {
+    this._diag.debug(`Unpatching langchain/chains.cjs@${moduleVersion}`);
+
     this._unwrap(moduleExports.RetrievalQAChain.prototype, "_call");
     this._unwrap(moduleExports.BaseChain.prototype, "call");
     return moduleExports;
   }
 
-  private unpatchAgentModule(moduleExports: any) {
+  private unpatchAgentModule(
+    moduleExports: typeof AgentsModule,
+    moduleVersion?: string,
+  ) {
+    this._diag.debug(`Unpatching langchain/agents.cjs@${moduleVersion}`);
+
     this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
     return moduleExports;
   }
 
-  private unpatchToolsModule(moduleExports: any) {
-    this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
+  private unpatchToolsModule(moduleExports: typeof ToolsModule) {
+    this._diag.debug(`Unpatching langchain/tools.cjs`);
+
+    this._unwrap(moduleExports.Tool.prototype, "call");
     return moduleExports;
   }
 

diff --git a/packages/instrumentation-llamaindex/src/instrumentation.ts b/packages/instrumentation-llamaindex/src/instrumentation.ts
@@ -40,6 +40,8 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
   }
 
   public manuallyInstrument(module: typeof llamaindex) {
+    this._diag.debug("Manually instrumenting llamaindex");
+
     this.patch(module);
   }
 
@@ -78,7 +80,9 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
     return retriever && (retriever as BaseRetriever).retrieve !== undefined;
   }
 
-  private patch(moduleExports: typeof llamaindex) {
+  private patch(moduleExports: typeof llamaindex, moduleVersion?: string) {
+    this._diag.debug(`Patching llamaindex@${moduleVersion}`);
+
     const customLLMInstrumentation = new CustomLLMInstrumentation(
       this._config,
       () => this.tracer, // this is on purpose. Tracer may change
@@ -158,7 +162,9 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private unpatch(moduleExports: typeof llamaindex) {
+  private unpatch(moduleExports: typeof llamaindex, moduleVersion?: string) {
+    this._diag.debug(`Unpatching llamaindex@${moduleVersion}`);
+
     this._unwrap(moduleExports.RetrieverQueryEngine.prototype, "query");
 
     for (const key in moduleExports) {

diff --git a/packages/instrumentation-openai/src/instrumentation.ts b/packages/instrumentation-openai/src/instrumentation.ts
@@ -56,6 +56,8 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
   }
 
   public manuallyInstrument(module: typeof openai.OpenAI) {
+    this._diag.debug(`Manually instrumenting openai`);
+
     // Old version of OpenAI API (v3.1.0)
     if ((module as any).OpenAIApi) {
       this._wrap(
@@ -92,7 +94,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
     return module;
   }
 
-  private patch(moduleExports: typeof openai) {
+  private patch(moduleExports: typeof openai, moduleVersion?: string) {
+    this._diag.debug(`Patching openai@${moduleVersion}`);
+
     // Old version of OpenAI API (v3.1.0)
     if ((moduleExports as any).OpenAIApi) {
       this._wrap(
@@ -120,7 +124,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private unpatch(moduleExports: typeof openai): void {
+  private unpatch(moduleExports: typeof openai, moduleVersion?: string): void {
+    this._diag.debug(`Unpatching openai@${moduleVersion}`);
+
     // Old version of OpenAI API (v3.1.0)
     if ((moduleExports as any).OpenAIApi) {
       this._unwrap(
@@ -171,8 +177,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            if (e) {
+              plugin._diag.error("OpenAI instrumentation: error", e);
+            }
+          },
         );
 
         if (

diff --git a/packages/instrumentation-pinecone/src/instrumentation.ts b/packages/instrumentation-pinecone/src/instrumentation.ts
@@ -47,7 +47,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
     return module;
   }
 
-  private patch(moduleExports: typeof pinecone) {
+  private patch(moduleExports: typeof pinecone, moduleVersion?: string) {
+    this._diag.debug(`Patching @pinecone-database/pinecone@${moduleVersion}`);
+
     this._wrap(
       moduleExports.Index.prototype,
       "query",
@@ -77,7 +79,12 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
     return moduleExports;
   }
 
-  private unpatch(moduleExports: typeof pinecone): void {
+  private unpatch(
+    moduleExports: typeof pinecone,
+    moduleVersion?: string,
+  ): void {
+    this._diag.debug(`Unpatching @pinecone-database/pinecone@${moduleVersion}`);
+
     this._unwrap(moduleExports.Index.prototype, "query");
     this._unwrap(moduleExports.Index.prototype, "upsert");
     this._unwrap(moduleExports.Index.prototype, "deleteAll");
@@ -86,6 +93,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
   }
 
   private genericWrapper(methodName: string, tracer: Tracer) {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const plugin = this;
+
     // eslint-disable-next-line @typescript-eslint/ban-types
     return (original: Function) => {
       return function method(this: any, ...args: unknown[]) {
@@ -98,8 +108,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            plugin._diag.error(`Error in Pinecone instrumentation`, e);
+          },
         );
         const wrappedPromise = execPromise
           .then((result: any) => {
@@ -125,6 +136,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
   }
 
   private queryWrapper(tracer: Tracer) {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const plugin = this;
+
     // eslint-disable-next-line @typescript-eslint/ban-types
     return (original: Function) => {
       return function method(this: any, ...args: unknown[]) {
@@ -164,8 +178,11 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
               return original.apply(this, args);
             });
           },
-          // eslint-disable-next-line @typescript-eslint/no-empty-function
-          () => {},
+          (e) => {
+            if (e) {
+              plugin._diag.error(`Error in Pinecone instrumentation`, e);
+            }
+          },
         );
         const wrappedPromise = execPromise
           .then((result: any) => {