228 lines
8.5 KiB
Python
Executable File
228 lines
8.5 KiB
Python
Executable File
from fastapi import APIRouter, HTTPException, Request, BackgroundTasks
|
|
from app.core.providers import get_provider_handler
|
|
from app.runtime import key_manager, dispatcher, stats_exporter
|
|
import httpx
|
|
import time
|
|
import logging
|
|
import traceback
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
@router.get("/v1/stats")
|
|
async def get_stats():
|
|
"""
|
|
Returns usage statistics for all API keys.
|
|
"""
|
|
stats = await dispatcher.get_all_key_stats()
|
|
return {"data": stats}
|
|
|
|
@router.get("/v1/usage")
|
|
async def get_usage():
|
|
"""
|
|
Returns overall usage summary.
|
|
"""
|
|
summary = await dispatcher.get_usage_summary()
|
|
return {"data": summary}
|
|
|
|
@router.post("/v1/stats/reset/rpd")
|
|
async def reset_today_rpd():
|
|
"""
|
|
Reset today's RPD counters in Redis and refresh dashboard output.
|
|
"""
|
|
day = dispatcher._get_day_str()
|
|
deleted_keys = await dispatcher.reset_today_rpd_usage()
|
|
stats_exporter.clear_daily_usage_day(day)
|
|
await stats_exporter.export_once()
|
|
return {
|
|
"success": True,
|
|
"day": day,
|
|
"deleted_keys": deleted_keys,
|
|
"message": f"Reset today's RPD counters for {deleted_keys} keys",
|
|
}
|
|
|
|
@router.post("/v1/keys/config/{config_id}/enable")
|
|
async def enable_config(config_id: str):
|
|
"""
|
|
Enable all keys for a specific config_id (auto-saves to config file).
|
|
"""
|
|
success = key_manager.set_config_enabled(config_id, True)
|
|
if success:
|
|
return {"success": True, "message": f"Config {config_id} enabled and saved"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Config {config_id} not found")
|
|
|
|
@router.post("/v1/keys/config/{config_id}/disable")
|
|
async def disable_config(config_id: str):
|
|
"""
|
|
Disable all keys for a specific config_id (auto-saves to config file).
|
|
"""
|
|
success = key_manager.set_config_enabled(config_id, False)
|
|
if success:
|
|
return {"success": True, "message": f"Config {config_id} disabled and saved"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Config {config_id} not found")
|
|
|
|
@router.post("/v1/keys/{key_id}/enable")
|
|
async def enable_key(key_id: str):
|
|
"""
|
|
Enable a specific key.
|
|
"""
|
|
success = key_manager.set_key_enabled(key_id, True)
|
|
if success:
|
|
return {"success": True, "message": f"Key {key_id} enabled"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Key {key_id} not found")
|
|
|
|
@router.post("/v1/keys/{key_id}/disable")
|
|
async def disable_key(key_id: str):
|
|
"""
|
|
Disable a specific key.
|
|
"""
|
|
success = key_manager.set_key_enabled(key_id, False)
|
|
if success:
|
|
return {"success": True, "message": f"Key {key_id} disabled"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Key {key_id} not found")
|
|
|
|
|
|
@router.post("/v1/keys/config/{config_id}/endpoint/{endpoint_idx}/enable")
|
|
async def enable_endpoint(config_id: str, endpoint_idx: int):
|
|
"""
|
|
Enable a specific endpoint for a config (auto-saves to config file).
|
|
"""
|
|
success = key_manager.set_endpoint_enabled(config_id, endpoint_idx, True)
|
|
if success:
|
|
return {"success": True, "message": f"Endpoint {endpoint_idx} of {config_id} enabled and saved"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Endpoint {endpoint_idx} of {config_id} not found")
|
|
|
|
|
|
@router.post("/v1/keys/config/{config_id}/endpoint/{endpoint_idx}/disable")
|
|
async def disable_endpoint(config_id: str, endpoint_idx: int):
|
|
"""
|
|
Disable a specific endpoint for a config (auto-saves to config file).
|
|
"""
|
|
success = key_manager.set_endpoint_enabled(config_id, endpoint_idx, False)
|
|
if success:
|
|
return {"success": True, "message": f"Endpoint {endpoint_idx} of {config_id} disabled and saved"}
|
|
else:
|
|
raise HTTPException(status_code=404, detail=f"Endpoint {endpoint_idx} of {config_id} not found")
|
|
|
|
@router.get("/v1/configs")
|
|
async def get_configs():
|
|
"""
|
|
Returns all config IDs.
|
|
"""
|
|
configs = []
|
|
for config_id in key_manager.get_config_ids():
|
|
keys = key_manager.get_keys_by_config(config_id)
|
|
if keys:
|
|
configs.append({
|
|
"config_id": config_id,
|
|
"model_name": keys[0].model_name,
|
|
"provider": keys[0].provider,
|
|
"key_count": len(keys),
|
|
"enabled": keys[0].enabled
|
|
})
|
|
return {"data": configs}
|
|
|
|
@router.post("/v1/chat/completions")
|
|
async def chat_completions(request: Request, background_tasks: BackgroundTasks):
|
|
"""
|
|
Unified entry point.
|
|
Expects JSON body with "model" field.
|
|
"""
|
|
try:
|
|
body = await request.json()
|
|
except:
|
|
logger.error(f"[Invalid JSON] detail={traceback.format_exc()}...")
|
|
raise HTTPException(status_code=400, detail="Invalid JSON")
|
|
|
|
model_name = body.get("model")
|
|
|
|
key, error_reason = await dispatcher.select_key(model_name)
|
|
if not key:
|
|
logger.error(f"[Key Not Found] model={model_name} | error={error_reason}...")
|
|
raise HTTPException(status_code=503, detail=error_reason)
|
|
|
|
if not model_name:
|
|
model_name = key.model_name
|
|
|
|
acquired = await dispatcher.acquire_lease(key)
|
|
if not acquired:
|
|
logger.error(f"[System busy (Concurrency Limit)] model={model_name} | provider={key.provider} | key={key.id}|owner={key.owner or 'N/A'}|key_prefix={key.key[:20] if key.key else 'N/A'}...")
|
|
raise HTTPException(status_code=503, detail="System busy (Concurrency Limit)")
|
|
|
|
try:
|
|
handler = get_provider_handler(key)
|
|
url = handler.get_url()
|
|
headers = handler.get_headers()
|
|
|
|
params = {k: v for k, v in body.items() if k not in ["model", "messages", "stream"]}
|
|
|
|
payload = handler.get_payload(body.get("messages", []), params)
|
|
|
|
timeout = httpx.Timeout(300.0, connect=10.0)
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
resp = await client.post(url, headers=headers, json=payload)
|
|
|
|
if resp.status_code != 200:
|
|
error_detail = resp.text[:500] if len(resp.text) > 500 else resp.text
|
|
logger.error(
|
|
f"[Request Failed] model={model_name} | key={key.id}|owner={key.owner or 'N/A'}|key_prefix={key.key[:20] if key.key else 'N/A'} | "
|
|
f"api_base={key.api_base} | "
|
|
f"status={resp.status_code} | response={error_detail}"
|
|
)
|
|
|
|
if resp.status_code == 429:
|
|
error_message = ""
|
|
try:
|
|
error_data = resp.json()
|
|
if isinstance(error_data, dict):
|
|
error_message = error_data.get("error", {}).get("message", "") or error_data.get("message", "")
|
|
except:
|
|
error_message = resp.text
|
|
await dispatcher.report_failure(key, is_rate_limit=True, error_message=error_message)
|
|
|
|
raise HTTPException(status_code=resp.status_code, detail=f"Provider Error: {error_detail}")
|
|
|
|
data = resp.json()
|
|
parsed_result = handler.parse_response(data)
|
|
|
|
total_tokens = parsed_result["usage"].get("total_tokens", 0)
|
|
background_tasks.add_task(dispatcher.record_usage, key, total_tokens)
|
|
|
|
response_data = {
|
|
"id": f"chatcmpl-{int(time.time())}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model_name,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": parsed_result["content"],
|
|
"reasoning_content": parsed_result.get("reasoning_content")
|
|
},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"usage": parsed_result["usage"]
|
|
}
|
|
|
|
return response_data
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
_key_info = f"{key.id}|owner={key.owner or 'N/A'}|key_prefix={key.key[:20] if key.key else 'N/A'}" if key else 'N/A'
|
|
logger.exception(
|
|
f"[Internal Error] model={model_name} | "
|
|
f"key={_key_info} | api_base={getattr(key, 'api_base', 'N/A') if key else 'N/A'} | error={str(e)}"
|
|
)
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
finally:
|
|
await dispatcher.release_lease(key)
|