88 AsyncStreamingCallbackT ,
99 ChatMessage ,
1010 ChatRole ,
11+ ComponentInfo ,
1112 StreamingChunk ,
1213 SyncStreamingCallbackT ,
1314 ToolCall ,
@@ -235,7 +236,9 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
235236
236237
237238# Bedrock streaming to Haystack util methods
238- def _convert_event_to_streaming_chunk (event : Dict [str , Any ], model : str ) -> StreamingChunk :
239+ def _convert_event_to_streaming_chunk (
240+ event : Dict [str , Any ], model : str , component_info : ComponentInfo
241+ ) -> StreamingChunk :
239242 """
240243 Convert a Bedrock streaming event to a Haystack StreamingChunk.
241244
@@ -244,6 +247,7 @@ def _convert_event_to_streaming_chunk(event: Dict[str, Any], model: str) -> Stre
244247
245248 :param event: Dictionary containing a Bedrock streaming event.
246249 :param model: The model ID used for generation, included in chunk metadata.
250+ :param component_info: ComponentInfo object
247251 :returns: StreamingChunk object containing the content and metadata extracted from the event.
248252 """
249253 # Initialize an empty StreamingChunk to return if no relevant event is found
@@ -358,6 +362,8 @@ def _convert_event_to_streaming_chunk(event: Dict[str, Any], model: str) -> Stre
358362 },
359363 )
360364
365+ streaming_chunk .component_info = component_info
366+
361367 return streaming_chunk
362368
363369
@@ -438,18 +444,20 @@ def _parse_streaming_response(
438444 response_stream : EventStream ,
439445 streaming_callback : SyncStreamingCallbackT ,
440446 model : str ,
447+ component_info : ComponentInfo ,
441448) -> List [ChatMessage ]:
442449 """
443450 Parse a streaming response from Bedrock.
444451
445452 :param response_stream: EventStream from Bedrock API
446453 :param streaming_callback: Callback for streaming chunks
447454 :param model: The model ID used for generation
455+ :param component_info: ComponentInfo object
448456 :return: List of ChatMessage objects
449457 """
450458 chunks : List [StreamingChunk ] = []
451459 for event in response_stream :
452- streaming_chunk = _convert_event_to_streaming_chunk (event = event , model = model )
460+ streaming_chunk = _convert_event_to_streaming_chunk (event = event , model = model , component_info = component_info )
453461 streaming_callback (streaming_chunk )
454462 chunks .append (streaming_chunk )
455463 replies = [_convert_streaming_chunks_to_chat_message (chunks = chunks )]
@@ -460,18 +468,20 @@ async def _parse_streaming_response_async(
460468 response_stream : EventStream ,
461469 streaming_callback : AsyncStreamingCallbackT ,
462470 model : str ,
471+ component_info : ComponentInfo ,
463472) -> List [ChatMessage ]:
464473 """
465474 Parse a streaming response from Bedrock.
466475
467476 :param response_stream: EventStream from Bedrock API
468477 :param streaming_callback: Callback for streaming chunks
469478 :param model: The model ID used for generation
479+ :param component_info: ComponentInfo object
470480 :return: List of ChatMessage objects
471481 """
472482 chunks : List [StreamingChunk ] = []
473483 async for event in response_stream :
474- streaming_chunk = _convert_event_to_streaming_chunk (event = event , model = model )
484+ streaming_chunk = _convert_event_to_streaming_chunk (event = event , model = model , component_info = component_info )
475485 await streaming_callback (streaming_chunk )
476486 chunks .append (streaming_chunk )
477487 replies = [_convert_streaming_chunks_to_chat_message (chunks = chunks )]
0 commit comments