feat: add bot token usage analyzer

server/auth/clients/auth0.py

@@ -67,7 +67,6 @@ async def get_user_info(self, request: Request) -> dict:
                 "picture": user_info.get("picture"),
                 "sub": user_info["sub"],
                 "sid": secrets.token_urlsafe(32),
-                "agreement_accepted": user_info.get("agreement_accepted"),
             }
             return data
         else:

server/auth/clients/local.py

@@ -34,7 +34,6 @@ async def get_user_info(self, user_id):
       "name": username,
       "picture": f"https://picsum.photos/seed/{seed}/100/100",
       "sid": secrets.token_urlsafe(32),
-      "agreement_accepted": False,
     }
 
 

server/auth/middleware.py

@@ -51,37 +51,37 @@ async def oauth(self, request: Request):
             return False
 
     async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
-        try:
-            if ENVIRONMENT == "development":
-                return await call_next(request)
-
-            # Auth 相关的直接放过
-            if request.url.path.startswith("/api/auth"):
-                return await call_next(request)
+
+        if ENVIRONMENT == "development":
+            return await call_next(request)
 
-            if request.url.path in ALLOW_LIST:
-                return await call_next(request)
+        # Auth 相关的直接放过
+        if request.url.path.startswith("/api/auth"):
+            return await call_next(request)
 
-            if await self.oauth(request=request):
-                return await call_next(request)
+        if request.url.path in ALLOW_LIST:
+            return await call_next(request)
 
-            # 获取 session 中的用户信息
-            user = request.session.get("user")
-            if not user:
-                raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Unauthorized")
+        if await self.oauth(request=request):
+            return await call_next(request)
 
-            if user['sub'].startswith("client|"):
-                if request.url.path in ANONYMOUS_USER_ALLOW_LIST:
-                    return await call_next(request)
-                else:
-                    # 如果没有用户信息，返回 401 Unauthorized 错误
-                    raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Anonymous User Not Allow")
+        # 获取 session 中的用户信息
+        user = request.session.get("user")
+        if not user:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Unauthorized")
 
-            return await call_next(request)
-        except HTTPException as e:
-            print(traceback.format_exception(e))
-            # 处理 HTTP 异常
-            return JSONResponse(status_code=e.status_code, content={"detail": e.detail})
-        except Exception as e:
-            # 处理其他异常
-            return JSONResponse(status_code=500, content={"detail": f"Internal Server Error: {e}"})
+        if user['sub'].startswith("client|"):
+            if request.url.path in ANONYMOUS_USER_ALLOW_LIST:
+                return await call_next(request)
+            else:
+                # 如果没有用户信息，返回 401 Unauthorized 错误
+                raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Anonymous User Not Allow")
+
+        return await call_next(request)
+        # except HTTPException as e:
+        #     print(traceback.format_exception(e))
+        #     # 处理 HTTP 异常
+        #     return JSONResponse(status_code=e.status_code, content={"detail": e.detail})
+        # except Exception as e:
+        #     # 处理其他异常
+        #     return JSONResponse(status_code=500, content={"detail": f"Internal Server Error: {e}"})

server/auth/router.py

@@ -43,10 +43,19 @@ async def logout(request: Request):
 @router.get("/callback")
 async def callback(request: Request, auth_client: BaseAuthClient = Depends(get_auth_client)):
     user_info = await auth_client.get_user_info(request)
+    profiles_dao = ProfilesDAO()
+    profile = profiles_dao.get_profile(user_info['id'])
+
     if user_info:
-        request.session["user"] = dict(user_info)
+        upsert_user = {
+            **user_info,
+            'agreement_accepted': profile['agreement_accepted'],
+            'is_admin': profile['is_admin'],
+        }
+
+        request.session["user"] = dict(upsert_user)
         supabase = get_client()
-        supabase.table("profiles").upsert(user_info).execute()
+        supabase.table("profiles").upsert(upsert_user).execute()
     return RedirectResponse(url=f"{WEB_LOGIN_SUCCESS_URL}", status_code=302)
 
 

server/auth/verify_admin.py

@@ -0,0 +1,11 @@
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, Request
+
+from auth.get_user_info import get_user
+from core.models.user import User
+
+async def verify_admin(request: Request, user: Annotated[User | None, Depends(get_user)] = None,):
+    if not user or not user.is_admin:
+      raise HTTPException(status_code=403, detail="Must Login")
+

server/core/dao/profilesDAO.py

@@ -11,6 +11,15 @@ def __init__(self):
         super().__init__()
         self.client = get_client()
 
+    def get_profile(self, user_id: str):
+        resp = (
+            self.client.table("profiles")
+            .select('id, is_admin, agreement_accepted')
+            .eq('id', user_id)
+            .execute()
+        )
+        return resp.data[0]
+
     def get_agreement_status(self, user_id: str):
 
         resp = (

server/core/models/user.py

@@ -9,6 +9,6 @@ class User(BaseModel):
     avatar: Optional[str] = None
     picture: Optional[str]
     agreement_token: Optional[bool] = False
-
+    is_admin: Optional[bool] = False
     anonymous: Optional[bool] = True
     access_token: Optional[str] = None

server/user/router.py

@@ -6,6 +6,7 @@
 
 from agent.llm import get_registered_llm_client, import_clients
 from auth.get_user_info import get_user_id
+from auth.verify_admin import verify_admin
 from core.service.user_llm_token import UserLLMTokenService, CreateUserLLMTokenVO, get_llm_token_service
 from core.service.user_token_usage import UserTokenUsageService, get_user_token_usage_service
 
@@ -67,23 +68,23 @@ def token_usage(
 ):
   return user_token_usage_service.usage_stats(user_id=user_id, start_date=start_date, end_date=end_date)
 
-@router.get("/llm_token_usages/analyzer")
+@router.get("/llm_token_usages/analyzer", dependencies=[Depends(verify_admin)])
 def token_usage_analyze(
   start_date: datetime = datetime.now() - timedelta(days=30),
   end_date: datetime = datetime.now(),
   user_token_usage_service: Annotated[UserTokenUsageService | None, Depends(get_user_token_usage_service)] = None,
 ):
   return user_token_usage_service.analyze_token_usage(start_date=start_date, end_date=end_date)
 
-@router.get("/llm_token_usages/top_bots")
+@router.get("/llm_token_usages/top_bots", dependencies=[Depends(verify_admin)])
 def top_used_bots(
   start_date: datetime = datetime.now() - timedelta(days=30),
   end_date: datetime = datetime.now(),
   user_token_usage_service: Annotated[UserTokenUsageService | None, Depends(get_user_token_usage_service)] = None,
 ):
   return user_token_usage_service.top_bots(start_date=start_date, end_date=end_date)
 
-@router.get("/llm_token_usages/top_users")
+@router.get("/llm_token_usages/top_users", dependencies=[Depends(verify_admin)])
 def top_used_users(
   start_date: datetime = datetime.now() - timedelta(days=30),
   end_date: datetime = datetime.now(),