demo / app.py
wkplhc's picture
Update app.py
37e2ec9 verified
import gradio as gr
import requests
import json
import base64
from io import BytesIO
from PIL import Image
import os
API_BASE = "https://www.gpt4novel.com/api/xiaoshuoai/ext/v1"
def get_models(api_key):
try:
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 "", []
card_file.change(load_card, card_file, [system_prompt, chatbot])
def connect_key(key):
models = get_models(key)
return gr.Dropdown(choices=models, value=models[0] if models else None)
api_key.submit(connect_key, api_key, model)
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.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)