PR feedback

Author: vblagoje

haystack/components/generators/chat/hugging_face_local.py

@@ -43,6 +43,35 @@
 )
 
 
+def default_tool_parser(text: str) -> Optional[List[ToolCall]]:
+    """
+    Default implementation for parsing tool calls from model output text.
+
+    Uses DEFAULT_TOOL_PATTERN to extract tool calls.
+
+    :param text: The text to parse for tool calls.
+    :returns: A list containing a single ToolCall if a valid tool call is found, None otherwise.
+    """
+    try:
+        match = re.search(DEFAULT_TOOL_PATTERN, text, re.DOTALL)
+    except re.error:
+        logger.warning("Invalid regex pattern for tool parsing: {pattern}", pattern=DEFAULT_TOOL_PATTERN)
+        return None
+
+    if not match:
+        return None
+
+    name = match.group(1) or match.group(3)
+    args_str = match.group(2) or match.group(4)
+
+    try:
+        arguments = json.loads(args_str)
+        return [ToolCall(tool_name=name, arguments=arguments)]
+    except json.JSONDecodeError:
+        logger.warning("Failed to parse tool call arguments: {args_str}", args_str=args_str)
+        return None
+
+
 @component
 class HuggingFaceLocalChatGenerator:
     """
@@ -93,7 +122,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         stop_words: Optional[List[str]] = None,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
         tools: Optional[List[Tool]] = None,
-        tool_pattern: Optional[Union[str, Callable[[str], Optional[List[ToolCall]]]]] = None,
+        tool_parsing_function: Optional[Callable[[str], Optional[List[ToolCall]]]] = None,
     ):
         """
         Initializes the HuggingFaceLocalChatGenerator component.
@@ -133,11 +162,9 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             In these cases, make sure your prompt has no stop words.
         :param streaming_callback: An optional callable for handling streaming responses.
         :param tools: A list of tools for which the model can prepare calls.
-        :param tool_pattern:
-            A pattern or callable to parse tool calls from model output.
-            If a string, it will be used as a regex pattern to extract ToolCall object.
-            If a callable, it should take a string and return a ToolCall object or None.
-            If None, a default pattern will be used.
+        :param tool_parsing_function:
+            A callable that takes a string and returns a list of ToolCall objects or None.
+            If None, the default_tool_parser will be used which extracts tool calls using a predefined pattern.
         """
         torch_and_transformers_import.check()
 
@@ -188,7 +215,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         generation_kwargs["stop_sequences"] = generation_kwargs.get("stop_sequences", [])
         generation_kwargs["stop_sequences"].extend(stop_words or [])
 
-        self.tool_pattern = tool_pattern or DEFAULT_TOOL_PATTERN
+        self.tool_parsing_function = tool_parsing_function or default_tool_parser
         self.huggingface_pipeline_kwargs = huggingface_pipeline_kwargs
         self.generation_kwargs = generation_kwargs
         self.chat_template = chat_template
@@ -228,6 +255,7 @@ def to_dict(self) -> Dict[str, Any]:
             token=self.token.to_dict() if self.token else None,
             chat_template=self.chat_template,
             tools=serialized_tools,
+            tool_parsing_function=serialize_callable(self.tool_parsing_function),
         )
 
         huggingface_pipeline_kwargs = serialization_dict["init_parameters"]["huggingface_pipeline_kwargs"]
@@ -254,6 +282,10 @@ def from_dict(cls, data: Dict[str, Any]) -> "HuggingFaceLocalChatGenerator":
         if serialized_callback_handler:
             data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
 
+        tool_parsing_function = init_params.get("tool_parsing_function")
+        if tool_parsing_function:
+            init_params["tool_parsing_function"] = deserialize_callable(tool_parsing_function)
+
         huggingface_pipeline_kwargs = init_params.get("huggingface_pipeline_kwargs", {})
         deserialize_hf_model_kwargs(huggingface_pipeline_kwargs)
         return default_from_dict(cls, data)
@@ -371,7 +403,7 @@ def create_message(  # pylint: disable=too-many-positional-arguments
         prompt_token_count = len(tokenizer.encode(prompt, add_special_tokens=False))
         total_tokens = prompt_token_count + completion_tokens
 
-        tool_calls = self._parse_tool_call(text) if parse_tool_calls else None
+        tool_calls = self.tool_parsing_function(text) if parse_tool_calls else None
 
         # Determine finish reason based on context
         if completion_tokens >= generation_kwargs.get("max_new_tokens", sys.maxsize):
@@ -392,7 +424,8 @@ def create_message(  # pylint: disable=too-many-positional-arguments
             },
         }
 
-        return ChatMessage.from_assistant(tool_calls=tool_calls, text=text, meta=meta)
+        # If tool calls are detected, don't include the text content since it contains the raw tool call format
+        return ChatMessage.from_assistant(tool_calls=tool_calls, text=None if tool_calls else text, meta=meta)
 
     def _validate_stop_words(self, stop_words: Optional[List[str]]) -> Optional[List[str]]:
         """
