amazon-bedrock-haystack 5.1.0__tar.gz → 5.2.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 (50) hide show
  1. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/CHANGELOG.md +6 -0
  2. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/PKG-INFO +2 -2
  3. amazon_bedrock_haystack-5.2.0/pydoc/config_docusaurus.yml +41 -0
  4. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/pyproject.toml +1 -1
  5. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +20 -18
  6. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_chat_generator.py +165 -1
  7. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/.gitignore +0 -0
  8. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/LICENSE.txt +0 -0
  9. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/README.md +0 -0
  10. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/examples/bedrock_ranker_example.py +0 -0
  11. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/examples/chatgenerator_example.py +0 -0
  12. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/examples/embedders_generator_with_rag_example.py +0 -0
  13. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/examples/s3_downloader_example.py +0 -0
  14. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/pydoc/config.yml +0 -0
  15. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/amazon_bedrock/__init__.py +0 -0
  16. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/amazon_bedrock/errors.py +0 -0
  17. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/amazon_bedrock/utils.py +0 -0
  18. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/py.typed +0 -0
  19. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/s3/__init__.py +0 -0
  20. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/s3/errors.py +0 -0
  21. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/common/s3/utils.py +0 -0
  22. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/downloaders/py.typed +0 -0
  23. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/downloaders/s3/__init__.py +0 -0
  24. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/downloaders/s3/s3_downloader.py +0 -0
  25. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/embedders/amazon_bedrock/__init__.py +0 -0
  26. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py +0 -0
  27. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/embedders/amazon_bedrock/document_image_embedder.py +0 -0
  28. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py +0 -0
  29. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  30. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/__init__.py +0 -0
  31. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py +0 -0
  32. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py +0 -0
  33. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/utils.py +0 -0
  34. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +0 -0
  35. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  36. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/rankers/amazon_bedrock/__init__.py +0 -0
  37. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/rankers/amazon_bedrock/ranker.py +0 -0
  38. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  39. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/__init__.py +0 -0
  40. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/conftest.py +0 -0
  41. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_chat_generator_utils.py +0 -0
  42. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_document_embedder.py +0 -0
  43. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_document_image_embedder.py +0 -0
  44. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_files/apple.jpg +0 -0
  45. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_files/haystack-logo.png +0 -0
  46. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_files/sample_pdf_1.pdf +0 -0
  47. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_generator.py +0 -0
  48. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_ranker.py +0 -0
  49. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_s3_downloader.py +0 -0
  50. {amazon_bedrock_haystack-5.1.0 → amazon_bedrock_haystack-5.2.0}/tests/test_text_embedder.py +0 -0
