-
Notifications
You must be signed in to change notification settings - Fork 184
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Update BedrockPromptDriver to use Converse API
- Loading branch information
1 parent
7a137e1
commit 64d5f2c
Showing
33 changed files
with
181 additions
and
1,244 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,55 +1,67 @@ | ||
| from __future__ import annotations | ||
| import json | ||
| from typing import TYPE_CHECKING, Any | ||
| from collections.abc import Iterator | ||
| from attrs import define, field, Factory | ||
| from griptape.drivers import BasePromptDriver | ||
| from griptape.artifacts import TextArtifact | ||
| from griptape.utils import import_optional_dependency | ||
| from .base_multi_model_prompt_driver import BaseMultiModelPromptDriver | ||
| from griptape.tokenizers import SimpleTokenizer, BaseTokenizer | ||
|
|
||
| if TYPE_CHECKING: | ||
| from griptape.utils import PromptStack | ||
| import boto3 | ||
|
|
||
|
|
||
| @define | ||
| class AmazonBedrockPromptDriver(BaseMultiModelPromptDriver): | ||
| class AmazonBedrockPromptDriver(BasePromptDriver): | ||
| session: boto3.Session = field(default=Factory(lambda: import_optional_dependency("boto3").Session()), kw_only=True) | ||
| bedrock_client: Any = field( | ||
| default=Factory(lambda self: self.session.client("bedrock-runtime"), takes_self=True), kw_only=True | ||
| ) | ||
| additional_model_request_fields: dict = field(default=Factory(dict), kw_only=True) | ||
| tokenizer: BaseTokenizer = field(default=Factory(lambda: SimpleTokenizer(characters_per_token=4)), kw_only=True) | ||
|
|
||
| def try_run(self, prompt_stack: PromptStack) -> TextArtifact: | ||
| model_input = self.prompt_model_driver.prompt_stack_to_model_input(prompt_stack) | ||
| payload = {**self.prompt_model_driver.prompt_stack_to_model_params(prompt_stack)} | ||
| if isinstance(model_input, dict): | ||
| payload.update(model_input) | ||
| response = self.bedrock_client.converse(**self._base_params(prompt_stack)) | ||
|
|
||
| response = self.bedrock_client.invoke_model( | ||
| modelId=self.model, contentType="application/json", accept="application/json", body=json.dumps(payload) | ||
| ) | ||
| output_message = response["output"]["message"] | ||
| output_content = output_message["content"][0]["text"] | ||
|
|
||
| response_body = response["body"].read() | ||
| return TextArtifact(output_content) | ||
|
|
||
| if response_body: | ||
| return self.prompt_model_driver.process_output(response_body) | ||
| def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]: | ||
| response = self.bedrock_client.converse_stream(**self._base_params(prompt_stack)) | ||
|
|
||
| stream = response.get("stream") | ||
| if stream is not None: | ||
| for event in stream: | ||
| if "contentBlockDelta" in event: | ||
| yield TextArtifact(event["contentBlockDelta"]["delta"]["text"]) | ||
| else: | ||
| raise Exception("model response is empty") | ||
|
|
||
| def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]: | ||
| model_input = self.prompt_model_driver.prompt_stack_to_model_input(prompt_stack) | ||
| payload = {**self.prompt_model_driver.prompt_stack_to_model_params(prompt_stack)} | ||
| if isinstance(model_input, dict): | ||
| payload.update(model_input) | ||
|
|
||
| response = self.bedrock_client.invoke_model_with_response_stream( | ||
| modelId=self.model, contentType="application/json", accept="application/json", body=json.dumps(payload) | ||
| ) | ||
|
|
||
| response_body = response["body"] | ||
| if response_body: | ||
| for chunk in response["body"]: | ||
| chunk_bytes = chunk["chunk"]["bytes"] | ||
| yield self.prompt_model_driver.process_output(chunk_bytes) | ||
| def _base_params(self, prompt_stack: PromptStack) -> dict: | ||
| system_messages = [ | ||
| {"text": input.content} for input in prompt_stack.inputs if input.is_system() and input.content | ||
| ] | ||
| messages = [ | ||
| {"role": self.__to_amazon_bedrock_role(input), "content": [{"text": input.content}]} | ||
| for input in prompt_stack.inputs | ||
| if not input.is_system() | ||
| ] | ||
|
|
||
| return { | ||
| "modelId": self.model, | ||
| "messages": messages, | ||
| "system": system_messages, | ||
| "inferenceConfig": {"temperature": self.temperature}, | ||
| "additionalModelRequestFields": self.additional_model_request_fields, | ||
| } | ||
|
|
||
| def __to_amazon_bedrock_role(self, prompt_input: PromptStack.Input) -> str: | ||
| if prompt_input.is_system(): | ||
| return "system" | ||
| elif prompt_input.is_assistant(): | ||
| return "assistant" | ||
| else: | ||
| raise Exception("model response is empty") | ||
| return "user" |
Oops, something went wrong.