@@ -410,36 +443,3 @@ def _validate_stop_words(self, stop_words: Optional[List[str]]) -> Optional[List
             return None
 
         return list(set(stop_words or []))
-
-    def _parse_tool_call(self, text: str) -> Optional[List[ToolCall]]:
-        """
-        Parse a tool call from model output text.
-
-        :param text: The text to parse for tool calls.
-        :returns: A ToolCall object if a valid tool call is found, None otherwise.
-        """
-        # if the tool pattern is a callable, call it with the text and return the result
-        if callable(self.tool_pattern):
-            return self.tool_pattern(text)
-
-        # if the tool pattern is a regex pattern, search for it in the text
-        try:
-            match = re.search(self.tool_pattern, text, re.DOTALL)
-        except re.error:
-            logger.warning("Invalid regex pattern for tool parsing: {pattern}", pattern=self.tool_pattern)
-            return None
-
-        if not match:
-            return None
-
-        # seem like most models are not producing tool ids, so we omit them
-        # and just use the tool name and arguments
-        name = match.group(1) or match.group(3)
-        args_str = match.group(2) or match.group(4)
-
-        try:
-            arguments = json.loads(args_str)
-            return [ToolCall(tool_name=name, arguments=arguments)]
-        except json.JSONDecodeError:
-            logger.warning("Failed to parse tool call arguments: {args_str}", args_str=args_str)
-            return None

test/components/generators/chat/test_hugging_face_local.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from unittest.mock import Mock, patch
+from typing import Optional, List
 
 from haystack.dataclasses.streaming_chunk import StreamingChunk
 import pytest
@@ -68,6 +69,11 @@ def tools():
     return [tool]
 
 
+def custom_tool_parser(text: str) -> Optional[List[ToolCall]]:
+    """Test implementation of a custom tool parser."""
+    return [ToolCall(tool_name="weather", arguments={"city": "Berlin"})]
+
+
 class TestHuggingFaceLocalChatGenerator:
     def test_initialize_with_valid_model_and_generation_parameters(self, model_info_mock):
         model = "HuggingFaceH4/zephyr-7b-alpha"
@@ -433,26 +439,38 @@ def test_run_with_tools_and_tool_response(self, model_info_mock, tools):
         assert "22°C" in message.text
         assert message.meta["finish_reason"] == "stop"
 
-    def test_run_with_invalid_tool_pattern(self, model_info_mock, tools):
+    def test_run_with_custom_tool_parser(self, model_info_mock, tools):
+        """Test that a custom tool parsing function works correctly."""
         generator = HuggingFaceLocalChatGenerator(
-            model="meta-llama/Llama-2-13b-chat-hf",
-            tools=tools,
-            tool_pattern=r"invalid[pattern",  # Invalid regex pattern
+            model="meta-llama/Llama-2-13b-chat-hf", tools=tools, tool_parsing_function=custom_tool_parser
         )
+        generator.pipeline = Mock(return_value=[{"generated_text": "Let me check the weather for you"}])
+        generator.pipeline.tokenizer = Mock()
+        generator.pipeline.tokenizer.encode.return_value = [1, 2, 3]
+        generator.pipeline.tokenizer.pad_token_id = 1
 
-        # Mock pipeline and tokenizer
-        mock_pipeline = Mock(return_value=[{"generated_text": '{"name": "weather", "arguments": {"city": "Paris"}}'}])
-        mock_tokenizer = Mock(spec=PreTrainedTokenizer)
-        mock_tokenizer.encode.return_value = ["some", "tokens"]
-        mock_tokenizer.pad_token_id = 100
-        mock_tokenizer.apply_chat_template.return_value = "test prompt"
-        mock_pipeline.tokenizer = mock_tokenizer
-        generator.pipeline = mock_pipeline
+        messages = [ChatMessage.from_user("What's the weather like in Berlin?")]
+        results = generator.run(messages=messages)
 
-        messages = [ChatMessage.from_user("What's the weather in Paris?")]
+        assert len(results["replies"]) == 1
+        assert len(results["replies"][0].tool_calls) == 1
+        assert results["replies"][0].tool_calls[0].tool_name == "weather"
+        assert results["replies"][0].tool_calls[0].arguments == {"city": "Berlin"}
+
+    def test_default_tool_parser(self, model_info_mock, tools):
+        """Test that the default tool parser works correctly with valid tool call format."""
+        generator = HuggingFaceLocalChatGenerator(model="meta-llama/Llama-2-13b-chat-hf", tools=tools)
+        generator.pipeline = Mock(
+            return_value=[{"generated_text": '{"name": "weather", "arguments": {"city": "Berlin"}}'}]
+        )
+        generator.pipeline.tokenizer = Mock()
+        generator.pipeline.tokenizer.encode.return_value = [1, 2, 3]
+        generator.pipeline.tokenizer.pad_token_id = 1
+
+        messages = [ChatMessage.from_user("What's the weather like in Berlin?")]
         results = generator.run(messages=messages)
 
         assert len(results["replies"]) == 1
-        message = results["replies"][0]
-        assert not message.tool_calls  # No tool calls due to invalid pattern
-        assert message.meta["finish_reason"] == "stop"
+        assert len(results["replies"][0].tool_calls) == 1
+        assert results["replies"][0].tool_calls[0].tool_name == "weather"
+        assert results["replies"][0].tool_calls[0].arguments == {"city": "Berlin"}