Update app.py

app.py

@@ -1,108 +1,87 @@
 import gradio as gr
 import requests
 import os
-import fitz
+import json
+import fitz  # PyMuPDF
 import re
 from huggingface_hub import HfApi, hf_hub_download
 
 # --- 核心配置 ---
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+# 建议在 Hugging Face 的 Settings -> Secrets 中设置这些变量
+OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
 HF_TOKEN = os.environ.get("HF_TOKEN")
-MODEL_ID = "gemini-1.5-flash"
 
-DATASETS = ["MGGroup/Treaties", "MGGroup/InvestmentGuide"]
+# 推荐使用 OpenRouter 的免费模型，稳定性比直连好
+# 备选: "deepseek/deepseek-chat:free" 或 "meta-llama/llama-3.3-70b-instruct:free"
+MODEL_ID = "google/gemini-2.0-flash-001" 
 
-# 顶部的专业描述（保持原样，符合您的要求）
-DESCRIPTION = """
-<div style="text-align: left; border-left: 4px solid #2196F3; padding-left: 15px; margin-bottom: 20px;">
-    <h3>MG TaxAI | 跨境财税合规实验室 (Beta)</h3>
-    <p>本系统依托 <b>MG 核心智库</b> 构建，旨在实现解析结果实时溯源至各国官方税收协定与法律文本。目前系统正处于<b>知识库全量装载阶段</b>，已优先上线核心业务国家的官方协定库。</p>
-    <p>我们正持续同步全球各主要经济体的国别投资税收指南及多税种年度税收报告。受限于测试版的数据填充进度，相关解析结果仅供专业参考。MG 团队正加速完善每一条咨询建议的合规证据链，以确保交付专家级的数字化合规支持。</p>
-    <hr style="border: 0; border-top: 1px solid #eee; margin: 10px 0;">
-    <p style="font-size: 0.85em; color: #666;">
-        <b>⚠️ AI 免责声明：</b><br>
-        本系统生成的内容由人工智能根据现有库文件分析得出，不构成正式的法律或税务建议。在使用本系统结果进行任何商业决策前，请务必咨询 MG Consult 专业团队。
-    </p>
-</div>
-"""
-
-def fetch_dataset_context(query):
-    if not HF_TOKEN: return ""
-    api = HfApi(token=HF_TOKEN)
-    combined_text = ""
-    keywords = [k for k in re.findall(r'[\u4e00-\u9fa5]{2,}|[a-zA-Z]{3,}', query) if k not in ["资料", "关于", "查询", "政策"]]
-    if not keywords: keywords = re.findall(r'[\u4e00-\u9fa5]+|[a-zA-Z]+', query)
-    
-    for repo in DATASETS:
-        try:
-            files = api.list_repo_files(repo_id=repo, repo_type="dataset")
-            matched = sorted([f for f in files if f.lower().endswith(".pdf") and any(k.lower() in f.lower() for k in keywords)])
-            for f_path in matched[:8]: 
-                temp_path = hf_hub_download(repo_id=repo, filename=f_path, repo_type="dataset", token=HF_TOKEN)
-                doc = fitz.open(temp_path)
-                combined_text += f"\n[Ref: {f_path}]\n" + "".join([page.get_text() for page in doc[:15]])
-                doc.close()
-                os.remove(temp_path)
-        except: continue
-    return combined_text[:12000]
-
-def find_local_context(query):
-    base_dir = "./treaties"
-    if not os.path.exists(base_dir): return ""
+# --- PDF 处理逻辑 ---
+def extract_text_from_pdf(pdf_path):
+    text = ""
     try:
-        keywords = re.findall(r'[\u4e00-\u9fa5]+|[a-zA-Z]+', query)
-        folders = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
-        selected = next((f for f in folders if any(k.lower() in f.lower() for k in keywords)), None)
-        if not selected: return ""
-        path = os.path.join(base_dir, selected)
-        texts = [f"\n[Local: {pdf}]\n" + "".join([p.get_text() for p in fitz.open(os.path.join(path, pdf))[:15]]) for pdf in sorted([f for f in os.listdir(path) if f.endswith(".pdf")])[:5]]
-        return "\n".join(texts)[:8000]
-    except: return ""
+        with fitz.open(pdf_path) as doc:
+            for page in doc:
+                text += page.get_text()
+    except Exception as e:
+        print(f"读取 PDF 出错: {e}")
+    return text
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    knowledge = fetch_dataset_context(message)
-    if not knowledge: knowledge = find_local_context(message)
-    
-    # --- 核心修改：在系统提示词中强制要求 AI 保持静默，不要自我介绍 ---
-    full_system_prompt = (
-        f"{system_message}\n\n"
-        "【重要约束】：\n"
-        "1. 严禁进行任何形式的自我介绍或身份说明（例如：不要说'我是MG的专家'、'你好'等）。\n"
-        "2. 严禁包含任何开场白，直接针对用户问题进入专业分析。\n"
-        "3. 必须严格基于以下参考资料进行回答。\n\n"
-        f"参考资料：\n{knowledge}"
-    )
+# --- OpenRouter API 调用逻辑 ---
+def ask_ai(user_query, context=""):
+    if not OPENROUTER_API_KEY:
+        return "错误：未配置 OPENROUTER_API_KEY。请在环境变量中设置。"
 
