128 lines
3.3 KiB
Python
128 lines
3.3 KiB
Python
import logging
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query, status
|
|
from pydantic import BaseModel
|
|
|
|
from app.api.deps import get_current_user
|
|
from app.models.user import User
|
|
from app.services.api_key_manager import APIKeyManager, KeySource, KeyStatus
|
|
from app.services.smart_router import ENGINE_COST_PROFILES
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
_key_manager: APIKeyManager = APIKeyManager()
|
|
|
|
|
|
def set_key_manager(mgr: APIKeyManager) -> None:
|
|
global _key_manager
|
|
_key_manager = mgr
|
|
|
|
|
|
def get_key_manager() -> APIKeyManager:
|
|
return _key_manager
|
|
|
|
|
|
class AddKeyRequest(BaseModel):
|
|
engine_type: str
|
|
api_key: str
|
|
source: str = "user"
|
|
|
|
|
|
class VerifyKeyRequest(BaseModel):
|
|
engine_type: str
|
|
|
|
|
|
@router.post("/")
|
|
async def add_key(
|
|
body: AddKeyRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
source = KeySource.USER if body.source == "user" else KeySource.SYSTEM
|
|
config = get_key_manager().add_key(
|
|
engine_type=body.engine_type,
|
|
api_key=body.api_key,
|
|
source=source,
|
|
user_id=str(current_user.id),
|
|
)
|
|
return {
|
|
"engine_type": config.engine_type,
|
|
"key_hint": config.key_hint,
|
|
"status": KeyStatus.ACTIVE.value,
|
|
}
|
|
|
|
|
|
@router.get("/")
|
|
async def list_keys(
|
|
engine_type: str | None = Query(None),
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
configs = get_key_manager().list_keys(engine_type=engine_type)
|
|
items = [
|
|
{
|
|
"engine_type": c.engine_type,
|
|
"key_hint": c.key_hint,
|
|
"source": c.key_source.value,
|
|
"status": c.status.value,
|
|
}
|
|
for c in configs
|
|
]
|
|
return {"items": items}
|
|
|
|
|
|
@router.post("/verify")
|
|
async def verify_key(
|
|
body: VerifyKeyRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
mgr = get_key_manager()
|
|
api_key = mgr.get_key(body.engine_type, user_id=str(current_user.id))
|
|
if not api_key:
|
|
api_key = mgr.get_key(body.engine_type)
|
|
if not api_key:
|
|
api_key = mgr.get_any_available_key(body.engine_type)
|
|
if not api_key:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"No API Key configured for {body.engine_type}",
|
|
)
|
|
key_status = await mgr.verify_key(body.engine_type, api_key)
|
|
return {
|
|
"engine_type": body.engine_type,
|
|
"status": key_status.value,
|
|
}
|
|
|
|
|
|
@router.get("/engines")
|
|
async def get_engines(
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
engines = [
|
|
{
|
|
"type": profile.engine_type,
|
|
"cost_tier": profile.cost_tier.value,
|
|
"has_free_tier": profile.has_free_tier,
|
|
"requires_own_key": profile.requires_own_key,
|
|
"input_price": profile.input_price_per_million,
|
|
"output_price": profile.output_price_per_million,
|
|
}
|
|
for profile in ENGINE_COST_PROFILES.values()
|
|
]
|
|
return {"engines": engines}
|
|
|
|
|
|
@router.delete("/{engine_type}/{key_hint}")
|
|
async def delete_key(
|
|
engine_type: str,
|
|
key_hint: str,
|
|
current_user: User = Depends(get_current_user),
|
|
):
|
|
deleted = get_key_manager().remove_key(engine_type, key_hint)
|
|
if not deleted:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail="API Key not found",
|
|
)
|
|
return {"deleted": True}
|