Tolerate malformed tool calls in Python SDK chat parsing (fixes #864)

sdk/python/src/openai/chat_client.py

@@ -18,11 +18,63 @@
                                                        CompletionCreateParamsStreaming
 from openai.types.chat import ChatCompletion
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from pydantic import ValidationError
 from typing import Any, Dict, Generator, List, Optional
 
 logger = logging.getLogger(__name__)
 
 
+def _tool_call_is_valid(tool_call: Any) -> bool:
+    """Return True if a tool call is well-formed enough for the OpenAI schema.
+
+    Validity is decided by the explicit ``type`` discriminator: a ``function`` tool
+    call must carry a non-empty ``function.name`` and a ``custom`` tool call must carry
+    a ``custom`` object. A missing or unknown discriminator is itself malformed, as is a
+    ``function`` call with only ``{"arguments": "null"}`` and no name.
+    """
+    if not isinstance(tool_call, dict):
+        return False
+    tool_call_type = tool_call.get("type")
+    if tool_call_type == "custom":
+        return isinstance(tool_call.get("custom"), dict)
+    if tool_call_type == "function":
+        function = tool_call.get("function")
+        return isinstance(function, dict) and bool(function.get("name"))
+    return False
+
+
+def _drop_malformed_tool_calls(payload: Any) -> None:
+    """Remove malformed tool calls from a chat completion payload in place.
+
+    Small models sometimes emit an invalid extra tool call (e.g. missing
+    ``function.name``). Dropping only the bad entries keeps the rest of the response
+    usable instead of failing the whole parse; a warning is logged for each drop.
+    """
+    if not isinstance(payload, dict):
+        return
+    for choice in payload.get("choices") or []:
+        if not isinstance(choice, dict):
+            continue
+        # Non-streaming responses carry tool calls under "message"; streaming under "delta".
+        for key in ("message", "delta"):
+            message = choice.get(key)
+            if not isinstance(message, dict):
+                continue
+            tool_calls = message.get("tool_calls")
+            if not isinstance(tool_calls, list):
+                continue
+            valid = [tc for tc in tool_calls if _tool_call_is_valid(tc)]
+            if len(valid) == len(tool_calls):
+                continue
+            for tc in tool_calls:
+                if not _tool_call_is_valid(tc):
+                    logger.warning("Dropping malformed tool call from Foundry Local response: %r", tc)
+            if valid:
+                message["tool_calls"] = valid
+            else:
+                message.pop("tool_calls", None)
+
+
 class ChatClientSettings:
     """Settings for chat completion requests.
 
@@ -216,7 +268,22 @@ def complete_chat(self, messages: List[ChatCompletionMessageParam], tools: Optio
         if response.error is not None:
             raise FoundryLocalException(f"Error during chat completion: {response.error}")
 
-        completion = ChatCompletion.model_validate_json(response.data)
+        try:
+            payload = json.loads(response.data)
+        except (TypeError, ValueError) as exc:
+            raise FoundryLocalException(
+                f"Foundry Local returned a non-JSON chat completion response: {response.data!r}"
+            ) from exc
+
+        _drop_malformed_tool_calls(payload)
+
+        try:
+            completion = ChatCompletion.model_validate(payload)
+        except ValidationError as exc:
+            raise FoundryLocalException(
+                "Failed to parse the chat completion response from Foundry Local. "
+                f"Raw response: {response.data!r}"
+            ) from exc
 
         return completion
 
@@ -227,18 +294,26 @@ def _stream_chunks(self, chat_request_json: str) -> Generator[ChatCompletionChun
         errors: List[Exception] = []
 
         def _on_chunk(response_str: str) -> None:
-            raw = json.loads(response_str)
-            # Foundry Local returns tool call chunks with "message.tool_calls" instead
-            # of the standard streaming "delta.tool_calls". Normalize to delta format
-            # so ChatCompletionChunk parses correctly.
-            for choice in raw.get("choices", []):
-                if "message" in choice and "delta" not in choice:
-                    msg = choice.pop("message")
-                    # ChoiceDeltaToolCall requires "index"; add if missing
-                    for i, tc in enumerate(msg.get("tool_calls", [])):
-                        tc.setdefault("index", i)
-                    choice["delta"] = msg
-            chunk_queue.put(ChatCompletionChunk.model_validate(raw))
+            try:
+                raw = json.loads(response_str)
+                # Foundry Local returns tool call chunks with "message.tool_calls" instead
+                # of the standard streaming "delta.tool_calls". Normalize to delta format
+                # so ChatCompletionChunk parses correctly.
+                for choice in raw.get("choices", []):
+                    if "message" in choice and "delta" not in choice:
+                        msg = choice.pop("message")
+                        # ChoiceDeltaToolCall requires "index"; add if missing
+                        for i, tc in enumerate(msg.get("tool_calls", [])):
+                            tc.setdefault("index", i)
+                        choice["delta"] = msg
+                _drop_malformed_tool_calls(raw)
+                chunk = ChatCompletionChunk.model_validate(raw)
+            except (ValidationError, ValueError, TypeError, AttributeError) as exc:
+                raise FoundryLocalException(
+                    "Failed to parse a streaming chat completion chunk from Foundry Local. "
+                    f"Raw chunk: {response_str!r}"
+                ) from exc
+            chunk_queue.put(chunk)
 
         def _run() -> None:
             try:

sdk/python/test/openai/test_chat_client_tool_calls.py

@@ -0,0 +1,154 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Unit tests for tolerant tool-call parsing in ChatClient.
+
+These cover the pure sanitization helpers used by ``complete_chat`` /
+``complete_streaming_chat`` and do not require a loaded model or native binary.
+"""
+
+from __future__ import annotations
+
+import json
+
+from openai.types.chat import ChatCompletion
+
+from foundry_local_sdk.openai.chat_client import (
+    _drop_malformed_tool_calls,
+    _tool_call_is_valid,
+)
+
+
+def _completion_with_tool_calls(tool_calls: list[dict]) -> dict:
+    return {
+        "id": "chatcmpl-x",
+        "object": "chat.completion",
+        "created": 0,
+        "model": "smollm3-3b",
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": "tool_calls",
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": tool_calls,
+                },
+            }
+        ],
+    }
+
+
+class TestToolCallValidation:
+    """Tests for ``_tool_call_is_valid``."""
+
+    def test_function_call_with_name_is_valid(self):
+        assert _tool_call_is_valid(
+            {"type": "function", "function": {"name": "f", "arguments": "{}"}}
+        )
+
+    def test_function_call_without_name_is_invalid(self):
+        # Exactly what smollm3-3b emitted: no name, arguments == "null".
+        assert not _tool_call_is_valid(
+            {"type": "function", "function": {"arguments": "null"}}
+        )
+
+    def test_custom_call_is_valid(self):
+        assert _tool_call_is_valid(
+            {"type": "custom", "custom": {"name": "c", "input": "x"}}
+        )
+
+    def test_function_call_with_extra_custom_key_is_still_valid(self):
+        # Dispatch on the explicit "type" discriminator: a valid function call must
+        # not be dropped merely because it also carries an extra "custom" key.
+        assert _tool_call_is_valid(
+            {"type": "function", "custom": None,
+             "function": {"name": "f", "arguments": "{}"}}
+        )
+
+    def test_missing_discriminator_is_invalid(self):
+        assert not _tool_call_is_valid(
+            {"function": {"name": "f", "arguments": "{}"}}
+        )
+
+    def test_unknown_discriminator_is_invalid(self):
+        assert not _tool_call_is_valid(
+            {"type": "mystery", "function": {"name": "f", "arguments": "{}"}}
+        )
+
+    def test_non_dict_is_invalid(self):
+        assert not _tool_call_is_valid("nope")
+        assert not _tool_call_is_valid(None)
+
+
+class TestDropMalformedToolCalls:
+    """Tests for ``_drop_malformed_tool_calls`` + downstream pydantic validation."""
+
+    def test_drops_only_malformed_entry_and_keeps_valid(self):
+        payload = _completion_with_tool_calls(
+            [
+                {"index": 0, "id": "a", "type": "function",
+                 "function": {"name": "get_weather", "arguments": "{}"}},
+                {"index": 1, "id": "b", "type": "function",
+                 "function": {"arguments": "null"}},
+            ]
+        )
+
+        _drop_malformed_tool_calls(payload)
+        completion = ChatCompletion.model_validate(payload)  # must not raise
+
+        tool_calls = completion.choices[0].message.tool_calls
+        assert tool_calls is not None
+        assert len(tool_calls) == 1
+        assert tool_calls[0].function.name == "get_weather"
+
+    def test_all_malformed_removes_tool_calls_entirely(self):
+        payload = _completion_with_tool_calls(
+            [{"index": 0, "id": "b", "type": "function", "function": {"arguments": "null"}}]
+        )
+
+        _drop_malformed_tool_calls(payload)
+        completion = ChatCompletion.model_validate(payload)  # must not raise
+
+        assert completion.choices[0].message.tool_calls is None
+
+    def test_valid_response_is_untouched(self):
+        payload = _completion_with_tool_calls(
+            [{"index": 0, "id": "a", "type": "function",
+              "function": {"name": "f", "arguments": "{}"}}]
+        )
+        expected = json.loads(json.dumps(payload))  # deep copy
+
+        _drop_malformed_tool_calls(payload)
+
+        assert payload == expected
+
+    def test_streaming_delta_tool_calls_are_sanitized(self):
+        payload = {
+            "id": "chatcmpl-x",
+            "object": "chat.completion.chunk",
+            "created": 0,
+            "model": "smollm3-3b",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "tool_calls",
+                    "delta": {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {"index": 0, "id": "a", "type": "function",
+                             "function": {"name": "f", "arguments": "{}"}},
+                            {"index": 1, "id": "b", "type": "function",
+                             "function": {"arguments": "null"}},
+                        ],
+                    },
+                }
+            ],
+        }
+
+        _drop_malformed_tool_calls(payload)
+
+        remaining = payload["choices"][0]["delta"]["tool_calls"]
+        assert len(remaining) == 1
+        assert remaining[0]["function"]["name"] == "f"