GML-2008 GML-2009 Azure Library Upgrades

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"
             ],

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"
             ],

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"],
         )
 

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,
         )