fix: fix Anthropic types + add py.typed (#1940)

* draft

* fix: Anthropic - fix types + add py.typed

* CI

* simplify

Author: anakin87

.github/workflows/anthropic.yml

@@ -50,11 +50,9 @@ jobs:
       - name: Install Hatch
         run: pip install --upgrade hatch
 
-    # TODO: Once this integration is properly typed, use hatch run test:types
-    # https://github.com/deepset-ai/haystack-core-integrations/issues/1771
       - name: Lint
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Run tests
         run: hatch run test:cov-retry

integrations/anthropic/pyproject.toml

@@ -66,18 +66,14 @@ integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
 
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = "mypy -p haystack_integrations.components.generators.anthropic {args}"
 
-# TODO: remove lint environment once this integration is properly typed
-# test environment should be used instead
-# https://github.com/deepset-ai/haystack-core-integrations/issues/1771
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
 
 [tool.black]
 target-version = ["py38"]
@@ -159,15 +155,6 @@ omit = ["*/tests/*", "*/__init__.py"]
 show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
-[[tool.mypy.overrides]]
-module = [
-  "anthropic.*",
-  "haystack.*",
-  "haystack_integrations.*",
-  "pytest.*",
-  "numpy.*",
-]
-ignore_missing_imports = true
 
 [tool.pytest.ini_options]
 addopts = "--strict-markers"

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

@@ -1,15 +1,13 @@
 import json
-from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Union
 
 from haystack import component, default_from_dict, default_to_dict, logging
-from haystack.dataclasses import (
+from haystack.dataclasses.chat_message import ChatMessage, ChatRole, ToolCall, ToolCallResult
+from haystack.dataclasses.streaming_chunk import (
     AsyncStreamingCallbackT,
-    ChatMessage,
-    ChatRole,
     StreamingCallbackT,
     StreamingChunk,
-    ToolCall,
-    ToolCallResult,
+    SyncStreamingCallbackT,
     select_streaming_callback,
 )
 from haystack.tools import (
@@ -19,98 +17,102 @@
     deserialize_tools_or_toolset_inplace,
     serialize_tools_or_toolset,
 )
-from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
+from haystack.utils.auth import Secret, deserialize_secrets_inplace
+from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
 from anthropic import Anthropic, AsyncAnthropic
+from anthropic.resources.messages.messages import Message, RawMessageStreamEvent, Stream
+from anthropic.types import MessageParam, TextBlockParam, ToolParam, ToolResultBlockParam, ToolUseBlockParam
 
 logger = logging.getLogger(__name__)
 
 
 def _update_anthropic_message_with_tool_call_results(
-    tool_call_results: List[ToolCallResult], anthropic_msg: Dict[str, Any]
+    tool_call_results: List[ToolCallResult],
+    content: List[Union[TextBlockParam, ToolUseBlockParam, ToolResultBlockParam]],
 ) -> None:
     """
-    Update an Anthropic message with tool call results.
+    Update an Anthropic message content list with tool call results.
 
     :param tool_call_results: The list of ToolCallResults to update the message with.
-    :param anthropic_msg: The Anthropic message to update.
+    :param content: The Anthropic message content list to update.
     """
-    if "content" not in anthropic_msg:
-        anthropic_msg["content"] = []
-
     for tool_call_result in tool_call_results:
         if tool_call_result.origin.id is None:
             msg = "`ToolCall` must have a non-null `id` attribute to be used with Anthropic."
             raise ValueError(msg)
-        anthropic_msg["content"].append(
-            {
-                "type": "tool_result",
-                "tool_use_id": tool_call_result.origin.id,
-                "content": [{"type": "text", "text": tool_call_result.result}],
-                "is_error": tool_call_result.error,
-            }
+
+        tool_result_block = ToolResultBlockParam(
+            type="tool_result",
+            tool_use_id=tool_call_result.origin.id,
+            content=[{"type": "text", "text": tool_call_result.result}],
+            is_error=tool_call_result.error,
         )
+        content.append(tool_result_block)
 
 
-def _convert_tool_calls_to_anthropic_format(tool_calls: List[ToolCall]) -> List[Dict[str, Any]]:
+def _convert_tool_calls_to_anthropic_format(tool_calls: List[ToolCall]) -> List[ToolUseBlockParam]:
     """
     Convert a list of tool calls to the format expected by Anthropic Chat API.
 
     :param tool_calls: The list of ToolCalls to convert.
-    :return: A list of dictionaries in the format expected by Anthropic API.
+    :return: A list of ToolUseBlockParam objects in the format expected by Anthropic API.
     """
     anthropic_tool_calls = []
     for tc in tool_calls:
         if tc.id is None:
             msg = "`ToolCall` must have a non-null `id` attribute to be used with Anthropic."
             raise ValueError(msg)
-        anthropic_tool_calls.append(
-            {
-                "type": "tool_use",
-                "id": tc.id,
-                "name": tc.tool_name,
-                "input": tc.arguments,
-            }
+
+        tool_use_block = ToolUseBlockParam(
+            type="tool_use",
+            id=tc.id,
+            name=tc.tool_name,
+            input=tc.arguments,
         )
+        anthropic_tool_calls.append(tool_use_block)
     return anthropic_tool_calls
 
 
 def _convert_messages_to_anthropic_format(
     messages: List[ChatMessage],
-) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+) -> Tuple[List[TextBlockParam], List[MessageParam]]:
     """
     Convert a list of messages to the format expected by Anthropic Chat API.
 
     :param messages: The list of ChatMessages to convert.
     :return: A tuple of two lists:
-        - A list of system message dictionaries in the format expected by Anthropic API.
-        - A list of non-system message dictionaries in the format expected by Anthropic API.
+        - A list of system message TextBlockParam objects in the format expected by Anthropic API.
+        - A list of non-system MessageParam objects in the format expected by Anthropic API.
     """
 
-    anthropic_system_messages = []
-    anthropic_non_system_messages = []
+    anthropic_system_messages: List[TextBlockParam] = []
+    anthropic_non_system_messages: List[MessageParam] = []
 
     i = 0
     while i < len(messages):
         message = messages[i]
 
-        # allow passing cache_control
-        cache_control = {"cache_control": message.meta.get("cache_control")} if "cache_control" in message.meta else {}
-
         # system messages have special format requirements for Anthropic API
         # they can have only type and text fields, and they need to be passed separately
         # to the Anthropic API endpoint
-        if message.is_from(ChatRole.SYSTEM):
-            anthropic_system_messages.append({"type": "text", "text": message.text, **cache_control})
+        if message.is_from(ChatRole.SYSTEM) and message.text:
+            sys_message = TextBlockParam(type="text", text=message.text)
+            if cache_control := message.meta.get("cache_control"):
+                sys_message["cache_control"] = cache_control
+            anthropic_system_messages.append(sys_message)
             i += 1
             continue
 
-        anthropic_msg: Dict[str, Any] = {"role": message._role.value, "content": [], **cache_control}
+        content: List[Union[TextBlockParam, ToolUseBlockParam, ToolResultBlockParam]] = []
 
         if message.texts and message.texts[0]:
-            anthropic_msg["content"].append({"type": "text", "text": message.texts[0]})
+            text_block = TextBlockParam(type="text", text=message.texts[0])
+            content.append(text_block)
+
         if message.tool_calls:
-            anthropic_msg["content"] += _convert_tool_calls_to_anthropic_format(message.tool_calls)
+            tool_use_blocks = _convert_tool_calls_to_anthropic_format(message.tool_calls)
+            content.extend(tool_use_blocks)
 
         if message.tool_call_results:
             results = message.tool_call_results.copy()
@@ -119,14 +121,20 @@ def _convert_messages_to_anthropic_format(
                 i += 1
                 results.extend(messages[i].tool_call_results)
 
-            _update_anthropic_message_with_tool_call_results(results, anthropic_msg)
-            anthropic_msg["role"] = "user"
+            _update_anthropic_message_with_tool_call_results(results, content)
 
-        if not anthropic_msg["content"]:
+        if not content:
             msg = "A `ChatMessage` must contain at least one `TextContent`, `ToolCall`, or `ToolCallResult`."
             raise ValueError(msg)
 
-        anthropic_non_system_messages.append(anthropic_msg)
+        # Anthropic only supports assistant and user roles in messages. User role is also used for tool messages.
+        # System messages are passed separately.
+        role: Union[Literal["assistant"], Literal["user"]] = "user"
+        if message._role == ChatRole.ASSISTANT:
+            role = "assistant"
+
+        anthropic_message = MessageParam(role=role, content=content)
+        anthropic_non_system_messages.append(anthropic_message)
         i += 1
 
     return anthropic_system_messages, anthropic_non_system_messages
@@ -340,11 +348,14 @@ def _convert_streaming_chunks_to_chat_message(
         for chunk in chunks:
             chunk_type = chunk.meta.get("type")
             if chunk_type == "content_block_start":
-                if chunk.meta.get("content_block", {}).get("type") == "tool_use":
-                    delta_block = chunk.meta.get("content_block")
+                content_block = chunk.meta.get("content_block")
+                if content_block is None:
+                    msg = "Invalid streaming chunk. Expected 'content_block' field."
+                    raise ValueError(msg)
+                if content_block.get("type") == "tool_use":
                     current_tool_call = {
-                        "id": delta_block.get("id"),
-                        "name": delta_block.get("name"),
+                        "id": content_block.get("id"),
+                        "name": content_block.get("name"),
                         "arguments": "",
                     }
             elif chunk_type == "content_block_delta":
@@ -388,21 +399,12 @@ def _convert_streaming_chunks_to_chat_message(
 
         return message
 
-    @staticmethod
-    def _remove_cache_control(message: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Removes the cache_control key from the message.
-        :param message: The message to remove the cache_control key from.
-        :returns: The message with the cache_control key removed.
-        """
-        return {k: v for k, v in message.items() if k != "cache_control"}
-
     def _prepare_request_params(
         self,
         messages: List[ChatMessage],
         generation_kwargs: Optional[Dict[str, Any]] = None,
         tools: Optional[Union[List[Tool], Toolset]] = None,
-    ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]:
+    ) -> Tuple[List[TextBlockParam], List[MessageParam], Dict[str, Any], List[ToolParam]]:
         """
         Prepare the parameters for the Anthropic API request.
 
@@ -433,8 +435,8 @@ def _prepare_request_params(
         # prompt caching
         extra_headers = generation_kwargs.get("extra_headers", {})
         prompt_caching_on = "anthropic-beta" in extra_headers and "prompt-caching" in extra_headers["anthropic-beta"]
-        has_cached_messages = any("cache_control" in m for m in system_messages) or any(
-            "cache_control" in m for m in non_system_messages
+        has_cached_messages = any(m.get("cache_control") is not None for m in system_messages) or any(
+            m.get("cache_control") is not None for m in non_system_messages
         )
         if has_cached_messages and not prompt_caching_on:
             # this avoids Anthropic errors when prompt caching is not enabled
@@ -443,32 +445,28 @@ def _prepare_request_params(
                 "Prompt caching is not enabled but you requested individual messages to be cached. "
                 "Messages will be sent to the API without prompt caching."
             )
-            system_messages = list(map(self._remove_cache_control, system_messages))
-            non_system_messages = list(map(self._remove_cache_control, non_system_messages))
+            for message in system_messages:
+                if message.get("cache_control"):
+                    del message["cache_control"]
 
         # tools management
         tools = tools or self.tools
         tools = list(tools) if isinstance(tools, Toolset) else tools
         _check_duplicate_tool_names(tools)  # handles Toolset as well
-        anthropic_tools = (
-            [
-                {
-                    "name": tool.name,
-                    "description": tool.description,
-                    "input_schema": tool.parameters,
-                }
-                for tool in tools
-            ]
-            if tools
-            else []
-        )
+
+        anthropic_tools: List[ToolParam] = []
+        if tools:
+            for tool in tools:
+                anthropic_tools.append(
+                    ToolParam(name=tool.name, description=tool.description, input_schema=tool.parameters)
+                )
 
         return system_messages, non_system_messages, generation_kwargs, anthropic_tools
 
     def _process_response(
         self,
-        response: Any,
-        streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
+        response: Union[Message, Stream[RawMessageStreamEvent]],
+        streaming_callback: Optional[SyncStreamingCallbackT] = None,
     ) -> Dict[str, List[ChatMessage]]:
         """
         Process the response from the Anthropic API.
@@ -478,8 +476,8 @@ def _process_response(
         :returns: A dictionary containing the processed response as a list of ChatMessage objects.
         """
         # workaround for https://github.com/DataDog/dd-trace-py/issues/12562
-        stream = streaming_callback is not None
-        if stream:
+        # we cannot use isinstance(Stream)
+        if not isinstance(response, Message):
             chunks: List[StreamingChunk] = []
             model: Optional[str] = None
             for chunk in response:
@@ -552,7 +550,7 @@ def run(
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
         tools: Optional[Union[List[Tool], Toolset]] = None,
-    ):
+    ) -> Dict[str, List[ChatMessage]]:
         """
         Invokes the Anthropic API with the given messages and generation kwargs.
 
@@ -584,7 +582,8 @@ def run(
             **generation_kwargs,
         )
 
-        return self._process_response(response, streaming_callback)
+        # select_streaming_callback returns a StreamingCallbackT, but we know it's SyncStreamingCallbackT
+        return self._process_response(response=response, streaming_callback=streaming_callback)  # type: ignore[arg-type]
 
     @component.output_types(replies=List[ChatMessage])
     async def run_async(
@@ -593,7 +592,7 @@ async def run_async(
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
         tools: Optional[Union[List[Tool], Toolset]] = None,
-    ):
+    ) -> Dict[str, List[ChatMessage]]:
         """
         Async version of the run method. Invokes the Anthropic API with the given messages and generation kwargs.
 
@@ -625,4 +624,5 @@ async def run_async(
             **generation_kwargs,
         )
 
-        return await self._process_response_async(response, streaming_callback)
+        # select_streaming_callback returns a StreamingCallbackT, but we know it's AsyncStreamingCallbackT
+        return await self._process_response_async(response, streaming_callback)  # type: ignore[arg-type]