1- import {
2- BedrockAgentRuntimeClient ,
3- InvokeAgentCommand ,
4- InvokeAgentCommandInput ,
5- InvokeAgentCommandOutput ,
6- } from '@aws-sdk/client-bedrock-agent-runtime' ;
1+ // Make imports optional with try/catch
2+ let BedrockAgentRuntimeClient : any ;
3+ let InvokeAgentCommand : any ;
4+ let InvokeAgentCommandInput : any ;
5+ let InvokeAgentCommandOutput : any ;
6+
7+ try {
8+ const bedrockModule = require ( '@aws-sdk/client-bedrock-agent-runtime' ) ;
9+ BedrockAgentRuntimeClient = bedrockModule . BedrockAgentRuntimeClient ;
10+ InvokeAgentCommand = bedrockModule . InvokeAgentCommand ;
11+ InvokeAgentCommandInput = bedrockModule . InvokeAgentCommandInput ;
12+ InvokeAgentCommandOutput = bedrockModule . InvokeAgentCommandOutput ;
13+ } catch ( error ) {
14+ // AWS SDK not available
15+ }
16+
717import { addChatCompletionStepToTrace } from '../tracing/tracer' ;
818
9- export function traceBedrockAgent ( client : BedrockAgentRuntimeClient ) : BedrockAgentRuntimeClient {
19+ export function traceBedrockAgent ( client : any ) : any {
20+ if ( ! BedrockAgentRuntimeClient || ! InvokeAgentCommand ) {
21+ throw new Error (
22+ 'AWS SDK for Bedrock Agent Runtime is not installed. Please install it with: npm install @aws-sdk/client-bedrock-agent-runtime' ,
23+ ) ;
24+ }
25+
1026 const originalSend = client . send . bind ( client ) ;
1127
12- client . send = async function ( this : BedrockAgentRuntimeClient , command : any , options ?: any ) : Promise < any > {
28+ client . send = async function ( this : any , command : any , options ?: any ) : Promise < any > {
1329 // Only trace InvokeAgentCommand
1430 if ( ! ( command instanceof InvokeAgentCommand ) ) {
1531 return originalSend ( command , options ) ;
1632 }
1733
1834 const startTime = performance . now ( ) ;
19- const input = command . input as InvokeAgentCommandInput ;
35+ const input = command . input ;
2036
2137 try {
2238 // Call the original send method
23- const response = ( await originalSend ( command , options ) ) as InvokeAgentCommandOutput ;
39+ const response = await originalSend ( command , options ) ;
2440
2541 if ( ! response . completion ) {
2642 throw new Error ( 'Completion is undefined' ) ;
@@ -46,7 +62,7 @@ export function traceBedrockAgent(client: BedrockAgentRuntimeClient): BedrockAge
4662// Create a traced completion that collects data while yielding original events
4763function createTracedCompletion (
4864 originalCompletion : AsyncIterable < any > ,
49- input : InvokeAgentCommandInput ,
65+ input : any ,
5066 startTime : number ,
5167) : AsyncIterable < any > {
5268 return {
@@ -64,7 +80,7 @@ function createTracedCompletion(
6480
6581 try {
6682 for await ( const chunkEvent of originalCompletion ) {
67- // YIELD FIRST - ensure user gets data immediately
83+ // Yield first - ensure user gets data immediately
6884 yield chunkEvent ;
6985
7086 // Then collect tracing data
@@ -173,7 +189,7 @@ function createTracedCompletion(
173189 } ;
174190}
175191
176- function extractInputs ( input : InvokeAgentCommandInput , traceData : any [ ] ) : Record < string , any > {
192+ function extractInputs ( input : any , traceData : any [ ] ) : Record < string , any > {
177193 const inputs : Record < string , any > = { } ;
178194
179195 // Build the prompt in OpenAI-compatible format
@@ -192,7 +208,7 @@ function extractInputs(input: InvokeAgentCommandInput, traceData: any[]): Record
192208 for ( const message of input . sessionState . conversationHistory . messages ) {
193209 const content =
194210 message . content ?
195- message . content . map ( ( block ) => ( 'text' in block ? block . text || '' : '' ) ) . join ( '' )
211+ message . content . map ( ( block : any ) => ( 'text' in block ? block . text || '' : '' ) ) . join ( '' )
196212 : '' ;
197213
198214 const role = message . role || 'user' ;
@@ -225,7 +241,7 @@ function extractInputs(input: InvokeAgentCommandInput, traceData: any[]): Record
225241 }
226242
227243 if ( input . sessionState ?. files && input . sessionState . files . length > 0 ) {
228- inputs [ 'files' ] = input . sessionState . files . map ( ( file ) => ( {
244+ inputs [ 'files' ] = input . sessionState . files . map ( ( file : any ) => ( {
229245 name : file . name ,
230246 useCase : file . useCase ,
231247 sourceType : file . source ?. sourceType ,
@@ -263,7 +279,7 @@ function extractReasoning(traceData: any[]): string[] | undefined {
263279 return reasoning . length > 0 ? reasoning : undefined ;
264280}
265281
266- function extractModelParameters ( input : InvokeAgentCommandInput ) : Record < string , any > {
282+ function extractModelParameters ( input : any ) : Record < string , any > {
267283 const params : Record < string , any > = { } ;
268284
269285 if ( input . enableTrace !== undefined ) {
0 commit comments