amazon-bedrock-haystack 3.11.0__tar.gz → 4.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/CHANGELOG.md +6 -0
  2. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/PKG-INFO +3 -3
  3. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/README.md +1 -1
  4. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/pyproject.toml +1 -1
  5. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/utils.py +118 -167
  6. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_chat_generator.py +15 -9
  7. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_chat_generator_utils.py +234 -236
  8. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/.gitignore +0 -0
  9. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/LICENSE.txt +0 -0
  10. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/bedrock_ranker_example.py +0 -0
  11. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/chatgenerator_example.py +0 -0
  12. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/embedders_generator_with_rag_example.py +0 -0
  13. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/pydoc/config.yml +0 -0
  14. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/__init__.py +0 -0
  15. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/errors.py +0 -0
  16. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/utils.py +0 -0
  17. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/py.typed +0 -0
  18. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/__init__.py +0 -0
  19. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py +0 -0
  20. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/document_image_embedder.py +0 -0
  21. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py +0 -0
  22. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  23. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/__init__.py +0 -0
  24. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py +0 -0
  25. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py +0 -0
  26. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +0 -0
  27. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +0 -0
  28. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  29. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/amazon_bedrock/__init__.py +0 -0
  30. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/amazon_bedrock/ranker.py +0 -0
  31. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  32. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/__init__.py +0 -0
  33. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/conftest.py +0 -0
  34. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_document_embedder.py +0 -0
  35. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_document_image_embedder.py +0 -0
  36. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/apple.jpg +0 -0
  37. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/haystack-logo.png +0 -0
  38. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/sample_pdf_1.pdf +0 -0
  39. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_generator.py +0 -0
  40. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_ranker.py +0 -0
  41. {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_text_embedder.py +0 -0
@@ -1,5 +1,11 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/amazon_bedrock-v3.11.0] - 2025-08-21
4
+
5
+ ### 🚀 Features
6
+
7
+ - Add `AmazonBedrockDocumentImageEmbedder` component (#2185)
8
+
3
9
  ## [integrations/amazon_bedrock-v3.10.0] - 2025-08-06
4
10
 
5
11
  ### 🚀 Features
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: amazon-bedrock-haystack
3
- Version: 3.11.0
3
+ Version: 4.0.0
4
4
  Summary: An integration of Amazon Bedrock as an AmazonBedrockGenerator component.
5
5
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme
6
6
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
21
21
  Requires-Python: >=3.9
22
22
  Requires-Dist: aioboto3>=14.0.0
23
23
  Requires-Dist: boto3>=1.28.57
24
- Requires-Dist: haystack-ai>=2.16.0
24
+ Requires-Dist: haystack-ai>=2.17.1
25
25
  Description-Content-Type: text/markdown
26
26
 
27
27
  # amazon-bedrock-haystack
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
31
31
 
32
32
  - [Integration page](https://haystack.deepset.ai/integrations/amazon-bedrock)
33
33
  - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/amazon_bedrock/CHANGELOG.md)
34
- -----
34
+ ---
35
35
 
36
36
  ## Contributing
37
37
 
@@ -5,7 +5,7 @@
5
5
 
6
6
  - [Integration page](https://haystack.deepset.ai/integrations/amazon-bedrock)
7
7
  - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/amazon_bedrock/CHANGELOG.md)
8
- -----
8
+ ---
9
9
 
10
10
  ## Contributing
11
11
 
@@ -23,7 +23,7 @@ classifiers = [
23
23
  "Programming Language :: Python :: Implementation :: CPython",
24
24
  "Programming Language :: Python :: Implementation :: PyPy",
25
25
  ]
26
- dependencies = ["haystack-ai>=2.16.0", "boto3>=1.28.57", "aioboto3>=14.0.0"]
26
+ dependencies = ["haystack-ai>=2.17.1", "boto3>=1.28.57", "aioboto3>=14.0.0"]
27
27
 
28
28
  [project.urls]
29
29
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme"
@@ -5,16 +5,20 @@ from typing import Any, Dict, List, Optional, Tuple
5
5
 
6
6
  from botocore.eventstream import EventStream
7
7
  from haystack import logging
8
+ from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message
8
9
  from haystack.dataclasses import (
9
10
  AsyncStreamingCallbackT,
10
11
  ChatMessage,
11
12
  ChatRole,
12
13
  ComponentInfo,
14
+ FinishReason,
13
15
  ImageContent,
16
+ ReasoningContent,
14
17
  StreamingChunk,
15
18
  SyncStreamingCallbackT,
16
19
  TextContent,
17
20
  ToolCall,
21
+ ToolCallDelta,
18
22
  )
19
23
  from haystack.tools import Tool
20
24
 
@@ -24,6 +28,16 @@ logger = logging.getLogger(__name__)
24
28
  # see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html for supported formats
25
29
  IMAGE_SUPPORTED_FORMATS = ["png", "jpeg", "gif", "webp"]
26
30
 
31
+ # see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_MessageStopEvent.html
32
+ FINISH_REASON_MAPPING: Dict[str, FinishReason] = {
33
+ "end_turn": "stop",
34
+ "stop_sequence": "stop",
35
+ "max_tokens": "length",
36
+ "guardrail_intervened": "content_filter",
37
+ "content_filtered": "content_filter",
38
+ "tool_use": "tool_calls",
39
+ }
40
+
27
41
 
28
42
  # Haystack to Bedrock util methods
29
43
  def _format_tools(tools: Optional[List[Tool]] = None) -> Optional[Dict[str, Any]]:
@@ -57,8 +71,8 @@ def _format_tool_call_message(tool_call_message: ChatMessage) -> Dict[str, Any]:
57
71
  content: List[Dict[str, Any]] = []
58
72
 
59
73
  # tool call messages can contain reasoning content
60
- if reasoning_contents := tool_call_message.meta.get("reasoning_contents"):
61
- content.extend(_format_reasoning_contents(reasoning_contents=reasoning_contents))
74
+ if reasoning_content := tool_call_message.reasoning:
75
+ content.extend(_format_reasoning_content(reasoning_content=reasoning_content))
62
76
 
63
77
  # Tool call message can contain text
64
78
  if tool_call_message.text:
@@ -162,16 +176,16 @@ def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]
162
176
  return [msg for _, msg in repaired_bedrock_formatted_messages]
