chore!: CohereChatGenerator - remove **kwargs init parameter + fix test (#2948)

* fix Cohere test

* simplify

* better client handling

* update haystack pin

* cheap model

Author: anakin87

integrations/cohere/pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.22.0", "cohere>=5.17.0"]
+dependencies = ["haystack-ai>=2.23.0", "cohere>=5.17.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere#readme"

integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py

@@ -22,7 +22,7 @@
     flatten_tools_or_toolsets,
     serialize_tools_or_toolset,
 )
-from haystack.utils import Secret, deserialize_secrets_inplace
+from haystack.utils import Secret
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 from httpx import AsyncClient as AsyncHTTPXClient
 from httpx import AsyncHTTPTransport, HTTPTransport
@@ -508,7 +508,6 @@ def __init__(
         *,
         timeout: float | None = None,
         max_retries: int | None = None,
-        **kwargs: Any,
     ):
         """
         Initialize the CohereChatGenerator instance.
@@ -537,23 +536,18 @@ def __init__(
         :param max_retries:
             Maximum number of retries to attempt for failed requests. If not set, it defaults to the default set by
             the Cohere client.
-        :param kwargs:
-            Additional generation parameters. These are merged into `generation_kwargs` for backward compatibility.
 
         """
         _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
 
         if not api_base_url:
             api_base_url = "https://api.cohere.com"
-        if generation_kwargs is None:
-            generation_kwargs = {}
-        if kwargs:
-            generation_kwargs = {**generation_kwargs, **kwargs}
+
         self.api_key = api_key
         self.model = model
         self.streaming_callback = streaming_callback
         self.api_base_url = api_base_url
-        self.generation_kwargs = generation_kwargs
+        self.generation_kwargs = generation_kwargs or {}
         self.tools = tools
         self.timeout = timeout
         self.max_retries = max_retries
@@ -565,14 +559,15 @@ def __init__(
         }
         if timeout is not None:
             client_kwargs["timeout"] = timeout
+
+        sync_kwargs = {**client_kwargs}
+        async_kwargs = {**client_kwargs}
         if max_retries is not None:
-            sync_httpx_client = HTTPXClient(transport=HTTPTransport(retries=max_retries))
-            async_httpx_client = AsyncHTTPXClient(transport=AsyncHTTPTransport(retries=max_retries))
-            self.client = ClientV2(**client_kwargs, httpx_client=sync_httpx_client)
-            self.async_client = AsyncClientV2(**client_kwargs, httpx_client=async_httpx_client)
-        else:
-            self.client = ClientV2(**client_kwargs)
-            self.async_client = AsyncClientV2(**client_kwargs)
+            sync_kwargs["httpx_client"] = HTTPXClient(transport=HTTPTransport(retries=max_retries))
+            async_kwargs["httpx_client"] = AsyncHTTPXClient(transport=AsyncHTTPTransport(retries=max_retries))
+
+        self.client = ClientV2(**sync_kwargs)
+        self.async_client = AsyncClientV2(**async_kwargs)
 
     def _get_telemetry_data(self) -> dict[str, Any]:
         """
@@ -593,7 +588,7 @@ def to_dict(self) -> dict[str, Any]:
             model=self.model,
             streaming_callback=callback_name,
             api_base_url=self.api_base_url,
-            api_key=self.api_key.to_dict(),
+            api_key=self.api_key,
             generation_kwargs=self.generation_kwargs,
             tools=serialize_tools_or_toolset(self.tools),
             timeout=self.timeout,
@@ -611,7 +606,6 @@ def from_dict(cls, data: dict[str, Any]) -> "CohereChatGenerator":
                Deserialized component.
         """
         init_params = data.get("init_parameters", {})
-        deserialize_secrets_inplace(init_params, ["api_key"])
         deserialize_tools_or_toolset_inplace(init_params, key="tools")
         serialized_callback_handler = init_params.get("streaming_callback")
         if serialized_callback_handler:

integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py

@@ -51,8 +51,21 @@ def __init__(
             You can check them in model's documentation.
         """
 
+        # from_dict deserialization, where `generation_kwargs` is in **kwargs
+        if "generation_kwargs" in kwargs:
+            generation_kwargs = kwargs.pop("generation_kwargs")
+        else:
+            # direct construction like `CohereGenerator(max_tokens=10)`
+            generation_kwargs = kwargs
+
         # Note we have to call super() like this because of the way components are dynamically built with the decorator
-        super(CohereGenerator, self).__init__(api_key, model, streaming_callback, api_base_url, None, **kwargs)  # noqa
+        super(CohereGenerator, self).__init__(  # noqa: UP008
+            api_key=api_key,
+            model=model,
+            streaming_callback=streaming_callback,
+            api_base_url=api_base_url,
+            generation_kwargs=generation_kwargs,
+        )
 
     @component.output_types(replies=list[str], meta=list[dict[str, Any]])
     def run(  # type: ignore[override] # due to incompatible signature with ChatGenerator

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, stream=True)
+        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,
         )
@@ -702,7 +702,7 @@ def test_pipeline_with_cohere_chat_generator(self):
         )
 
         pipeline = Pipeline()
-        pipeline.add_component("generator", CohereChatGenerator(model="command-r-08-2024", tools=[weather_tool]))
+        pipeline.add_component("generator", CohereChatGenerator(model="command-r7b-12-2024", tools=[weather_tool]))
         pipeline.add_component("tool_invoker", ToolInvoker(tools=[weather_tool]))
 
         pipeline.connect("generator", "tool_invoker")
@@ -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,
         )

integrations/cohere/tests/test_generator.py

@@ -91,12 +91,10 @@ def test_from_dict(self, monkeypatch):
             "type": "haystack_integrations.components.generators.cohere.generator.CohereGenerator",
             "init_parameters": {
                 "model": "command-a-03-2025",
-                "max_tokens": 10,
+                "generation_kwargs": {"max_tokens": 10, "some_test_param": "test-params"},
                 "api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
-                "some_test_param": "test-params",
                 "api_base_url": "test-base-url",
                 "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
-                "tools": None,
             },
         }
         component: CohereGenerator = CohereGenerator.from_dict(data)