Added .configure() for passing configs instead of using ctor.

Author: rohit-rptless

aisuitealt/client.py

@@ -2,12 +2,13 @@
 
 
 class Client:
-    def __init__(self, provider_configs: dict):
+    def __init__(self, provider_configs: dict = {}):
         """
         Initialize the client with provider configurations. 
         Use the ProviderFactory to create provider instances.
         """
         self.providers = {}
+        self.provider_configs = provider_configs
         for provider_key, config in provider_configs.items():
             # Check if the provider key is a valid ProviderNames enum
             if not isinstance(provider_key, ProviderNames):
@@ -17,6 +18,20 @@ def __init__(self, provider_configs: dict):
 
         self._chat = None
 
+    def configure(self, provider_configs: dict = None):
+        """
+        Configure the client with provider configurations.
+        """
+        if provider_configs is None:
+            return
+ 
+        self.provider_configs.update(provider_configs)
+
+        for provider_key, config in self.provider_configs.items():
+            if not isinstance(provider_key, ProviderNames):
+                raise ValueError(f"Provider {provider_key} is not a valid ProviderNames enum")
+            self.providers[provider_key.value] = ProviderFactory.create_provider(provider_key, config)
+
     @property
     def chat(self):
         """Return the chat API interface."""
@@ -39,18 +54,33 @@ class Completions:
     def __init__(self, client: 'Client'):
         self.client = client
 
-    def create(self, model: str, messages: list):
-        """Create chat completion based on the model."""
+    def create(self, model: str, messages: list, **kwargs):
+        """
+        Create chat completion based on the model, messages, and any extra arguments.
+        """
+        # Check that correct format is used
+        if ':' not in model:
+            raise ValueError(f"Invalid model format. Expected 'provider:model', got '{model}'")
+ 
         # Extract the provider key from the model identifier, e.g., "aws-bedrock:model-name"
         provider_key, model_name = model.split(":", 1)
 
-        # Use the correct provider instance created by the factory
+        if provider_key not in ProviderNames._value2member_map_:
+            raise ValueError(f"Provider {provider_key} is not a valid ProviderNames enum")
+
+        if provider_key not in self.client.providers:
+            config = {}
+            if provider_key in self.client.provider_configs:
+                config = self.client.provider_configs[provider_key]
+            self.client.providers[provider_key] = ProviderFactory.create_provider(ProviderNames(provider_key), config)
+ 
         provider = self.client.providers.get(provider_key)
         if not provider:
-            # Add the providers to the ValueError
-            raise ValueError(f"Provider {provider_key} is not present in the client. Here are the providers: {self.client.providers}")
+            raise ValueError(f"Could not load provider for {provider_key}.")
 
         # Delegate the chat completion to the correct provider's implementation
-        return provider.chat_completions_create(model_name, messages)
+        # Any additional arguments will be passed to the provider's implementation.
+        # Eg: max_tokens, temperature, etc.
+        return provider.chat_completions_create(model_name, messages, **kwargs)
 
 

aisuitealt/providers/aws_bedrock_provider.py

@@ -2,21 +2,21 @@
 from provider import Provider, LLMError
 
 class AWSBedrockProvider(Provider):
-    def __init__(self, access_key, secret_key, session_token, region):
-        self.client = AnthropicBedrock(
-            aws_access_key=access_key,
-            aws_secret_key=secret_key,
-            aws_session_token=session_token,
-            aws_region=region
-        )
+    def __init__(self, **config):
+        """
+        Initialize the AWS Bedrock provider with the given configuration.
+        Pass the entire configuration dictionary to the Anthropic Bedrock client constructor.
+        """
+        # Anthropic Bedrock client will use the default AWS credential providers, such as
+        # using ~/.aws/credentials or the "AWS_SECRET_ACCESS_KEY" and "AWS_ACCESS_KEY_ID" environment variables.
+        # Any overrides from the user is passed to the constructor.
+        self.client = AnthropicBedrock(**config)
 
-    def chat_completions_create(self, model, messages):
-        try:
-            response = self.client.messages.create(
-                model=model,
-                max_tokens=256,
-                messages=messages
-            )
-            return response['choices'][0]['message']['content']
-        except Exception as e:
-            raise LLMError(f"AWS Bedrock API error: {str(e)}")
+    def chat_completions_create(self, model, messages, **kwargs):
+        # Any exception raised by Anthropic will be returned to the caller.
+        # Maybe we should catch them and raise a custom LLMError.
+        return self.client.messages.create(
+            model=model,
+            messages=messages,
+            **kwargs # Pass any additional arguments to the Anthropic API. Eg: max_tokens.
+        )

aisuitealt/providers/openai_provider.py

@@ -1,16 +1,30 @@
 import openai
+import os
 from provider import Provider, LLMError
 
 class OpenAIProvider(Provider):
-    def __init__(self, api_key):
-        openai.api_key = api_key
+    def __init__(self, **config):
+        """
+        Initialize the OpenAI provider with the given configuration.
+        Pass the entire configuration dictionary to the OpenAI client constructor.
+        """
+        # Ensure API key is provided either in config or via environment variable
+        config.setdefault('api_key', os.getenv('OPENAI_API_KEY'))
+        if not config['api_key']:
+            raise ValueError("OpenAI API key is missing. Please provide it in the config or set the OPENAI_API_KEY environment variable.")
 
-    def chat_completions_create(self, model, messages):
-        try:
-            response = openai.ChatCompletion.create(
-                model=model,
-                messages=messages
-            )
-            return response['choices'][0]['message']['content']
-        except Exception as e:
-            raise LLMError(f"OpenAI API error: {str(e)}")
+        # NOTE: We could choose to remove above lines for api_key since OpenAI will automatically
+        # infer certain values from the environment variables.
+        # Eg: OPENAI_API_KEY, OPENAI_ORG_ID, OPENAI_PROJECT_ID, OPENAI_BASE_URL, etc.
+
+        # Pass the entire config to the OpenAI client constructor
+        self.client = openai.OpenAI(**config)
+
+    def chat_completions_create(self, model, messages, **kwargs):
+        # Any exception raised by OpenAI will be returned to the caller.
+        # Maybe we should catch them and raise a custom LLMError.
+        return self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            **kwargs # Pass any additional arguments to the OpenAI API
+        )

aisuitealt/tests/test_client.py

@@ -18,15 +18,16 @@ def test_client_chat_completions(self, mock_bedrock, mock_openai):
                 'api_key': 'test_openai_api_key'
             },
             ProviderNames.AWS_BEDROCK: {
-                'access_key': 'test_aws_access_key', 
-                'secret_key': 'test_aws_secret_key',
-                'session_token': 'test_aws_session_token',
-                'region': 'us-west-2'
+                'aws_access_key': 'test_aws_access_key',
+                'aws_secret_key': 'test_aws_secret_key',
+                'aws_session_token': 'test_aws_session_token',
+                'aws_region': 'us-west-2'
             }
         }
 
         # Initialize the client
-        client = Client(provider_configs)
+        client = Client()
+        client.configure(provider_configs)
 
         # Test OpenAI model
         open_ai_model = ProviderNames.OPENAI + ":" + "gpt-4o"