163
177
 
164
178
 
165
- def _format_reasoning_contents(reasoning_contents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
179
+ def _format_reasoning_content(reasoning_content: ReasoningContent) -> List[Dict[str, Any]]:
166
180
  """
167
- Format reasoning contents to match Bedrock's expected structure.
181
+ Format ReasoningContent to match Bedrock's expected structure.
168
182
 
169
- :param reasoning_contents: List of reasoning content dictionaries from Haystack ChatMessage metadata.
183
+ :param reasoning_content: ReasoningContent object containing reasoning contents to format.
170
184
  :returns: List of formatted reasoning content dictionaries for Bedrock.
171
185
  """
172
186
  formatted_contents = []
173
- for reasoning_content in reasoning_contents:
174
- formatted_content = {"reasoningContent": reasoning_content["reasoning_content"]}
187
+ for content in reasoning_content.extra.get("reasoning_contents", []):
188
+ formatted_content = {"reasoningContent": content["reasoning_content"]}
175
189
  if reasoning_text := formatted_content["reasoningContent"].pop("reasoning_text", None):
176
190
  formatted_content["reasoningContent"]["reasoningText"] = reasoning_text
177
191
  if redacted_content := formatted_content["reasoningContent"].pop("redacted_content", None):
@@ -192,8 +206,8 @@ def _format_text_image_message(message: ChatMessage) -> Dict[str, Any]:
192
206
 
193
207
  bedrock_content_blocks: List[Dict[str, Any]] = []
194
208
  # Add reasoning content if available as the first content block
195
- if message.meta.get("reasoning_contents"):
196
- bedrock_content_blocks.extend(_format_reasoning_contents(reasoning_contents=message.meta["reasoning_contents"]))
209
+ if message.reasoning:
210
+ bedrock_content_blocks.extend(_format_reasoning_content(reasoning_content=message.reasoning))
197
211
 
198
212
  for part in content_parts:
199
213
  if isinstance(part, TextContent):
@@ -269,7 +283,7 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
269
283
  base_meta = {
270
284
  "model": model,
271
285
  "index": 0,
272
- "finish_reason": response_body.get("stopReason"),
286
+ "finish_reason": FINISH_REASON_MAPPING.get(response_body.get("stopReason", "")),
273
287
  "usage": {
274
288
  # OpenAI's format for usage for cross ChatGenerator compatibility
275
289
  "prompt_tokens": response_body.get("usage", {}).get("inputTokens", 0),
@@ -303,11 +317,26 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
303
317
  reasoning_content["redacted_content"] = reasoning_content.pop("redactedContent")
304
318
  reasoning_contents.append({"reasoning_content": reasoning_content})
305
319
 
306
- # If reasoning contents were found, add them to the base meta
307
- base_meta.update({"reasoning_contents": reasoning_contents})
320
+ reasoning_text = ""
321
+ for content in reasoning_contents:
322
+ if "reasoning_text" in content["reasoning_content"]:
323
+ reasoning_text += content["reasoning_content"]["reasoning_text"]["text"]
324
+ elif "redacted_content" in content["reasoning_content"]:
325
+ reasoning_text += "[REDACTED]"
308
326
 
309
327
  # Create a single ChatMessage with combined text and tool calls
310
- replies.append(ChatMessage.from_assistant(" ".join(text_content), tool_calls=tool_calls, meta=base_meta))
328
+ replies.append(
329
+ ChatMessage.from_assistant(
330
+ " ".join(text_content),
331
+ tool_calls=tool_calls,
332
+ meta=base_meta,
333
+ reasoning=ReasoningContent(
334
+ reasoning_text=reasoning_text, extra={"reasoning_contents": reasoning_contents}
335
+ )
336
+ if reasoning_contents
337
+ else None,
338
+ )
339
+ )
311
340
 
312
341
  return replies
313
342
 
@@ -328,9 +357,8 @@ def _convert_event_to_streaming_chunk(
328
357
  """
329
358
  # Initialize an empty StreamingChunk to return if no relevant event is found
330
359
  # (e.g. for messageStart and contentBlockStop)
331
- streaming_chunk = StreamingChunk(
332
- content="", meta={"model": model, "received_at": datetime.now(timezone.utc).isoformat()}
333
- )
360
+ base_meta = {"model": model, "received_at": datetime.now(timezone.utc).isoformat()}
361
+ streaming_chunk = StreamingChunk(content="", meta=base_meta)
334
362
 
335
363
  if "contentBlockStart" in event:
336
364
  # contentBlockStart always has the key "contentBlockIndex"
@@ -340,26 +368,15 @@ def _convert_event_to_streaming_chunk(
340
368
  tool_start = block_start["start"]["toolUse"]
341
369
  streaming_chunk = StreamingChunk(
342
370
  content="",
343
- meta={
344
- "model": model,
345
- # This is always 0 b/c it represents the choice index
346
- "index": 0,
347
- # We follow the same format used in the OpenAIChatGenerator
348
- "tool_calls": [ # Optional[List[ChoiceDeltaToolCall]]
349
- {
350
- "index": block_idx, # int
351
- "id": tool_start["toolUseId"], # Optional[str]
352
- "function": { # Optional[ChoiceDeltaToolCallFunction]
353
- # Will accumulate deltas as string
354
- "arguments": "", # Optional[str]
355
- "name": tool_start["name"], # Optional[str]
356
- },
357
- "type": "function", # Optional[Literal["function"]]
358
- }
359
- ],
360
- "finish_reason": None,
361
- "received_at": datetime.now(timezone.utc).isoformat(),
362
- },
371
+ index=block_idx,
372
+ tool_calls=[
373
+ ToolCallDelta(
374
+ index=block_idx,
375
+ id=tool_start["toolUseId"],
376
+ tool_name=tool_start["name"],
377
+ )
378
+ ],
379
+ meta=base_meta,
363
380
  )
364
381
 
365
382
  elif "contentBlockDelta" in event:
@@ -370,39 +387,22 @@ def _convert_event_to_streaming_chunk(
370
387
  if "text" in delta:
371
388
  streaming_chunk = StreamingChunk(
372
389
  content=delta["text"],
373
- meta={
374
- "model": model,
375
- # This is always 0 b/c it represents the choice index
376
- "index": 0,
377
- "tool_calls": None,
378
- "finish_reason": None,
379
- "received_at": datetime.now(timezone.utc).isoformat(),
380
- },
390
+ index=block_idx,
391
+ meta=base_meta,
381
392
  )
382
393
  # This only occurs when accumulating the arguments for a toolUse
383
394
  # The content_block for this tool should already exist at this point
384
395
  elif "toolUse" in delta:
385
396
  streaming_chunk = StreamingChunk(
386
397
  content="",
387
- meta={
388
- "model": model,
389
- # This is always 0 b/c it represents the choice index
390
- "index": 0,
391
- "tool_calls": [ # Optional[List[ChoiceDeltaToolCall]]
392
- {
393
- "index": block_idx, # int
394
- "id": None, # Optional[str]
395
- "function": { # Optional[ChoiceDeltaToolCallFunction]
396
- # Will accumulate deltas as string
397
- "arguments": delta["toolUse"].get("input", ""), # Optional[str]
398
- "name": None, # Optional[str]
399
- },
400
- "type": "function", # Optional[Literal["function"]]
401
- }
402
- ],
403
- "finish_reason": None,
404
- "received_at": datetime.now(timezone.utc).isoformat(),
405
- },
398
+ index=block_idx,
399
+ tool_calls=[
400
+ ToolCallDelta(
401
+ index=block_idx,
402
+ arguments=delta["toolUse"].get("input", ""),
403
+ )
404
+ ],
405
+ meta=base_meta,
406
406
  )
407
407
  # This is for accumulating reasoning content deltas
408
408
  elif "reasoningContent" in delta:
@@ -411,28 +411,19 @@ def _convert_event_to_streaming_chunk(
411
411
  reasoning_content["redacted_content"] = reasoning_content.pop("redactedContent")
412
412
  streaming_chunk = StreamingChunk(
413
413
  content="",
414
+ index=block_idx,
414
415
  meta={
415
- "model": model,
416
- "index": 0,
417
- "tool_calls": None,
418
- "finish_reason": None,
419
- "received_at": datetime.now(timezone.utc).isoformat(),
416
+ **base_meta,
420
417
  "reasoning_contents": [{"index": block_idx, "reasoning_content": reasoning_content}],
421
418
  },
422
419
  )
423
420
 
424
421
  elif "messageStop" in event:
425
- finish_reason = event["messageStop"].get("stopReason")
422
+ finish_reason = FINISH_REASON_MAPPING.get(event["messageStop"].get("stopReason"))
426
423
  streaming_chunk = StreamingChunk(
427
424
  content="",
428
- meta={
429
- "model": model,
430
- # This is always 0 b/c it represents the choice index
431
- "index": 0,
432
- "tool_calls": None,
433
- "finish_reason": finish_reason,
434
- "received_at": datetime.now(timezone.utc).isoformat(),
435
- },
425
+ finish_reason=finish_reason,
426
+ meta=base_meta,
436
427
  )
437
428
 
438
429
  elif "metadata" in event and "usage" in event["metadata"]:
@@ -440,12 +431,7 @@ def _convert_event_to_streaming_chunk(
440
431
  streaming_chunk = StreamingChunk(
441
432
  content="",
442
433
  meta={
443
- "model": model,
444
- # This is always 0 b/c it represents the choice index
445
- "index": 0,
446
- "tool_calls": None,
447
- "finish_reason": None,
448
- "received_at": datetime.now(timezone.utc).isoformat(),
434
+ **base_meta,
449
435
  "usage": {
450
436
  "prompt_tokens": metadata["usage"].get("inputTokens", 0),
451
437
  "completion_tokens": metadata["usage"].get("outputTokens", 0),
@@ -459,7 +445,7 @@ def _convert_event_to_streaming_chunk(
459
445
  return streaming_chunk
460
446
 
461
447
 
462
- def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str, Any]]:
448
+ def _process_reasoning_contents(chunks: List[StreamingChunk]) -> Optional[ReasoningContent]:
463
449
  """
464
450
  Process reasoning contents from a list of StreamingChunk objects into the Bedrock expected format.
465
451
 
@@ -491,6 +477,8 @@ def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str,
491
477
  )
492
478
  if redacted_content:
493
479
  formatted_reasoning_contents.append({"reasoning_content": {"redacted_content": redacted_content}})
480
+
481
+ # Reset accumulators for new group
494
482
  reasoning_text = ""
495
483
  reasoning_signature = None
496
484
  redacted_content = None
@@ -516,85 +504,22 @@ def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str,
516
504
  if redacted_content:
517
505
  formatted_reasoning_contents.append({"reasoning_content": {"redacted_content": redacted_content}})
518
506
 
519
- return formatted_reasoning_contents
520
-
521
-
522
- def _convert_streaming_chunks_to_chat_message(chunks: List[StreamingChunk]) -> ChatMessage:
523
- """
524
- Converts a list of streaming chunks into a ChatMessage object.
525
-
526
- The function processes streaming chunks to build a ChatMessage object, including extracting and constructing
527
- tool calls, managing metadata such as model type, finish reason, and usage information.
528
- The tool call processing handles accumulating data across the chunks and attempts to parse JSON-formatted
529
- arguments for tool calls.
530
-
531
- :param chunks: A list of StreamingChunk objects representing parts of the assistant's response.
532
-
533
- :returns:
534
- A ChatMessage object constructed from the streaming chunks, containing the aggregated text, processed tool
535
- calls, and metadata.
536
- """
537
- # Join all text content from the chunks
538
- text = "".join([chunk.content for chunk in chunks])
539
-
540
- # If reasoning content is present in any chunk, accumulate it
541
- reasoning_contents = _process_reasoning_contents(chunks=chunks)
542
-
543
- # Process tool calls if present in any chunk
544
- tool_calls = []
545
- tool_call_data: Dict[int, Dict[str, str]] = {} # Track tool calls by index
546
- for chunk_payload in chunks:
547
- tool_calls_meta = chunk_payload.meta.get("tool_calls")
548
- if tool_calls_meta is not None:
549
- for delta in tool_calls_meta:
550
- # We use the index of the tool call to track it across chunks since the ID is not always provided
551
- if delta["index"] not in tool_call_data:
552
- tool_call_data[delta["index"]] = {"id": "", "name": "", "arguments": ""}
553
-
554
- # Save the ID if present
555
- if delta.get("id"):
556
- tool_call_data[delta["index"]]["id"] = delta["id"]
557
-
558
- if delta.get("function"):
559
- if delta["function"].get("name"):
560
- tool_call_data[delta["index"]]["name"] += delta["function"]["name"]
561
- if delta["function"].get("arguments"):
562
- tool_call_data[delta["index"]]["arguments"] += delta["function"]["arguments"]
563
-
564
- # Convert accumulated tool call data into ToolCall objects
565
- for call_data in tool_call_data.values():
566
- try:
567
- arguments = json.loads(call_data.get("arguments", "{}")) if call_data.get("arguments") else {}
568
- tool_calls.append(ToolCall(id=call_data["id"], tool_name=call_data["name"], arguments=arguments))
569
- except json.JSONDecodeError:
570
- logger.warning(
571
- "Amazon Bedrock returned a malformed JSON string for tool call arguments. This tool call will be "
572
- "skipped. Tool call ID: {tool_id}, Tool name: {tool_name}, Arguments: {tool_arguments}",
573
- tool_id=call_data["id"],
574
- tool_name=call_data["name"],
575
- tool_arguments=call_data["arguments"],
576
- )
577
-
578
- # finish_reason can appear in different places so we look for the last one
579
- finish_reasons = [
580
- chunk.meta.get("finish_reason") for chunk in chunks if chunk.meta.get("finish_reason") is not None
581
- ]
582
- finish_reason = finish_reasons[-1] if finish_reasons else None
583
-
584
- # usage is usually last but we look for it as well
585
- usages = [chunk.meta.get("usage") for chunk in chunks if chunk.meta.get("usage") is not None]
586
- usage = usages[-1] if usages else None
587
-
588
- meta = {
589
- "model": chunks[-1].meta["model"],
590
- "index": 0,
591
- "finish_reason": finish_reason,
592
- "completion_start_time": chunks[0].meta.get("received_at"), # first chunk received
593
- "usage": usage,
594
- "reasoning_contents": reasoning_contents,
595
- }
596
-
597
- return ChatMessage.from_assistant(text=text or None, tool_calls=tool_calls, meta=meta)
507
+ # Combine all reasoning texts into a single string for the main reasoning_text field
508
+ final_reasoning_text = ""
509
+ for content in formatted_reasoning_contents:
510
+ if "reasoning_text" in content["reasoning_content"]:
511
+ # mypy somehow thinks that content["reasoning_content"]["reasoning_text"]["text"] can be of type None
512
+ final_reasoning_text += content["reasoning_content"]["reasoning_text"]["text"] # type: ignore[operator]
513
+ elif "redacted_content" in content["reasoning_content"]:
514
+ final_reasoning_text += "[REDACTED]"
515
+
516
+ return (
517
+ ReasoningContent(
518
+ reasoning_text=final_reasoning_text, extra={"reasoning_contents": formatted_reasoning_contents}
519
+ )
520
+ if formatted_reasoning_contents
521
+ else None
522
+ )
598
523
 
599
524
 
600
525
  def _parse_streaming_response(
@@ -612,13 +537,26 @@ def _parse_streaming_response(
612
537
  :param component_info: ComponentInfo object
613
538
  :return: List of ChatMessage objects
614
539
  """
540
+ content_block_idxs = set()
615
541
  chunks: List[StreamingChunk] = []
616
542
  for event in response_stream:
617
543
  streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model, component_info=component_info)
544
+ content_block_idx = streaming_chunk.index
545
+ if content_block_idx is not None and content_block_idx not in content_block_idxs:
546
+ streaming_chunk.start = True
547
+ content_block_idxs.add(content_block_idx)
618
548
  streaming_callback(streaming_chunk)
619
549
  chunks.append(streaming_chunk)
620
- replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
621
- return replies
550
+ reply = _convert_streaming_chunks_to_chat_message(chunks=chunks)
551
+ reasoning_content = _process_reasoning_contents(chunks=chunks)
552
+ reply = ChatMessage.from_assistant(
553
+ text=reply.text,
554
+ meta=reply.meta,
555
+ name=reply.name,
556
+ tool_calls=reply.tool_calls,
557
+ reasoning=reasoning_content,
558
+ )
559
+ return [reply]
622
560
 
623
561
 
624
562
  async def _parse_streaming_response_async(
@@ -636,10 +574,23 @@ async def _parse_streaming_response_async(
636
574
  :param component_info: ComponentInfo object
637
575
  :return: List of ChatMessage objects
638
576
  """
577
+ content_block_idxs = set()
639
578
  chunks: List[StreamingChunk] = []
640
579
  async for event in response_stream:
641
580
  streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model, component_info=component_info)
581
+ content_block_idx = streaming_chunk.index
582
+ if content_block_idx is not None and content_block_idx not in content_block_idxs:
583
+ streaming_chunk.start = True
584
+ content_block_idxs.add(content_block_idx)
642
585
  await streaming_callback(streaming_chunk)
643
586
  chunks.append(streaming_chunk)
644
- replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
645
- return replies
587
+ reply = _convert_streaming_chunks_to_chat_message(chunks=chunks)
588
+ reasoning_content = _process_reasoning_contents(chunks=chunks)
589
+ reply = ChatMessage.from_assistant(
590
+ text=reply.text,
591
+ meta=reply.meta,
592
+ name=reply.name,
593
+ tool_calls=reply.tool_calls,
594
+ reasoning=reasoning_content,
595
+ )
596
+ return [reply]
@@ -392,7 +392,7 @@ class TestAmazonBedrockChatGeneratorInference:
392
392
  assert tool_call.id, "Tool call does not contain value for 'id' key"
393
393
  assert tool_call.tool_name == "weather"
394
394
  assert tool_call.arguments["city"] in ["Paris", "Berlin"]
395
- assert tool_call_message.meta["finish_reason"] == "tool_use"
395
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
396
396
 
397
397
  # Mock the response we'd get from ToolInvoker
398
398
  tool_result_messages = [
@@ -438,12 +438,14 @@ class TestAmazonBedrockChatGeneratorInference:
438
438
  assert isinstance(tool_call_message, ChatMessage), "Tool message is not a ChatMessage instance"
439
439
  assert ChatMessage.is_from(tool_call_message, ChatRole.ASSISTANT), "Tool message is not from the assistant"
440
440
 
441
+ assert tool_call_message.reasoning is not None, "Tool message does not contain reasoning"
442
+
441
443
  tool_calls = tool_call_message.tool_calls
442
444
  assert len(tool_calls) == 1
443
445
  assert tool_calls[0].id, "Tool call does not contain value for 'id' key"
444
446
  assert tool_calls[0].tool_name == "weather"
445
447
  assert tool_calls[0].arguments["city"] == "Paris"
446
- assert tool_call_message.meta["finish_reason"] == "tool_use"
448
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
447
449
 
448
450
  # Mock the response we'd get from ToolInvoker
449
451
  tool_result_messages = [
@@ -489,12 +491,14 @@ class TestAmazonBedrockChatGeneratorInference:
489
491
  assert isinstance(tool_call_message, ChatMessage), "Tool message is not a ChatMessage instance"
490
492
  assert ChatMessage.is_from(tool_call_message, ChatRole.ASSISTANT), "Tool message is not from the assistant"
491
493
 
494
+ assert tool_call_message.reasoning is not None, "Tool message does not contain reasoning"
495
+
492
496
  tool_calls = tool_call_message.tool_calls
493
497
  assert len(tool_calls) == 1
494
498
  assert tool_calls[0].id, "Tool call does not contain value for 'id' key"
495
499
  assert tool_calls[0].tool_name == "weather"
496
500
  assert tool_calls[0].arguments["city"] == "Paris"
497
- assert tool_call_message.meta["finish_reason"] == "tool_use"
501
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
498
502
 
499
503
  # Mock the response we'd get from ToolInvoker
500
504
  tool_result_messages = [
@@ -533,7 +537,8 @@ class TestAmazonBedrockChatGeneratorInference:
533
537
 
534
538
  assert len(results["replies"]) > 0, "No replies received"
535
539
  assert isinstance(
536
- results["replies"][0].meta["reasoning_contents"][0]["reasoning_content"]["redacted_content"], bytes
540
+ results["replies"][0].reasoning.extra["reasoning_contents"][0]["reasoning_content"]["redacted_content"],
541
+ bytes,
537
542
  )
538
543
 
539
544
  def test_live_run_with_redacted_thinking_streaming(self, tools):
@@ -560,7 +565,8 @@ class TestAmazonBedrockChatGeneratorInference:
560
565
 
561
566
  assert len(results["replies"]) > 0, "No replies received"
562
567
  assert isinstance(
563
- results["replies"][0].meta["reasoning_contents"][0]["reasoning_content"]["redacted_content"], bytes
568
+ results["replies"][0].reasoning.extra["reasoning_contents"][0]["reasoning_content"]["redacted_content"],
569
+ bytes,
564
570
  )
565
571
 
566
572
  @pytest.mark.parametrize("model_name", STREAMING_TOOL_MODELS)
@@ -591,7 +597,7 @@ class TestAmazonBedrockChatGeneratorInference:
591
597
  assert tool_call.id, "Tool call does not contain value for 'id' key"
592
598
  assert tool_call.tool_name == "weather"
593
599
  assert tool_call.arguments["city"] in ["Paris", "Berlin"]
594
- assert tool_call_message.meta["finish_reason"] == "tool_use"
600
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
595
601
 
596
602
  # Mock the response we'd get from ToolInvoker
597
603
  tool_result_messages = [
@@ -633,7 +639,7 @@ class TestAmazonBedrockChatGeneratorInference:
633
639
  assert tool_call.id is not None
634
640
  assert tool_call.tool_name == "hello_world"
635
641
  assert tool_call.arguments == {}
636
- assert message.meta["finish_reason"] == "tool_use"
642
+ assert message.meta["finish_reason"] == "tool_calls"
637
643
 
638
644
  new_messages = [
639
645
  *initial_messages,
@@ -724,7 +730,7 @@ class TestAmazonBedrockChatGeneratorAsyncInference:
724
730
  assert tool_call.id, "Tool call does not contain value for 'id' key"
725
731
  assert tool_call.tool_name == "weather"
726
732
  assert tool_call.arguments["city"] in ["Paris", "Berlin"]
727
- assert tool_call_message.meta["finish_reason"] == "tool_use"
733
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
728
734
 
729
735
  # Mock the response we'd get from ToolInvoker
730
736
  tool_result_messages = [
@@ -800,7 +806,7 @@ class TestAmazonBedrockChatGeneratorAsyncInference:
800
806
  assert tool_call.id, "Tool call does not contain value for 'id' key"
801
807
  assert tool_call.tool_name == "weather"
802
808
  assert tool_call.arguments["city"] in ["Paris", "Berlin"]
803
- assert tool_call_message.meta["finish_reason"] == "tool_use"
809
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
804
810
 
805
811
  # Mock the response we'd get from ToolInvoker
806
812
  tool_result_messages = [