90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
from ...utils.http_client_async import AsyncHttpClient
|
|
from ...utils.error_handler import handle_response_error
|
|
from ...types import ConcurrencyCheck, CreditUsage, TokenUsage, CreditUsageHistoricalResponse, TokenUsageHistoricalResponse, QueueStatusResponse
|
|
|
|
|
|
async def get_concurrency(client: AsyncHttpClient) -> ConcurrencyCheck:
|
|
resp = await client.get("/v2/concurrency-check")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get concurrency")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
data = body.get("data", body)
|
|
return ConcurrencyCheck(
|
|
concurrency=data.get("concurrency"),
|
|
max_concurrency=data.get("maxConcurrency", data.get("max_concurrency")),
|
|
)
|
|
|
|
|
|
async def get_credit_usage(client: AsyncHttpClient) -> CreditUsage:
|
|
resp = await client.get("/v2/team/credit-usage")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get credit usage")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
data = body.get("data", body)
|
|
return CreditUsage(
|
|
remaining_credits=data.get("remainingCredits", data.get("remaining_credits", 0)),
|
|
plan_credits=data.get("planCredits", data.get("plan_credits")),
|
|
billing_period_start=data.get("billingPeriodStart", data.get("billing_period_start")),
|
|
billing_period_end=data.get("billingPeriodEnd", data.get("billing_period_end")),
|
|
)
|
|
|
|
|
|
async def get_token_usage(client: AsyncHttpClient) -> TokenUsage:
|
|
resp = await client.get("/v2/team/token-usage")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get token usage")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
data = body.get("data", body)
|
|
return TokenUsage(
|
|
remaining_tokens=data.get("remainingTokens", data.get("remaining_tokens", 0)),
|
|
plan_tokens=data.get("planTokens", data.get("plan_tokens")),
|
|
billing_period_start=data.get("billingPeriodStart", data.get("billing_period_start")),
|
|
billing_period_end=data.get("billingPeriodEnd", data.get("billing_period_end")),
|
|
)
|
|
|
|
|
|
async def get_queue_status(client: AsyncHttpClient) -> QueueStatusResponse:
|
|
resp = await client.get("/v2/team/queue-status")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get queue status")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
data = body.get("data", body)
|
|
return QueueStatusResponse(
|
|
jobs_in_queue=data.get("jobsInQueue", 0),
|
|
active_jobs_in_queue=data.get("activeJobsInQueue", 0),
|
|
waiting_jobs_in_queue=data.get("waitingJobsInQueue", 0),
|
|
max_concurrency=data.get("maxConcurrency", 0),
|
|
most_recent_success=data.get("mostRecentSuccess", None),
|
|
)
|
|
|
|
|
|
async def get_credit_usage_historical(client: AsyncHttpClient, by_api_key: bool = False) -> CreditUsageHistoricalResponse:
|
|
query = "?byApiKey=true" if by_api_key else ""
|
|
resp = await client.get(f"/v2/team/credit-usage/historical{query}")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get credit usage historical")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
return CreditUsageHistoricalResponse(**body)
|
|
|
|
|
|
async def get_token_usage_historical(client: AsyncHttpClient, by_api_key: bool = False) -> TokenUsageHistoricalResponse:
|
|
query = "?byApiKey=true" if by_api_key else ""
|
|
resp = await client.get(f"/v2/team/token-usage/historical{query}")
|
|
if resp.status_code >= 400:
|
|
handle_response_error(resp, "get token usage historical")
|
|
body = resp.json()
|
|
if not body.get("success"):
|
|
raise Exception(body.get("error", "Unknown error"))
|
|
return TokenUsageHistoricalResponse(**body)
|
|
|