@@ -1,5 +1,11 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/amazon_bedrock-v5.1.0] - 2025-09-29
4
+
5
+ ### 🚀 Features
6
+
7
+ - S3Downloader - add `s3_key_generation_function` param to customize S3 key generation (#2343)
8
+
3
9
  ## [integrations/amazon_bedrock-v5.0.0] - 2025-09-22
4
10
 
5
11
  ### 🧹 Chores
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: amazon-bedrock-haystack
3
- Version: 5.1.0
3
+ Version: 5.2.0
4
4
  Summary: An integration of AWS S3 and Bedrock as a Downloader and Generator components.
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.17.1
24
+ Requires-Dist: haystack-ai>=2.19.0
25
25
  Description-Content-Type: text/markdown
26
26
 
27
27
  # amazon-bedrock-haystack
@@ -0,0 +1,41 @@
1
+ loaders:
2
+ - ignore_when_discovered:
3
+ - __init__
4
+ modules:
5
+ - haystack_integrations.common.amazon_bedrock.errors
6
+ - haystack_integrations.components.embedders.amazon_bedrock.document_embedder
7
+ - haystack_integrations.components.embedders.amazon_bedrock.text_embedder
8
+ - haystack_integrations.components.embedders.amazon_bedrock.document_image_embedder
9
+ - haystack_integrations.components.generators.amazon_bedrock.generator
10
+ - haystack_integrations.components.generators.amazon_bedrock.adapters
11
+ - haystack_integrations.common.amazon_bedrock.errors
12
+ - haystack_integrations.components.generators.amazon_bedrock.chat.chat_generator
13
+ - haystack_integrations.components.rankers.amazon_bedrock.ranker
14
+ - haystack_integrations.components.downloaders.s3.s3_downloader
15
+ - haystack_integrations.common.s3.utils
16
+ - haystack_integrations.common.s3.errors
17
+ search_path:
18
+ - ../src
19
+ type: haystack_pydoc_tools.loaders.CustomPythonLoader
20
+ processors:
21
+ - do_not_filter_modules: false
22
+ documented_only: true
23
+ expression: null
24
+ skip_empty_modules: true
25
+ type: filter
26
+ - expression: name not in ['BedrockRanker']
27
+ type: filter
28
+ - type: smart
29
+ - type: crossref
30
+ renderer:
31
+ description: Amazon Bedrock integration for Haystack
32
+ id: integrations-amazon-bedrock
33
+ markdown:
34
+ add_member_class_prefix: false
35
+ add_method_class_prefix: true
36
+ classdef_code_block: false
37
+ descriptive_class_title: false
38
+ descriptive_module_title: true
39
+ filename: amazon_bedrock.md
40
+ title: Amazon Bedrock
41
+ type: haystack_pydoc_tools.renderers.DocusaurusRenderer
@@ -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.17.1", "boto3>=1.28.57", "aioboto3>=14.0.0"]
26
+ dependencies = ["haystack-ai>=2.19.0", "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"
@@ -1,4 +1,4 @@
1
- from typing import Any, Dict, List, Optional, Tuple, Union
1
+ from typing import Any, Dict, List, Optional, Tuple
2
2
 
3
3
  import aioboto3
4
4
  from botocore.config import Config
@@ -7,10 +7,10 @@ from botocore.exceptions import ClientError
7
7
  from haystack import component, default_from_dict, default_to_dict, logging
8
8
  from haystack.dataclasses import ChatMessage, ComponentInfo, StreamingCallbackT, select_streaming_callback
9
9
  from haystack.tools import (
10
- Tool,
11
- Toolset,
10
+ ToolsType,
12
11
  _check_duplicate_tool_names,
13
12
  deserialize_tools_or_toolset_inplace,
13
+ flatten_tools_or_toolsets,
14
14
  serialize_tools_or_toolset,
15
15
  )
16
16
  from haystack.utils.auth import Secret, deserialize_secrets_inplace
@@ -157,7 +157,7 @@ class AmazonBedrockChatGenerator:
157
157
  generation_kwargs: Optional[Dict[str, Any]] = None,
158
158
  streaming_callback: Optional[StreamingCallbackT] = None,
159
159
  boto3_config: Optional[Dict[str, Any]] = None,
160
- tools: Optional[Union[List[Tool], Toolset]] = None,
160
+ tools: Optional[ToolsType] = None,
161
161
  *,
162
162
  guardrail_config: Optional[Dict[str, str]] = None,
163
163
  ) -> None:
@@ -187,7 +187,8 @@ class AmazonBedrockChatGenerator:
187
187
  [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk) object and switches
188
188
  the streaming mode on.
189
189
  :param boto3_config: The configuration for the boto3 client.
190
- :param tools: A list of Tool objects or a Toolset that the model can use. Each tool should have a unique name.
190
+ :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
191
+ Each tool should have a unique name.
191
192
  :param guardrail_config: Optional configuration for a guardrail that has been created in Amazon Bedrock.
192
193
  This must be provided as a dictionary matching either
193
194
  [GuardrailConfiguration](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GuardrailConfiguration.html).
@@ -218,7 +219,7 @@ class AmazonBedrockChatGenerator:
218
219
  self.streaming_callback = streaming_callback
219
220
  self.boto3_config = boto3_config
220
221
 
221
- _check_duplicate_tool_names(list(tools or [])) # handles Toolset as well
222
+ _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
222
223
  self.tools = tools
223
224
 
224
225
  _validate_guardrail_config(guardrail_config=guardrail_config, streaming=streaming_callback is not None)
@@ -342,7 +343,7 @@ class AmazonBedrockChatGenerator:
342
343
  messages: List[ChatMessage],
343
344
  streaming_callback: Optional[StreamingCallbackT] = None,
344
345
  generation_kwargs: Optional[Dict[str, Any]] = None,
345
- tools: Optional[Union[List[Tool], Toolset]] = None,
346
+ tools: Optional[ToolsType] = None,
346
347
  requires_async: bool = False,
347
348
  ) -> Tuple[Dict[str, Any], Optional[StreamingCallbackT]]:
