cheap model

Author: anakin87

integrations/cohere/tests/test_chat_generator.py

@@ -499,7 +499,7 @@ def test_run_image(self):
 class TestCohereChatGeneratorInference:
     def test_live_run(self):
         chat_messages = [ChatMessage.from_user("What's the capital of France")]
-        component = CohereChatGenerator(generation_kwargs={"temperature": 0.8})
+        component = CohereChatGenerator(model="command-r7b-12-2024", generation_kwargs={"temperature": 0.8})
         results = component.run(chat_messages)
         assert len(results["replies"]) == 1
         message: ChatMessage = results["replies"][0]
@@ -525,7 +525,7 @@ def __call__(self, chunk: StreamingChunk) -> None:
                 self.responses += chunk.content if chunk.content else ""
 
         callback = Callback()
-        component = CohereChatGenerator(streaming_callback=callback)
+        component = CohereChatGenerator(model="command-r7b-12-2024", streaming_callback=callback)
         results = component.run([ChatMessage.from_user("What's the capital of France? answer in a word")])
 
         assert len(results["replies"]) == 1
@@ -559,7 +559,7 @@ def test_tools_use_old_way(self):
                 },
             }
         ]
-        client = CohereChatGenerator()
+        client = CohereChatGenerator(model="command-r7b-12-2024")
         response = client.run(
             messages=[ChatMessage.from_user("What is the current price of AAPL?")],
             generation_kwargs={"tools": tools_schema},
@@ -595,7 +595,7 @@ def test_tools_use_with_tools(self):
             function=stock_price,
         )
         initial_messages = [ChatMessage.from_user("What is the current price of AAPL?")]
-        client = CohereChatGenerator()
+        client = CohereChatGenerator(model="command-r7b-12-2024")
         response = client.run(
             messages=initial_messages,
             tools=[stock_price_tool],
@@ -650,7 +650,7 @@ def test_live_run_with_tools_streaming(self):
 
         initial_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
         component = CohereChatGenerator(
-            # Cohere's model that supports tools
+            model="command-r7b-12-2024",
             tools=[weather_tool],
             streaming_callback=print_streaming_chunk,
         )
@@ -787,7 +787,7 @@ def test_live_run_with_mixed_tools(self):
         initial_messages = [
             ChatMessage.from_user("What's the weather like in Paris and what is the population of Berlin?")
         ]
-        component = CohereChatGenerator(model="command-r-08-2024", tools=mixed_tools)
+        component = CohereChatGenerator(model="command-r7b-12-2024", tools=mixed_tools)
         results = component.run(messages=initial_messages)
 
         assert len(results["replies"]) > 0, "No replies received"

integrations/cohere/tests/test_chat_generator_async.py

@@ -25,7 +25,7 @@ def stock_price(ticker: str):
 class TestCohereChatGeneratorAsyncInference:
     async def test_live_run_async(self):
         chat_messages = [ChatMessage.from_user("What's the capital of France")]
-        component = CohereChatGenerator(generation_kwargs={"temperature": 0.8})
+        component = CohereChatGenerator(model="command-r7b-12-2024", generation_kwargs={"temperature": 0.8})
         results = await component.run_async(chat_messages)
         assert len(results["replies"]) == 1
         message: ChatMessage = results["replies"][0]
@@ -44,7 +44,7 @@ async def callback(chunk: StreamingChunk):
             counter += 1
             responses += chunk.content if chunk.content else ""
 
-        component = CohereChatGenerator(streaming_callback=callback)
+        component = CohereChatGenerator(model="command-r7b-12-2024", streaming_callback=callback)
         results = await component.run_async([ChatMessage.from_user("What's the capital of France? answer in a word")])
 
         assert len(results["replies"]) == 1
@@ -74,7 +74,7 @@ async def test_tools_use_with_tools_async(self):
             function=stock_price,
         )
         initial_messages = [ChatMessage.from_user("What is the current price of AAPL?")]
-        client = CohereChatGenerator()
+        client = CohereChatGenerator(model="command-r7b-12-2024")
         response = await client.run_async(
             messages=initial_messages,
             tools=[stock_price_tool],
@@ -137,7 +137,7 @@ async def print_streaming_chunk_async(chunk: StreamingChunk) -> None:
 
         initial_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
         component = CohereChatGenerator(
-            # Cohere's model that supports tools
+            model="command-r7b-12-2024",
             tools=[weather_tool],
             streaming_callback=print_streaming_chunk_async,
         )