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.
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/CHANGELOG.md +6 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/PKG-INFO +3 -3
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/README.md +1 -1
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/pyproject.toml +1 -1
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/utils.py +118 -167
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_chat_generator.py +15 -9
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_chat_generator_utils.py +234 -236
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/.gitignore +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/LICENSE.txt +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/bedrock_ranker_example.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/chatgenerator_example.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/examples/embedders_generator_with_rag_example.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/pydoc/config.yml +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/__init__.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/errors.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/amazon_bedrock/utils.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/common/py.typed +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/__init__.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py +0 -0
- {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
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/__init__.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py +0 -0
- {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
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/amazon_bedrock/__init__.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/amazon_bedrock/ranker.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/__init__.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/conftest.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_document_embedder.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_document_image_embedder.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/apple.jpg +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/haystack-logo.png +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_files/sample_pdf_1.pdf +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_generator.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_ranker.py +0 -0
- {amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_text_embedder.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: amazon-bedrock-haystack
|
|
3
|
-
Version:
|
|
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.
|
|
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
|
|
|
@@ -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.
|
|
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
|
|
61
|
-
content.extend(
|
|
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
|
|
179
|
+
def _format_reasoning_content(reasoning_content: ReasoningContent) -> List[Dict[str, Any]]:
|
|
166
180
|
"""
|
|
167
|
-
Format
|
|
181
|
+
Format ReasoningContent to match Bedrock's expected structure.
|
|
168
182
|
|
|
169
|
-
:param
|
|
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
|
|
174
|
-
formatted_content = {"reasoningContent":
|
|
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.
|
|
196
|
-
bedrock_content_blocks.extend(
|
|
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
|
-
|
|
307
|
-
|
|
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(
|
|
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
|
-
|
|
332
|
-
|
|
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
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
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
|
-
|
|
374
|
-
|
|
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
|
-
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
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
|
-
|
|
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
|
-
|
|
429
|
-
|
|
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
|
-
|
|
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]) ->
|
|
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
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
|
|
529
|
-
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
534
|
-
|
|
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
|
-
|
|
621
|
-
|
|
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
|
-
|
|
645
|
-
|
|
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]
|
{amazon_bedrock_haystack-3.11.0 → amazon_bedrock_haystack-4.0.0}/tests/test_chat_generator.py
RENAMED
|
@@ -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"] == "
|
|
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"] == "
|
|
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"] == "
|
|
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].
|
|
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].
|
|
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"] == "
|
|
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"] == "
|
|
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"] == "
|
|
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"] == "
|
|
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 = [
|