Update app.py

app.py

@@ -1,209 +1,88 @@
 import gradio as gr
-import aiohttp
+import requests
 import json
 import base64
 from io import BytesIO
 from PIL import Image
 import os
 
-# ====================== 配置 ======================
-DEFAULT_API_BASE = "https://www.gpt4novel.com/api/xiaoshuoai/ext/v1/chat/completions"
-MODELS_CACHE = None
-
-# ====================== 获取模型列表 ======================
-async def fetch_available_models(api_key: str):
-    global MODELS_CACHE
-    if MODELS_CACHE is not None:
-        return MODELS_CACHE
-
-    url = DEFAULT_API_BASE.replace("/chat/completions", "/models")
-    headers = {"Authorization": f"Bearer {api_key}"}
-
-    async with aiohttp.ClientSession() as session:
-        async with session.get(url, headers=headers, timeout=15) as resp:
-            if resp.status != 200:
-                text = await resp.text()
-                raise gr.Error(f"获取模型列表失败: {resp.status} - {text}")
-            data = await resp.json()
-            models = [m["id"] for m in data.get("data", []) if "nalang" in m["id"].lower()]
-            MODELS_CACHE = models
-            return models
-
-# ====================== 流式聊天核心函数 ======================
-async def stream_chat(messages, api_key, model, temperature=0.7, max_tokens=800):
-    if not api_key.strip():
-        yield "请先在设置中输入API Key并连接"
-        return
-
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-
-    payload = {
-        "model": model,
-        "messages": messages,
-        "stream": True,
-        "temperature": temperature,
-        "max_tokens": max_tokens,
-        "top_p": 0.35,
-        "repetition_penalty": 1.05
-    }
-
-    full_response = ""
-    async with aiohttp.ClientSession() as session:
-        async with session.post(DEFAULT_API_BASE, headers=headers, json=payload) as resp:
-            if resp.status != 200:
-                text = await resp.text()
-                yield f"请求失败: {resp.status} - {text}"
-                return
-
-            async for line in resp.content:
-                line = line.decode('utf-8').strip()
-                if line.startswith("data: "):
-                    data_str = line[6:]
-                    if data_str == "[DONE]":
-                        break
-                    try:
-                        data = json.loads(data_str)
-                        delta = data["choices"][0]["delta"]
-                        if "content" in delta:
-                            content = delta["content"]
-                            full_response += content
-                            yield full_response
-                    except:
-                        pass
-
-# ====================== 角色卡处理 ======================
-def extract_role_card(file):
-    if file is None:
-        return "未上传角色卡", None, None
+API_BASE = "https://www.gpt4novel.com/api/xiaoshuoai/ext/v1"
 
+def get_models(api_key):
     try:
