diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py index 12cb9f2..15900dd 100644 --- a/common/embeddings/embedding_services.py +++ b/common/embeddings/embedding_services.py @@ -141,9 +141,9 @@ class AzureOpenAI_Ada002(EmbeddingModel): def __init__(self, config): super().__init__(config, model_name=config.get("model_name", "OpenAI ada-002")) - from langchain.embeddings import AzureOpenAIEmbeddings + from langchain_openai import AzureOpenAIEmbeddings - self.embeddings = AzureOpenAIEmbeddings(deployment=config["azure_deployment"]) + self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, dimensions=self.dimensions, deployment=config["azure_deployment"]) class OpenAI_Embedding(EmbeddingModel): @@ -180,14 +180,22 @@ class AWS_Bedrock_Embedding(EmbeddingModel): """AWS Bedrock Embedding Model""" def __init__(self, config): - import boto3 + import boto3, botocore from langchain_aws import BedrockEmbeddings - super().__init__(config=config, model_name=config["model_name"]) + super().__init__(config=config, model_name=config.get("model_name", "amazon.titan-embed-text-v1")) + + boto3_config = config.get("boto3_config", {}) + client_config = botocore.config.Config( + max_pool_connections=boto3_config.get("max_pool_connections", 20), + read_timeout=boto3_config.get("read_timeout", 300), + retries={"max_attempts": boto3_config.get("retries", 5)}, + ) client = boto3.client( "bedrock-runtime", - region_name="us-east-1", + region_name=config.get("region_name", "us-east-1"), + config=client_config, aws_access_key_id=config["authentication_configuration"][ "AWS_ACCESS_KEY_ID" ], diff --git a/common/llm_services/aws_bedrock_service.py b/common/llm_services/aws_bedrock_service.py index 29348ea..a4eb05f 100644 --- a/common/llm_services/aws_bedrock_service.py +++ b/common/llm_services/aws_bedrock_service.py @@ -13,7 +13,7 @@ # limitations under the License. import os -import boto3 +import boto3, botocore from langchain_aws import ChatBedrock import logging from common.llm_services import LLM_Model @@ -27,9 +27,18 @@ class AWSBedrock(LLM_Model): def __init__(self, config): super().__init__(config) model_name = config["llm_model"] + + boto3_config = config.get("boto3_config", {}) + client_config = botocore.config.Config( + max_pool_connections=boto3_config.get("max_pool_connections", 20), + read_timeout=boto3_config.get("read_timeout", 300), + retries={"max_attempts": boto3_config.get("retries", 5)}, + ) + client = boto3.client( "bedrock-runtime", region_name=config.get("region_name", "us-east-1"), + config=client_config, aws_access_key_id=config["authentication_configuration"][ "AWS_ACCESS_KEY_ID" ], diff --git a/common/llm_services/azure_openai_service.py b/common/llm_services/azure_openai_service.py index 2094125..e4dc6f5 100644 --- a/common/llm_services/azure_openai_service.py +++ b/common/llm_services/azure_openai_service.py @@ -14,13 +14,14 @@ def __init__(self, config): os.environ[auth_detail] = config["authentication_configuration"][ auth_detail ] - from langchain.chat_models.azure_openai import AzureChatOpenAI + from langchain_openai import AzureChatOpenAI model_name = config["llm_model"] self.llm = AzureChatOpenAI( azure_deployment=config["azure_deployment"], openai_api_version=config["openai_api_version"], model_name=config["llm_model"], + max_tokens=config.get("token_limit"), temperature=config["model_kwargs"]["temperature"], ) diff --git a/common/llm_services/google_genai_service.py b/common/llm_services/google_genai_service.py index 883e670..9c9727b 100644 --- a/common/llm_services/google_genai_service.py +++ b/common/llm_services/google_genai_service.py @@ -36,7 +36,7 @@ def __init__(self, config): self.llm = ChatGoogleGenerativeAI( temperature=config["model_kwargs"]["temperature"], model=model_name, - max_tokens=None, + max_tokens=config.get("token_limit"), timeout=None, max_retries=2, )