1+ import logging
2+ from aixplain .utils import config
3+ from aixplain .utils .file_utils import _request_with_retry
14from aixplain .modules import Model
25from datetime import datetime
36from typing import Dict , List , Optional , Text , Union
@@ -23,6 +26,22 @@ def __init__(
2326 self .model = ModelFactory .get (model )
2427
2528
29+ class APIKeyUsageLimit :
30+ def __init__ (self , request_count : int , request_count_limit : int , token_count : int , token_count_limit : int ):
31+ """Get the usage limits of an API key
32+
33+ Args:
34+ request_count (int): number of requests made
35+ request_count_limit (int): limit of requests
36+ token_count (int): number of tokens used
37+ token_count_limit (int): limit of tokens
38+ """
39+ self .request_count = request_count
40+ self .request_count_limit = request_count_limit
41+ self .token_count = token_count
42+ self .token_count_limit = token_count_limit
43+
44+
2645class APIKey :
2746 def __init__ (
2847 self ,
@@ -54,7 +73,7 @@ def __init__(
5473 token_per_day = asset_limit ["tpd" ],
5574 request_per_minute = asset_limit ["rpm" ],
5675 request_per_day = asset_limit ["rpd" ],
57- model = asset_limit ["model " ],
76+ model = asset_limit ["assetId " ],
5877 )
5978 self .expires_at = expires_at
6079 self .access_key = access_key
@@ -110,17 +129,13 @@ def to_dict(self) -> Dict:
110129 "tpd" : asset_limit .token_per_day ,
111130 "rpm" : asset_limit .request_per_minute ,
112131 "rpd" : asset_limit .request_per_day ,
113- "model " : asset_limit .model .id ,
132+ "assetId " : asset_limit .model .id ,
114133 }
115134 )
116135 return payload
117136
118137 def delete (self ) -> None :
119138 """Delete an API key by its ID"""
120- import logging
121- from aixplain .utils import config
122- from aixplain .utils .file_utils import _request_with_retry
123-
124139 try :
125140 url = f"{ config .BACKEND_URL } { self .id }
126141 headers = {"Authorization" : f"Token { config .TEAM_API_KEY } , "Content-Type" : "application/json" }
@@ -132,3 +147,28 @@ def delete(self) -> None:
132147 message = "API Key Deletion Error: Make sure the API Key exists and you are the owner."
133148 logging .error (message )
134149 raise Exception (f"{ message } )
150+
151+ def get_usage (self , asset_id : Optional [Text ] = None ) -> APIKeyUsageLimit :
152+ """Get the usage limits of an API key"""
153+ try :
154+ url = f"{ config .BACKEND_URL } { self .id }
155+ headers = {"Authorization" : f"Token { config .TEAM_API_KEY } , "Content-Type" : "application/json" }
156+ logging .info (f"Start service for GET API Key Usage - { url } { headers } )
157+ if asset_id is not None :
158+ url += f"?assetId={ asset_id }
159+ r = _request_with_retry ("GET" , url , headers = headers )
160+ resp = r .json ()
161+ except Exception :
162+ message = "API Key Usage Error: Make sure the API Key exists and you are the owner."
163+ logging .error (message )
164+ raise Exception (f"{ message } )
165+
166+ if 200 <= r .status_code < 300 :
167+ return APIKeyUsageLimit (
168+ request_count = resp ["requestCount" ],
169+ request_count_limit = resp ["requestCountLimit" ],
170+ token_count = resp ["tokenCount" ],
171+ token_count_limit = resp ["tokenCountLimit" ],
172+ )
173+ else :
174+ raise Exception (f"API Key Usage Error: Failed to get usage. Error: { str (resp )} )
0 commit comments