chore: enable ANN ruff ruleset for cohere integration

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: julian-risch

integrations/cohere/CHANGELOG.md

@@ -1,12 +1,5 @@
 # Changelog
 
-## [integrations/cohere-v8.0.1] - 2026-03-17
-
-### 🐛 Bug Fixes
-
-- Fix the conversion of cohere chunks to Haystack streaming chunks (#2968)
-
-
 ## [integrations/cohere-v8.0.0] - 2026-03-11
 
 ### 🚀 Features

integrations/cohere/pyproject.toml

@@ -84,6 +84,7 @@ line-length = 120
 [tool.ruff.lint]
 select = [
   "A",
+  "ANN",
   "ARG",
   "B",
   "C",
@@ -111,6 +112,7 @@ select = [
 ignore = [
   # Allow non-abstract empty methods in abstract base classes
   "B027",
+  "ANN401",  # Allow Any - used legitimately for dynamic types and SDK boundaries
   # Ignore checks for possible passwords
   "S105",
   "S106",
@@ -134,7 +136,7 @@ ban-relative-imports = "parents"
 
 [tool.ruff.lint.per-file-ignores]
 # Tests can use magic values, assertions, and relative imports
-"tests/**/*" = ["PLR2004", "S101", "TID252"]
+"tests/**/*" = ["PLR2004", "S101", "TID252", "ANN"]
 
 [tool.coverage.run]
 source = ["haystack_integrations"]

integrations/cohere/src/haystack_integrations/components/embedders/cohere/document_embedder.py

@@ -49,7 +49,7 @@ def __init__(
         meta_fields_to_embed: list[str] | None = None,
         embedding_separator: str = "\n",
         embedding_type: EmbeddingTypes | None = None,
-    ):
+    ) -> None:
         """
         :param api_key: the Cohere API key.
         :param model: the name of the model to use. Supported Models are:

integrations/cohere/src/haystack_integrations/components/embedders/cohere/embedding_types.py

@@ -21,7 +21,7 @@ class EmbeddingTypes(Enum):
     BINARY = "binary"
     UBINARY = "ubinary"
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.value
 
     @staticmethod

integrations/cohere/src/haystack_integrations/components/embedders/cohere/text_embedder.py

@@ -41,7 +41,7 @@ def __init__(
         truncate: str = "END",
         timeout: float = 120.0,
         embedding_type: EmbeddingTypes | None = None,
-    ):
+    ) -> None:
         """
         :param api_key: the Cohere API key.
         :param model: the name of the model to use. Supported Models are:

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

@@ -38,13 +38,6 @@
 ImageFormat = Literal["image/png", "image/jpeg", "image/webp", "image/gif"]
 IMAGE_SUPPORTED_FORMATS: list[ImageFormat] = list(get_args(ImageFormat))
 
-FINISH_REASON_MAPPING: dict[str, FinishReason] = {
-    "COMPLETE": "stop",
-    "STOP_SEQUENCE": "stop",
-    "MAX_TOKENS": "length",
-    "TOOL_CALL": "tool_calls",
-}
-
 
 def _format_tool(tool: Tool) -> dict[str, Any]:
     """
@@ -58,11 +51,17 @@ def _format_tool(tool: Tool) -> dict[str, Any]:
     """
     return {
         "type": "function",
-        "function": {"name": tool.name, "description": tool.description, "parameters": tool.parameters},
+        "function": {
+            "name": tool.name,
+            "description": tool.description,
+            "parameters": tool.parameters,
+        },
     }
 
 
-def _format_message(message: ChatMessage) -> dict[str, Any]:
+def _format_message(
+    message: ChatMessage,
+) -> dict[str, Any]:
     """
     Formats a Haystack ChatMessage into Cohere's chat format.
 
@@ -103,10 +102,17 @@ def _format_message(message: ChatMessage) -> dict[str, Any]:
                 {
                     "id": tool_call.id,
                     "type": "function",
-                    "function": {"name": tool_call.tool_name, "arguments": json.dumps(tool_call.arguments)},
+                    "function": {
+                        "name": tool_call.tool_name,
+                        "arguments": json.dumps(tool_call.arguments),
+                    },
                 }
             )
-        return {"role": "assistant", "tool_calls": tool_calls, "tool_plan": message.text if message.text else ""}
+        return {
+            "role": "assistant",
+            "tool_calls": tool_calls,
+            "tool_plan": message.text if message.text else "",
+        }
 
     if message.role.value == "user":
         if not message.images and not message.text:
@@ -169,43 +175,42 @@ def _parse_response(chat_response: ChatResponse, model: str) -> ChatMessage:
     :param model: The name of the model that generated the response.
     :return: A Haystack ChatMessage containing the formatted response.
     """
-    # Extract text content from the response
-    text_content = ""
-    if chat_response.message.content:
-        for content_item in chat_response.message.content:
-            if content_item.type == "text":
-                text_content = content_item.text
-
-    # Extract tool calls if present in the response
-    tool_calls = None
     if chat_response.message.tool_calls:
         tool_calls = []
         for tc in chat_response.message.tool_calls:
             if tc.function and tc.function.name and tc.function.arguments and isinstance(tc.function.arguments, str):
                 tool_calls.append(
-                    ToolCall(id=tc.id, tool_name=tc.function.name, arguments=json.loads(tc.function.arguments))
+                    ToolCall(
+                        id=tc.id,
+                        tool_name=tc.function.name,
+                        arguments=json.loads(tc.function.arguments),
+                    )
                 )
-        # If a tool plan is provided we use that as our text content over the default text content
-        text_content = chat_response.message.tool_plan or text_content
-
-    # Create metadata for the message
-    resolved_finish_reason = None
-    if chat_response.finish_reason:
-        resolved_finish_reason = FINISH_REASON_MAPPING.get(chat_response.finish_reason, chat_response.finish_reason)
-    base_meta = {
-        "model": model,
-        "index": 0,
-        "finish_reason": resolved_finish_reason,
-        "citations": chat_response.message.citations,
-    }
+
+        # Create message with tool plan as text and tool calls in the format Haystack expects
+        tool_plan = chat_response.message.tool_plan or ""
+        message = ChatMessage.from_assistant(text=tool_plan, tool_calls=tool_calls)
+    elif chat_response.message.content and hasattr(chat_response.message.content[0], "text"):
+        message = ChatMessage.from_assistant(chat_response.message.content[0].text)
+    else:
+        # Handle the case where neither tool_calls nor content exists
+        logger.warning(f"Received empty response from Cohere API: {chat_response.message}")
+        message = ChatMessage.from_assistant("")
+
     # In V2, token usage is part of the response object, not the message
+    message._meta.update(
+        {
+            "model": model,
+            "index": 0,
+            "finish_reason": chat_response.finish_reason,
+            "citations": chat_response.message.citations,
+        }
+    )
     if chat_response.usage and chat_response.usage.billed_units:
-        base_meta["usage"] = {
+        message._meta["usage"] = {
             "prompt_tokens": chat_response.usage.billed_units.input_tokens,
             "completion_tokens": chat_response.usage.billed_units.output_tokens,
         }
-
-    message = ChatMessage.from_assistant(text=text_content, tool_calls=tool_calls, meta=base_meta)
     return message
 
 
@@ -214,7 +219,6 @@ def _convert_cohere_chunk_to_streaming_chunk(
     model: str,
     component_info: ComponentInfo | None = None,
     global_index: int = 0,
-    previous_original_chunks: list[StreamedChatResponseV2] | None = None,
 ) -> StreamingChunk:
     """
     Converts a Cohere streaming response chunk to a StreamingChunk.
@@ -233,6 +237,12 @@ def _convert_cohere_chunk_to_streaming_chunk(
     :returns:
         A StreamingChunk object representing the content of the chunk from the Cohere API.
     """
+    finish_reason_mapping: dict[str, FinishReason] = {
+        "COMPLETE": "stop",
+        "MAX_TOKENS": "length",
+        "TOOL_CALLS": "tool_calls",
+    }
+
     # Initialize default values
     content = ""
     index = global_index
@@ -244,23 +254,24 @@ def _convert_cohere_chunk_to_streaming_chunk(
     if chunk.type == "content-delta" and chunk.delta and chunk.delta.message:
         if chunk.delta.message and chunk.delta.message.content and chunk.delta.message.content.text is not None:
             content = chunk.delta.message.content.text
-        # If the previous chunk is a content-start chunk, we set start to True for the first content-delta chunk
-        if previous_original_chunks and previous_original_chunks[-1].type == "content-start":
-            start = True
 
     elif chunk.type == "tool-plan-delta" and chunk.delta and chunk.delta.message:
         if chunk.delta.message and chunk.delta.message.tool_plan is not None:
             content = chunk.delta.message.tool_plan
-        # If the previous chunk is a message-start chunk, we set start to True for the first tool-plan-delta chunk
-        if previous_original_chunks and previous_original_chunks[-1].type == "message-start":
-            start = True
 
     elif chunk.type == "tool-call-start" and chunk.delta and chunk.delta.message:
         if chunk.delta.message and chunk.delta.message.tool_calls:
             tool_call = chunk.delta.message.tool_calls
             function = tool_call.function
             if function is not None and function.name is not None:
-                tool_calls = [ToolCallDelta(index=global_index, id=tool_call.id, tool_name=function.name)]
+                tool_calls = [
+                    ToolCallDelta(
+                        index=global_index,
+                        id=tool_call.id,
+                        tool_name=function.name,
+                        arguments=None,
+                    )
+                ]
                 start = True  # This starts a tool call
                 if tool_call.id is not None:
                     meta["tool_call_id"] = tool_call.id
@@ -273,11 +284,21 @@ def _convert_cohere_chunk_to_streaming_chunk(
             and chunk.delta.message.tool_calls.function.arguments is not None
         ):
             arguments = chunk.delta.message.tool_calls.function.arguments
-            tool_calls = [ToolCallDelta(index=global_index, arguments=arguments)]
+            tool_calls = [
+                ToolCallDelta(
+                    index=global_index,
+                    tool_name=None,
+                    arguments=arguments,
+                )
+            ]
+
+    elif chunk.type == "tool-call-end":
+        # Tool call end doesn't have content, just signals completion
+        start = True
 
     elif chunk.type == "message-end":
         finish_reason_raw = getattr(chunk.delta, "finish_reason", None)
-        finish_reason = FINISH_REASON_MAPPING.get(finish_reason_raw) if finish_reason_raw else None
+        finish_reason = finish_reason_mapping.get(finish_reason_raw) if finish_reason_raw else None
 
         # The Cohere API is subject to changes in how usage data is returned. We try to support both dict and objects.
         usage_data = getattr(chunk.delta, "usage", None)
@@ -325,7 +346,6 @@ def _parse_streaming_response(
 
     Loops through each stream object from Cohere and converts it into a StreamingChunk.
     """
-    original_chunks: list[StreamedChatResponseV2] = []
     chunks: list[StreamingChunk] = []
     global_index = 0
 
@@ -338,10 +358,11 @@ def _parse_streaming_response(
             component_info=component_info,
             model=model,
             global_index=global_index,
-            previous_original_chunks=original_chunks,
         )
 
-        original_chunks.append(chunk)
+        if not streaming_chunk:
+            continue
+
         chunks.append(streaming_chunk)
         streaming_callback(streaming_chunk)
 
@@ -357,7 +378,6 @@ async def _parse_async_streaming_response(
     """
     Parses Cohere's async streaming chat response into a Haystack ChatMessage.
     """
-    original_chunks: list[StreamedChatResponseV2] = []
     chunks: list[StreamingChunk] = []
     global_index = 0
 
@@ -366,14 +386,11 @@ async def _parse_async_streaming_response(
             global_index += 1
 
         streaming_chunk = _convert_cohere_chunk_to_streaming_chunk(
-            chunk=chunk,
-            component_info=component_info,
-            model=model,
-            global_index=global_index,
-            previous_original_chunks=original_chunks,
+            chunk=chunk, component_info=component_info, model=model, global_index=global_index
         )
+        if not streaming_chunk:
+            continue
 
-        original_chunks.append(chunk)
         chunks.append(streaming_chunk)
         await streaming_callback(streaming_chunk)
 
@@ -491,7 +508,7 @@ def __init__(
         *,
         timeout: float | None = None,
         max_retries: int | None = None,
-    ):
+    ) -> None:
         """
         Initialize the CohereChatGenerator instance.
 
@@ -621,7 +638,10 @@ def run(
         """
 
         # update generation kwargs by merging with the generation kwargs passed to the run method
-        generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
+        generation_kwargs = {
+            **self.generation_kwargs,
+            **(generation_kwargs or {}),
+        }
 
         # Handle tools
         tools = tools or self.tools
@@ -685,7 +705,10 @@ async def run_async(
         """
 
         # update generation kwargs by merging with the generation kwargs passed to the run method
-        generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
+        generation_kwargs = {
+            **self.generation_kwargs,
+            **(generation_kwargs or {}),
+        }
 
         # Handle tools
         tools = tools or self.tools

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

@@ -37,7 +37,7 @@ def __init__(
         streaming_callback: Callable | None = None,
         api_base_url: str | None = None,
         **kwargs: Any,
-    ):
+    ) -> None:
         """
         Instantiates a `CohereGenerator` component.
 

integrations/cohere/src/haystack_integrations/components/rankers/cohere/ranker.py

@@ -41,7 +41,7 @@ def __init__(
         meta_fields_to_embed: list[str] | None = None,
         meta_data_separator: str = "\n",
         max_tokens_per_doc: int = 4096,
-    ):
+    ) -> None:
         """
         Creates an instance of the 'CohereRanker'.
 

integrations/cohere/tests/test_chat_generator.py

@@ -462,7 +462,6 @@ def test_run_image(self):
         mock_response = MagicMock()
         mock_response.message.content = [MagicMock()]
         mock_response.message.content[0].text = "This is a test image response"
-        mock_response.message.content[0].type = "text"
         mock_response.message.tool_calls = None
         mock_response.finish_reason = "COMPLETE"
         mock_response.usage = None