Update app.py

app.py

@@ -9,149 +9,168 @@ from fastapi import FastAPI, Request, HTTPException, Depends, status
 from fastapi.responses import StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
-# --- 1. 初始化 FastAPI 与 安全验证 ---
+# --- 1. 配置与全局状态 ---
 app = FastAPI()
 security = HTTPBearer()
 
-# 建议在 HF Secret 中设置 ACCESS_TOKEN
 ACCESS_TOKEN = os.getenv("ACCESS_TOKEN", "sk-admin-123456")
 API_KEY = "AIzaSyAZaD22Mzi9HkTcW3ErNxRA_sNEFolLBCA"
 SUPPORTED_MODELS = ["auto", "gpt-5-mini", "gemini-3-flash", "pro"]
 
-async def verify_token(auth: HTTPAuthorizationCredentials = Depends(security)):
-    if auth.credentials != ACCESS_TOKEN:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid API Key",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    return auth.credentials
+# 全局缓存：存储当前可用的 Token 及其余额
+class TokenCache:
+    def __init__(self):
+        self.token = ""
+        self.balance = 0
+
+cache = TokenCache()
 
-# --- 2. 核心后端逻辑 ---
-def get_jwt():
-    url = f"https://identitytoolkit.googleapis.com/v1/accounts:signUp?key={API_KEY}"
+# --- 2. 核心逻辑函数 ---
+
+def fetch_new_account():
+    """注册新账号并初始化缓存"""
+    reg_url = f"https://identitytoolkit.googleapis.com/v1/accounts:signUp?key={API_KEY}"
     try:
-        resp = requests.post(url, json={"returnSecureToken": True}, headers={"origin": "https://www.aidocmaker.com"}, timeout=10)
-        return resp.json().get("idToken")
+        # 注册
+        reg_resp = requests.post(reg_url, json={"returnSecureToken": True}, headers={"origin": "https://www.aidocmaker.com"}, timeout=10).json()
+        token = reg_resp.get("idToken")
+        if not token: return None
+        
+        # 查询初始余额
+        bal_url = "https://level2labs-prod--adm-agent-helper-user-get-credits.modal.run/"
+        bal_resp = requests.post(bal_url, headers={"Authorization": f"Bearer {token}"}, timeout=10).json()
+        
+        cache.token = token
+        cache.balance = bal_resp.get("premium_quota", 0)
+        return token
     except:
         return None
 
-def generate_id():
-    return ''.join(random.choice('0123456789abcdef') for _ in range(32))
+def get_valid_token():
+    """获取有效 Token：如果有缓存且有余额则复用，否则获取新的"""
+    if cache.token and cache.balance > 0:
+        return cache.token
+    return fetch_new_account()
+
+def sync_balance():
+    """同步当前 Token 的实际余额"""
+    if not cache.token: return 0
+    url = "https://level2labs-prod--adm-agent-helper-user-get-credits.modal.run/"
+    try:
+        data = requests.post(url, headers={"Authorization": f"Bearer {cache.token}"}).json()
+        cache.balance = data.get("premium_quota", 0)
+        return cache.balance
+    except:
+        cache.balance = 0
+        return 0
 
-# --- 3. OpenAI 兼容接口 (/v1/...) ---
+# --- 3. OpenAI 兼容接口 ---
+
+async def verify_api_token(auth: HTTPAuthorizationCredentials = Depends(security)):
+    if auth.credentials != ACCESS_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+    return auth.credentials
 
 @app.get("/v1/models")
-async def list_models(token: str = Depends(verify_token)):
-    models_data = [{"id": m, "object": "model", "created": int(time.time()), "owned_by": "aidocmaker"} for m in SUPPORTED_MODELS]
-    return {"object": "list", "data": models_data}
+async def list_models(token: str = Depends(verify_api_token)):
+    return {"object": "list", "data": [{"id": m, "object": "model"} for m in SUPPORTED_MODELS]}
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: Request, token: str = Depends(verify_token)):
+async def chat_completions(request: Request, token: str = Depends(verify_api_token)):
     body = await request.json()
-    messages = body.get("messages", [])
     model = body.get("model", "auto")
