Bedrock titan prompt model driver

BedrockTitanPromptModelDriver

Bases: BasePromptModelDriver

Source code in griptape/drivers/prompt_model/bedrock_titan_prompt_model_driver.py

@define
class BedrockTitanPromptModelDriver(BasePromptModelDriver):
    top_p: float = field(default=0.9, kw_only=True, metadata={"serializable": True})
    _tokenizer: BedrockTitanTokenizer = field(default=None, kw_only=True)
    prompt_driver: Optional[AmazonBedrockPromptDriver] = field(default=None, kw_only=True)

    @property
    def tokenizer(self) -> BedrockTitanTokenizer:
        """Returns the tokenizer for this driver.

        We need to pass the `session` field from the Prompt Driver to the
        Tokenizer. However, the Prompt Driver is not initialized until after
        the Prompt Model Driver is initialized. To resolve this, we make the `tokenizer`
        field a @property that is only initialized when it is first accessed.
        This ensures that by the time we need to initialize the Tokenizer, the
        Prompt Driver has already been initialized.

        See this thread for more information: https://github.com/griptape-ai/griptape/issues/244

        Returns:
            BedrockTitanTokenizer: The tokenizer for this driver.
        """
        if self._tokenizer:
            return self._tokenizer
        else:
            self._tokenizer = BedrockTitanTokenizer(model=self.prompt_driver.model)
            return self._tokenizer

    def prompt_stack_to_model_input(self, prompt_stack: PromptStack) -> dict:
        prompt_lines = []

        for i in prompt_stack.inputs:
            if i.is_user():
                prompt_lines.append(f"User: {i.content}")
            elif i.is_assistant():
                prompt_lines.append(f"Bot: {i.content}")
            elif i.is_system():
                prompt_lines.append(f"Instructions: {i.content}")
            else:
                prompt_lines.append(i.content)
        prompt_lines.append("Bot:")

        prompt = "\n\n".join(prompt_lines)

        return {"inputText": prompt}

    def prompt_stack_to_model_params(self, prompt_stack: PromptStack) -> dict:
        prompt = self.prompt_stack_to_model_input(prompt_stack)["inputText"]

        return {
            "textGenerationConfig": {
                "maxTokenCount": self.prompt_driver.max_output_tokens(prompt),
                "stopSequences": self.tokenizer.stop_sequences,
                "temperature": self.prompt_driver.temperature,
                "topP": self.top_p,
            }
        }

    def process_output(self, output: list[dict] | str | bytes) -> TextArtifact:
        # When streaming, the response body comes back as bytes.
        if isinstance(output, str) or isinstance(output, bytes):
            if isinstance(output, bytes):
                output = output.decode()

            body = json.loads(output)

            if self.prompt_driver.stream:
                return TextArtifact(body["outputText"])
            else:
                return TextArtifact(body["results"][0]["outputText"])
        else:
            raise ValueError("output must be an instance of 'str' or 'bytes'")

prompt_driver: Optional[AmazonBedrockPromptDriver] = field(default=None, kw_only=True) class-attribute instance-attribute

tokenizer: BedrockTitanTokenizer property

Returns the tokenizer for this driver.

We need to pass the session field from the Prompt Driver to the Tokenizer. However, the Prompt Driver is not initialized until after the Prompt Model Driver is initialized. To resolve this, we make the tokenizer field a @property that is only initialized when it is first accessed. This ensures that by the time we need to initialize the Tokenizer, the Prompt Driver has already been initialized.

See this thread for more information: https://github.com/griptape-ai/griptape/issues/244

Returns:

Name	Type	Description
BedrockTitanTokenizer	BedrockTitanTokenizer	The tokenizer for this driver.

top_p: float = field(default=0.9, kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

process_output(output)

Source code in griptape/drivers/prompt_model/bedrock_titan_prompt_model_driver.py

def process_output(self, output: list[dict] | str | bytes) -> TextArtifact:
    # When streaming, the response body comes back as bytes.
    if isinstance(output, str) or isinstance(output, bytes):
        if isinstance(output, bytes):
            output = output.decode()

        body = json.loads(output)

        if self.prompt_driver.stream:
            return TextArtifact(body["outputText"])
        else:
            return TextArtifact(body["results"][0]["outputText"])
    else:
        raise ValueError("output must be an instance of 'str' or 'bytes'")

prompt_stack_to_model_input(prompt_stack)

Source code in griptape/drivers/prompt_model/bedrock_titan_prompt_model_driver.py

def prompt_stack_to_model_input(self, prompt_stack: PromptStack) -> dict:
    prompt_lines = []

    for i in prompt_stack.inputs:
        if i.is_user():
            prompt_lines.append(f"User: {i.content}")
        elif i.is_assistant():
            prompt_lines.append(f"Bot: {i.content}")
        elif i.is_system():
            prompt_lines.append(f"Instructions: {i.content}")
        else:
            prompt_lines.append(i.content)
    prompt_lines.append("Bot:")

    prompt = "\n\n".join(prompt_lines)

    return {"inputText": prompt}

prompt_stack_to_model_params(prompt_stack)

Source code in griptape/drivers/prompt_model/bedrock_titan_prompt_model_driver.py

def prompt_stack_to_model_params(self, prompt_stack: PromptStack) -> dict:
    prompt = self.prompt_stack_to_model_input(prompt_stack)["inputText"]

    return {
        "textGenerationConfig": {
            "maxTokenCount": self.prompt_driver.max_output_tokens(prompt),
            "stopSequences": self.tokenizer.stop_sequences,
            "temperature": self.prompt_driver.temperature,
            "topP": self.top_p,
        }
    }