Add support for AWS Bedrock LLM Provider

docs/source/extend/adding-an-llm-provider.md

@@ -112,6 +112,37 @@ Similar to the registration function for the provider, the client registration f
 In the above example, the `ChatOpenAI` class is imported lazily, allowing for the client to be registered without importing the client class until it is needed. Thus, improving performance and startup times.
 :::
 
+## Test the Combination of LLM Provider and Client
+
+After implementing a new LLM provider, it's important to verify that it works correctly with all existing LLM clients. This can be done by writing integration tests. Here's an example of how to test the integration between the NIM LLM provider and the LangChain framework:
+
+```python
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_nim_langchain_agent():
+    """
+    Test NIM LLM with LangChain agent. Requires NVIDIA_API_KEY to be set.
+    """
+
+    prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant."), ("human", "{input}")])
+
+    llm_config = NIMModelConfig(model_name="meta/llama-3.1-70b-instruct", temperature=0.0)
+
+    async with WorkflowBuilder() as builder:
+        await builder.add_llm("nim_llm", llm_config)
+        llm = await builder.get_llm("nim_llm", wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+        agent = prompt | llm
+
+        response = await agent.ainvoke({"input": "What is 1+2?"})
+        assert isinstance(response, AIMessage)
+        assert response.content is not None
+        assert isinstance(response.content, str)
+        assert "3" in response.content.lower()
+```
+
+Note: Since this test requires an API key, it's marked with `@pytest.mark.integration` to exclude it from CI runs. However, these tests are valuable for maintaining and verifying the functionality of LLM providers and their client integrations.
+
 ## Packaging the Provider and Client
 
 The provider and client will need to be bundled into a Python package, which in turn will be registered with AIQ toolkit as a [plugin](../extend/plugins.md). In the `pyproject.toml` file of the package the `project.entry-points.'aiq.components'` section, defines a Python module as the entry point of the plugin. Details on how this is defined are found in the [Entry Point](../extend/plugins.md#entry-point) section of the plugins document. By convention, the entry point module is named `register.py`, but this is not a requirement.

docs/source/extend/integrating-aws-bedrock-models.md

@@ -0,0 +1,49 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# Integrating AWS Bedrock Models to AIQ Toolkit Workflow
+
+To integrate AWS Bedrock models into your AIQ Toolkit workflow, follow these steps:
+
+1. **Prerequisites**:
+   - Set up AWS credentials by configuring `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`. For detailed setup instructions, refer to the [AWS Bedrock setup guide](https://docs.aws.amazon.com/bedrock/latest/userguide/setting-up.html)
+
+1. **Configuration**:
+   Add the AWS Bedrock LLM configuration to your workflow config file. Make sure the `region_name` matches the region of your `AWS` account, and the `credentials_profile_name` matches the field in your credential file. Here's an example:
+
+```yaml
+llms:
+  aws_bedrock_llm:
+    _type: aws_bedrock
+    model_name: meta.llama3-3-70b-instruct-v1:0
+    temperature: 0.0
+    max_tokens: 1024
+    region_name: us-east-2
+    credentials_profile_name: default
+```
+
+3. **Usage in Workflow**:
+   Reference the AWS Bedrock LLM in your workflow configuration:
+
+```yaml
+workflow:
+  _type: react_agent
+  llm_name: aws_bedrock_llm
+  # ... other workflow configurations
+```
+
+The AWS Bedrock integration supports various models and configurations, allowing you to leverage AWS's managed LLM services within your AIQ Toolkit workflows.

packages/aiqtoolkit_langchain/pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
   # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to aiq packages.
   # Keep sorted!!!
   "aiqtoolkit~=1.2",
+  "langchain-aws~=0.2.1",
   "langchain-core~=0.3.7",
   "langchain-nvidia-ai-endpoints~=0.3.5",
   "langchain-milvus~=0.1.5",

packages/aiqtoolkit_langchain/src/aiq/plugins/langchain/llm.py

@@ -16,6 +16,7 @@
 from aiq.builder.builder import Builder
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.cli.register_workflow import register_llm_client
+from aiq.llm.aws_bedrock_llm import AWSBedrockModelConfig
 from aiq.llm.nim_llm import NIMModelConfig
 from aiq.llm.openai_llm import OpenAIModelConfig
 
@@ -34,3 +35,11 @@ async def openai_langchain(llm_config: OpenAIModelConfig, builder: Builder):
     from langchain_openai import ChatOpenAI
 
     yield ChatOpenAI(**llm_config.model_dump(exclude={"type"}, by_alias=True))
+
+
+@register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def aws_bedrock_langchain(llm_config: AWSBedrockModelConfig, builder: Builder):
+
+    from langchain_aws import ChatBedrockConverse
+
+    yield ChatBedrockConverse(**llm_config.model_dump(exclude={"type", "context_size"}, by_alias=True))

packages/aiqtoolkit_llama_index/pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
   # error
   "llama-index-core==0.12.21",
   "llama-index-embeddings-nvidia==0.3.1",
+  "llama-index-llms-bedrock==0.3.8",
   "llama-index-llms-nvidia==0.3.1",
   "llama-index-readers-file==0.4.4",
   "llama-index==0.12.21",

packages/aiqtoolkit_llama_index/src/aiq/plugins/llama_index/llm.py

@@ -16,6 +16,7 @@
 from aiq.builder.builder import Builder
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.cli.register_workflow import register_llm_client
+from aiq.llm.aws_bedrock_llm import AWSBedrockModelConfig
 from aiq.llm.nim_llm import NIMModelConfig
 from aiq.llm.openai_llm import OpenAIModelConfig
 
@@ -47,7 +48,16 @@ async def openai_llama_index(llm_config: OpenAIModelConfig, builder: Builder):
 
     llm = OpenAI(**kwargs)
 
-    # Disable content blocks
-    llm.supports_content_blocks = False
+    yield llm
+
+
+@register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def aws_bedrock_llama_index(llm_config: AWSBedrockModelConfig, builder: Builder):
+
+    from llama_index.llms.bedrock import Bedrock
+
+    kwargs = llm_config.model_dump(exclude={"type", "max_tokens"}, by_alias=True)
+
+    llm = Bedrock(**kwargs)
 
     yield llm

src/aiq/llm/aws_bedrock_llm.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.llm import LLMProviderInfo
+from aiq.cli.register_workflow import register_llm_provider
+from aiq.data_models.llm import LLMBaseConfig
+
+
+class AWSBedrockModelConfig(LLMBaseConfig, name="aws_bedrock"):
+    """An AWS Bedrock llm provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    # Completion parameters
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The model name for the hosted AWS Bedrock.")
+    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
+    max_tokens: int | None = Field(default=1024,
+                                   description="Maximum number of tokens to generate."
+                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
+    context_size: int | None = Field(default=1024,
+                                     description="Maximum number of tokens to generate."
+                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+
+    # Client parameters
+    region_name: str | None = Field(default="None", description="AWS region to use.")
+    base_url: str | None = Field(
+        default=None, description="Bedrock endpoint to use. Needed if you don't want to default to us-east-1 endpoint.")
+    credentials_profile_name: str | None = Field(
+        default=None, description="The name of the profile in the ~/.aws/credentials or ~/.aws/config files.")
+
+
+@register_llm_provider(config_type=AWSBedrockModelConfig)
+async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, builder: Builder):
+
+    yield LLMProviderInfo(config=llm_config, description="A AWS Bedrock model for use with an LLM client.")

src/aiq/llm/register.py

@@ -20,3 +20,4 @@
 # Import any providers which need to be automatically registered here
 from . import nim_llm
 from . import openai_llm
+from . import aws_bedrock_llm

tests/aiq/llm_providers/test_langchain_agents.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+from langchain_core.messages import AIMessage
+from langchain_core.prompts import ChatPromptTemplate
+
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.builder.workflow_builder import WorkflowBuilder
+from aiq.llm.aws_bedrock_llm import AWSBedrockModelConfig
+from aiq.llm.nim_llm import NIMModelConfig
+from aiq.llm.openai_llm import OpenAIModelConfig
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_nim_langchain_agent():
+    """
+    Test NIM LLM with LangChain agent. Requires NVIDIA_API_KEY to be set.
+    """
+
+    prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant."), ("human", "{input}")])
+
+    llm_config = NIMModelConfig(model_name="meta/llama-3.1-70b-instruct", temperature=0.0)
+
+    async with WorkflowBuilder() as builder:
+        await builder.add_llm("nim_llm", llm_config)
+        llm = await builder.get_llm("nim_llm", wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+        agent = prompt | llm
+
+        response = await agent.ainvoke({"input": "What is 1+2?"})
+        assert isinstance(response, AIMessage)
+        assert response.content is not None
+        assert isinstance(response.content, str)
+        assert "3" in response.content.lower()
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_openai_langchain_agent():
+    """
+    Test OpenAI LLM with LangChain agent. Requires OPENAI_API_KEY to be set.
+    """
+    prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant."), ("human", "{input}")])
+
+    llm_config = OpenAIModelConfig(model_name="gpt-3.5-turbo", temperature=0.0)
+
+    async with WorkflowBuilder() as builder:
+        await builder.add_llm("openai_llm", llm_config)
+        llm = await builder.get_llm("openai_llm", wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+        agent = prompt | llm
+
+        response = await agent.ainvoke({"input": "What is 1+2?"})
+        assert isinstance(response, AIMessage)
+        assert response.content is not None
+        assert isinstance(response.content, str)
+        assert "3" in response.content.lower()
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_aws_bedrock_langchain_agent():
+    """
+    Test AWS Bedrock LLM with LangChain agent.
+    Requires AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to be set.
+    See https://docs.aws.amazon.com/bedrock/latest/userguide/setting-up.html for more information.
+    """
+    prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant."), ("human", "{input}")])
+
+    llm_config = AWSBedrockModelConfig(model_name="meta.llama3-3-70b-instruct-v1:0",
+                                       temperature=0.0,
+                                       region_name="us-east-2",
+                                       max_tokens=1024)
+
+    async with WorkflowBuilder() as builder:
+        await builder.add_llm("aws_bedrock_llm", llm_config)
+        llm = await builder.get_llm("aws_bedrock_llm", wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+        agent = prompt | llm
+
+        response = await agent.ainvoke({"input": "What is 1+2?"})
+        assert isinstance(response, AIMessage)
+        assert response.content is not None
+        assert isinstance(response.content, str)
+        assert "3" in response.content.lower()