feat: bedrock converse

integrations/amazon_bedrock/CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## [unreleased]
 
+### 🚀 Features
+
+- New generator for the bedrock converse API (#977)
+
+## [unreleased]
+
 ### 🐛 Bug Fixes
 
 - *(Bedrock)* Allow tools kwargs for AWS Bedrock Claude model (#976)

integrations/amazon_bedrock/examples/converse_generator_example.py

@@ -0,0 +1,46 @@
+from haystack import Pipeline
+
+from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockConverseGenerator
+from haystack_integrations.components.generators.amazon_bedrock.converse.utils import ConverseMessage, ToolConfig
+
+
+def get_current_weather(location: str, unit: str = "celsius") -> str:
+    """Get the current weather in a given location"""
+    return f"The weather in {location} is 22 degrees {unit}."
+
+
+def get_current_time(timezone: str) -> str:
+    """Get the current time in a given timezone"""
+    return f"The current time in {timezone} is 14:30."
+
+
+generator = AmazonBedrockConverseGenerator(
+    model="anthropic.claude-3-5-sonnet-20240620-v1:0",
+    streaming_callback=print,
+)
+
+# Create ToolConfig from functions
+tool_config = ToolConfig.from_functions([get_current_weather, get_current_time])
+
+tool_config_dict = tool_config.to_dict()
+
+
+pipeline = Pipeline()
+pipeline.add_component("generator", generator)
+
+result = pipeline.run(
+    data={
+        "generator": {
+            "inference_config": {
+                "temperature": 0.1,
+                "maxTokens": 256,
+                "topP": 0.1,
+                "stopSequences": ["\\n"],
+            },
+            "messages": [
+                ConverseMessage.from_user(["What's the weather like in Paris and what time is it in New York?"]),
+            ],
+            "tool_config": tool_config_dict,
+        },
+    },
+)

integrations/amazon_bedrock/pydoc/config.yml

@@ -7,6 +7,7 @@ loaders:
       "haystack_integrations.common.amazon_bedrock.errors",
       "haystack_integrations.components.generators.amazon_bedrock.handlers",
       "haystack_integrations.components.generators.amazon_bedrock.chat.chat_generator",
+      "haystack_integrations.components.generators.amazon_bedrock.converse.converse_generator",
       "haystack_integrations.components.embedders.amazon_bedrock.text_embedder",
       "haystack_integrations.components.embedders.amazon_bedrock.document_embedder",
     ]

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/__init__.py

@@ -2,6 +2,16 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from .chat.chat_generator import AmazonBedrockChatGenerator
+from .converse.capabilities import MODEL_CAPABILITIES
+from .converse.converse_generator import AmazonBedrockConverseGenerator
+from .converse.utils import ConverseMessage, ToolConfig
 from .generator import AmazonBedrockGenerator
 
-__all__ = ["AmazonBedrockGenerator", "AmazonBedrockChatGenerator"]
+__all__ = [
+    "AmazonBedrockGenerator",
+    "AmazonBedrockChatGenerator",
+    "AmazonBedrockConverseGenerator",
+    "ConverseMessage",
+    "ToolConfig",
+    "MODEL_CAPABILITIES",
+]

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/converse/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <[email protected]>
+#
+# SPDX-License-Identifier: Apache-2.0

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/converse/capabilities.py

@@ -0,0 +1,109 @@
+from enum import Enum, auto
+
+
+class ModelCapability(Enum):
+    CONVERSE = auto()
+    CONVERSE_STREAM = auto()
+    SYSTEM_PROMPTS = auto()
+    DOCUMENT_CHAT = auto()
+    VISION = auto()
+    TOOL_USE = auto()
+    STREAMING_TOOL_USE = auto()
+    GUARDRAILS = auto()
+
+
+# https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html
+
+MODEL_CAPABILITIES = {
+    "ai21.jamba-instruct-.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+    },
+    "ai21.j2-.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.GUARDRAILS,
+    },
+    "amazon.titan-text-.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "anthropic.claude-v2.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "anthropic.claude-3.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.VISION,
+        ModelCapability.TOOL_USE,
+        ModelCapability.STREAMING_TOOL_USE,
+        ModelCapability.GUARDRAILS,
+    },
+    "cohere.command-text.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "cohere.command-light.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.GUARDRAILS,
+    },
+    "cohere.command-r.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.TOOL_USE,
+    },
+    "meta.llama2.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "meta.llama3.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "meta.llama3-1.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.TOOL_USE,
+        ModelCapability.GUARDRAILS,
+    },
+    "mistral.mistral-.*-instruct": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.GUARDRAILS,
+    },
+    "mistral.mistral-large.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.DOCUMENT_CHAT,
+        ModelCapability.TOOL_USE,
+        ModelCapability.GUARDRAILS,
+    },
+    "mistral.mistral-small.*": {
+        ModelCapability.CONVERSE,
+        ModelCapability.CONVERSE_STREAM,
+        ModelCapability.SYSTEM_PROMPTS,
+        ModelCapability.TOOL_USE,
+        ModelCapability.GUARDRAILS,
+    },
+}