-    stream = body.get("stream", False)
-    content = messages[-1]["content"] if messages else ""
-
-    jwt = get_jwt()
-    if not jwt: raise HTTPException(status_code=500, detail="Failed to get JWT")
+    content = body.get("messages", [])[-1]["content"] if body.get("messages") else ""
+    
+    # 自动获取/刷新 Token
+    jwt = get_valid_token()
+    if not jwt: raise HTTPException(status_code=500, detail="Provider Error")
 
     url = "https://level2labs-prod--adm-agent-send-chat-message.modal.run/"
-    headers = {"Authorization": f"Bearer {jwt}"}
     form_data = {
-        "content": (None, content),
-        "model": (None, model),
-        "conversation_id": (None, generate_id()),
-        "message_id": (None, generate_id()),
-        "max_mode": (None, "true")
+        "content": (None, content), "model": (None, model),
+        "conversation_id": (None, f"conv-{uuid.uuid4()}"),
+        "message_id": (None, f"msg-{uuid.uuid4()}"), "max_mode": (None, "true")
     }
 
-    target_resp = requests.post(url, files=form_data, headers=headers, stream=True, timeout=60)
+    target_resp = requests.post(url, files=form_data, headers={"Authorization": f"Bearer {jwt}"}, stream=True)
 
     def generate():
-        chat_id = f"chatcmpl-{uuid.uuid4()}"
         for line in target_resp.iter_lines():
             if line:
                 try:
-                    raw_data = json.loads(line.decode("utf-8"))
-                    chunk = {
-                        "id": chat_id, "object": "chat.completion.chunk", "created": int(time.time()),
-                        "model": model, "choices": [{"index": 0, "delta": {"content": raw_data.get("data", "")}, "finish_reason": None}]
-                    }
-                    yield f"data: {json.dumps(chunk)}\n\n"
+                    raw = json.loads(line.decode("utf-8"))
+                    yield f"data: {json.dumps({'choices': [{'delta': {'content': raw.get('data', '')}}]})}\n\n"
                 except: continue
         yield "data: [DONE]\n\n"
+        # 响应结束后异步同步一下余额（简单处理，实际可根据响应次数扣减）
+        sync_balance()
 
-    if stream:
-        return StreamingResponse(generate(), media_type="text/event-stream")
-    else:
-        full_text = "".join([json.loads(l.decode("utf-8")).get("data", "") for l in target_resp.iter_lines() if l])
-        return {
-            "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(time.time()),
-            "model": model, "choices": [{"index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": "stop"}]
-        }
+    return StreamingResponse(generate(), media_type="text/event-stream")
 
-# --- 4. Gradio 管理界面逻辑 ---
+# --- 4. Gradio UI (带门禁验证) ---
 
-def ui_check_balance(token):
-    if not token: return "请先获取 Token"
-    url = "https://level2labs-prod--adm-agent-helper-user-get-credits.modal.run/"
-    headers = {"Authorization": f"Bearer {token}"}
-    try:
-        data = requests.post(url, headers=headers).json()
-        return f"当前账号余额: {data.get('premium_quota', 0)}"
-    except: return "查询失败"
+with gr.Blocks(title="AI Doc Maker 管理系统") as demo:
+    # 登录界面
+    with gr.Column(visible=True) as login_panel:
+        gr.Markdown("## 🔐 访问受限\n请输入系统的 ACCESS_TOKEN 以解锁管理面板。")
+        pwd_input = gr.Textbox(label="ACCESS_TOKEN", type="password", placeholder="输入 sk-xxxx...")
+        login_btn = gr.Button("解锁系统", variant="primary")
+        login_msg = gr.Markdown()
 
