Files
ai-gateway/app/api/admin_backend.py

121 lines
3.4 KiB
Python

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from app.core.database import get_db
from app.core.security import get_current_user
from app.models.module import Module
import secrets
from pydantic import BaseModel
from datetime import datetime
from typing import List, Optional
router = APIRouter()
class ModuleCreate(BaseModel):
name: str
program: Optional[str] = None
lob: Optional[str] = None
job_code: Optional[str] = None
class ModuleResponse(BaseModel):
id: int
name: str
secret_key: str
program: Optional[str]
lob: Optional[str]
job_code: Optional[str]
is_active: bool
created_at: datetime
last_rotated_at: datetime
ingress_tokens: int
egress_tokens: int
total_tokens: int
class Config:
from_attributes = True
@router.post("/modules", response_model=ModuleResponse)
def create_module(
module_in: ModuleCreate,
db: Session = Depends(get_db),
current_user: str = Depends(get_current_user)
):
db_module = db.query(Module).filter(Module.name == module_in.name).first()
if db_module:
raise HTTPException(status_code=400, detail="Module already exists")
new_key = secrets.token_hex(32)
db_module = Module(
name=module_in.name,
secret_key=new_key,
program=module_in.program,
lob=module_in.lob,
job_code=module_in.job_code
)
db.add(db_module)
db.commit()
db.refresh(db_module)
return db_module
@router.get("/modules", response_model=List[ModuleResponse])
def get_modules(
include_archived: bool = False,
db: Session = Depends(get_db),
current_user: str = Depends(get_current_user)
):
query = db.query(Module)
if not include_archived:
query = query.filter(Module.is_active == True)
return query.all()
@router.post("/modules/{module_id}/rotate", response_model=ModuleResponse)
def rotate_module_key(
module_id: int,
db: Session = Depends(get_db),
current_user: str = Depends(get_current_user)
):
db_module = db.query(Module).filter(Module.id == module_id).first()
if not db_module:
raise HTTPException(status_code=404, detail="Module not found")
db_module.secret_key = secrets.token_hex(32)
db_module.last_rotated_at = datetime.utcnow()
db.commit()
db.refresh(db_module)
return db_module
@router.delete("/modules/{module_id}")
def delete_module(
module_id: int,
hard_delete: bool = False,
db: Session = Depends(get_db),
current_user: str = Depends(get_current_user)
):
db_module = db.query(Module).filter(Module.id == module_id).first()
if not db_module:
raise HTTPException(status_code=404, detail="Module not found")
if hard_delete:
db.delete(db_module)
message = "Module permanently deleted"
else:
db_module.is_active = False
message = "Module archived"
db.commit()
return {"status": "success", "message": message}
@router.post("/modules/{module_id}/restore", response_model=ModuleResponse)
def restore_module(
module_id: int,
db: Session = Depends(get_db),
current_user: str = Depends(get_current_user)
):
db_module = db.query(Module).filter(Module.id == module_id).first()
if not db_module:
raise HTTPException(status_code=404, detail="Module not found")
db_module.is_active = True
db.commit()
db.refresh(db_module)
return db_module