348
349
  """
@@ -358,7 +359,8 @@ class AmazonBedrockChatGenerator:
358
359
  - `stopSequences`: List of stop sequences to stop generation.
359
360
  - `temperature`: Sampling temperature.
360
361
  - `topP`: Nucleus sampling parameter.
361
- :param tools: Optional list of Tool objects or a Toolset that the model can use.
362
+ :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
363
+ Each tool should have a unique name.
362
364
  :param requires_async: Boolean flag to indicate if an async-compatible streaming callback function is needed.
363
365
 
364
366
  :returns:
@@ -380,14 +382,12 @@ class AmazonBedrockChatGenerator:
380
382
 
381
383
  # Handle tools - either toolConfig or Haystack Tool objects but not both
382
384
  tools = tools or self.tools
383
- _check_duplicate_tool_names(list(tools or []))
385
+ flattened_tools = flatten_tools_or_toolsets(tools)
386
+ _check_duplicate_tool_names(flattened_tools)
384
387
  tool_config = merged_kwargs.pop("toolConfig", None)
385
- if tools:
386
- # Convert Toolset to list if needed
387
- if isinstance(tools, Toolset):
388
- tools = list(tools)
388
+ if flattened_tools:
389
389
  # Format Haystack tools to Bedrock format
390
- tool_config = _format_tools(tools)
390
+ tool_config = _format_tools(flattened_tools)
391
391
 
392
392
  # Any remaining kwargs go to additionalModelRequestFields
393
393
  additional_fields = merged_kwargs if merged_kwargs else None
@@ -425,7 +425,7 @@ class AmazonBedrockChatGenerator:
425
425
  messages: List[ChatMessage],
426
426
  streaming_callback: Optional[StreamingCallbackT] = None,
427
427
  generation_kwargs: Optional[Dict[str, Any]] = None,
428
- tools: Optional[Union[List[Tool], Toolset]] = None,
428
+ tools: Optional[ToolsType] = None,
429
429
  ) -> Dict[str, List[ChatMessage]]:
430
430
  """
431
431
  Executes a synchronous inference call to the Amazon Bedrock model using the Converse API.
@@ -439,7 +439,8 @@ class AmazonBedrockChatGenerator:
439
439
  - `stopSequences`: List of stop sequences to stop generation.
440
440
  - `temperature`: Sampling temperature.
441
441
  - `topP`: Nucleus sampling parameter.
442
- :param tools: Optional list of Tools that the model may call during execution.
442
+ :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
443
+ Each tool should have a unique name.
443
444
 
444
445
  :returns:
445
446
  A dictionary containing the model-generated replies under the `"replies"` key.
@@ -485,7 +486,7 @@ class AmazonBedrockChatGenerator:
485
486
  messages: List[ChatMessage],
486
487
  streaming_callback: Optional[StreamingCallbackT] = None,
487
488
  generation_kwargs: Optional[Dict[str, Any]] = None,
488
- tools: Optional[Union[List[Tool], Toolset]] = None,
489
+ tools: Optional[ToolsType] = None,
489
490
  ) -> Dict[str, List[ChatMessage]]:
