Add support for AWS Bedrock LLM Provider (#238)

- Add support for AWS Bedrock LLM provider. The provider can be accessed by type `aws_bedrock`, and it is compatible with both `langchain` and `llamaindex`
- Added unit tests for existing combinations of existing LLM providers and LLM frameworks. The unit tests are skipped by default, since each of them needs credential to interact with LLM models, but it should be updated and checked each time we add/modify the LLM provider/client to make sure they are still working
- Fixed OpenAI + `llamaindex` LLM client that is not working

Closes [AIQ-1213](https://jirasw.nvidia.com/browse/AIQ-1213)

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AIQToolkit/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Yuchen Zhang (https://github.com/yczhang-nv)
  - David Gardner (https://github.com/dagardner-nv)
  - https://github.com/liamy-nv
  - Matthew Penn (https://github.com/mpenn)
  - Anuradha Karuppiah (https://github.com/AnuradhaKaruppiah)
  - Ayush Thakur (https://github.com/ayulockin)
  - Soumili Nandi (https://github.com/soumilinandi)
  - Eric Evans II (https://github.com/ericevans-nv)
  - https://github.com/hsin-c
  - Zac Wang (https://github.com/zac-wang-nv)
  - Hritik Raj (https://github.com/Hritik003)
  - Victor Yudin (https://github.com/VictorYudin)
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)
  - Michael Demoret (https://github.com/mdemoret-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: #238

Author: yczhang-nv

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

@@ -112,6 +112,36 @@ 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
+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 necessary 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,61 @@
+<!--
+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.
+-->
+
+# AWS Bedrock Integration
+
+The Agent Intelligence Toolkit supports integration with multiple LLM providers, including AWS Bedrock. This documentation provides a comprehensive guide on how to integrate AWS Bedrock models into your AIQ Toolkit workflow. To view the full list of supported LLM providers, run `aiq info components -t llm_provider`.
+
+
+## Configuration
+
+### Prerequisites
+Before integrating AWS Bedrock, ensure you have:
+- 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)
+
+### Example 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:
+
+```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
+```
+
+### Configurable Options
+* `model_name`: The name of the AWS Bedrock model to use (required)
+* `temperature`: Controls randomness in the output (0.0 to 1.0, default: 0.0)
+* `max_tokens`: Maximum number of tokens to generate (must be > 0, default: 1024)
+* `context_size`: Maximum number of tokens for context (must be > 0, default: 1024, required for LlamaIndex)
+* `region_name`: AWS region where your Bedrock service is hosted (default: "None")
+* `base_url`: Custom Bedrock endpoint URL (default: None, needed if you don't want to use the default us-east-1 endpoint)
+* `credentials_profile_name`: AWS credentials profile name from ~/.aws/credentials or ~/.aws/config files (default: None)
+
+## 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
+```

docs/source/index.md

@@ -108,6 +108,7 @@ Adding a Custom Evaluator <./extend/custom-evaluator.md>
 ./extend/adding-a-retriever.md
 ./extend/memory.md
 Adding an LLM Provider <./extend/adding-an-llm-provider.md>
+Integrating AWS Bedrock Models <./extend/integrating-aws-bedrock-models.md>
 ```
 
 ```{toctree}

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,56 @@
+# 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, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")
+    max_tokens: int | None = Field(default=1024,
+                                   gt=0,
+                                   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,
+                                     gt=0,
+                                     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

@@ -18,5 +18,6 @@
 # isort:skip_file
 
 # Import any providers which need to be automatically registered here
+from . import aws_bedrock_llm
 from . import nim_llm
 from . import openai_llm

tests/aiq/llm_providers/test_langchain_agents.py

@@ -0,0 +1,95 @@
+# 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
+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
+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
+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()