Update test_ollama.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
This commit is contained in:
Kurt Heiden
2024-10-28 21:59:36 -06:00
committed by GitHub
parent b23f732d55
commit 83d430a310
+15 -11
View File
@@ -48,17 +48,21 @@ class TestOllama(unittest.TestCase):
self.assertIsInstance(result, sm.models.Message)
def test_structure_response(self):
class Poem(BaseModel):
title: str
content: str
with self.assertRaises(NotImplementedError):
data_obj = sm.generate_data(
prompt="Write a poem about love",
llm_provider="ollama",
llm_model="llama3.2",
response_model=Poem)
self.assertIsNotNone(data_obj)
self.assertIsInstance(data_obj, Poem)
class Poem(BaseModel):
title: str
content: str
# Test for NotImplementedError
with self.assertRaises(NotImplementedError):
sm.generate_data(
prompt="Write a poem about love",
llm_provider="ollama",
llm_model="llama3.2",
response_model=Poem)
# If implementation is added later, uncomment and modify:
# data_obj = sm.generate_data(...)
# self.assertIsNotNone(data_obj)
# self.assertIsInstance(data_obj, Poem)
if __name__ == '__main__':