litellm-vector-store/app/routers/admin.py

from fastapi import APIRouter, Depends, HTTPException
from app.auth import verify_api_key
from app.database import get_db
import httpx
import os

router = APIRouter()

LITELLM_URL = os.getenv("LITELLM_PROXY_URL", "http://litellm:4000")
MASTER_KEY  = os.getenv("LITELLM_MASTER_KEY")
ADMIN_IDS   = os.getenv("ADMIN_USER_IDS", "").split(",")


# --- Admin-Check ---
async def require_admin(user: dict = Depends(verify_api_key)):
    if user["user_id"] not in ADMIN_IDS:
        raise HTTPException(403, "Admin-Zugriff erforderlich")
    return user


# --- Stats ---
@router.get("/stats")
async def get_stats(
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    stats = await db.fetchrow(
        """SELECT
               (SELECT COUNT(*) FROM vector_stores)       AS total_stores,
               (SELECT COUNT(*) FROM documents)           AS total_documents,
               (SELECT COUNT(DISTINCT owner_user_id)
                FROM vector_stores)                       AS total_users,
               (SELECT COUNT(*) FROM store_permissions)   AS total_permissions"""
    )
    return dict(stats)


# --- User Endpoints ---
@router.get("/users")
async def list_users(
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    rows = await db.fetch(
        """SELECT
               owner_user_id  AS user_id,
               COUNT(id)      AS store_count,
               MAX(created_at) AS last_activity
           FROM vector_stores
           GROUP BY owner_user_id
           ORDER BY last_activity DESC"""
    )
    return [dict(r) for r in rows]


@router.get("/users/{user_id}/stores")
async def get_user_stores(
    user_id: str,
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    rows = await db.fetch(
        """SELECT
               vs.id,
               vs.name,
               vs.created_at,
               COUNT(d.id) AS document_count
           FROM vector_stores vs
           LEFT JOIN documents d ON d.store_id = vs.id
           WHERE vs.owner_user_id = $1
           GROUP BY vs.id, vs.name, vs.created_at""",
        user_id
    )
    return [dict(r) for r in rows]


@router.delete("/users/{user_id}/stores/{store_id}")
async def admin_delete_store(
    user_id: str,
    store_id: str,
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    deleted = await db.fetchval(
        """DELETE FROM vector_stores
           WHERE id = $1 AND owner_user_id = $2
           RETURNING id""",
        store_id, user_id
    )
    if not deleted:
        raise HTTPException(404, "Store nicht gefunden")
    return {"deleted": str(deleted)}


# --- Permission Endpoints ---
@router.get("/stores/{store_id}/permissions")
async def get_permissions(
    store_id: str,
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    rows = await db.fetch(
        """SELECT user_id, permission, created_at
           FROM store_permissions
           WHERE store_id = $1""",
        store_id
    )
    return [dict(r) for r in rows]


@router.post("/stores/{store_id}/permissions")
async def grant_permission(
    store_id: str,
    user_id: str,
    permission: str = "read",
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    if permission not in ("read", "write", "admin"):
        raise HTTPException(400, "Ungültige Permission: read, write oder admin")

    await db.execute(
        """INSERT INTO store_permissions (store_id, user_id, permission)
           VALUES ($1, $2, $3)
           ON CONFLICT (store_id, user_id)
           DO UPDATE SET permission = $3""",
        store_id, user_id, permission
    )
    return {"granted": permission, "user_id": user_id}


@router.delete("/stores/{store_id}/permissions/{user_id}")
async def revoke_permission(
    store_id: str,
    user_id: str,
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    await db.execute(
        "DELETE FROM store_permissions WHERE store_id=$1 AND user_id=$2",
        store_id, user_id
    )
    return {"revoked": user_id}


# --- Key Management ---
@router.post("/users/{user_id}/rotate-key")
async def rotate_key(
    user_id: str,
    admin=Depends(require_admin),
    db=Depends(get_db)
):
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            f"{LITELLM_URL}/key/generate",
            headers={"Authorization": f"Bearer {MASTER_KEY}"},
            json={
                "user_id":   user_id,
                "key_alias": f"{user_id}-rotated"
            }
        )
    if resp.status_code != 200:
        raise HTTPException(500, "Key-Rotation fehlgeschlagen")

    store_count = await db.fetchval(
        "SELECT COUNT(*) FROM vector_stores WHERE owner_user_id=$1",
        user_id
    )
    return {
        "new_key":          resp.json()["key"],
        "user_id":          user_id,
        "stores_preserved": store_count
    }

@router.get("/verify")
async def verify_admin(
    admin=Depends(require_admin),
):
    """
    Prüft ob der API Key Admin-Rechte hat.
    Gibt 200 zurück wenn Admin, 403 wenn nicht.
    """
    return {
        "admin":   True,
        "user_id": admin["user_id"],
    }