test(mlx): add async tool test

Author: diego-coder

libs/community/langchain_community/chat_models/mlx.py

@@ -99,7 +99,8 @@ def _generate(
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        llm_input = self._to_chat_prompt(messages)
+        tools = kwargs.pop("tools", None)
+        llm_input = self._to_chat_prompt(messages, tools=tools)
         llm_result = self.llm._generate(
             prompts=[llm_input], stop=stop, run_manager=run_manager, **kwargs
         )
@@ -112,7 +113,8 @@ async def _agenerate(
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        llm_input = self._to_chat_prompt(messages)
+        tools = kwargs.pop("tools", None)
+        llm_input = self._to_chat_prompt(messages, tools=tools)
         llm_result = await self.llm._agenerate(
             prompts=[llm_input], stop=stop, run_manager=run_manager, **kwargs
         )
@@ -123,8 +125,17 @@ def _to_chat_prompt(
         messages: List[BaseMessage],
         tokenize: bool = False,
         return_tensors: Optional[str] = None,
+        tools: Sequence[dict] | None = None,
     ) -> str:
-        """Convert a list of messages into a prompt format expected by wrapped LLM."""
+        """Convert messages to the prompt format expected by the wrapped LLM.
+
+        Args:
+            messages: Chat messages to include in the prompt.
+            tokenize: Whether to return token IDs instead of text.
+            return_tensors: Framework for returned tensors when ``tokenize`` is
+                True.
+            tools: Optional tool definitions to include in the prompt.
+        """
         if not messages:
             raise ValueError("At least one HumanMessage must be provided!")
 
@@ -137,6 +148,7 @@ def _to_chat_prompt(
             tokenize=tokenize,
             add_generation_prompt=True,
             return_tensors=return_tensors,
+            tools=tools,
         )
 
     def _to_chatml_format(self, message: BaseMessage) -> dict:

libs/community/tests/unit_tests/chat_models/test_mlx.py

@@ -2,6 +2,42 @@
 
 from importlib import import_module
 
+import pytest
+from langchain_core.messages import HumanMessage
+
+from langchain_community.chat_models.mlx import ChatMLX
+
+
+class _FakeTokenizer:
+    def __init__(self) -> None:
+        self.tools = None
+
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        return_tensors=None,
+        tools=None,
+    ) -> str:
+        self.tools = tools
+        return "prompt"
+
+
+class _FakeLLM:
+    def __init__(self) -> None:
+        self.tokenizer = _FakeTokenizer()
+
+    def _generate(self, prompts, stop=None, run_manager=None, **kwargs):
+        class _Res:
+            generations = [[type("G", (), {"text": "", "generation_info": {}})]]
+            llm_output = {}
+
+        return _Res()
+
+    async def _agenerate(self, prompts, stop=None, run_manager=None, **kwargs):
+        return self._generate(prompts, stop=stop, run_manager=run_manager, **kwargs)
+
 
 def test_import_class() -> None:
     """Test that the class can be imported."""
@@ -10,3 +46,38 @@ def test_import_class() -> None:
 
     module = import_module(module_name)
     assert hasattr(module, class_name)
+
+
+def test_generate_passes_tools_to_tokenizer() -> None:
+    llm = _FakeLLM()
+    chat = ChatMLX(llm=llm)
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "foo",
+                "description": "",
+                "parameters": {"type": "object", "properties": {}},
+            },
+        }
+    ]
+    chat._generate([HumanMessage(content="hi")], tools=tools)
+    assert llm.tokenizer.tools == tools
+
+
+@pytest.mark.asyncio
+async def test_agenerate_passes_tools_to_tokenizer() -> None:
+    llm = _FakeLLM()
+    chat = ChatMLX(llm=llm)
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "foo",
+                "description": "",
+                "parameters": {"type": "object", "properties": {}},
+            },
+        }
+    ]
+    await chat._agenerate([HumanMessage(content="hi")], tools=tools)
+    assert llm.tokenizer.tools == tools