mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-23 11:53:08 -05:00
Alan Hagedorn
d7a704c0d9
WIP Cleanup 🧹 Remove existing rate limiting logic Cleanup 🧼 Undo nit Cleanup 🧽 Move db constants (avoids circular import) WIP WIP Cleanup Lint Resolve alembic conflict Fix mypy Add backfill to migration Update comment Make unauthenticated users still adhere to global limits Use Depends Remove enum from table Update migration error handling + deletion Address verbal feedback, cleanup urls, minor nits
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 danswer.auth.users import current_admin_user
|
|
from danswer.db.engine import get_session
|
|
from danswer.db.models import User
|
|
from danswer.server.query_and_chat.token_limit import any_rate_limit_exists
|
|
from danswer.server.token_rate_limits.models import TokenRateLimitArgs
|
|
from danswer.server.token_rate_limits.models import TokenRateLimitDisplay
|
|
from ee.danswer.db.token_limit import delete_token_rate_limit
|
|
from ee.danswer.db.token_limit import fetch_all_global_token_rate_limits
|
|
from ee.danswer.db.token_limit import insert_global_token_rate_limit
|
|
from ee.danswer.db.token_limit import update_token_rate_limit
|
|
|
|
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,
|
|
)
|