490
491
  """
491
492
  Executes an asynchronous inference call to the Amazon Bedrock model using the Converse API.
@@ -499,7 +500,8 @@ class AmazonBedrockChatGenerator:
499
500
  - `stopSequences`: List of stop sequences to stop generation.
500
501
  - `temperature`: Sampling temperature.
501
502
  - `topP`: Nucleus sampling parameter.
502
- :param tools: Optional list of Tool objects or a Toolset that the model can use.
503
+ :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
504
+ Each tool should have a unique name.
503
505
 
504
506
  :returns:
505
507
  A dictionary containing the model-generated replies under the `"replies"` key.
@@ -6,7 +6,7 @@ from haystack import Pipeline
6
6
  from haystack.components.generators.utils import print_streaming_chunk
7
7
  from haystack.components.tools import ToolInvoker
8
8
  from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk, ToolCall
9
- from haystack.tools import Tool
9
+ from haystack.tools import Tool, Toolset
10
10
 
11
11
  from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
12
12
 
@@ -59,6 +59,11 @@ def weather(city: str):
59
59
  return f"The weather in {city} is sunny and 32°C"
60
60
 
61
61
 
62
+ def population(city: str):
63
+ """Get population for a given city."""
64
+ return f"The population of {city} is 2.2 million"
65
+
66
+
62
67
  @pytest.fixture
63
68
  def tools():
64
69
  tool_parameters = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
@@ -71,6 +76,25 @@ def tools():
71
76
  return [tool]
72
77
 
73
78
 
79
+ @pytest.fixture
80
+ def mixed_tools():
81
+ """Fixture that returns a mixed list of Tool and Toolset."""
82
+ weather_tool = Tool(
83
+ name="weather",
84
+ description="useful to determine the weather in a given location",
85
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
86
+ function=weather,
87
+ )
88
+ population_tool = Tool(
89
+ name="population",
90
+ description="useful to determine the population of a given location",
91
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
92
+ function=population,
93
+ )
94
+ toolset = Toolset([population_tool])
95
+ return [weather_tool, toolset]
96
+
97
+
74
98
  @pytest.fixture
75
99
  def chat_messages():
76
100
  messages = [
@@ -305,6 +329,86 @@ class TestAmazonBedrockChatGenerator:
305
329
  assert request_params["messages"] == [{"content": [{"text": "What's the capital of France?"}], "role": "user"}]
306
330
  assert request_params["guardrailConfig"] == {"guardrailIdentifier": "test", "guardrailVersion": "test"}
307
331
 
332
+ def test_init_with_mixed_tools(self, mock_boto3_session, set_env_variables):
333
+ def tool_fn(city: str) -> str:
334
+ return city
335
+
336
+ weather_tool = Tool(
337
+ name="weather",
338
+ description="Weather lookup",
339
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
340
+ function=tool_fn,
341
+ )
342
+ population_tool = Tool(
343
+ name="population",
344
+ description="Population lookup",
345
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
346
+ function=tool_fn,
347
+ )
348
+ toolset = Toolset([population_tool])
349
+
350
+ generator = AmazonBedrockChatGenerator(
351
+ model="anthropic.claude-3-5-sonnet-20240620-v1:0",
352
+ tools=[weather_tool, toolset],
353
+ )
354
+
355
+ assert generator.tools == [weather_tool, toolset]
356
+
357
+ def test_prepare_request_params_with_mixed_tools(self, mock_boto3_session, set_env_variables):
358
+ def tool_fn(city: str) -> str:
359
+ return city
360
+
361
+ weather_tool = Tool(
362
+ name="weather",
363
+ description="Weather lookup",
364
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
365
+ function=tool_fn,
366
+ )
367
+ population_tool = Tool(
368
+ name="population",
369
+ description="Population lookup",
370
+ parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
371
+ function=tool_fn,
372
+ )
373
+ toolset = Toolset([population_tool])
374
+
375
+ generator = AmazonBedrockChatGenerator(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
376
+ request_params, _ = generator._prepare_request_params(
377
+ messages=[ChatMessage.from_user("What's the capital of France?")],
378
+ tools=[weather_tool, toolset],
379
+ )
380
+
381
+ assert request_params["toolConfig"] == {
382
+ "tools": [
383
+ {
384
+ "toolSpec": {
385
+ "name": "weather",
386
+ "description": "Weather lookup",
387
+ "inputSchema": {
388
+ "json": {
389
+ "type": "object",
390
+ "properties": {"city": {"type": "string"}},
391
+ "required": ["city"],
392
+ }
393
+ },
394
+ }
395
+ },
396
+ {
397
+ "toolSpec": {
398
+ "name": "population",
399
+ "description": "Population lookup",
400
+ "inputSchema": {
401
+ "json": {
402
+ "type": "object",
403
+ "properties": {"city": {"type": "string"}},
404
+ "required": ["city"],
405
+ }
406
+ },
407
+ }
408
+ },
409
+ ]
410
+ }
411
+
308
412
 
309
413
  # In the CI, those tests are skipped if AWS Authentication fails
310
414
  @pytest.mark.integration
@@ -423,6 +527,66 @@ class TestAmazonBedrockChatGeneratorInference:
423
527
  assert "paris" in final_message.text.lower()
424
528
  assert "berlin" in final_message.text.lower()
425
529
 
530
+ @pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_TOOLS)
531
+ def test_live_run_with_mixed_tools(self, model_name, mixed_tools):
532
+ """
533
+ Integration test that verifies AmazonBedrockChatGenerator works with mixed Tool and Toolset.
534
+ This tests that the LLM can correctly invoke tools from both a standalone Tool and a Toolset.
535
+ """
536
+ initial_messages = [
537
+ ChatMessage.from_user("What's the weather like in Paris and what is the population of Berlin?")
538
+ ]
539
+ component = AmazonBedrockChatGenerator(model=model_name, tools=mixed_tools)
540
+ results = component.run(messages=initial_messages)
541
+
542
+ assert len(results["replies"]) > 0, "No replies received"
543
+
544
+ # Find the message with tool calls
545
+ tool_call_message = None
546
+ for message in results["replies"]:
547
+ if message.tool_calls:
548
+ tool_call_message = message
549
+ break
550
+
551
+ assert tool_call_message is not None, "No message with tool call found"
552
+ assert isinstance(tool_call_message, ChatMessage), "Tool message is not a ChatMessage instance"
553
+ assert ChatMessage.is_from(tool_call_message, ChatRole.ASSISTANT), "Tool message is not from the assistant"
554
+
555
+ tool_calls = tool_call_message.tool_calls
556
+ assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}"
557
+
558
+ # Verify we got calls to both weather and population tools
559
+ tool_names = {tc.tool_name for tc in tool_calls}
560
+ assert "weather" in tool_names, "Expected 'weather' tool call"
561
+ assert "population" in tool_names, "Expected 'population' tool call"
562
+
563
+ # Verify tool call details
564
+ for tool_call in tool_calls:
565
+ assert tool_call.id, "Tool call does not contain value for 'id' key"
566
+ assert tool_call.tool_name in ["weather", "population"]
567
+ assert "city" in tool_call.arguments
568
+ assert tool_call.arguments["city"] in ["Paris", "Berlin"]
569
+ assert tool_call_message.meta["finish_reason"] == "tool_calls"
570
+
571
+ # Mock the response we'd get from ToolInvoker
572
+ tool_result_messages = []
573
+ for tool_call in tool_calls:
574
+ if tool_call.tool_name == "weather":
575
+ result = "The weather in Paris is sunny and 32°C"
576
+ else: # population
577
+ result = "The population of Berlin is 2.2 million"
578
+ tool_result_messages.append(ChatMessage.from_tool(tool_result=result, origin=tool_call))
579
+
580
+ new_messages = [*initial_messages, tool_call_message, *tool_result_messages]
581
+ results = component.run(new_messages)
582
+
583
+ assert len(results["replies"]) == 1
584
+ final_message = results["replies"][0]
585
+ assert not final_message.tool_call
586
+ assert len(final_message.text) > 0
587
+ assert "paris" in final_message.text.lower()
588
+ assert "berlin" in final_message.text.lower()
589
+
426
590
  @pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_THINKING)
427
591
  def test_live_run_with_tool_call_and_thinking(self, model_name, tools):
428
592
  initial_messages = [ChatMessage.from_user("What's the weather like in Paris?")]