From dae36ab0e127a198746e901a12bc2215534c702d Mon Sep 17 00:00:00 2001 From: Abhi Date: Sat, 9 Aug 2025 19:54:21 -0400 Subject: [PATCH] [Feature] Adding AWS as a provider to the models Current models file contains only OpenAI and Azure authentication. Adding AWS auth allows users to select from wide range of models available on Bedrock --- environment.yml | 1 + .../notebooks/03_tableau_datasource_qa.ipynb | 1 + experimental/utilities/models.py | 33 +++++++++++++++++++ pkg/langchain_tableau/utilities/models.py | 33 +++++++++++++++++++ 4 files changed, 68 insertions(+) diff --git a/environment.yml b/environment.yml index f9deb50..b460bc2 100644 --- a/environment.yml +++ b/environment.yml @@ -24,4 +24,5 @@ dependencies: - "pinecone-client[grpc]==6.0.0" - tavily-python==0.5.3 - build==1.2.2 + - langchain-aws==0.2.27 diff --git a/experimental/notebooks/03_tableau_datasource_qa.ipynb b/experimental/notebooks/03_tableau_datasource_qa.ipynb index 381b03b..2a746f2 100644 --- a/experimental/notebooks/03_tableau_datasource_qa.ipynb +++ b/experimental/notebooks/03_tableau_datasource_qa.ipynb @@ -191,6 +191,7 @@ " tableau_user=tableau_user,\n", " datasource_luid=datasource_luid,\n", " tooling_llm_model=tooling_model\n", + " model_provider=openai\n", ")\n", "\n", "# add the tool to a List to give to the agent\n", diff --git a/experimental/utilities/models.py b/experimental/utilities/models.py index f08e5a2..35e0235 100644 --- a/experimental/utilities/models.py +++ b/experimental/utilities/models.py @@ -3,6 +3,7 @@ from langchain_openai import ChatOpenAI, AzureChatOpenAI, OpenAIEmbeddings, AzureOpenAIEmbeddings from langchain.chat_models.base import BaseChatModel from langchain.embeddings.base import Embeddings +from langchain_aws import ChatBedrockConverse, BedrockEmbeddings def select_model(provider: str = "openai", model_name: str = "gpt-4o-mini", temperature: float = 0.2) -> BaseChatModel: @@ -15,6 +16,23 @@ def select_model(provider: str = "openai", model_name: str = "gpt-4o-mini", temp model_name=model_name, temperature=temperature ) + elif provider == "aws": + auth_type = 'profile' if os.environ.get("aws_cred_profile",None) is None else 'creds' + if auth_type == 'profile': + return ChatBedrockConverse( + model=model_name, + temperature=temperature, + credentials_profile_name = os.environ.get("aws_cred_profile") + ) + else: + return ChatBedrockConverse( + model=model_name, + temperature=temperature, + aws_access_key_id=os.environ.get("aws_access_key_id"), + aws_secret_access_key=os.environ.get("aws_secret_access_key"), + aws_session_token=os.environ.get("aws_session_token") + ) + else: # default to OpenAI return ChatOpenAI( model_name=model_name, @@ -32,6 +50,21 @@ def select_embeddings(provider: str = "openai", model_name: str = "text-embeddin openai_api_key=os.environ.get("AZURE_OPENAI_API_KEY"), model=model_name ) + elif provider == "aws": + auth_type = 'profile' if os.environ.get("aws_cred_profile",None) is None else 'creds' + if auth_type == 'profile': + + return BedrockEmbeddings( + model=model_name, + credentials_profile_name = os.environ.get("aws_cred_profile") + ) + else: + return BedrockEmbeddings( + model=model_name, + aws_access_key_id=os.environ.get("aws_access_key_id"), + aws_secret_access_key=os.environ.get("aws_secret_access_key"), + aws_session_token=os.environ.get("aws_session_token") + ) else: # default to OpenAI return OpenAIEmbeddings( model=model_name, diff --git a/pkg/langchain_tableau/utilities/models.py b/pkg/langchain_tableau/utilities/models.py index f08e5a2..35e0235 100644 --- a/pkg/langchain_tableau/utilities/models.py +++ b/pkg/langchain_tableau/utilities/models.py @@ -3,6 +3,7 @@ from langchain_openai import ChatOpenAI, AzureChatOpenAI, OpenAIEmbeddings, AzureOpenAIEmbeddings from langchain.chat_models.base import BaseChatModel from langchain.embeddings.base import Embeddings +from langchain_aws import ChatBedrockConverse, BedrockEmbeddings def select_model(provider: str = "openai", model_name: str = "gpt-4o-mini", temperature: float = 0.2) -> BaseChatModel: @@ -15,6 +16,23 @@ def select_model(provider: str = "openai", model_name: str = "gpt-4o-mini", temp model_name=model_name, temperature=temperature ) + elif provider == "aws": + auth_type = 'profile' if os.environ.get("aws_cred_profile",None) is None else 'creds' + if auth_type == 'profile': + return ChatBedrockConverse( + model=model_name, + temperature=temperature, + credentials_profile_name = os.environ.get("aws_cred_profile") + ) + else: + return ChatBedrockConverse( + model=model_name, + temperature=temperature, + aws_access_key_id=os.environ.get("aws_access_key_id"), + aws_secret_access_key=os.environ.get("aws_secret_access_key"), + aws_session_token=os.environ.get("aws_session_token") + ) + else: # default to OpenAI return ChatOpenAI( model_name=model_name, @@ -32,6 +50,21 @@ def select_embeddings(provider: str = "openai", model_name: str = "text-embeddin openai_api_key=os.environ.get("AZURE_OPENAI_API_KEY"), model=model_name ) + elif provider == "aws": + auth_type = 'profile' if os.environ.get("aws_cred_profile",None) is None else 'creds' + if auth_type == 'profile': + + return BedrockEmbeddings( + model=model_name, + credentials_profile_name = os.environ.get("aws_cred_profile") + ) + else: + return BedrockEmbeddings( + model=model_name, + aws_access_key_id=os.environ.get("aws_access_key_id"), + aws_secret_access_key=os.environ.get("aws_secret_access_key"), + aws_session_token=os.environ.get("aws_session_token") + ) else: # default to OpenAI return OpenAIEmbeddings( model=model_name,