diff --git a/examples/function_call_example.py b/examples/function_call_example.py index 0257baa..72247b5 100644 --- a/examples/function_call_example.py +++ b/examples/function_call_example.py @@ -41,7 +41,7 @@ def parse_function_call(model_response, messages): tools=tools, ) print(response.choices[0].message) - messages.append(response.choices[0].message.model_dump()) + messages.append(response.choices[0].message) messages = [] tools = [ @@ -104,7 +104,7 @@ response = client.chat.completions.create( tools=tools, ) print(response.choices[0].message) -messages.append(response.choices[0].message.model_dump()) +messages.append(response.choices[0].message) parse_function_call(response, messages) @@ -115,6 +115,6 @@ response = client.chat.completions.create( tools=tools, ) print(response.choices[0].message) -messages.append(response.choices[0].message.model_dump()) +messages.append(response.choices[0].message) parse_function_call(response, messages) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 469c3bb..8b556ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "zai-sdk" -version = "0.0.3" +version = "0.0.3.1" description = "A SDK library for accessing big model apis from Z.ai" authors = ["Z.ai"] readme = "README.md" diff --git a/src/zai/_version.py b/src/zai/_version.py index 9098f97..f38c874 100644 --- a/src/zai/_version.py +++ b/src/zai/_version.py @@ -1,2 +1,2 @@ __title__ = 'Z.ai' -__version__ = '0.0.3' +__version__ = '0.0.3.1' diff --git a/src/zai/api_resource/chat/completions.py b/src/zai/api_resource/chat/completions.py index 9aa0bb8..9a503f2 100644 --- a/src/zai/api_resource/chat/completions.py +++ b/src/zai/api_resource/chat/completions.py @@ -18,6 +18,7 @@ from zai.core import ( make_request_options, maybe_transform, ) +from zai.core._base_models import BaseModel from zai.types.chat.chat_completion import Completion from zai.types.chat.chat_completion_chunk import ChatCompletionChunk from zai.types.chat.code_geex import code_geex_params @@ -112,7 +113,9 @@ class Completions(BaseAPI): logger.debug(f'temperature:{temperature}, top_p:{top_p}') if isinstance(messages, List): for item in messages: - if item.get('content'): + if isinstance(item, BaseModel) and hasattr(item, 'content'): + item.content = drop_prefix_image_data(item.content) + elif isinstance(item, dict) and item.get('content'): item['content'] = drop_prefix_image_data(item['content']) body = deepcopy_minimal( diff --git a/src/zai/core/_base_models.py b/src/zai/core/_base_models.py index 859060b..1cce184 100644 --- a/src/zai/core/_base_models.py +++ b/src/zai/core/_base_models.py @@ -160,6 +160,64 @@ class BaseModel(pydantic.BaseModel): warnings=warnings, ) + def get(self, key: str, default: Any = None) -> Any: + """Get the value of an attribute by name, with an optional default value. + + This method allows you to access model attributes by their string name, + similar to how dict.get() works. + + Args: + key: The name of the attribute to get + default: The value to return if the attribute doesn't exist or is None. + Defaults to None. + + Returns: + The value of the attribute if it exists, otherwise the default value. + + Examples: + >>> model = MyModel(name="test", age=25) + >>> model.get("name") # Returns "test" + >>> model.get("nonexistent") # Returns None + >>> model.get("nonexistent", "default") # Returns "default" + """ + try: + value = getattr(self, key) + return value if value is not None else default + except AttributeError: + return default + + def __json__(self) -> dict[str, Any]: + """Custom JSON serialization method. + + This method is called by JSON encoders that support the __json__ protocol, + making BaseModel objects directly serializable without requiring model_dump(). + + Returns: + A dictionary representation of the model suitable for JSON serialization. + """ + return self.model_dump(by_alias=True, exclude_unset=True) + + def __reduce_ex__(self, protocol): + """Support for pickle serialization by returning the dict representation.""" + return (self.__class__.model_validate, (self.model_dump(by_alias=True, exclude_unset=True),)) + + def __iter__(self): + """Make BaseModel iterable to support dict() conversion.""" + data = self.model_dump(by_alias=True, exclude_unset=True) + return iter(data.items()) + + def keys(self): + """Return keys for dict-like interface.""" + return self.model_dump(by_alias=True, exclude_unset=True).keys() + + def values(self): + """Return values for dict-like interface.""" + return self.model_dump(by_alias=True, exclude_unset=True).values() + + def items(self): + """Return items for dict-like interface.""" + return self.model_dump(by_alias=True, exclude_unset=True).items() + @override def __str__(self) -> str: # mypy complains about an invalid self arg diff --git a/src/zai/core/_http_client.py b/src/zai/core/_http_client.py index 8586a32..080f8b2 100644 --- a/src/zai/core/_http_client.py +++ b/src/zai/core/_http_client.py @@ -65,6 +65,7 @@ from ._request_opt import FinalRequestOptions, UserRequestInput from ._response import APIResponse, BaseAPIResponse, extract_response_type from ._streaming import StreamResponse from ._utils import flatten, is_given, is_mapping +from ._json_encoder import json_dumps log: logging.Logger = logging.getLogger(__name__) @@ -355,6 +356,9 @@ class HttpClient: else: raise RuntimeError(f'Unexpected JSON data type, {type(json_data)}, cannot merge with `extra_body`') + # Convert BaseModel objects to dicts before passing to httpx + json_data = self._prepare_json_data(json_data) + content_type = headers.get('Content-Type') # multipart/form-data; boundary=---abc-- if headers.get('Content-Type') == 'multipart/form-data': @@ -377,6 +381,24 @@ class HttpClient: **kwargs, ) + def _prepare_json_data(self, json_data: Any) -> Any: + """Prepare JSON data for httpx by converting BaseModel objects to dicts.""" + from ._base_models import BaseModel + + if json_data is None: + return None + + if isinstance(json_data, BaseModel): + return json_data.model_dump(by_alias=True, exclude_unset=True) + + if isinstance(json_data, list): + return [self._prepare_json_data(item) for item in json_data] + + if isinstance(json_data, dict): + return {key: self._prepare_json_data(value) for key, value in json_data.items()} + + return json_data + def _object_to_formfata(self, key: str, value: Data | Mapping[object, object]) -> list[tuple[str, str]]: items = [] diff --git a/src/zai/core/_json_encoder.py b/src/zai/core/_json_encoder.py new file mode 100644 index 0000000..a228fa8 --- /dev/null +++ b/src/zai/core/_json_encoder.py @@ -0,0 +1,44 @@ +"""JSON encoding utilities for BaseModel objects.""" + +import json +from typing import Any + +from ._base_models import BaseModel + + +class ZAIJSONEncoder(json.JSONEncoder): + """Custom JSON encoder that handles BaseModel objects.""" + + def default(self, obj: Any) -> Any: + """Override default method to handle BaseModel objects.""" + if isinstance(obj, BaseModel): + return obj.model_dump(by_alias=True, exclude_unset=True) + return super().default(obj) + + +def json_dumps(obj: Any, **kwargs) -> str: + """ + JSON dumps with support for BaseModel objects. + + Args: + obj: Object to serialize + **kwargs: Additional arguments to pass to json.dumps() + + Returns: + JSON string representation + """ + return json.dumps(obj, cls=ZAIJSONEncoder, **kwargs) + + +def json_loads(s: str, **kwargs) -> Any: + """ + JSON loads with consistent interface. + + Args: + s: JSON string to deserialize + **kwargs: Additional arguments to pass to json.loads() + + Returns: + Deserialized object + """ + return json.loads(s, **kwargs) diff --git a/src/zai/types/chat/chat_completion_chunk.py b/src/zai/types/chat/chat_completion_chunk.py index 5bf19fc..4c4ba43 100644 --- a/src/zai/types/chat/chat_completion_chunk.py +++ b/src/zai/types/chat/chat_completion_chunk.py @@ -94,6 +94,26 @@ class Choice(BaseModel): finish_reason: Optional[str] = None index: int +class PromptTokensDetails(BaseModel): + """ + Detailed breakdown of token usage for the input prompt + + Attributes: + cached_tokens: Number of tokens reused from cache + """ + + cached_tokens: int + + +class CompletionTokensDetails(BaseModel): + """ + Detailed breakdown of token usage for the model completion + + Attributes: + reasoning_tokens: Number of tokens used for reasoning steps + """ + + reasoning_tokens: int class CompletionUsage(BaseModel): """ @@ -106,7 +126,9 @@ class CompletionUsage(BaseModel): """ prompt_tokens: int + prompt_tokens_details: Optional[PromptTokensDetails] = None completion_tokens: int + completion_tokens_details: Optional[CompletionTokensDetails] = None total_tokens: int diff --git a/tests/integration_tests/asr.wav b/tests/integration_tests/asr.wav new file mode 100644 index 0000000..e29ac13 Binary files /dev/null and b/tests/integration_tests/asr.wav differ diff --git a/tests/integration_tests/asr1.wav b/tests/integration_tests/asr1.wav deleted file mode 100644 index ecff0a3..0000000 Binary files a/tests/integration_tests/asr1.wav and /dev/null differ diff --git a/tests/integration_tests/test_audio.py b/tests/integration_tests/test_audio.py index a0bb2f4..e143571 100644 --- a/tests/integration_tests/test_audio.py +++ b/tests/integration_tests/test_audio.py @@ -30,9 +30,9 @@ def test_audio_speech(logging_conf): def test_audio_customization(logging_conf): logging.config.dictConfig(logging_conf) client = ZaiClient() # Fill in your own API Key - with open(Path(__file__).parent / 'asr1.wav', 'rb') as file: + with open(Path(__file__).parent / 'asr.wav', 'rb') as file: try: - speech_file_path = Path(__file__).parent / 'asr1.wav' + speech_file_path = Path(__file__).parent / 'asr.wav' response = client.audio.customization( model='cogtts', input='Hello, welcome to Z.ai Open Platform', diff --git a/tests/integration_tests/test_client_cleanup_integration.py b/tests/integration_tests/test_client_cleanup_integration.py deleted file mode 100644 index fe9c392..0000000 --- a/tests/integration_tests/test_client_cleanup_integration.py +++ /dev/null @@ -1,267 +0,0 @@ -""" -Integration tests for client cleanup functionality -""" - -import gc -import os -import sys -import time -import unittest -from unittest.mock import patch - -# Add src to path for testing -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../src')) - -from zai import ZaiClient, ZhipuAiClient - - -class TestClientCleanupIntegration(unittest.TestCase): - """Integration tests for client cleanup functionality""" - - def setUp(self): - """Set up test environment""" - # Set a dummy API key for testing - os.environ['ZAI_API_KEY'] = 'test-api-key' - - def tearDown(self): - """Clean up test environment""" - if 'ZAI_API_KEY' in os.environ: - del os.environ['ZAI_API_KEY'] - - def test_context_manager_with_api_calls(self): - """Test context manager with actual API calls (mocked)""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - with ZaiClient() as client: - # Simulate API call - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate some API operations - # (In real usage, these would be actual API calls) - pass - - # Client should be closed after context manager exits - # After our fix, the client should still be open because we check is_closed - # and don't close if it's already closed - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_explicit_close_with_api_calls(self): - """Test explicit close with actual API calls (mocked)""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - client = ZaiClient() - try: - # Simulate API call - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate some API operations - # (In real usage, these would be actual API calls) - pass - finally: - client.close() - - # Client should be closed - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_multiple_clients_in_sequence(self): - """Test multiple clients created and closed in sequence""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - # Create and close multiple clients - for i in range(3): - with ZaiClient() as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - - # Verify that close was called for each client - # Our fix allows multiple close calls, so the count might be higher - self.assertGreaterEqual(mock_client.close.call_count, 3) - - def test_client_reuse_after_close(self): - """Test reusing a client after it has been closed""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - client = ZaiClient() - client.close() - self.assertFalse(client.is_closed()) - - # Try to close again - should not raise exception - client.close() - self.assertFalse(client.is_closed()) - - def test_zhipu_client_cleanup_integration(self): - """Test ZhipuAiClient cleanup integration""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - with ZhipuAiClient() as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call - pass - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_client_with_custom_http_client_integration(self): - """Test client with custom HTTP client integration""" - import httpx - - # Create a real httpx client for testing - custom_client = httpx.Client() - - with ZaiClient(http_client=custom_client) as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call - pass - - # Client should be closed after context manager exits - # The client is actually closed in this case - self.assertTrue(client.is_closed()) - - # Custom client should be closed by our client - self.assertTrue(custom_client.is_closed) - - # No need to clean up custom client as it's already closed - - def test_client_cleanup_with_exception_handling(self): - """Test client cleanup when exceptions occur during API calls""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - try: - with ZaiClient() as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate an exception during API call - raise Exception("Simulated API error") - except Exception: - # Exception should be caught and client should still be closed - pass - - # Client should be closed even after exception - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_client_cleanup_with_timeout(self): - """Test client cleanup with timeout scenarios""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - # Create client with custom timeout - import httpx - custom_timeout = httpx.Timeout(timeout=30.0, connect=5.0) - - with ZaiClient(timeout=custom_timeout) as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call with timeout - pass - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_client_cleanup_with_retries(self): - """Test client cleanup with retry scenarios""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - # Create client with custom retry settings - with ZaiClient(max_retries=5) as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call with retries - pass - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_client_cleanup_with_custom_headers(self): - """Test client cleanup with custom headers""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - # Create client with custom headers - custom_headers = {'X-Custom-Header': 'test-value'} - - with ZaiClient(custom_headers=custom_headers) as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call - pass - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - def test_client_cleanup_with_source_channel(self): - """Test client cleanup with source channel""" - with patch('httpx.Client') as mock_client_class: - # Mock the httpx client - mock_client = mock_client_class.return_value - mock_client.is_closed = False - mock_client.close = unittest.mock.Mock() - - # Create client with source channel - with ZaiClient(source_channel='test-channel') as client: - self.assertIsNotNone(client) - self.assertFalse(client.is_closed()) - - # Simulate API call - pass - - # Client should be closed after context manager exits - self.assertFalse(client.is_closed()) - mock_client.close.assert_called_once() - - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/integration_tests/test_file.py b/tests/integration_tests/test_file.py index f69a071..382be42 100644 --- a/tests/integration_tests/test_file.py +++ b/tests/integration_tests/test_file.py @@ -13,7 +13,6 @@ from zai import ZaiClient @pytest.fixture(scope='class') def test_server(): class SharedData: - client = ZaiClient() file_id1 = None file_id2 = None @@ -26,7 +25,10 @@ class TestZaiClientFileServer: def test_files(self, test_server, test_file_path): try: - result = test_server.client.files.create( + import os + print('ZAI_API_KEY in test_files:', os.environ.get('ZAI_API_KEY')) + client = ZaiClient() + result = client.files.create( file=open(os.path.join(test_file_path, 'demo.jsonl'), 'rb'), purpose='fine-tune', ) @@ -42,7 +44,10 @@ class TestZaiClientFileServer: def test_files_validation(self, test_server, test_file_path): try: - result = test_server.client.files.create( + import os + print('ZAI_API_KEY in test_files_validation:', os.environ.get('ZAI_API_KEY')) + client = ZaiClient() + result = client.files.create( file=open(os.path.join(test_file_path, 'demo.jsonl'), 'rb'), purpose='fine-tune', ) @@ -59,7 +64,10 @@ class TestZaiClientFileServer: def test_files_list(self, test_server): try: - list = test_server.client.files.list() + import os + print('ZAI_API_KEY in test_files_list:', os.environ.get('ZAI_API_KEY')) + client = ZaiClient() + list = client.files.list() print(list) except zai.core._errors.APIRequestFailedError as err: @@ -71,13 +79,16 @@ class TestZaiClientFileServer: def test_delete_files(self, test_server): try: + import os + print('ZAI_API_KEY in test_delete_files:', os.environ.get('ZAI_API_KEY')) + client = ZaiClient() # Only delete files if they were successfully created if test_server.file_id1: - delete1 = test_server.client.files.delete(file_id=test_server.file_id1) + delete1 = client.files.delete(file_id=test_server.file_id1) print(delete1) if test_server.file_id2: - delete2 = test_server.client.files.delete(file_id=test_server.file_id2) + delete2 = client.files.delete(file_id=test_server.file_id2) print(delete2) except zai.core._errors.APIRequestFailedError as err: diff --git a/tests/unit_tests/test_client.py b/tests/unit_tests/test_client.py new file mode 100644 index 0000000..3c3cfa4 --- /dev/null +++ b/tests/unit_tests/test_client.py @@ -0,0 +1,23 @@ +def test_completions_image_url_data_image(monkeypatch): + from zai import ZaiClient + client = ZaiClient() + # 构造 messages，包含 image_url 且 url 以 data:image/ 开头 + base64_str = 'abc123==' + messages = [ + { + 'role': 'user', + 'content': [ + {'type': 'text', 'text': 'What is in this image?'}, + {'type': 'image_url', 'image_url': {'url': f'data:image/png;base64,{base64_str}'}}, + ] + } + ] + # mock client.post 方法只返回 body 以便断言 + def fake_post(*args, **kwargs): + return kwargs.get('body', {}) + monkeypatch.setattr(client, 'post', fake_post) + result = client.chat.completions.create(model='glm-4', messages=messages) + # 验证 image_url 的 url 字段已去除前缀 + content = result['messages'][0]['content'] + image_url = [x for x in content if x.get('type') == 'image_url'][0] + assert image_url['image_url']['url'] == base64_str \ No newline at end of file diff --git a/tests/unit_tests/test_json_serialization.py b/tests/unit_tests/test_json_serialization.py new file mode 100644 index 0000000..e911194 --- /dev/null +++ b/tests/unit_tests/test_json_serialization.py @@ -0,0 +1,166 @@ +"""Test JSON serialization for BaseModel objects.""" + +import json +import pytest +from typing import Optional, List + +from zai.core._base_models import BaseModel +from zai.core._json_encoder import ZAIJSONEncoder, json_dumps +from zai.core._http_client import HttpClient + + +class ModelForTest(BaseModel): + """Test model for JSON serialization.""" + name: str + age: Optional[int] = None + tags: Optional[List[str]] = None + + +class MessageForTest(BaseModel): + """Test message model similar to CompletionMessage.""" + content: Optional[str] = None + role: str + tool_calls: Optional[List[dict]] = None + + +def test_base_model_dict_conversion(): + """Test that BaseModel can be converted to dict.""" + model = ModelForTest(name="test", age=25, tags=["tag1", "tag2"]) + + # Test model_dump() method + model_dict = model.model_dump() + assert model_dict == {"name": "test", "age": 25, "tags": ["tag1", "tag2"]} + + # Test keys/values/items methods + assert list(model.keys()) == ["name", "age", "tags"] + assert list(model.values()) == ["test", 25, ["tag1", "tag2"]] + assert list(model.items()) == [("name", "test"), ("age", 25), ("tags", ["tag1", "tag2"])] + + +def test_zai_json_encoder(): + """Test custom JSON encoder for BaseModel objects.""" + model = ModelForTest(name="test", age=25) + + # Test with custom encoder + json_str = json.dumps(model, cls=ZAIJSONEncoder) + result = json.loads(json_str) + assert result == {"name": "test", "age": 25} + + # Test with nested BaseModel + message = MessageForTest(content="hello", role="user") + nested_data = {"message": message, "other": "data"} + + json_str = json.dumps(nested_data, cls=ZAIJSONEncoder) + result = json.loads(json_str) + assert result == { + "message": {"content": "hello", "role": "user"}, + "other": "data" + } + + +def test_json_dumps_utility(): + """Test the json_dumps utility function.""" + model = ModelForTest(name="test", age=25) + + json_str = json_dumps(model) + result = json.loads(json_str) + assert result == {"name": "test", "age": 25} + + +def test_http_client_json_preparation(): + """Test that HttpClient can prepare JSON data with BaseModel objects.""" + # We'll create a minimal HttpClient to test the _prepare_json_data method + import httpx + from zai.core._http_client import HttpClient + + # Create a minimal client for testing + client = HttpClient( + base_url="https://test.com", + version="v1", + _strict_response_validation=True, + timeout=10.0 + ) + + # Test with BaseModel + model = ModelForTest(name="test", age=25) + result = client._prepare_json_data(model) + assert result == {"name": "test", "age": 25} + + # Test with list containing BaseModel + models = [ModelForTest(name="test1", age=25), ModelForTest(name="test2", age=30)] + result = client._prepare_json_data(models) + expected = [ + {"name": "test1", "age": 25}, + {"name": "test2", "age": 30} + ] + assert result == expected + + # Test with dict containing BaseModel + data = { + "message": MessageForTest(content="hello", role="user"), + "count": 1 + } + result = client._prepare_json_data(data) + expected = { + "message": {"content": "hello", "role": "user"}, + "count": 1 + } + assert result == expected + + # Test with None + assert client._prepare_json_data(None) is None + + # Test with regular data + regular_data = {"key": "value", "number": 42} + assert client._prepare_json_data(regular_data) == regular_data + + +def test_completion_message_like_serialization(): + """Test serialization for CompletionMessage-like objects.""" + # Simulate a CompletionMessage-like object + message = MessageForTest( + content="Hello, how can I help you?", + role="assistant", + tool_calls=[{"id": "call_123", "function": {"name": "test_func"}}] + ) + + # Test that it can be serialized + json_str = json_dumps(message) + result = json.loads(json_str) + + expected = { + "content": "Hello, how can I help you?", + "role": "assistant", + "tool_calls": [{"id": "call_123", "function": {"name": "test_func"}}] + } + assert result == expected + + # Test in a messages list (similar to the original error case) + messages = [ + {"role": "user", "content": "Help me"}, + message # This is a BaseModel object + ] + + # This should work now with our custom encoder + json_str = json_dumps(messages) + result = json.loads(json_str) + + expected = [ + {"role": "user", "content": "Help me"}, + { + "content": "Hello, how can I help you?", + "role": "assistant", + "tool_calls": [{"id": "call_123", "function": {"name": "test_func"}}] + } + ] + assert result == expected + + +if __name__ == "__main__": + # Run basic tests + test_base_model_dict_conversion() + test_zai_json_encoder() + test_json_dumps_utility() + test_http_client_json_preparation() + test_completion_message_like_serialization() + print("All tests passed!")