Minor change to avoid using .value each time

andrewyng · Sep 5, 2024 · 6b312eb · 6b312eb
1 parent 406e644
commit 6b312eb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/aisuitealt/provider.py b/aisuitealt/provider.py
@@ -15,7 +15,7 @@ def chat_completions_create(self, model, messages):
         pass
 
 
-class ProviderNames(Enum):
+class ProviderNames(str, Enum):
     OPENAI = 'openai'
     AWS_BEDROCK = 'aws-bedrock'
 

diff --git a/aisuitealt/tests/test_client.py b/aisuitealt/tests/test_client.py
@@ -29,15 +29,17 @@ def test_client_chat_completions(self, mock_bedrock, mock_openai):
         client = Client(provider_configs)
 
         # Test OpenAI model
-        openai_response = client.chat.completions.create(ProviderNames.OPENAI.value + ":" + "gpt-4o", [
+        open_ai_model = ProviderNames.OPENAI + ":" + "gpt-4o"
+        openai_response = client.chat.completions.create(open_ai_model, [
             {"role": "system", "content": "You are a helpful assistant."},
             {"role": "user", "content": "Who won the world series in 2020?"}
         ])
         self.assertEqual(openai_response, "OpenAI Response")
         mock_openai.assert_called_once()
 
         # Test AWS Bedrock model
-        bedrock_response = client.chat.completions.create(ProviderNames.AWS_BEDROCK.value + ":" + "claude-v3", [
+        bedrock_model = ProviderNames.AWS_BEDROCK + ":" + "claude-v3"
+        bedrock_response = client.chat.completions.create(bedrock_model, [
             {"role": "user", "content": "Hello, world!"}
         ])
         self.assertEqual(bedrock_response, "AWS Bedrock Response")