Gordonlim tlm chat typing fix (#107)

Author: gordon-lim

pyproject.toml

@@ -48,7 +48,7 @@ extra-dependencies = [
   "pytest-asyncio",
   "python-dotenv",
   "tiktoken",
-  "openai<=1.99.1",
+  "openai",
 ]
 [tool.hatch.envs.types.scripts]
 check = "mypy --strict --install-types --non-interactive {args:src/cleanlab_tlm tests}"
@@ -60,7 +60,7 @@ extra-dependencies = [
   "python-dotenv",
   "pytest-asyncio",
   "tiktoken",
-  "openai<=1.99.1",
+  "openai",
 ]
 
 [tool.hatch.envs.hatch-test.env-vars]

src/cleanlab_tlm/utils/chat.py

@@ -339,7 +339,8 @@ def _form_prompt_chat_completions_api(
 
     # Only return content directly if there's a single user message AND no tools
     if len(messages) == 1 and messages[0].get("role") == _USER_ROLE and (tools is None or len(tools) == 0):
-        return output + str(messages[0]["content"])
+        first_msg = cast(dict[str, Any], messages[0])
+        return output + str(first_msg["content"])
 
     # Warn if the last message is an assistant message with tool calls
     if messages and (messages[-1].get("role") == _ASSISTANT_ROLE or "tool_calls" in messages[-1]):
@@ -358,22 +359,26 @@ def _form_prompt_chat_completions_api(
         if msg["role"] == _ASSISTANT_ROLE:
             output += _ASSISTANT_PREFIX
             # Handle content if present
-            if msg.get("content"):
-                output += f"{msg['content']}\n\n"
+            content_value = cast(Optional[str], msg.get("content"))
+            if content_value:
+                output += f"{content_value}\n\n"
             # Handle tool calls if present
             if "tool_calls" in msg:
                 for tool_call in msg["tool_calls"]:
-                    call_id = tool_call["id"]
-                    function_names[call_id] = tool_call["function"]["name"]
-                    # Format function call as JSON within XML tags, now including call_id
-                    function_call = {
-                        "name": tool_call["function"]["name"],
-                        "arguments": json.loads(tool_call["function"]["arguments"])
-                        if tool_call["function"]["arguments"]
-                        else {},
-                        "call_id": call_id,
-                    }
-                    output += f"{_TOOL_CALL_TAG_START}\n{json.dumps(function_call, indent=2)}\n{_TOOL_CALL_TAG_END}\n\n"
+                    if tool_call["type"] == "function":
+                        call_id = tool_call["id"]
+                        function_names[call_id] = tool_call["function"]["name"]
+                        # Format function call as JSON within XML tags, now including call_id
+                        function_call = {
+                            "name": tool_call["function"]["name"],
+                            "arguments": json.loads(tool_call["function"]["arguments"])
+                            if tool_call["function"]["arguments"]
+                            else {},
+                            "call_id": call_id,
+                        }
+                        output += (
+                            f"{_TOOL_CALL_TAG_START}\n{json.dumps(function_call, indent=2)}\n{_TOOL_CALL_TAG_END}\n\n"
+                        )
         elif msg["role"] == _TOOL_ROLE:
             # Handle tool responses
             output += _TOOL_PREFIX
@@ -506,13 +511,19 @@ def form_response_string_chat_completions_api(
     """
     response_dict = _response_to_dict(response)
     content = response_dict.get("content") or ""
-    tool_calls = response_dict.get("tool_calls")
+    tool_calls = cast(Optional[list[dict[str, Any]]], response_dict.get("tool_calls"))
     if tool_calls is not None:
         try:
-            tool_calls_str = "\n".join(
-                f"{_TOOL_CALL_TAG_START}\n{json.dumps({'name': call['function']['name'], 'arguments': json.loads(call['function']['arguments']) if call['function']['arguments'] else {}}, indent=2)}\n{_TOOL_CALL_TAG_END}"
-                for call in tool_calls
-            )
+            rendered_calls: list[str] = []
+            for call in tool_calls:
+                function_dict = call["function"]
+                name = cast(str, function_dict["name"])
+                args_str = cast(Optional[str], function_dict.get("arguments"))
+                args_obj = json.loads(args_str) if args_str else {}
+                rendered_calls.append(
+                    f"{_TOOL_CALL_TAG_START}\n{json.dumps({'name': name, 'arguments': args_obj}, indent=2)}\n{_TOOL_CALL_TAG_END}"
+                )
+            tool_calls_str = "\n".join(rendered_calls)
             return f"{content}\n{tool_calls_str}".strip() if content else tool_calls_str
         except (KeyError, TypeError, json.JSONDecodeError) as e:
             # Log the error but continue with just the content

tests/openai_compat.py

@@ -0,0 +1,34 @@
+"""Compatibility layer for OpenAI SDK tool call types across versions.
+
+This module exposes two names for tests to import:
+- ChatCompletionMessageToolCall: the class representing a function tool call
+    In previous versions, this was called `ChatCompletionMessageToolCall`,
+    but that has now become a union for both Custom Tool Calls and Function Tool Calls,
+    which have different schemas.
+    This is a shim to allow tests to work with both types.
+- Function: the inner function payload model
+    The Function model is only used for Function Tool Calls, which are currently supported
+    by this package.
+
+Works with SDKs that expose either the legacy
+`chat_completion_message_tool_call.ChatCompletionMessageToolCall` or the
+newer `chat_completion_message_function_tool_call.ChatCompletionMessageFunctionToolCall`.
+"""
+
+from __future__ import annotations
+
+try:  # OpenAI SDK >= 1.99.2
+    from openai.types.chat.chat_completion_message_function_tool_call import (
+        ChatCompletionMessageFunctionToolCall as ChatCompletionMessageToolCall,
+    )
+    from openai.types.chat.chat_completion_message_function_tool_call import (
+        Function,
+    )
+except Exception:  # OpenAI SDK <= 1.99.1
+    import importlib
+
+    _legacy = importlib.import_module("openai.types.chat.chat_completion_message_tool_call")
+    ChatCompletionMessageToolCall = _legacy.ChatCompletionMessageToolCall  # type: ignore
+    Function = _legacy.Function  # type: ignore
+
+__all__ = ["ChatCompletionMessageToolCall", "Function"]

tests/test_chat.py

@@ -3,7 +3,6 @@
 import pytest
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
-from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 
 from cleanlab_tlm.utils.chat import (
     _form_prompt_chat_completions_api,
@@ -12,6 +11,7 @@
     form_response_string_chat_completions,
     form_response_string_chat_completions_api,
 )
+from tests.openai_compat import ChatCompletionMessageToolCall, Function
 
 if TYPE_CHECKING:
     from openai.types.chat import ChatCompletionMessageParam

tests/test_chat_completions.py

@@ -4,7 +4,6 @@
 import pytest
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
-from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 from openai.types.completion_usage import (
     CompletionTokensDetails,
     CompletionUsage,
@@ -17,6 +16,7 @@
 from cleanlab_tlm.utils.chat_completions import TLMChatCompletion
 from tests.conftest import make_text_unique
 from tests.constants import TEST_PROMPT, TEST_RESPONSE
+from tests.openai_compat import ChatCompletionMessageToolCall, Function
 from tests.test_get_trustworthiness_score import is_trustworthiness_score_json_format
 
 test_prompt = make_text_unique(TEST_PROMPT)