@traccia2/sdk 0.0.5 → 0.0.7

package/src/integrations/langchain-callback.ts

@@ -6,31 +6,7 @@
 
 import { ISpan } from '../types';
 import { getTracer } from '../auto';
-
-// Try to import BaseCallbackHandler from different LangChain versions
-let BaseCallbackHandler: any = class {
-  ignoreAgent = false;
-  ignoreChain = false;
-  ignoreLLM = false;
-  ignoreRetriever = false;
-  ignoreToolUse = false;
-};
-
-try {
-  // Try LangChain 0.2.x+
-  // eslint-disable-next-line @typescript-eslint/no-var-requires
-  const { BaseCallbackHandler: Handler } = require('@langchain/core/callbacks/base');
-  BaseCallbackHandler = Handler;
-} catch {
-  try {
-    // Try older LangChain versions
-    // eslint-disable-next-line @typescript-eslint/no-var-requires
-    const { BaseCallbackHandler: Handler } = require('langchain/callbacks');
-    BaseCallbackHandler = Handler;
-  } catch {
-    // Fallback to basic implementation - user likely has LangChain installed
-  }
-}
+import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
 
 /**
  * LangChain Callback Handler for Traccia SDK.
@@ -41,17 +17,19 @@ try {
  *
  * @example
  * import { ChatOpenAI } from '@langchain/openai';
- * import { TraciaCallbackHandler } from '@traccia/sdk/integrations';
+ * import { TracciaCallbackHandler } from '@traccia/sdk/integrations';
  *
- * const handler = new TraciaCallbackHandler();
+ * const handler = new TracciaCallbackHandler();
  * const model = new ChatOpenAI({ callbacks: [handler] });
  *
  * const response = await model.invoke({ input: 'Hello!' });
  * // Automatically traced with spans for LLM calls, tokens, latency, etc.
  */
-export class TraciaCallbackHandler extends BaseCallbackHandler {
+export class TracciaCallbackHandler extends BaseCallbackHandler {
+  name = 'TracciaCallbackHandler';
   private tracer = getTracer('langchain');
   private spanStack: Map<string, ISpan> = new Map();
+  private streamingStartTimes: Record<string, Date> = {};
 
   /**
    * Extract model name from LLM instance, checking multiple property locations.
@@ -125,23 +103,47 @@ export class TraciaCallbackHandler extends BaseCallbackHandler {
     if (span) {
       try {
         // Try multiple ways to get token usage
-        // OpenAI format
+        // OpenAI format and new @langchain/core format
         const tokenUsage = 
           output?.llmOutput?.token_usage || 
+          output?.llmOutput?.tokenUsage ||
           output?.token_usage || 
           output?.metadata?.token_usage;
 
         if (tokenUsage) {
-          if (tokenUsage.prompt_tokens !== undefined) {
-            span.setAttribute('llm.tokens.prompt', tokenUsage.prompt_tokens);
+          // Handle standard token counts
+          const promptTokens = tokenUsage.prompt_tokens ?? tokenUsage.promptTokens;
+          const completionTokens = tokenUsage.completion_tokens ?? tokenUsage.completionTokens;
+          const totalTokens = tokenUsage.total_tokens ?? tokenUsage.totalTokens;
+
+          if (promptTokens !== undefined) {
+            span.setAttribute('llm.tokens.prompt', promptTokens);
           }
-          if (tokenUsage.completion_tokens !== undefined) {
-            span.setAttribute('llm.tokens.completion', tokenUsage.completion_tokens);
+          if (completionTokens !== undefined) {
+            span.setAttribute('llm.tokens.completion', completionTokens);
           }
-          const totalTokens = (tokenUsage.prompt_tokens || 0) + (tokenUsage.completion_tokens || 0);
-          if (totalTokens > 0) {
+          if (totalTokens !== undefined) {
             span.setAttribute('llm.tokens.total', totalTokens);
           }
+
+          // Handle detailed token breakdown for models like GPT-4o vision
+          // input_token_details contains breakdown of prompt token usage
+          if (tokenUsage.input_token_details && typeof tokenUsage.input_token_details === 'object') {
+            for (const [key, value] of Object.entries(tokenUsage.input_token_details)) {
+              if (typeof value === 'number') {
+                span.setAttribute(`llm.tokens.input_${key}`, value);
+              }
+            }
+          }
+
+          // output_token_details contains breakdown of completion token usage
+          if (tokenUsage.output_token_details && typeof tokenUsage.output_token_details === 'object') {
+            for (const [key, value] of Object.entries(tokenUsage.output_token_details)) {
+              if (typeof value === 'number') {
+                span.setAttribute(`llm.tokens.output_${key}`, value);
+              }
+            }
+          }
         }
 
         // Capture output text length
@@ -187,6 +189,29 @@ export class TraciaCallbackHandler extends BaseCallbackHandler {
     }
   }
 
+  /**
+   * Handle LLM new token - called when a new token is generated during streaming.
+   * Tracks first token latency and token count for streaming scenarios.
+   */
+  public async handleLLMNewToken(
+    _token: string,
+    _idx?: any,
+    runId?: string
+  ): Promise<void> {
+    if (runId && !(runId in this.streamingStartTimes)) {
+      // Record the time of the first streaming token
+      this.streamingStartTimes[runId] = new Date();
+      const span = this.spanStack.get(runId);
+      if (span) {
+        try {
+          span.setAttribute('stream.first_token_generated', true);
+        } catch (error) {
+          // Silently fail
+        }
+      }
+    }
+  }
+
   /**
    * Handle chain start - called when a chain begins execution.
    */
@@ -360,6 +385,14 @@ export class TraciaCallbackHandler extends BaseCallbackHandler {
     return this.handleLLMError(error, runId);
   }
 
+  public async onLLMNewToken(
+    _token: string,
+    idx?: any,
+    runId?: string
+  ): Promise<void> {
+    return this.handleLLMNewToken(_token, idx, runId);
+  }
+
   public async onChainStart(
     chain: any,
     inputs: any,