fix: proper tracer fetcing in llamaindex and langchain

Author: nirga

packages/instrumentation-langchain/src/instrumentation.ts

@@ -85,15 +85,15 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
       moduleExports.RetrievalQAChain.prototype,
       "_call",
       workflowWrapper(
-        this.tracer,
+        () => this.tracer,
         this._shouldSendPrompts(),
         "retrieval_qa.workflow",
       ),
     );
     this._wrap(
       moduleExports.BaseChain.prototype,
       "call",
-      taskWrapper(this.tracer, this._shouldSendPrompts()),
+      taskWrapper(() => this.tracer, this._shouldSendPrompts()),
     );
     return moduleExports;
   }
@@ -109,7 +109,7 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
       moduleExports.AgentExecutor.prototype,
       "_call",
       workflowWrapper(
-        this.tracer,
+        () => this.tracer,
         this._shouldSendPrompts(),
         "langchain.agent",
       ),
@@ -127,7 +127,7 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
     this._wrap(
       moduleExports.Tool.prototype,
       "call",
-      taskWrapper(this.tracer, this._shouldSendPrompts()),
+      taskWrapper(() => this.tracer, this._shouldSendPrompts()),
     );
     return moduleExports;
   }

packages/instrumentation-langchain/src/utils.ts

@@ -6,15 +6,15 @@ import {
 } from "@traceloop/ai-semantic-conventions";
 
 export function genericWrapper(
-  tracer: Tracer,
+  tracer: () => Tracer,
   shouldSendPrompts: boolean,
   spanKind: TraceloopSpanKindValues,
   spanName?: string,
 ) {
   // eslint-disable-next-line @typescript-eslint/ban-types
   return (original: Function) => {
     return function method(this: any, ...args: unknown[]) {
-      const span = tracer.startSpan(
+      const span = tracer().startSpan(
         spanName || `langchain.${spanKind}.${this.constructor.name}`,
       );
       span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, spanKind);
@@ -97,7 +97,7 @@ export function genericWrapper(
 }
 
 export function taskWrapper(
-  tracer: Tracer,
+  tracer: () => Tracer,
   shouldSendPrompts: boolean,
   spanName?: string,
 ) {
@@ -110,7 +110,7 @@ export function taskWrapper(
 }
 
 export function workflowWrapper(
-  tracer: Tracer,
+  tracer: () => Tracer,
   shouldSendPrompts: boolean,
   spanName: string,
 ) {

packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts

@@ -25,9 +25,9 @@ type AsyncResponseType =
 
 export class CustomLLMInstrumentation {
   private config: LlamaIndexInstrumentationConfig;
-  private tracer: Tracer;
+  private tracer: () => Tracer;
 
-  constructor(config: LlamaIndexInstrumentationConfig, tracer: Tracer) {
+  constructor(config: LlamaIndexInstrumentationConfig, tracer: () => Tracer) {
     this.config = config;
     this.tracer = tracer;
   }
@@ -42,10 +42,11 @@ export class CustomLLMInstrumentation {
         const messages = params?.messages;
         const streaming = params?.stream;
 
-        const span = plugin.tracer.startSpan(
-          `llamaindex.${lodash.snakeCase(className)}.chat`,
-          { kind: SpanKind.CLIENT },
-        );
+        const span = plugin
+          .tracer()
+          .startSpan(`llamaindex.${lodash.snakeCase(className)}.chat`, {
+            kind: SpanKind.CLIENT,
+          });
 
         span.setAttribute(SpanAttributes.LLM_VENDOR, className);
         span.setAttribute(

packages/instrumentation-llamaindex/src/instrumentation.ts

@@ -91,7 +91,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
 
     const customLLMInstrumentation = new CustomLLMInstrumentation(
       this._config,
-      this.tracer,
+      () => this.tracer, // this is on purpose. Tracer may change
     );
 
     this._wrap(
@@ -101,7 +101,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
         moduleExports.RetrieverQueryEngine.name,
         "query",
         TraceloopSpanKindValues.WORKFLOW,
-        this.tracer,
+        () => this.tracer,
         shouldSendPrompts(this._config),
       ),
     );
@@ -113,7 +113,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
         moduleExports.ContextChatEngine.name,
         "chat",
         TraceloopSpanKindValues.WORKFLOW,
-        this.tracer,
+        () => this.tracer,
         shouldSendPrompts(this._config),
       ),
     );
@@ -134,7 +134,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             cls.name,
             "getQueryEmbedding",
             TraceloopSpanKindValues.TASK,
-            this.tracer,
+            () => this.tracer,
             shouldSendPrompts(this._config),
           ),
         );
@@ -146,7 +146,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             cls.name,
             "synthesize",
             TraceloopSpanKindValues.TASK,
-            this.tracer,
+            () => this.tracer,
             shouldSendPrompts(this._config),
           ),
         );
@@ -158,7 +158,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
             cls.name,
             "retrieve",
             TraceloopSpanKindValues.TASK,
-            this.tracer,
+            () => this.tracer,
             shouldSendPrompts(this._config),
           ),
         );

packages/instrumentation-llamaindex/src/utils.ts

@@ -44,14 +44,14 @@ export function genericWrapper(
   className: string,
   methodName: string,
   kind: TraceloopSpanKindValues,
-  tracer: Tracer,
+  tracer: () => Tracer,
   shouldSendPrompts: boolean,
 ) {
   // eslint-disable-next-line @typescript-eslint/ban-types
   return (original: Function) => {
     return function method(this: any, ...args: unknown[]) {
       const name = `${lodash.snakeCase(className)}.${lodash.snakeCase(methodName)}`;
-      const span = tracer.startSpan(`${name}`);
+      const span = tracer().startSpan(`${name}`, {}, context.active());
       span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, kind);
 
       if (kind === TraceloopSpanKindValues.WORKFLOW) {

packages/traceloop-sdk/src/lib/node-server-sdk.ts

@@ -1,3 +1,5 @@
+import { initInstrumentations } from "./tracing";
+
 export * from "./errors";
 export { InitializeOptions } from "./interfaces";
 export { initialize } from "./configuration";
@@ -6,3 +8,5 @@ export * from "./tracing/decorators";
 export * from "./tracing/association";
 export * from "./tracing/score";
 export * from "./prompts";
+
+initInstrumentations();