Refactor message parameter in Anthropic provider

This commit is contained in:
2024-10-31 16:33:01 -04:00
parent 48ac97f070
commit febf5473d5
3 changed files with 36 additions and 2 deletions
+28
View File
@@ -0,0 +1,28 @@
import pytest
from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama
from pydantic import BaseModel
class ResponseModel(BaseModel):
result: int
@pytest.mark.parametrize(
"provider_cls",
[
Anthropic,
Gemini,
OpenAI,
Groq,
Ollama,
],
)
def test_generate_data(provider_cls):
provider = provider_cls()
prompt = "What is 2+2?"
data = provider.structured_response(prompt=prompt, response_model=ResponseModel)
assert isinstance(data, ResponseModel)
assert type(data.result) == int