mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-15 22:30:38 +02:00
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
from fastapi import APIRouter
|
|
from fastapi import Depends
|
|
from sqlalchemy.orm import Session
|
|
|
|
from onyx.auth.users import current_admin_user
|
|
from onyx.db.engine import get_session
|
|
from onyx.db.models import User
|
|
from onyx.db.token_limit import delete_token_rate_limit
|
|
from onyx.db.token_limit import fetch_all_global_token_rate_limits
|
|
from onyx.db.token_limit import insert_global_token_rate_limit
|
|
from onyx.db.token_limit import update_token_rate_limit
|
|
from onyx.server.query_and_chat.token_limit import any_rate_limit_exists
|
|
from onyx.server.token_rate_limits.models import TokenRateLimitArgs
|
|
from onyx.server.token_rate_limits.models import TokenRateLimitDisplay
|
|
|
|
router = APIRouter(prefix="/admin/token-rate-limits")
|
|
|
|
|
|
"""
|
|
Global Token Limit Settings
|
|
"""
|
|
|
|
|
|
@router.get("/global")
|
|
def get_global_token_limit_settings(
|
|
_: User | None = Depends(current_admin_user),
|
|
db_session: Session = Depends(get_session),
|
|
) -> list[TokenRateLimitDisplay]:
|
|
return [
|
|
TokenRateLimitDisplay.from_db(token_rate_limit)
|
|
for token_rate_limit in fetch_all_global_token_rate_limits(db_session)
|
|
]
|
|
|
|
|
|
@router.post("/global")
|
|
def create_global_token_limit_settings(
|
|
token_limit_settings: TokenRateLimitArgs,
|
|
_: User | None = Depends(current_admin_user),
|
|
db_session: Session = Depends(get_session),
|
|
) -> TokenRateLimitDisplay:
|
|
rate_limit_display = TokenRateLimitDisplay.from_db(
|
|
insert_global_token_rate_limit(db_session, token_limit_settings)
|
|
)
|
|
# clear cache in case this was the first rate limit created
|
|
any_rate_limit_exists.cache_clear()
|
|
return rate_limit_display
|
|
|
|
|
|
"""
|
|
General Token Limit Settings
|
|
"""
|
|
|
|
|
|
@router.put("/rate-limit/{token_rate_limit_id}")
|
|
def update_token_limit_settings(
|
|
token_rate_limit_id: int,
|
|
token_limit_settings: TokenRateLimitArgs,
|
|
_: User | None = Depends(current_admin_user),
|
|
db_session: Session = Depends(get_session),
|
|
) -> TokenRateLimitDisplay:
|
|
return TokenRateLimitDisplay.from_db(
|
|
update_token_rate_limit(
|
|
db_session=db_session,
|
|
token_rate_limit_id=token_rate_limit_id,
|
|
token_rate_limit_settings=token_limit_settings,
|
|
)
|
|
)
|
|
|
|
|
|
@router.delete("/rate-limit/{token_rate_limit_id}")
|
|
def delete_token_limit_settings(
|
|
token_rate_limit_id: int,
|
|
_: User | None = Depends(current_admin_user),
|
|
db_session: Session = Depends(get_session),
|
|
) -> None:
|
|
return delete_token_rate_limit(
|
|
db_session=db_session,
|
|
token_rate_limit_id=token_rate_limit_id,
|
|
)
|