325 lines
11 KiB
Python
325 lines
11 KiB
Python
from fastapi import FastAPI, Request, HTTPException, Depends, Form
|
||
from fastapi.responses import HTMLResponse, RedirectResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.templating import Jinja2Templates
|
||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials, HTTPBasic, HTTPBasicCredentials
|
||
from starlette.middleware.sessions import SessionMiddleware
|
||
import os
|
||
from loguru import logger
|
||
import sys
|
||
import uvicorn
|
||
import json
|
||
import aiosqlite
|
||
import time
|
||
from datetime import datetime, timedelta
|
||
from typing import Optional, List, Dict, Any
|
||
from contextlib import asynccontextmanager
|
||
from pydantic import BaseModel
|
||
import secrets
|
||
|
||
listen_port = os.getenv("LISTEN_PORT", "8000")
|
||
listen_host = os.getenv("LISTEN_HOST", "0.0.0.0")
|
||
config_path = os.getenv("CONFIG_PATH", "./config.json")
|
||
database_path = os.getenv("DATABASE_PATH", "./data.sqlite3")
|
||
logging_level = os.getenv("LOGGING_LEVEL", "TRACE")
|
||
|
||
# Configure loguru properly - remove default handler and add new one with correct level
|
||
logger.remove() # Remove the default handler
|
||
logger.add(sys.stderr, level=logging_level) # Add new handler with specified level
|
||
|
||
security_bearer = HTTPBearer()
|
||
security_basic = HTTPBasic()
|
||
|
||
class APICallRecord(BaseModel):
|
||
timestamp: int
|
||
model_id: str
|
||
user_email: str
|
||
input_tokens: int
|
||
output_tokens: int
|
||
cost_usd: float
|
||
|
||
def normalize_path(path: str) -> str:
|
||
"""
|
||
处理路径，支持绝对路径、用户目录和相对路径
|
||
|
||
Args:
|
||
path: 输入路径
|
||
|
||
Returns:
|
||
处理后的路径
|
||
"""
|
||
if path.startswith("/") or path.startswith("~"):
|
||
# 绝对路径，如果是~开头则展开为用户目录
|
||
return os.path.expanduser(path)
|
||
elif path.startswith("./"):
|
||
# 相对于脚本本身的相对路径
|
||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||
return os.path.join(script_dir, path[2:]) # 去掉'./'前缀
|
||
else:
|
||
# 其他情况保持原样（相对于pwd的相对路径）
|
||
return path
|
||
|
||
# 处理config_path和database_path路径
|
||
config_path = normalize_path(config_path)
|
||
database_path = normalize_path(database_path)
|
||
|
||
with open(config_path, "r") as f:
|
||
config = json.load(f)
|
||
|
||
admin_api_key = config.get("admin_api_key", "")
|
||
report_api_key = config.get("report_api_key", "")
|
||
session_secret_key = config.get("session_secret_key", secrets.token_hex(32))
|
||
|
||
async def init_database():
|
||
"""Initialize the database and create tables if they don't exist"""
|
||
async with aiosqlite.connect(database_path) as db:
|
||
await db.execute("""
|
||
CREATE TABLE IF NOT EXISTS api_calls (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
timestamp INTEGER NOT NULL,
|
||
model_id TEXT NOT NULL,
|
||
user_email TEXT NOT NULL,
|
||
input_tokens INTEGER NOT NULL,
|
||
output_tokens INTEGER NOT NULL,
|
||
cost_usd REAL NOT NULL,
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
await db.commit()
|
||
logger.info("Database initialized successfully")
|
||
|
||
async def verify_admin_session(request: Request):
|
||
"""Verify admin session for dashboard access"""
|
||
if not request.session.get("authenticated"):
|
||
raise HTTPException(
|
||
status_code=302,
|
||
headers={"Location": "/login"}
|
||
)
|
||
return True
|
||
|
||
async def verify_report_key(credentials: HTTPAuthorizationCredentials = Depends(security_bearer)):
|
||
"""Verify report API key for API access using Bearer authentication"""
|
||
if credentials.credentials != report_api_key:
|
||
raise HTTPException(status_code=403, detail="Invalid report API key")
|
||
return credentials.credentials
|
||
|
||
async def get_24h_stats():
|
||
"""Get statistics for the past 24 hours"""
|
||
twenty_four_hours_ago = int(time.time()) - (24 * 60 * 60)
|
||
|
||
async with aiosqlite.connect(database_path) as db:
|
||
# Total chats (assuming each API call is a chat)
|
||
cursor = await db.execute(
|
||
"SELECT COUNT(*) FROM api_calls WHERE timestamp >= ?",
|
||
(twenty_four_hours_ago,)
|
||
)
|
||
total_chats = (await cursor.fetchone())[0]
|
||
|
||
# Total tokens
|
||
cursor = await db.execute(
|
||
"SELECT SUM(input_tokens + output_tokens) FROM api_calls WHERE timestamp >= ?",
|
||
(twenty_four_hours_ago,)
|
||
)
|
||
result = await cursor.fetchone()
|
||
total_tokens = result[0] if result[0] is not None else 0
|
||
|
||
# Total cost
|
||
cursor = await db.execute(
|
||
"SELECT SUM(cost_usd) FROM api_calls WHERE timestamp >= ?",
|
||
(twenty_four_hours_ago,)
|
||
)
|
||
result = await cursor.fetchone()
|
||
total_cost = result[0] if result[0] is not None else 0.0
|
||
|
||
return {
|
||
"total_chats": total_chats,
|
||
"total_tokens": total_tokens,
|
||
"total_cost": round(total_cost, 6)
|
||
}
|
||
|
||
async def get_recent_logs(page: int = 1, page_size: int = 100):
|
||
"""Get recent API call logs with pagination"""
|
||
offset = (page - 1) * page_size
|
||
|
||
async with aiosqlite.connect(database_path) as db:
|
||
# Get total count for pagination
|
||
cursor = await db.execute("SELECT COUNT(*) FROM api_calls")
|
||
total_count = (await cursor.fetchone())[0]
|
||
|
||
# Get paginated results
|
||
cursor = await db.execute("""
|
||
SELECT timestamp, model_id, user_email, input_tokens, output_tokens, cost_usd
|
||
FROM api_calls
|
||
ORDER BY timestamp DESC
|
||
LIMIT ? OFFSET ?
|
||
""", (page_size, offset))
|
||
|
||
logs = []
|
||
async for row in cursor:
|
||
logs.append({
|
||
"timestamp": row[0],
|
||
"model_id": row[1],
|
||
"user_email": row[2],
|
||
"input_tokens": row[3],
|
||
"output_tokens": row[4],
|
||
"total_tokens": row[3] + row[4],
|
||
"cost_usd": round(row[5], 6)
|
||
})
|
||
|
||
total_pages = (total_count + page_size - 1) // page_size # Ceiling division
|
||
|
||
return {
|
||
"logs": logs,
|
||
"pagination": {
|
||
"current_page": page,
|
||
"total_pages": total_pages,
|
||
"total_count": total_count,
|
||
"page_size": page_size,
|
||
"has_prev": page > 1,
|
||
"has_next": page < total_pages
|
||
}
|
||
}
|
||
|
||
async def get_users_report():
|
||
"""Get user consumption statistics"""
|
||
async with aiosqlite.connect(database_path) as db:
|
||
cursor = await db.execute("""
|
||
SELECT
|
||
user_email,
|
||
COUNT(*) as total_calls,
|
||
SUM(input_tokens) as total_input_tokens,
|
||
SUM(output_tokens) as total_output_tokens,
|
||
SUM(input_tokens + output_tokens) as total_tokens,
|
||
SUM(cost_usd) as total_cost
|
||
FROM api_calls
|
||
GROUP BY user_email
|
||
ORDER BY total_cost DESC
|
||
""")
|
||
|
||
users = []
|
||
async for row in cursor:
|
||
users.append({
|
||
"email": row[0],
|
||
"total_calls": row[1],
|
||
"total_input_tokens": row[2],
|
||
"total_output_tokens": row[3],
|
||
"total_tokens": row[4],
|
||
"total_cost": round(row[5], 6)
|
||
})
|
||
|
||
return users
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
"""Lifespan context manager for startup and shutdown events"""
|
||
# Startup
|
||
await init_database()
|
||
logger.info("Application started")
|
||
|
||
yield
|
||
|
||
# Shutdown
|
||
logger.info("Application shutting down")
|
||
|
||
app = FastAPI(lifespan=lifespan)
|
||
|
||
# Add session middleware
|
||
app.add_middleware(SessionMiddleware, secret_key=session_secret_key)
|
||
|
||
# Set up templates
|
||
templates = Jinja2Templates(directory="templates")
|
||
|
||
# Routes
|
||
@app.get("/")
|
||
async def root(request: Request):
|
||
"""Redirect to dashboard or login"""
|
||
if request.session.get("authenticated"):
|
||
return RedirectResponse(url="/dashboard", status_code=302)
|
||
else:
|
||
return RedirectResponse(url="/login", status_code=302)
|
||
|
||
@app.get("/dashboard", response_class=HTMLResponse)
|
||
async def dashboard(request: Request, _: bool = Depends(verify_admin_session)):
|
||
"""Main dashboard page"""
|
||
stats = await get_24h_stats()
|
||
return templates.TemplateResponse("dashboard.html", {
|
||
"request": request,
|
||
"stats": stats
|
||
})
|
||
|
||
@app.get("/dashboard/logs", response_class=HTMLResponse)
|
||
async def dashboard_logs(request: Request, page: int = 1, _: bool = Depends(verify_admin_session)):
|
||
"""Dashboard logs page"""
|
||
if page < 1:
|
||
page = 1
|
||
|
||
result = await get_recent_logs(page=page, page_size=100)
|
||
return templates.TemplateResponse("logs.html", {
|
||
"request": request,
|
||
"logs": result["logs"],
|
||
"pagination": result["pagination"]
|
||
})
|
||
|
||
@app.get("/dashboard/users_report", response_class=HTMLResponse)
|
||
async def dashboard_users_report(request: Request, _: bool = Depends(verify_admin_session)):
|
||
"""Dashboard users report page"""
|
||
users = await get_users_report()
|
||
return templates.TemplateResponse("users_report.html", {
|
||
"request": request,
|
||
"users": users
|
||
})
|
||
|
||
@app.post("/api/record_api_call")
|
||
async def record_api_call(
|
||
record: APICallRecord,
|
||
_: str = Depends(verify_report_key)
|
||
):
|
||
"""Record an API call"""
|
||
try:
|
||
async with aiosqlite.connect(database_path) as db:
|
||
await db.execute("""
|
||
INSERT INTO api_calls (timestamp, model_id, user_email, input_tokens, output_tokens, cost_usd)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""", (
|
||
record.timestamp,
|
||
record.model_id,
|
||
record.user_email,
|
||
record.input_tokens,
|
||
record.output_tokens,
|
||
record.cost_usd
|
||
))
|
||
await db.commit()
|
||
|
||
logger.info(f"Recorded API call: {record.user_email} - {record.model_id} - ${record.cost_usd}")
|
||
return {"status": "success", "message": "API call recorded successfully"}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error recording API call: {e}")
|
||
raise HTTPException(status_code=500, detail="Failed to record API call")
|
||
|
||
@app.get("/login", response_class=HTMLResponse)
|
||
async def login_page(request: Request):
|
||
"""Login page"""
|
||
return templates.TemplateResponse("login.html", {"request": request})
|
||
|
||
@app.post("/login")
|
||
async def login(request: Request, password: str = Form(...)):
|
||
"""Handle login form submission"""
|
||
if secrets.compare_digest(password, admin_api_key):
|
||
request.session["authenticated"] = True
|
||
return RedirectResponse(url="/dashboard", status_code=302)
|
||
else:
|
||
return templates.TemplateResponse("login.html", {
|
||
"request": request,
|
||
"error": "Invalid password"
|
||
})
|
||
|
||
@app.get("/logout")
|
||
async def logout(request: Request):
|
||
"""Logout and clear session"""
|
||
request.session.clear()
|
||
return RedirectResponse(url="/login", status_code=302)
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host=listen_host, port=int(listen_port))
|