Files
simplemind/tests/test_generate_data.py
T
Siddhesh Agarwal 3a7383425f sorted imports
2024-11-01 11:09:54 +05:30

29 lines
579 B
Python

from pydantic import BaseModel
import pytest
from simplemind.providers import Anthropic, Gemini, Groq, Ollama, OpenAI
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 isinstance(data.result, int)