feat: Add async JWT support with session_name (#184)

cozepy/auth/__init__.py

@@ -110,6 +110,21 @@ def _refresh_access_token(self, refresh_token: str, secret: str = "") -> OAuthTo
         }
         return self._requester.request("post", url, False, OAuthToken, headers=headers, body=body)
 
+    def _gen_jwt(self, public_key_id: str, private_key: str, ttl: int, session_name: Optional[str] = None):
+        now = int(time.time())
+        header = {"alg": "RS256", "typ": "JWT", "kid": public_key_id}
+        payload = {
+            "iss": self._client_id,
+            "aud": self._api_endpoint,
+            "iat": now,
+            "exp": now + ttl,
+            "jti": random_hex(16),
+        }
+        if session_name:
+            payload["session_name"] = session_name
+        s = jwt.encode(header, payload, private_key)
+        return s.decode("utf-8")
+
     async def _arefresh_access_token(self, refresh_token: str, secret: str = "") -> OAuthToken:
         url = f"{self._base_url}/api/permission/oauth2/token"
         headers = {"Authorization": f"Bearer {secret}"} if secret else {}
@@ -285,7 +300,7 @@ def get_access_token(
         :param scope:
         :param session_name: Isolate different sub-resources under the same jwt account
         """
-        jwt_token = self._gen_jwt(3600, session_name)
+        jwt_token = self._gen_jwt(self._public_key_id, self._private_key, 3600, session_name)
         url = f"{self._base_url}/api/permission/oauth2/token"
         headers = {"Authorization": f"Bearer {jwt_token}"}
         body = {
@@ -295,21 +310,6 @@ def get_access_token(
         }
         return self._requester.request("post", url, False, OAuthToken, headers=headers, body=body)
 
-    def _gen_jwt(self, ttl: int, session_name: Optional[str] = None):
-        now = int(time.time())
-        header = {"alg": "RS256", "typ": "JWT", "kid": self._public_key_id}
-        payload = {
-            "iss": self._client_id,
-            "aud": self._api_endpoint,
-            "iat": now,
-            "exp": now + ttl,
-            "jti": random_hex(16),
-        }
-        if session_name:
-            payload["session_name"] = session_name
-        s = jwt.encode(header, payload, self._private_key)
-        return s.decode("utf-8")
-
 
 class AsyncJWTOAuthApp(OAuthApp):
     """
@@ -331,11 +331,16 @@ def __init__(self, client_id: str, private_key: str, public_key_id: str, base_ur
         self._public_key_id = public_key_id
         super().__init__(client_id, base_url, www_base_url="")
 
-    async def get_access_token(self, ttl: int, scope: Optional[Scope] = None) -> OAuthToken:
+    async def get_access_token(
+        self, ttl: int, scope: Optional[Scope] = None, session_name: Optional[str] = None
+    ) -> OAuthToken:
         """
         Get the token by jwt with jwt auth flow.
+        :param ttl:
+        :param scope:
+        :param session_name: Isolate different sub-resources under the same jwt account
         """
-        jwt_token = self._gen_jwt(3600)
+        jwt_token = self._gen_jwt(self._public_key_id, self._private_key, 3600, session_name)
         url = f"{self._base_url}/api/permission/oauth2/token"
         headers = {"Authorization": f"Bearer {jwt_token}"}
         body = {
@@ -345,19 +350,6 @@ async def get_access_token(self, ttl: int, scope: Optional[Scope] = None) -> OAu
         }
         return await self._requester.arequest("post", url, False, OAuthToken, headers=headers, body=body)
 
-    def _gen_jwt(self, ttl: int):
-        now = int(time.time())
-        header = {"alg": "RS256", "typ": "JWT", "kid": self._public_key_id}
-        payload = {
-            "iss": self._client_id,
-            "aud": self._api_endpoint,
-            "iat": now,
-            "exp": now + ttl,
-            "jti": random_hex(16),
-        }
-        s = jwt.encode(header, payload, self._private_key)
-        return s.decode("utf-8")
-
 
 class PKCEOAuthApp(OAuthApp):
     """

examples/files_upload.py

@@ -0,0 +1,22 @@
+"""
+This example describes how to upload files.
+"""
+
+import os
+import sys
+
+from cozepy import COZE_COM_BASE_URL, Coze, Stream, TokenAuth, WorkflowEvent, WorkflowEventType  # noqa
+
+# Get an access_token through personal access token or oauth.
+coze_api_token = os.getenv("COZE_API_TOKEN")
+# The default access is api.coze.com, but if you need to access api.coze.cn,
+# please use base_url to configure the api endpoint to access
+coze_api_base = os.getenv("COZE_API_BASE") or COZE_COM_BASE_URL
+
+# Init the Coze client through the access_token.
+coze = Coze(auth=TokenAuth(token=coze_api_token), base_url=coze_api_base)
+
+file_path = sys.argv[1] if len(sys.argv) > 1 else "/path/image.jpg"
+
+file = coze.files.upload(file=file_path)
+print(f"uploaded file: {file.id}\n   {file}")