From e31d8510a9a6cc85d5c1d2d80c4972d2234c433b Mon Sep 17 00:00:00 2001 From: Dmytro Nikolaiev Date: Thu, 12 Dec 2024 11:04:00 -0800 Subject: [PATCH] Update model used in test --- tests/integration/llm/test_groq_llm.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/integration/llm/test_groq_llm.py b/tests/integration/llm/test_groq_llm.py index a61d01ee..028f4cd8 100644 --- a/tests/integration/llm/test_groq_llm.py +++ b/tests/integration/llm/test_groq_llm.py @@ -10,14 +10,14 @@ class TestGroqLLM(unittest.TestCase): @staticmethod - def get_gemma_7b(): + def get_gemma(): dotenv.load_dotenv() - with OsEnviron("GROQ_LLM_MODEL", "gemma-7b-it"): + with OsEnviron("GROQ_LLM_MODEL", "gemma2-9b-it"): return GroqLLM.from_env() def test_message(self): messages = [LLMMessage.user_message("what is the capital of France?")] - llm = self.get_gemma_7b() + llm = self.get_gemma() result = llm.post_chat_request(LLMContext.empty(), messages) assert "Paris" in result.choices[0] @@ -29,15 +29,15 @@ def test_message(self): def test_consumptions(self): messages = [LLMMessage.user_message("Hello how are you?")] - llm = self.get_gemma_7b() + llm = self.get_gemma() result = llm.post_chat_request(LLMContext.empty(), messages) assert len(result.consumptions) == 12 # call, duration, 3 token kinds, 3 cost kinds and 4 groq duration for consumption in result.consumptions: - assert consumption.kind.startswith("gemma-7b-it") + assert consumption.kind.startswith("gemma2-9b-it") def test_max_tokens_param(self): - llm = self.get_gemma_7b() + llm = self.get_gemma() llm.configuration.temperature.set(0.8) llm.configuration.max_tokens.set(7) @@ -47,7 +47,7 @@ def test_max_tokens_param(self): def test_json_mode(self): messages = [LLMMessage.user_message("Output a JSON object with the data about RPG character.")] - llm = self.get_gemma_7b() + llm = self.get_gemma() result = llm.post_chat_request(LLMContext.empty(), messages, response_format={"type": "json_object"}) data = json.loads(result.first_choice)