Fix hatch run test:types

Author: sjrl

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

@@ -6,7 +6,7 @@
 import json
 import os
 from datetime import datetime, timezone
-from typing import Any
+from typing import Any, Dict, List, Optional, Union
 
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
@@ -177,7 +177,9 @@ def from_dict(cls, data: dict[str, Any]) -> WatsonxChatGenerator:
         return default_from_dict(cls, data)
 
     @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | None = None, stream: bool = False):
+    def run(
+        self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | None = None, stream: bool = False
+    ) -> Dict[str, List[ChatMessage]]:
         """
         Generate chat completions synchronously.
 
@@ -198,7 +200,7 @@ def run(self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | N
     @component.output_types(replies=list[ChatMessage])
     async def run_async(
         self, messages: list[ChatMessage], generation_kwargs: dict[str, Any] | None = None, stream: bool = False
-    ):
+    ) -> Dict[str, List[ChatMessage]]:
         """
         Generate chat completions asynchronously.
 
@@ -222,6 +224,7 @@ def _prepare_api_call(
         merged_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
 
         watsonx_messages = []
+        content: Union[Optional[str], Dict[str, Any]]
         for msg in messages:
             if msg.is_from("user"):
                 content = msg.text
@@ -347,7 +350,10 @@ def _process_response(self, response: dict[str, Any]) -> dict[str, Any]:
                             ToolCall(id=tc.get("id"), tool_name=tc["function"]["name"], arguments=arguments)
                         )
                     except json.JSONDecodeError:
-                        logger.warning("Failed to parse tool call arguments: %s", tc["function"]["arguments"])
+                        logger.warning(
+                            "Failed to parse tool call arguments: {tool_call_args}",
+                            tool_call_args=tc["function"]["arguments"]
+                        )
 
         return {
             "replies": [

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

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Dict, List, Union
 
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import StreamingChunk
@@ -181,7 +181,7 @@ def run(
         generation_kwargs: dict[str, Any] | None = None,
         guardrails: bool = False,
         stream: bool = False,
-    ):
+    ) -> Dict[str, Union[List[str],  List[Dict[str, Any]], List[StreamingChunk]]]:
         """
         Generate text using the watsonx.ai model.