Update app.py

app.py

@@ -9,16 +9,15 @@ from fastapi import FastAPI, Request, HTTPException, Depends, status
 from fastapi.responses import StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
+# --- 1. 初始化 FastAPI 与 安全验证 ---
 app = FastAPI()
 security = HTTPBearer()
 
-# --- 配置与模型列表 ---
+# 建议在 HF Secret 中设置 ACCESS_TOKEN
 ACCESS_TOKEN = os.getenv("ACCESS_TOKEN", "sk-admin-123456")
 API_KEY = "AIzaSyAZaD22Mzi9HkTcW3ErNxRA_sNEFolLBCA"
-# 原 Go 代码注释中提到的模型列表
 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(
@@ -28,7 +27,7 @@ async def verify_token(auth: HTTPAuthorizationCredentials = Depends(security)):
         )
     return auth.credentials
 
-# --- 核心逻辑函数 ---
+# --- 2. 核心后端逻辑 ---
 def get_jwt():
     url = f"https://identitytoolkit.googleapis.com/v1/accounts:signUp?key={API_KEY}"
     try:
@@ -40,40 +39,29 @@ def get_jwt():
 def generate_id():
     return ''.join(random.choice('0123456789abcdef') for _ in range(32))
 
-# --- OpenAI 兼容接口 ---
+# --- 3. OpenAI 兼容接口 (/v1/...) ---
 
-# 1. 获取模型列表接口
 @app.get("/v1/models")
 async def list_models(token: str = Depends(verify_token)):
-    models_list = []
-    for m in SUPPORTED_MODELS:
-        models_list.append({
-            "id": m,
-            "object": "model",
-            "created": int(time.time()),
-            "owned_by": "aidocmaker"
-        })
-    return {"object": "list", "data": models_list}
-
-# 2. 对话接口
+    models_data = [{"id": m, "object": "model", "created": int(time.time()), "owned_by": "aidocmaker"} for m in SUPPORTED_MODELS]
+    return {"object": "list", "data": models_data}
+
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request, token: str = Depends(verify_token)):
     body = await request.json()
     messages = body.get("messages", [])
-    # 使用请求中指定的模型，如果没传则默认 auto
-    requested_model = body.get("model", "auto") 
+    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")
+    if not jwt: raise HTTPException(status_code=500, detail="Failed to get JWT")
 
     url = "https://level2labs-prod--adm-agent-send-chat-message.modal.run/"
     headers = {"Authorization": f"Bearer {jwt}"}
     form_data = {
         "content": (None, content),
-        "model": (None, requested_model),
+        "model": (None, model),
         "conversation_id": (None, generate_id()),
         "message_id": (None, generate_id()),
         "max_mode": (None, "true")
@@ -83,44 +71,87 @@ async def chat_completions(request: Request, token: str = Depends(verify_token))
 
     def generate():
         chat_id = f"chatcmpl-{uuid.uuid4()}"
-        created_time = int(time.time())
         for line in target_resp.iter_lines():
             if line:
                 try:
                     raw_data = json.loads(line.decode("utf-8"))
-                    delta_content = raw_data.get("data", "")
                     chunk = {
-                        "id": chat_id,
-                        "object": "chat.completion.chunk",
-                        "created": created_time,
-                        "model": requested_model,
-                        "choices": [{"index": 0, "delta": {"content": delta_content}, "finish_reason": None}]
+                        "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"
-                except:
-                    continue
+                except: continue
         yield "data: [DONE]\n\n"
 
     if stream:
         return StreamingResponse(generate(), media_type="text/event-stream")
     else:
-        full_text = ""
-        for line in target_resp.iter_lines():
-            if line:
-                try:
-                    full_text += json.loads(line.decode("utf-8")).get("data", "")
-                except: continue
+        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": requested_model,
-            "choices": [{"index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": "stop"}]
+            "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"}]
         }
 
-# --- Gradio UI 逻辑 (略，保持之前的实现即可) ---
-# ... 这里放入之前提供的 demo 代码 ...
+# --- 4. Gradio 管理界面逻辑 ---
+
+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 "查询失败"
 
+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)
+        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="语音播放")
+
+    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])
+
+# 挂载到根路径，这样你打开 Space 首页就能看到面板
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":