Fix missing maxTokens in AmazonBedrockPromptDriver (#1123)

griptape-ai · Aug 29, 2024 · ab25735 · ab25735
1 parent ba47112
commit ab25735
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Parsing streaming response with some OpenAi compatible services.
 - Issue in `PromptSummaryEngine` if there are no artifacts during recursive summarization.
 - Issue in `GooglePromptDriver` using Tools with no schema.
+- Missing `maxTokens` inference parameter in `AmazonBedrockPromptDriver`.
 - Incorrect model in `OpenAiDriverConfig`'s `text_to_speech_driver`.
 
 **Note**: This release includes breaking changes. Please refer to the [Migration Guide](./MIGRATION.md#030x-to-031x) for details.

diff --git a/griptape/drivers/prompt/amazon_bedrock_prompt_driver.py b/griptape/drivers/prompt/amazon_bedrock_prompt_driver.py
@@ -98,7 +98,7 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             "modelId": self.model,
             "messages": messages,
             "system": system_messages,
-            "inferenceConfig": {"temperature": self.temperature},
+            "inferenceConfig": {"temperature": self.temperature, "maxTokens": self.max_tokens},
             "additionalModelRequestFields": self.additional_model_request_fields,
             **(
                 {"toolConfig": {"tools": self.__to_bedrock_tools(prompt_stack.tools), "toolChoice": self.tool_choice}}

diff --git a/tests/unit/drivers/prompt/test_amazon_bedrock_prompt_driver.py b/tests/unit/drivers/prompt/test_amazon_bedrock_prompt_driver.py
@@ -344,7 +344,7 @@ def test_try_run(self, mock_converse, prompt_stack, messages, use_native_tools):
         mock_converse.assert_called_once_with(
             modelId=driver.model,
             messages=messages,
-            inferenceConfig={"temperature": driver.temperature},
+            inferenceConfig={"temperature": driver.temperature, "maxTokens": driver.max_tokens},
             additionalModelRequestFields={},
             **({"system": [{"text": "system-input"}]} if prompt_stack.system_messages else {"system": []}),
             **(
@@ -376,7 +376,7 @@ def test_try_stream_run(self, mock_converse_stream, prompt_stack, messages, use_
         mock_converse_stream.assert_called_once_with(
             modelId=driver.model,
             messages=messages,
-            inferenceConfig={"temperature": driver.temperature},
+            inferenceConfig={"temperature": driver.temperature, "maxTokens": driver.max_tokens},
             additionalModelRequestFields={},
             **({"system": [{"text": "system-input"}]} if prompt_stack.system_messages else {"system": []}),
             **(