-        if file.name.lower().endswith('.json'):
-            with open(file.name, 'r', encoding='utf-8') as f:
-                card = json.load(f)
-        elif file.name.lower().endswith('.png'):
-            img = Image.open(file.name)
-            if "parameters" not in img.info:
-                return "这张PNG不包含JSON角色卡信息", None, None
-            json_str = img.info["parameters"]
-            card = json.loads(json_str)
-        else:
-            return "不支持的文件格式", None, None
-
-        name = card.get("name", "未知角色")
-        desc = card.get("description", "")
-        first_msg = card.get("first_mes", "你好呀～")
-        personality = card.get("personality", "")
-
-        system_prompt = f"""你现在是{ name }，以下是你的角色描述：
-{ desc }
-{ personality }
-
-请完全沉浸在角色中，用第一人称回复，语气、性格要符合设定。"""
-
-        messages = [{"role": "system", "content": system_prompt}]
-        if first_msg.strip():
-            messages.append({"role": "assistant", "content": first_msg})
-
-        preview = f"角色名：{name}\n\n描述：{desc[:200]}...\n\n首句：{first_msg[:100]}..."
-
-        return preview, messages, name
-
-    except Exception as e:
-        return f"解析角色卡失败：{str(e)}", None, None
-
-# ====================== Gradio界面 ======================
-with gr.Blocks(title="小说AI聊天室", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 小说AI聊天室")
-
-    api_key = gr.State("")
-    current_model = gr.State("nalang-xl-0826-10k")
-    chat_history = gr.State([])
-    role_messages = gr.State([])
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            chatbot = gr.Chatbot(height=650, show_label=False, bubble_full_width=False)
-            msg = gr.Textbox(placeholder="在这里输入消息...", container=False, lines=3)
-            with gr.Row():
-                submit_btn = gr.Button("发送", variant="primary")
-                clear_btn = gr.Button("清空对话")
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 设置")
-            api_key_input = gr.Textbox(label="API Key", type="password", placeholder="sk-xxx...")
-            connect_btn = gr.Button("连接", variant="primary")
-            model_dropdown = gr.Dropdown(label="选择模型", choices=[], interactive=False)
-            status = gr.Markdown("")
-
-            gr.Markdown("### 角色卡（支持JSON / 带JSON的PNG）")
-            role_upload = gr.File(label="上传角色卡", file_types=[".json", ".png"])
-            role_preview = gr.Markdown()
-            use_role_btn = gr.Button("使用该角色聊天", variant="secondary")
-
-    # ------------------ 连接API ------------------
-    async def connect_api(key):
-        if not key.strip():
-            return "API Key不能为空", [], False
-        try:
-            models = await fetch_available_models(key)
-            return "连接成功！可选择模型", models, True
-        except Exception as e:
-            return f"连接失败：{str(e)}", [], False
-
-    connect_btn.click(
-        connect_api,
-        inputs=api_key_input,
-        outputs=[status, model_dropdown, model_dropdown.interactive]
-    ).then(lambda key: key, inputs=api_key_input, outputs=api_key)
-
-    # ------------------ 选择模型 ------------------
-    model_dropdown.change(
-        lambda m: m,
-        inputs=model_dropdown,
-        outputs=current_model
+        r = requests.get(f"{API_BASE}/models", headers={"Authorization": f"Bearer {api_key}"})
+        models = [m["id"] for m in r.json()["data"]]
+        return sorted(models, reverse=True)
+    except:
+        return ["nalang-xl-0826-10k"]
+
+def extract_card_from_png(file_obj):
+    if not file_obj.name.endswith(".png"):
+        return None
+    img = Image.open(file_obj.name)
+    data = img.info.get("chara")
+    if data:
+        decoded = base64.b64decode(data).decode()
+        return json.loads(decoded)
+    return None
+
+def chat(message, history, model, api_key, system_prompt):
+    messages = [{"role": "system", "content": system_prompt or "你是一个有帮助的AI助手。"}]
+    for h in history:
+        messages.append({"role": "user", "content": h[0]})
+        if h[1]: messages.append({"role": "assistant", "content": h[1]})
+    messages.append({"role": "user", "content": message})
+
+    resp = requests.post(
+        f"{API_BASE}/chat/completions",
+        json={"model": model, "messages": messages, "stream": True, "temperature": 0.7, "max_tokens": 800},
+        headers={"Authorization": f"Bearer {api_key}"},
+        stream=True,
+        timeout=60
     )
+    resp.raise_for_status()
+    for line in resp.iter_lines():
+        if line and line.startswith(b"data: "):
+            try:
+                data = json.loads(line[6:].decode())
+                if token := data["choices"][0]["delta"].get("content"):
+                    yield token
+            except:
+                continue
+
+with gr.Blocks() as demo:
+    gr.HTML(open("index.html").read())
+    chatbot = gr.Chatbot(height="70vh")
+    with gr.Row():
+        msg = gr.Textbox(scale=8, placeholder="发送消息...", container=False)
+        send = gr.Button("发送", scale=1)
+    model = gr.Dropdown(choices=[], label="模型")
+    api_key = gr.Textbox(label="API Key", type="password")
+    card_file = gr.File(label="导入角色卡 (.json 或 .png)", file_types=[".json", ".png"])
+
+    system_prompt = gr.Textbox(label="系统提示词（角色卡自动填充）", lines=4, value="你是一个有帮助的AI助手。")
+
+    def load_card(file):
+        if not file: return "", []
+        if str(file).endswith(".png"):
+            card = extract_card_from_png(file)
+        else:
+            card = json.load(open(file.name, encoding="utf-8"))
+        if card:
+            prompt = card.get("data", card.get("system_prompt", card.get("description", "")))
+            name = card.get("name", card.get("data", {}).get("name", "角色"))
+            return prompt, [[None, f"已加载角色卡：{name}"]]
+        return "", []
 
-    # ------------------ 普通聊天 ------------------
-    def user_submit(message, history, role_msgs):
-        if not message.strip():
-            return "", history
-        history.append((message, None))
-        return "", history
-
-    async def bot_response(history, api_key_val, model_val, role_msgs):
-        messages = role_msgs + [{"role": "user" if i % 2 == 0 else "assistant", "content": msg} for i, (msg, _) in enumerate(history)]
-        messages = messages[:-1]  # 去掉最后一个None
-
-        response = ""
-        async for chunk in stream_chat(messages, api_key_val, model_val):
-            response = chunk
-            history[-1] = (history[-1][0], response)
-            yield history
-
-    submit_btn.click(user_submit, [msg, chatbot, role_messages], [msg, chatbot])\
-              .then(bot_response, [chatbot, api_key, current_model, role_messages], chatbot)
-
-    msg.submit(user_submit, [msg, chatbot, role_messages], [msg, chatbot])\
-       .then(bot_response, [chatbot, api_key, current_model, role_messages], chatbot)
-
-    # ------------------ 加载角色卡 ------------------
-    role_upload.upload(extract_role_card, inputs=role_upload, outputs=[role_preview, role_messages, gr.State()])
+    card_file.change(load_card, card_file, [system_prompt, chatbot])
 
-    def start_role_chat(role_msgs, history):
-        if not role_msgs:
-            return history
-        # 清空当前对话，加载角色首句
-        new_history = []
-        if len(role_msgs) > 1 and role_msgs[1]["role"] == "assistant":
-            new_history.append((None, role_msgs[1]["content"]))
-        return new_history
+    def connect_key(key):
+        models = get_models(key)
+        return gr.Dropdown(choices=models, value=models[0] if models else None)
 
-    use_role_btn.click(start_role_chat, [role_messages, chatbot], chatbot)
+    api_key.submit(connect_key, api_key, model)
 
-    # ------------------ 清空 ------------------
-    clear_btn.click(lambda: ([], []), None, [chatbot, role_messages])
+    send.click(chat, [msg, chatbot, model, api_key, system_prompt], chatbot).then(lambda: "", None, msg)
+    msg.submit(chat, [msg, chatbot, model, api_key, system_prompt], chatbot).then(lambda: "", None, msg)
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)