diff --git a/sdk/python/src/openai/chat_client.py b/sdk/python/src/openai/chat_client.py index 0b0d58bcd..14219705c 100644 --- a/sdk/python/src/openai/chat_client.py +++ b/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: diff --git a/sdk/python/test/openai/test_chat_client_tool_calls.py b/sdk/python/test/openai/test_chat_client_tool_calls.py new file mode 100644 index 000000000..6972d0197 --- /dev/null +++ b/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" diff --git a/sdk_v2/python/src/foundry_local_sdk/openai/chat_client.py b/sdk_v2/python/src/foundry_local_sdk/openai/chat_client.py index bcd7a9778..2063a3727 100644 --- a/sdk_v2/python/src/foundry_local_sdk/openai/chat_client.py +++ b/sdk_v2/python/src/foundry_local_sdk/openai/chat_client.py @@ -6,15 +6,72 @@ from __future__ import annotations import json +import logging from typing import TYPE_CHECKING, Any, Generator from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion_chunk import ChatCompletionChunk +from pydantic import ValidationError + +from foundry_local_sdk.exception import FoundryLocalException if TYPE_CHECKING: from foundry_local_sdk.imodel import IModel +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. @@ -221,7 +278,22 @@ def complete_chat( request_json = self._build_request_json(messages, streaming=False, tools=tools) response_json = self._run_native_request(request_json) - return ChatCompletion.model_validate_json(response_json) + try: + payload = json.loads(response_json) + except (TypeError, ValueError) as exc: + raise FoundryLocalException( + f"Foundry Local returned a non-JSON chat completion response: {response_json!r}" + ) from exc + + _drop_malformed_tool_calls(payload) + + try: + return ChatCompletion.model_validate(payload) + except ValidationError as exc: + raise FoundryLocalException( + "Failed to parse the chat completion response from Foundry Local. " + f"Raw response: {response_json!r}" + ) from exc def complete_streaming_chat( self, @@ -265,16 +337,24 @@ def complete_streaming_chat( request.add_item(TextItem(request_json, TextItemType.OPENAI_JSON)) for item in session.process_streaming_request(request): # Each item is a TextItem(OPENAI_JSON) — parse and normalize. - raw = json.loads(item.text) - - # Foundry Local streams tool calls under "message" instead of the - # standard "delta". Normalize to "delta" so ChatCompletionChunk parses. - for choice in raw.get("choices", []): - if "message" in choice and "delta" not in choice: - msg_obj = choice.pop("message") - # ChoiceDeltaToolCall requires "index"; add if absent. - for i, tc in enumerate(msg_obj.get("tool_calls", [])): - tc.setdefault("index", i) - choice["delta"] = msg_obj - - yield ChatCompletionChunk.model_validate(raw) + try: + raw = json.loads(item.text) + + # Foundry Local streams tool calls under "message" instead of the + # standard "delta". Normalize to "delta" so ChatCompletionChunk parses. + for choice in raw.get("choices", []): + if "message" in choice and "delta" not in choice: + msg_obj = choice.pop("message") + # ChoiceDeltaToolCall requires "index"; add if absent. + for i, tc in enumerate(msg_obj.get("tool_calls", [])): + tc.setdefault("index", i) + choice["delta"] = msg_obj + + _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: {item.text!r}" + ) from exc + yield chunk diff --git a/sdk_v2/python/test/unit/test_chat_client_tool_calls.py b/sdk_v2/python/test/unit/test_chat_client_tool_calls.py new file mode 100644 index 000000000..c1601b56f --- /dev/null +++ b/sdk_v2/python/test/unit/test_chat_client_tool_calls.py @@ -0,0 +1,161 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unit tests for tolerant tool-call parsing in the SDK v2 ChatClient. + +These cover the pure sanitization helpers used by ``complete_chat`` / +``complete_streaming_chat`` and do not require a native session or loaded model. +""" +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("openai") + +from openai.types.chat import ChatCompletion # noqa: E402 +from openai.types.chat.chat_completion_chunk import ChatCompletionChunk # noqa: E402 + +from foundry_local_sdk.openai.chat_client import ( # noqa: E402 + _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_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_custom_call_is_valid(self): + assert _tool_call_is_valid( + {"type": "custom", "custom": {"name": "c", "input": "x"}} + ) + + def test_missing_discriminator_is_invalid(self): + # No "type" — the discriminator is missing, so the entry is malformed. + 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) + chunk = ChatCompletionChunk.model_validate(payload) # must not raise + + remaining = chunk.choices[0].delta.tool_calls + assert remaining is not None + assert len(remaining) == 1 + assert remaining[0].function.name == "f"