增加apikey与token认证逻辑 (#31)

Author: glide-the

zhipuai/_client.py

@@ -17,6 +17,7 @@
 class ZhipuAI(HttpClient):
     chat: api_resource.chat.Chat
     api_key: str
+    _disable_token_cache: bool = True
 
     def __init__(
             self,
@@ -26,13 +27,15 @@ def __init__(
             timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
             max_retries: int = ZHIPUAI_DEFAULT_MAX_RETRIES,
             http_client: httpx.Client | None = None,
-            custom_headers: Mapping[str, str] | None = None
+            custom_headers: Mapping[str, str] | None = None,
+            disable_token_cache: bool = True
     ) -> None:
         if api_key is None:
             api_key = os.environ.get("ZHIPUAI_API_KEY")
         if api_key is None:
             raise ZhipuAIError("未提供api_key，请通过参数或环境变量提供")
         self.api_key = api_key
+        self._disable_token_cache = disable_token_cache
 
         if base_url is None:
             base_url = os.environ.get("ZHIPUAI_BASE_URL")
@@ -57,8 +60,10 @@ def __init__(
     @override
     def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key
-        # return {"Authorization": f"{_jwt_token.generate_token(api_key)}"}
-        return {"Authorization": f"{api_key}"}
+        if self._disable_token_cache:
+            return {"Authorization": f"{api_key}"}
+        else:
+            return {"Authorization": f"{_jwt_token.generate_token(api_key)}"}
 
     def __del__(self) -> None:
         if (not hasattr(self, "_has_custom_http_client")

zhipuai/core/_jwt_token.py

@@ -4,9 +4,11 @@
 import cachetools.func
 import jwt
 
-API_TOKEN_TTL_SECONDS = 3 * 60
+# 缓存时间 3分钟
+CACHE_TTL_SECONDS = 3 * 60
 
-CACHE_TTL_SECONDS = API_TOKEN_TTL_SECONDS - 30
+# token 有效期比缓存时间 多30秒
+API_TOKEN_TTL_SECONDS = CACHE_TTL_SECONDS + 30
 
 
 @cachetools.func.ttl_cache(maxsize=10, ttl=CACHE_TTL_SECONDS)