fix: client close graceful (#22)

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

Author: wellenzheng

examples/example_types.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import Optional, Dict
+
+Message = dict[str, str]  # keys role, content
+MessageList = list[Message]
+
+__all__ = ['Message', 'MessageList', 'Templates', 'SamplerBase', 'EvalResult', 'SingleEvalResult', 'Eval']
+
+Templates = {
+    'base': "{task_template}",
+
+    'meta-chat': "[INST] {task_template} [/INST]",
+
+    'vicuna-chat': "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {task_template} ASSISTANT:",
+
+    'lwm-chat': "You are a helpful assistant. USER: {task_template} ASSISTANT: ",
+
+    'command-r-chat': "<BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>{task_template}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+
+    'chatglm-chat': "[gMASK]sop<|user|> \n {task_template}<|assistant|> \n ",
+    
+    'glm-4-chat': "[gMASK]<sop><|user|>\n{task_template}<|assistant|>",
+    
+    'tgi-glm-4-chat': "<|user|>\n{task_template}<|assistant|>",
+    
+    'RWKV': "User: hi\n\nAssistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it\n\nUser: {task_template}\n\nAssistant:",
+}
+
+class SamplerBase:
+    """
+    Base class for defining a sampling model, which can be evaluated,
+    or used as part of the grading process.
+    """
+
+    def __call__(self, message_list: MessageList) -> str:
+        raise NotImplementedError
+
+
+@dataclass
+class EvalResult:
+    """
+    Result of running an evaluation (usually consisting of many samples)
+    """
+
+    score: Optional[float] = None   # top-line metric
+    metrics: Optional[Dict[str, float]] = None  # other metrics
+
+
+@dataclass
+class SingleEvalResult:
+    """
+    Result of evaluating a single sample
+    """
+
+    score: Optional[float] = None  # top-line metric
+    metrics: Dict[str, float] = field(default_factory=dict)  # other metrics with default empty dict
+
+
+class Eval:
+    """
+    Base class for defining an evaluation.
+    """
+
+    def __call__(self, sampler: SamplerBase) -> EvalResult:
+        raise NotImplementedError

examples/glm4_5_thinking.py

@@ -0,0 +1,113 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import time
+import traceback
+from typing import Optional
+
+from example_types import MessageList, SamplerBase
+from zai import ZhipuAiClient
+
+
+class ZaiSampler(SamplerBase):
+    """
+    Sample from TGI's completion API
+    """
+
+    def __init__(
+        self,
+        model: str = "glm-4.5",
+        api_key: str = '',
+        system_message: Optional[str] = None,
+        temperature: float = 0.0,
+        max_tokens: int = 4096,
+        stream: bool = False,
+    ):
+        self.system_message = system_message
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.model = model
+        self.client = ZhipuAiClient(api_key=api_key)
+        self.stream = stream
+        
+    def get_resp(self, message_list):
+        for _ in range(3):
+            try:
+                chat_completion = self.client.chat.completions.create(
+                    messages=message_list,
+                    model=self.model,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    max_tokens=self.max_tokens
+                )
+                output = chat_completion.choices[0].message.content
+                return output
+            except Exception as e:
+                print(f"Exception: {e}\nTraceback: {traceback.format_exc()}")
+                time.sleep(1)
+                continue
+        print(f"failed, last exception: {e if 'e' in locals() else ''}")
+        return ''
+
+
+    def get_resp_stream(self, message_list, top_p=-1, temperature=-1):
+        temperature = temperature if temperature > 0 else self.temperature
+        top_p = top_p if top_p > 0 else 0.95
+        final = ''
+        for _ in range(200):
+            try:
+                chat_completion_res = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=message_list,
+                    thinking={
+                        "type": "enabled",
+                    },
+                    stream=True,
+                    max_tokens=self.max_tokens,
+                    temperature=temperature
+                )
+                for chunk in chat_completion_res:
+                    if chunk.choices[0].delta.content:
+                        final += chunk.choices[0].delta.content
+                break
+            except Exception as e:
+                final = ""
+                print(f"Exception: {e}\nTraceback: {traceback.format_exc()}")
+                time.sleep(5)
+                continue
+            
+        if final == '':
+            print(f"failed in get_resp for 50 times, last exception: {e if 'e' in locals() else ''}")
+            return ''
+        
+        content = ''
+        if '</think>' in final:
+            content = final.split("</think>")[-1].strip()
+            if not content:
+                content = final[-512:].strip()
+        else:
+            content = final[-512:].strip()
+        
+        return content
+    
+    def __call__(self, message_list: MessageList, top_p=0.95, temperature=0.6) -> str:
+        if self.system_message:      
+            message_list = [
+                {
+                    "role": "system", "content": self.system_message
+                }
+            ] + message_list
+
+        if not self.stream:
+            return self.get_resp(message_list, top_p, temperature)
+        else:
+            return self.get_resp_stream(message_list, top_p, temperature)
+        
+
+if __name__ == "__main__":
+    client = ZaiSampler(model="glm-4.5", api_key=os.getenv("ZAI_API_KEY"), stream=True)
+    messages = [
+        {"role": "user", "content": "Hi?"},
+    ]
+    response = client(messages)
+    print(response)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zai-sdk"
-version = "0.0.1"
+version = "0.0.2"
 description = "A SDK library for accessing big model apis from Z.ai"
 authors = ["Z.ai"]
 readme = "README.md"

src/zai/_client.py

@@ -204,7 +204,13 @@ def __del__(self) -> None:
 		if self._has_custom_http_client:
 			return
 
-		self.close()
+		try:
+			# Check if client is still valid before closing
+			if hasattr(self, '_client') and self._client is not None:
+				self.close()
+		except Exception:
+			# Ignore any exceptions during cleanup to avoid masking the original error
+			pass
 
 
 class ZaiClient(BaseClient):

src/zai/_version.py

@@ -1,2 +1,2 @@
 __title__ = 'Z.ai'
-__version__ = '0.0.1'
+__version__ = '0.0.2'

src/zai/core/_http_client.py

@@ -481,7 +481,12 @@ def is_closed(self) -> bool:
 		return self._client.is_closed
 
 	def close(self):
-		self._client.close()
+		try:
+			if hasattr(self, '_client') and self._client is not None and not self._client.is_closed:
+				self._client.close()
+		except Exception:
+			# Ignore any exceptions during cleanup to avoid masking the original error
+			pass
 
 	def __enter__(self):
 		return self