Test cases for Providers.

Test case checks for translation into correct
set of messages and additional parameters
being passed to the respective SDK calls.

Author: rohit-rptless

aisuitealt/provider.py

@@ -18,6 +18,7 @@ def chat_completions_create(self, model, messages):
 class ProviderNames(str, Enum):
     OPENAI = 'openai'
     AWS_BEDROCK = 'aws-bedrock'
+    ANTHROPIC = 'anthropic'
 
 
 class ProviderFactory:
@@ -26,6 +27,7 @@ class ProviderFactory:
     _provider_modules = {
         ProviderNames.OPENAI: 'providers.openai_provider',
         ProviderNames.AWS_BEDROCK: 'providers.aws_bedrock_provider',
+        ProviderNames.ANTHROPIC: 'providers.anthropic_provider',
     }
 
     @classmethod
@@ -54,4 +56,5 @@ def _get_provider_class_name(provider_key):
         return {
             ProviderNames.OPENAI: 'OpenAIProvider',
             ProviderNames.AWS_BEDROCK: 'AWSBedrockProvider',
+            ProviderNames.ANTHROPIC: 'AnthropicProvider',
         }[provider_key]

aisuitealt/providers/aws_bedrock_provider.py

@@ -9,8 +9,8 @@ def __init__(self, **config):
         """
         # 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 = boto3.client("bedrock-runtime", **config)
+        # It does not like parameters passed to the constructor.
+        self.client = boto3.client("bedrock-runtime")
         # Maintain a list of Inference Parameters which Bedrock supports.
         # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html
         self.inference_parameters = ['maxTokens', 'temperature', 'topP', 'stopSequences']

aisuitealt/tests/test_anthropic_provider.py

@@ -0,0 +1,45 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from client import Client
+from provider import ProviderNames
+
+class TestAnthropicProvider(unittest.TestCase):
+
+  @patch('anthropic.Anthropic')
+  def test_anthropic_chat_completions(self, mock_anthropic):
+      """Test that correct parameters are passed to Anthropic's chat completions API."""
+      # Mocking the Anthropic client
+      mock_anthropic_client = MagicMock()
+      mock_anthropic.return_value = mock_anthropic_client
+
+      # Mock response from the Anthropic API
+      mock_anthropic_client.messages.create.return_value = {
+          "role": "assistant",
+          "content": "Test response from Anthropic"
+      }
+
+      client = Client()
+      client.configure({
+          ProviderNames.ANTHROPIC: {
+              'api_key': 'test_anthropic_api_key'
+          }
+      })
+
+      messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Describe cloud computing."}
+        ]
+        
+      response = client.chat.completions.create(ProviderNames.ANTHROPIC + ":" + "claude-v3", messages)
+
+      # Assert that the converse API was called with the correct parameters
+      mock_anthropic_client.messages.create.assert_called_once_with(
+          model="claude-v3",
+          system="You are a helpful assistant.",
+          messages=[
+              {"role": "user", "content": "Describe cloud computing."}
+          ]
+      )
+
+      # Check that the response is normalized correctly
+      self.assertEqual(response['choices'][0]['message']['content'], "Test response from Anthropic")

aisuitealt/tests/test_aws_bedrock_provider.py

@@ -0,0 +1,55 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from client import Client
+from provider import ProviderNames
+
+class TestAWSBedrockProvider(unittest.TestCase):
+
+    @patch('boto3.client')
+    def test_aws_bedrock_converse_parameters(self, mock_boto_client):
+        """Test that correct parameters are passed to AWS Bedrock's converse API."""
+        # Mocking the Bedrock client
+        mock_bedrock_client = MagicMock()
+        mock_boto_client.return_value = mock_bedrock_client
+
+        # Mock response from the Bedrock API
+        mock_bedrock_client.converse.return_value = {
+            "output": {
+                "message": {
+                    "content": "Test response from AWS Bedrock"
+                }
+            }
+        }
+
+        client = Client()
+        # No need to call client.configure() as boto3 uses environment variables.
+
+        # Call the client for AWS Bedrock
+        bedrock_model = ProviderNames.AWS_BEDROCK + ":" + "claude-v3"
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Describe cloud computing."}
+        ]
+        response = client.chat.completions.create(bedrock_model, messages, temperature=0.7, maxTokens=100, top_k=200)
+
+        # Assert that the converse API was called with the correct parameters
+        mock_bedrock_client.converse.assert_called_once_with(
+            model="claude-v3",
+            messages=[
+                {"role": "user", "content": "Describe cloud computing."}
+            ],
+            system=["You are a helpful assistant."],
+            inferenceConfig={
+                'temperature': 0.7,
+                'maxTokens': 100
+            },
+            additionalModelRequestFields={
+                'top_k': 200
+            }
+        )
+
+        # Check that the response is normalized correctly
+        self.assertEqual(response['choices'][0]['message']['content'], "Test response from AWS Bedrock")
+
+if __name__ == '__main__':
+    unittest.main()