fixed app.py

app.py

@@ -14,8 +14,11 @@ collection = load_collection()
 encoder = load_encoder()
 reranker = CrossEncoder("BAAI/bge-reranker-large")
 
-def build_rag_prompt(query, context):
-    prompt = f"""已知资料如下：
+
+def build_rag_prompt(query, context, system_message):
+    prompt = f"""{system_message}
+
+已知资料如下：
 {context}
 
 用户提问：{query}
@@ -23,53 +26,86 @@ def build_rag_prompt(query, context):
     return prompt
 
 
-def answer_fn(query, history=None):
-    query_vec = encode_query(encoder, query)
+def respond(
+        message,
+        history: list[dict[str, str]],
+        system_message,
+        max_tokens,
+        temperature,
+        top_p,
+):
+    """
+    message: 当前输入内容
+    history: [{"role": "user", "content": ...}, {"role": "assistant", "content": ...}, ...]
+    system_message: 自定义 System Prompt
+    """
+    chat_history = [
+        {"role": "system", "content": system_message.strip() or "你是BangDream知识问答助手, 只能基于提供资料作答。"}
+    ]
+    chat_history.extend(history)
+    chat_history.append({"role": "user", "content": message})
+
+    query_vec = encode_query(encoder, message)
     results = retrieve_docs(collection, query_vec, top_k=30)
-    reranked = query_rerank(reranker, query, results, top_n=10)
+    reranked = query_rerank(reranker, message, results, top_n=10)
     deduped = dedup_by_chapter_event(reranked, max_per_group=1)
     expanded_results = expand_with_neighbors(deduped[:3], collection)
-
     context = expanded_results[0][0] if expanded_results else ""
-    rag_prompt = build_rag_prompt(query, context)
-    system_prompt = "你是BangDream知识问答助手, 也就是邦学家. 只能基于提供的资料内容作答。"
 
-    response = client.chat.completions.create(
+    rag_prompt = build_rag_prompt(message, context, system_message)
+    messages = [
+        {"role": "system", "content": system_message.strip() or "你是BangDream知识问答助手, 只能基于提供资料作答。"},
+        {"role": "user", "content": rag_prompt}
+    ]
+
+    response = ""
+    stream = client.chat.completions.create(
         model="gpt-4o",
-        messages=[
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": rag_prompt}
-        ],
-        temperature=0.2,
-        max_tokens=512,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        top_p=top_p,
+        stream=True
     )
-    answer = response.choices[0].message.content.strip()
+    for chunk in stream:
+        delta = getattr(chunk.choices[0].delta, "content", None)
+        if delta:
+            response += delta
+            yield response
 
-    references = ""
-    for idx, (doc, score, meta) in enumerate(expanded_results, 1):
-        chapter = meta.get("chapterTitle", "UnknownChapter")
-        event = meta.get("eventName", "UnknownEvent")
-        references += f"\n--- reference: {idx} (chapter: {chapter}, event: {event}, score={score:.4f}) ---\n"
-        references += doc[:300] + "...\n"
 
-    return answer, references
+# ========== Gradio ChatInterface with extra sidebar inputs ==========
+chatbot = gr.ChatInterface(
+    respond,
+    type="messages",
+    additional_inputs=[
+        gr.Textbox(value="你是BangDream知识问答助手, 只能基于提供资料内容作答。", label="System message"),
+        gr.Slider(minimum=64, maximum=1024, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.2, step=0.05, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
+    examples=[
+        ["乐奈为什么喜欢吉他？"],
+        ["LOCK和CHU²第一次见面是什么情节？"],
+        ["谁是RAS的初代成员？"],
+    ],
+    description="输入你关于BangDream的问题，邦学家会基于资料库为你检索并作答",
+    title="Dr-Bang RAG QA Chatbot"
+)
 
-# Gradio UI
 with gr.Blocks(title="Dr-Bang RAG QA") as demo:
-    gr.Markdown("# Dr-Bang RAG QA\n\n基于BangDream知识库的RAG问答系统。")
-    with gr.Row():
-        chatbot = gr.ChatInterface(
-            fn=answer_fn,
-            title="Dr-Bang RAG Chat",
-            description="输入你的BangDream问题，AI助手会基于资料库为你检索并作答。",
-            examples=[
-                ["乐奈为什么喜欢吉他？"],
-                ["LOCK和CHU²第一次见面是什么情节？"],
-                ["谁是RAS的初代成员？"],
-            ],
-            outputs=[
-                gr.Textbox(label="Answer", lines=6, interactive=False),
-                gr.Textbox(label="Reference", lines=8, interactive=False)
-            ]
+    with gr.Sidebar():
+        gr.Markdown(
+            "## Dr-Bang QA\n\n"
         )
-demo.launch()
+        gr.LoginButton()
+    chatbot.render()
+
+if __name__ == "__main__":
+    demo.launch()