-def ui_chat(token, content):
-    if not token: yield "请先获取 Token", ""; return
-    url = "https://level2labs-prod--adm-agent-send-chat-message.modal.run/"
-    headers = {"Authorization": f"Bearer {token}"}
-    data = {"content": (None, content), "model": (None, "gemini-3-flash"), "conversation_id": (None, generate_id()), "message_id": (None, generate_id()), "max_mode": (None, "true")}
-    full_text = ""
-    try:
-        resp = requests.post(url, files=data, headers=headers, stream=True)
+    # 主管理界面 (初始隐藏)
+    with gr.Column(visible=False) as main_panel:
+        gr.Markdown("# 🚀 AI Doc Maker 后台管理")
+        
+        with gr.Row():
+            with gr.Column():
+                token_view = gr.Textbox(label="当前活跃 Token (缓存中)", interactive=False)
+                balance_view = gr.Textbox(label="剩余可用额度", interactive=False)
+                refresh_btn = gr.Button("强制刷新账号/同步余额")
+            
+            with gr.Column():
+                chat_input = gr.Textbox(label="测试对话", placeholder="输入消息...")
+                chat_btn = gr.Button("发送测试", variant="primary")
+                chat_output = gr.Textbox(label="AI 回复", interactive=False)
+
+    # --- UI 逻辑绑定 ---
+    def handle_login(entered_token):
+        if entered_token == ACCESS_TOKEN:
+            # 获取当前可用信息
+            t = get_valid_token()
+            b = sync_balance()
+            return gr.update(visible=False), gr.update(visible=True), t, b, ""
+        else:
+            return gr.update(visible=True), gr.update(visible=False), "", "", "❌ 密码错误，请重试！"
+
+    login_btn.click(
+        handle_login, 
+        inputs=pwd_input, 
+        outputs=[login_panel, main_panel, token_view, balance_view, login_msg]
+    )
+
+    def manual_refresh():
+        # 强制获取新账号
+        t = fetch_new_account()
+        b = sync_balance()
+        return t, b
+
+    refresh_btn.click(manual_refresh, outputs=[token_view, balance_view])
+
+    def ui_chat_test(msg):
+        jwt = get_valid_token()
+        url = "https://level2labs-prod--adm-agent-send-chat-message.modal.run/"
+        data = {"content": (None, msg), "model": (None, "auto"), "conversation_id": (None, "ui-test"), "message_id": (None, "ui-test"), "max_mode": (None, "true")}
+        full_text = ""
+        resp = requests.post(url, files=data, headers={"Authorization": f"Bearer {jwt}"}, stream=True)
         for line in resp.iter_lines():
             if line:
-                raw = json.loads(line.decode("utf-8"))
-                full_text += raw.get("data", "")
-                yield full_text, json.dumps(raw, ensure_ascii=False)
-    except Exception as e: yield f"错误: {e}", ""
-
-def ui_tts(token, prompt):
-    if not token: return "请先获取 Token", None
-    url = "https://level2labs-prod--adm-agent-audio-create-audio-with-assistant.modal.run/"
-    headers = {"Authorization": f"Bearer {token}"}
-    try:
-        resp = requests.post(url, files={"prompt": (None, prompt), "voice_id": (None, "clear"), "speed": (None, "1")}, headers=headers).json()
-        audio_url = requests.get(f"https://level2labs-prod--adm-agent-audio-get-audio-playback-url.modal.run/?name={resp.get('name')}", headers=headers).json().get("url")
-        return json.dumps(resp, indent=2), audio_url
-    except: return "生成失败", None
-
-# --- 5. 构建与挂载 Gradio ---
-
-with gr.Blocks(title="AI Doc Maker 管理面板") as demo:
-    gr.Markdown("# 🚀 AI Doc Maker 管理面板")
-    with gr.Group():
-        btn_reg = gr.Button("第一步：创建/刷新 Token", variant="primary")
-        token_display = gr.Textbox(label="当前 IdToken", interactive=True)
-    with gr.Group():
-        btn_balance = gr.Button("查询余额")
-        balance_display = gr.Textbox(label="余额详情", interactive=False)
-    with gr.Group():
-        msg_input = gr.Textbox(label="输入消息", placeholder="a joke")
-        with gr.Row():
-            btn_msg = gr.Button("发送消息", variant="primary")
-            btn_tts = gr.Button("生成语音")
-        reply_out = gr.Textbox(label="AI 回复")
-        audio_out = gr.Audio(label="语音播放")
+                full_text += json.loads(line.decode("utf-8")).get("data", "")
+                yield full_text
+        sync_balance() # 更新余额显示
 
-    btn_reg.click(get_jwt, outputs=token_display)
-    btn_balance.click(ui_check_balance, inputs=token_display, outputs=balance_display)
-    btn_msg.click(ui_chat, inputs=[token_display, msg_input], outputs=[reply_out, gr.State()])
-    btn_tts.click(ui_tts, inputs=[token_display, msg_input], outputs=[gr.State(), audio_out])
+    chat_btn.click(ui_chat_test, inputs=chat_input, outputs=chat_output)
 
-# 挂载到根路径，这样你打开 Space 首页就能看到面板
+# 挂载
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":