-    messages = [{"role": "system", "content": full_system_prompt}]
-    for u, a in history:
-        if u: messages.append({"role": "user", "content": u})
-        if a: messages.append({"role": "assistant", "content": a})
-    messages.append({"role": "user", "content": message})
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "HTTP-Referer": "https://mgconsult.net", # 可选
+        "X-Title": "TaxAI Assistant", # 可选
+        "Content-Type": "application/json"
+    }
+    
+    # 构造 Prompt，将 PDF 内容作为背景知识
+    full_prompt = f"背景知识：\n{context}\n\n问题：\n{user_query}" if context else user_query
+    
+    payload = {
+        "model": MODEL_ID,
+        "messages": [
+            {"role": "system", "content": "你是一个专业的国际税务与贸易合规专家，请基于提供的背景知识回答问题。"},
+            {"role": "user", "content": full_prompt}
+        ],
+        "temperature": 0.7
+    }
 
-    url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL_ID}:generateContent?key={GEMINI_API_KEY}"
     try:
-        res = requests.post(url, json={"contents": [{"role":"user" if m["role"] in ["user","system"] else "model", "parts":[{"text":m["content"]}]} for m in messages]}, timeout=60)
-        reply = res.json()['candidates'][0]['content']['parts'][0]['text']
-    except:
-        reply = "系统繁忙，请稍后再试。"
+        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
+        if response.status_code == 200:
+            result = response.json()
+            return result['choices'][0]['message']['content']
+        elif response.status_code == 429:
+            return "系统繁忙 (429)：OpenRouter 的免费额度也暂时达到上限，请等待一分钟再试。"
+        else:
+            return f"API 报错: {response.status_code} - {response.text}"
+    except Exception as e:
+        return f"请求失败: {str(e)}"
 
-    # --- 唯一的身份开场白在这里定义，由代码控制，不给 AI 发挥空间 ---
-    header = "您好，我是 **MG Consult** 的国际税收 AI 专家。很高兴为您提供专业咨询。\n\n---\n\n"
-    yield header + reply
+# --- Gradio 界面处理 ---
+def chat_handler(message, history):
+    # 这里可以添加你之前的数据集检索逻辑 (RAG)
+    # 目前演示直接调用 AI
+    response = ask_ai(message)
+    return response
 
-# 界面
-demo = gr.ChatInterface(
-    fn=respond,
-    description=DESCRIPTION,
-    theme="soft",
-    css=".gradio-container {max-width: 950px !important} .description {margin-bottom: 20px}",
-    additional_inputs=[
-        gr.Textbox(value="你代表 MG Consult，是国际税收专家。请严格基于参考资料提供深度分析。请直接进入正文，严禁自我介绍。", label="系统指令"),
-        gr.Slider(512, 4096, 2048, label="回复长度限制"),
-        gr.Slider(0, 1, 0.05, label="严谨度 (Temperature)"),
-        gr.Slider(0, 1, 0.95, label="采样率 (Top-p)"),
-    ],
-)
+# --- 构建 Gradio UI ---
+with gr.Blocks(title="MG Consulting TaxAI") as demo:
+    gr.Markdown("# 🌍 MG Consulting 国际税务 AI 助手")
+    gr.Markdown("基于 OpenRouter 引擎，支持国际税收协定与投资指南查询。")
+    
+    chatbot = gr.ChatInterface(
+        fn=chat_handler,
+        examples=["美国个人所得税税率是多少？", "如何查询丹麦的税收协定？"],
+        cache_examples=False,
+    )
 
+# --- 启动应用 ---
 if __name__ == "__main__":
-    demo.launch()  # 前面至少要有 4 个空格或 1 个 Tab
+    # 这里的缩进非常重要，确保 demo.launch() 在 if 语句块内
+    demo.launch()