feat: modeldump and usage support (#25)

Co-authored-by: zhengweijun <[email protected]>

Author: wellenzheng

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 @@ 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)
 
 parse_function_call(response, messages)
 
@@ -115,6 +115,6 @@ 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)
 
 parse_function_call(response, messages)

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"

src/zai/_version.py

@@ -1,2 +1,2 @@
 __title__ = 'Z.ai'
-__version__ = '0.0.3'
+__version__ = '0.0.3.1'

src/zai/api_resource/chat/completions.py

@@ -18,6 +18,7 @@
 	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 @@ def create(
 		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(

src/zai/core/_base_models.py

@@ -160,6 +160,64 @@ def to_json(
 			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

src/zai/core/_http_client.py

@@ -65,6 +65,7 @@
 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 @@ def _build_request(self, options: FinalRequestOptions) -> httpx.Request:
 			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 @@ def _build_request(self, options: FinalRequestOptions) -> httpx.Request:
 			**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 = []
 

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)

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