Skip to content

Commit 6b312eb

Browse files
committed
Minor change to avoid using .value each time
1 parent 406e644 commit 6b312eb

File tree

2 files changed

+5
-3
lines changed

2 files changed

+5
-3
lines changed

aisuitealt/provider.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def chat_completions_create(self, model, messages):
1515
pass
1616

1717

18-
class ProviderNames(Enum):
18+
class ProviderNames(str, Enum):
1919
OPENAI = 'openai'
2020
AWS_BEDROCK = 'aws-bedrock'
2121

aisuitealt/tests/test_client.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,15 +29,17 @@ def test_client_chat_completions(self, mock_bedrock, mock_openai):
2929
client = Client(provider_configs)
3030

3131
# Test OpenAI model
32-
openai_response = client.chat.completions.create(ProviderNames.OPENAI.value + ":" + "gpt-4o", [
32+
open_ai_model = ProviderNames.OPENAI + ":" + "gpt-4o"
33+
openai_response = client.chat.completions.create(open_ai_model, [
3334
{"role": "system", "content": "You are a helpful assistant."},
3435
{"role": "user", "content": "Who won the world series in 2020?"}
3536
])
3637
self.assertEqual(openai_response, "OpenAI Response")
3738
mock_openai.assert_called_once()
3839

3940
# Test AWS Bedrock model
40-
bedrock_response = client.chat.completions.create(ProviderNames.AWS_BEDROCK.value + ":" + "claude-v3", [
41+
bedrock_model = ProviderNames.AWS_BEDROCK + ":" + "claude-v3"
42+
bedrock_response = client.chat.completions.create(bedrock_model, [
4143
{"role": "user", "content": "Hello, world!"}
4244
])
4345
self.assertEqual(bedrock_response, "AWS Bedrock Response")

0 commit comments

Comments
 (0)