feat: add Bedrock agent tracer

gustavocidornelas · stainless-app[bot] · commit f0083a7cd081 · 2025-08-07T20:44:04.000Z
diff --git a/examples/bedrock-agent-tracing.ts b/examples/bedrock-agent-tracing.ts
@@ -0,0 +1,83 @@
+import { BedrockAgentRuntimeClient, InvokeAgentCommand } from '@aws-sdk/client-bedrock-agent-runtime';
+import { traceBedrockAgent } from '../src/lib/integrations/bedrockAgentTracer';
+
+/**
+ * Simple example to test the Bedrock agent tracer
+ * Replace the agent IDs with your actual Bedrock agent configuration
+ */
+export const simpleBedrockExample = async () => {
+  // Step 0: Make sure to set OPENLAYER_API_KEY and OPENLAYER_INFERENCE_PIPELINE_ID as environment variables
+
+  // Step 1: Create the Bedrock client
+  const client = new BedrockAgentRuntimeClient({
+    region: 'us-west-2', // Updated to match your agent's region
+  });
+
+  // Step 2: Apply Openlayer tracing
+  const tracedClient = traceBedrockAgent(client);
+
+  // Step 3: Configure your agent details
+  // Update the example with your actual agent details
+  const agentId = 'YOUR_AGENT_ID'; // Your actual agent ID
+  const agentAliasId = 'YOUR_ALIAS_ID'; // You'll need to create/find this
+  const sessionId = `test-session-${Date.now()}`;
+
+  // Step 4: Create the command
+  const command = new InvokeAgentCommand({
+    agentId: agentId,
+    agentAliasId: agentAliasId, // You'll get this after creating an alias
+    sessionId: `test-session-${Date.now()}`,
+    inputText: 'What is the capital of the moon?',
+    enableTrace: true,
+  });
+
+  try {
+    console.log('🤖 Sending request to Bedrock agent...');
+
+    // Step 5: Send the command and process the response
+    const response = await tracedClient.send(command);
+
+    if (!response.completion) {
+      throw new Error('No completion received');
+    }
+
+    let finalOutput = '';
+
+    // Step 6: Process the streaming response
+    for await (const event of response.completion) {
+      if (event.chunk?.bytes) {
+        const text = new TextDecoder('utf-8').decode(event.chunk.bytes);
+        finalOutput += text;
+        process.stdout.write(text); // Real-time output
+      }
+    }
+
+    console.log('\n\n✅ Request completed successfully!');
+    console.log('📊 Check your Openlayer dashboard for trace data');
+
+    return {
+      sessionId,
+      output: finalOutput,
+      success: true,
+    };
+  } catch (error) {
+    console.error('❌ Error:', error);
+    return {
+      sessionId,
+      output: '',
+      success: false,
+      error: error instanceof Error ? error.message : String(error),
+    };
+  }
+};
+
+// Run the example if this file is executed directly
+if (require.main === module) {
+  simpleBedrockExample()
+    .then((result) => {
+      console.log('\n📋 Final result:', result);
+    })
+    .catch((error) => {
+      console.error('💥 Failed to run example:', error);
+    });
+}
diff --git a/src/lib/integrations/bedrockAgentTracer.ts b/src/lib/integrations/bedrockAgentTracer.ts
@@ -0,0 +1,290 @@
+import {
+  BedrockAgentRuntimeClient,
+  InvokeAgentCommand,
+  InvokeAgentCommandInput,
+  InvokeAgentCommandOutput,
+} from '@aws-sdk/client-bedrock-agent-runtime';
+import { addChatCompletionStepToTrace } from '../tracing/tracer';
+
+export function traceBedrockAgent(client: BedrockAgentRuntimeClient): BedrockAgentRuntimeClient {
+  const originalSend = client.send.bind(client);
+
+  client.send = async function (this: BedrockAgentRuntimeClient, command: any, options?: any): Promise<any> {
+    // Only trace InvokeAgentCommand
+    if (!(command instanceof InvokeAgentCommand)) {
+      return originalSend(command, options);
+    }
+
+    const startTime = performance.now();
+    const input = command.input as InvokeAgentCommandInput;
+
+    try {
+      // Call the original send method
+      const response = (await originalSend(command, options)) as InvokeAgentCommandOutput;
+
+      if (!response.completion) {
+        throw new Error('Completion is undefined');
+      }
+
+      // Create a traced async iterator that preserves the original
+      const tracedCompletion = createTracedCompletion(response.completion, input, startTime);
+
+      // Return the response with the traced completion
+      return {
+        ...response,
+        completion: tracedCompletion,
+      };
+    } catch (error) {
+      console.error('Failed to trace the Bedrock agent invocation with Openlayer', error);
+      throw error;
+    }
+  };
+
+  return client;
+}
+
+// Create a traced completion that collects data while yielding original events
+function createTracedCompletion(
+  originalCompletion: AsyncIterable<any>,
+  input: InvokeAgentCommandInput,
+  startTime: number,
+): AsyncIterable<any> {
+  return {
+    async *[Symbol.asyncIterator]() {
+      let firstTokenTime: number | undefined;
+      let totalTokens = 0;
+      let promptTokens = 0;
+      let completionTokens = 0;
+      let collectedOutput = '';
+      const rawOutputChunks: any[] = [];
+      let agentModel: string | null = null;
+      let citations: any[] = [];
+      let traceData: any[] = [];
+      let chunkCount = 0;
+
+      try {
+        for await (const chunkEvent of originalCompletion) {
+          // YIELD FIRST - ensure user gets data immediately
+          yield chunkEvent;
+
+          // Then collect tracing data
+          if (chunkCount === 0) {
+            firstTokenTime = performance.now();
+          }
+          chunkCount++;
+
+          // Handle chunk events
+          if (chunkEvent.chunk) {
+            const chunk = chunkEvent.chunk;
+            rawOutputChunks.push(chunk);
+
+            if (chunk.bytes) {
+              const decodedResponse = new TextDecoder('utf-8').decode(chunk.bytes);
+              collectedOutput += decodedResponse;
+              completionTokens += 1;
+            }
+
+            if (chunk.attribution && chunk.attribution.citations) {
+              citations.push(...chunk.attribution.citations);
+            }
+          }
+
+          // Handle trace events
+          if (chunkEvent.trace) {
+            traceData.push(chunkEvent.trace);
+
+            if (chunkEvent.trace.trace) {
+              const trace = chunkEvent.trace.trace;
+
+              // Extract tokens and model info
+              if (
+                'orchestrationTrace' in trace &&
+                trace.orchestrationTrace?.modelInvocationOutput?.metadata?.usage
+              ) {
+                const usage = trace.orchestrationTrace.modelInvocationOutput.metadata.usage;
+                promptTokens += usage.inputTokens || 0;
+                completionTokens += usage.outputTokens || 0;
+              }
+
+              if (
+                'orchestrationTrace' in trace &&
+                trace.orchestrationTrace?.modelInvocationInput?.foundationModel
+              ) {
+                agentModel = trace.orchestrationTrace.modelInvocationInput.foundationModel;
+              }
+            }
+          }
+        }
+
+        // After the stream is complete, send trace data
+        const endTime = performance.now();
+        totalTokens = promptTokens + completionTokens;
+
+        // Send trace data to Openlayer
+        const inputs = extractInputs(input, traceData);
+        const metadata: Record<string, any> = {
+          agentId: input.agentId,
+          agentAliasId: input.agentAliasId,
+          sessionId: input.sessionId,
+          timeToFirstToken: firstTokenTime ? firstTokenTime - startTime : null,
+        };
+
+        if (citations.length > 0) {
+          metadata['citations'] = citations;
+        }
+
+        const reasoning = extractReasoning(traceData);
+        if (reasoning && reasoning.length > 0) {
+          metadata['reasoning'] = reasoning;
+        }
+
+        if (input.sessionState) {
+          metadata['sessionState'] = {
+            hasSessionAttributes: !!input.sessionState.sessionAttributes,
+            hasPromptSessionAttributes: !!input.sessionState.promptSessionAttributes,
+            hasFiles: !!input.sessionState.files && input.sessionState.files.length > 0,
+            hasKnowledgeBaseConfigurations:
+              !!input.sessionState.knowledgeBaseConfigurations &&
+              input.sessionState.knowledgeBaseConfigurations.length > 0,
+          };
+        }
+
+        const traceStepData = {
+          name: 'AWS Bedrock Agent Invocation',
+          inputs: inputs,
+          output: collectedOutput,
+          latency: endTime - startTime,
+          tokens: totalTokens > 0 ? totalTokens : null,
+          promptTokens: promptTokens > 0 ? promptTokens : null,
+          completionTokens: completionTokens > 0 ? completionTokens : null,
+          model: agentModel || `${input.agentId}:${input.agentAliasId}`,
+          modelParameters: extractModelParameters(input),
+          rawOutput: JSON.stringify(rawOutputChunks, null, 2),
+          metadata: metadata,
+          provider: 'Bedrock',
+        };
+
+        addChatCompletionStepToTrace(traceStepData);
+      } catch (error) {
+        console.error('Error in traced completion:', error);
+        // Don't rethrow - we don't want tracing errors to break the user's stream
+      }
+    },
+  };
+}
+
+function extractInputs(input: InvokeAgentCommandInput, traceData: any[]): Record<string, any> {
+  const inputs: Record<string, any> = {};
+
+  // Build the prompt in OpenAI-compatible format
+  const prompt: Array<{ role: string; content: string }> = [];
+
+  // Add the main user message
+  if (input.inputText) {
+    prompt.push({
+      role: 'user',
+      content: input.inputText,
+    });
+  }
+
+  // Add conversation history if present
+  if (input.sessionState?.conversationHistory?.messages) {
+    for (const message of input.sessionState.conversationHistory.messages) {
+      const content =
+        message.content ?
+          message.content.map((block) => ('text' in block ? block.text || '' : '')).join('')
+        : '';
+
+      const role = message.role || 'user';
+
+      prompt.unshift({
+        role: role,
+        content: content,
+      });
+    }
+  }
+
+  // Extract system prompt from trace data if available
+  const systemPrompt = extractSystemPrompt(traceData);
+  if (systemPrompt) {
+    prompt.unshift({
+      role: 'system',
+      content: systemPrompt,
+    });
+  }
+
+  inputs['prompt'] = prompt;
+
+  // Add additional context as separate fields
+  if (input.sessionState?.sessionAttributes) {
+    inputs['sessionAttributes'] = input.sessionState.sessionAttributes;
+  }
+
+  if (input.sessionState?.promptSessionAttributes) {
+    inputs['promptSessionAttributes'] = input.sessionState.promptSessionAttributes;
+  }
+
+  if (input.sessionState?.files && input.sessionState.files.length > 0) {
+    inputs['files'] = input.sessionState.files.map((file) => ({
+      name: file.name,
+      useCase: file.useCase,
+      sourceType: file.source?.sourceType,
+    }));
+  }
+
+  return inputs;
+}
+
+function extractSystemPrompt(traceData: any[]): string | null {
+  for (const trace of traceData) {
+    if (trace.trace?.orchestrationTrace?.modelInvocationInput?.text) {
+      try {
+        const parsed = JSON.parse(trace.trace.orchestrationTrace.modelInvocationInput.text);
+        if (parsed.system) {
+          return parsed.system;
+        }
+      } catch (e) {
+        // If parsing fails, continue
+      }
+    }
+  }
+  return null;
+}
+
+function extractReasoning(traceData: any[]): string[] | undefined {
+  const reasoning: string[] = [];
+
+  for (const trace of traceData) {
+    if (trace.trace?.orchestrationTrace?.rationale?.text) {
+      reasoning.push(trace.trace.orchestrationTrace.rationale.text);
+    }
+  }
+
+  return reasoning.length > 0 ? reasoning : undefined;
+}
+
+function extractModelParameters(input: InvokeAgentCommandInput): Record<string, any> {
+  const params: Record<string, any> = {};
+
+  if (input.enableTrace !== undefined) {
+    params['enableTrace'] = input.enableTrace;
+  }
+
+  if (input.endSession !== undefined) {
+    params['endSession'] = input.endSession;
+  }
+
+  if (input.bedrockModelConfigurations) {
+    params['bedrockModelConfigurations'] = input.bedrockModelConfigurations;
+  }
+
+  if (input.streamingConfigurations) {
+    params['streamingConfigurations'] = input.streamingConfigurations;
+  }
+
+  if (input.promptCreationConfigurations) {
+    params['promptCreationConfigurations'] = input.promptCreationConfigurations;
+  }
+
+  return params;
+}
