app.py

import os
import re
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import gradio as gr
import time
import numpy as np
import io
from PIL import Image
from dotenv import load_dotenv
load_dotenv()
import easyocr
import faiss
import random
import requests
from openai import OpenAI

# 全局会话存储
session_data = {}  # session_id -> {"chunks": [...], "index": faiss_index, "history": [...]}
reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)

def embed_texts(texts):
    # 这里可替换为真实嵌入接口调用
    return [np.random.rand(768).tolist() for _ in texts]

def safe_call_model(prompt, session_id, model_list, history):
    for model in model_list:
        try:
            if model == "openai":
                client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
                response = client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=[
                        {"role": "system", "content": "你是一个多模态文档理解助手。"},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=0.5
                )
                answer = response.choices[0].message.content
                return {
                    "answer": answer,
                    "model": "gpt-3.5-turbo",
                    "elapsed": 1.2,
                    "cost": 0.002
                }
            elif model == "claude":
                print("当前使用的ANTHROPIC_API_KEY:", os.getenv("ANTHROPIC_API_KEY"))
                headers = {
                    "x-api-key": os.getenv("ANTHROPIC_API_KEY"),
                    "content-type": "application/json",
                    "anthropic-version": "2023-06-01"
                }
                payload = {
                    "model": "claude-3-sonnet-20240229",
                    "max_tokens": 1024,
                    "temperature": 0.5,
                    "messages": [{"role": "user", "content": prompt}]
                }
                response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload)

                if response.status_code == 429:
                    print("⚠️ Claude API 限速，2秒后重试一次...")
                    time.sleep(2)
                    response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload)

                if response.status_code != 200:
                    raise Exception(f"Claude API 请求失败：{response.status_code} {response.text}")

                data = response.json()
                if "content" in data and isinstance(data["content"], list):
                    answer = data["content"][0].get("text", "⚠️ Claude返回内容结构异常")
                else:
                    raise Exception(f"Claude响应格式无效：{data}")

                return {
                    "answer": answer,
                    "model": "claude-3-sonnet-20240229",
                    "elapsed": 1.2,
                    "cost": 0.002
                }
            elif model == "deepseek":
                headers = {
                    "Authorization": f"Bearer {os.getenv('DEEPSEEK_API_KEY')}",
                    "Content-Type": "application/json"
                }
                payload = {
                    "model": "deepseek-chat",
                    "messages": [
                        {"role": "system", "content": "你是一个多模态文档理解助手。"},
                        {"role": "user", "content": prompt}
                    ],
                    "temperature": 0.5
                }
                response = requests.post(
                    os.getenv("DEEPSEEK_BASE_URL"),
                    headers=headers,
                    json=payload
                )
                if response.status_code != 200:
                    raise Exception(f"DeepSeek API 请求失败：{response.status_code} {response.text}")
                answer = response.json()["choices"][0]["message"]["content"]
                return {
                    "answer": answer,
                    "model": "deepseek-chat",
                    "elapsed": 1.2,
                    "cost": 0.002
                }
            elif model == "text-only":
                # 本地简单文本模拟回复
                answer = f"[text-only模式] 你的输入是：{prompt}"
                return {
                    "answer": answer,
                    "model": "text-only",
                    "elapsed": 0,
                    "cost": 0
                }
        except Exception as e:
            print(f"❌ {model} 调用失败：{e}")
            continue
    print("="*30)
    return {
        "answer": "❌ 所有模型均调用失败，请检查网络或API Key。",
        "model": "none",
        "elapsed": 0,
        "cost": 0
    }

def extract_text_from_image(pil_image):
    result = reader.readtext(np.array(pil_image), detail=0)
    return [line.strip() for line in result if len(line.strip()) > 5]

def build_index(vectors):
    dim = len(vectors[0])
    index = faiss.IndexFlatL2(dim)
    index.add(np.array(vectors).astype("float32"))
    return index

def search_index(index, query_vec, top_k=5):
    D, I = index.search(np.array([query_vec]).astype("float32"), top_k)
    return I[0].tolist()

def fetch_related_papers(query):
    return [
        {"title": "Document Understanding with OCR and LLMs", "url": "https://arxiv.org/abs/2403.01234"},
        {"title": "Multimodal Large Models on Visual Text", "url": "https://arxiv.org/abs/2402.05678"}
    ]

def clean_text(text: str) -> str:
    # 去除多余空行、控制字符，简单纠正断句
    text = re.sub(r'\n\s*\n', '\n', text)  # 连续空行合并
    text = text.replace('\r', '')
    return text.strip()

def chunk_text(text: str, max_chunk_size=1000) -> list:
    paragraphs = text.split('\n')
    chunks = []
    current_chunk = []
    current_len = 0
    for para in paragraphs:
        para_len = len(para)
        if current_len + para_len > max_chunk_size and current_chunk:
            chunks.append('\n'.join(current_chunk))
            current_chunk = [para]
            current_len = para_len
        else:
            current_chunk.append(para)
            current_len += para_len
    if current_chunk:
        chunks.append('\n'.join(current_chunk))
    return chunks

def model_call_func(prompt: str) -> str:
    # 简易同步调用示例，只调用 Claude，实际可改成 safe_call_model 同步包装
    system_prompt = (
        "你是一名专业的文档分析师，具备深厚的行业背景知识，"
        "能够根据用户上传的文档内容，结合上下文，进行深入的分析、总结和解读。"
        "即使内容简略，也要尝试推断合理信息，给出高质量的答案。"
        "请避免只字面翻译，要结合行业理解。"
    )
    headers = {
        "x-api-key": os.getenv("ANTHROPIC_API_KEY"),
        "content-type": "application/json",
        "anthropic-version": "2023-06-01"
    }
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    payload = {
        "model": "claude-3-sonnet-20240229",
        "max_tokens": 512,
        "temperature": 0.5,
        "messages": messages
    }
    response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload)
    if response.status_code != 200:
        print(f"摘要调用失败：{response.status_code} {response.text}")
        return ""
    data = response.json()
    if "content" in data and isinstance(data["content"], list):
        return data["content"][0].get("text", "")
    return ""

def summarize_chunks(chunks: list, model_call_func) -> str:
    summaries = []
    for chunk in chunks:
        prompt = f"请对以下内容进行简明扼要的总结：\n\n{chunk}"
        result = model_call_func(prompt)
        summaries.append(result)
    return "\n".join(summaries)

def handle_file_upload(file, session_id):
    if file is None:
        return "❌ 请上传文件", [], False

    ext = os.path.splitext(file.name)[1].lower()
    if ext in [".jpg", ".jpeg", ".png"]:
        img = Image.open(file.name).convert("RGB")
        chunks = extract_text_from_image(img)
        gallery = [img]
    elif ext == ".pdf":
        from PyPDF2 import PdfReader
        reader = PdfReader(file.name)
        text = "\n".join(page.extract_text() or "" for page in reader.pages)
        chunks = text.split("\n\n")
        gallery = []
    elif ext == ".docx":
        import docx
        doc = docx.Document(file.name)
        text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
        chunks = text.split("\n\n")
        gallery = []
    elif ext == ".txt":
        with open(file.name, "r", encoding="utf-8") as f:
            text = f.read()
        chunks = text.split("\n\n")
        gallery = []
    elif ext == ".md":
        with open(file.name, "r", encoding="utf-8") as f:
            text = f.read()
        chunks = text.split("\n\n")
        gallery = []
    elif ext == ".csv":
        import pandas as pd
        df = pd.read_csv(file.name)
        text = df.to_string()
        chunks = text.split("\n\n")
        gallery = []
    elif ext == ".xlsx":
        import pandas as pd
        df = pd.read_excel(file.name, sheet_name=None)
        text = "\n\n".join([df[sheet].to_string() for sheet in df])
        chunks = text.split("\n\n")
        gallery = []
    else:
        return "❌ 不支持的文件格式", [], False

    if not chunks:
        return "❌ 未提取到有效文本内容", gallery, False

    # 清洗文本合并所有块
    raw_text = "\n\n".join(chunks)
    cleaned_text = clean_text(raw_text)

    # 分段拆分
    chunk_list = chunk_text(cleaned_text, max_chunk_size=1500)

    # 生成摘要
    summary_text = summarize_chunks(chunk_list, model_call_func)

    # 将摘要插入chunks首位，方便后续检索提升相关性
    chunks.insert(0, "[摘要]\n" + summary_text)

    vectors = embed_texts(chunks)
    index = build_index(vectors)
    session_data[session_id] = {"chunks": chunks, "index": index, "history": []}
    return "✅ 识别成功，已建立索引", gallery, True

def handle_user_query(query, session_id, model_name, chat_history, image_uploaded, task_mode="default"):
    if task_mode == "summary":
        query = "请对上传的内容进行摘要提炼。"
    elif task_mode == "translation":
        query = "请将上传的内容翻译成中文。"
    elif task_mode == "qa-extract":
        query = "请从上传内容中提取结构化的问答对。"

    # 文本模式优先跳过OCR索引
    if model_name == "text-only":
        history = session_data.get(session_id, {}).get("history", [])
        prompt = f"""你是一个多模态文档理解助手，能够总结上传的文件、图像或纯文本内容。请根据以下用户输入的问题进行回答。

用户问题：{query}
"""
    else:
        if image_uploaded and session_id in session_data:
            chunks = session_data[session_id]["chunks"]
            index = session_data[session_id]["index"]
            history = session_data[session_id]["history"]
            q_vec = embed_texts([query])[0]
            top_ids = search_index(index, q_vec)
            context = "\n".join([chunks[i] for i in top_ids])
            MAX_PROMPT_LEN = 4000
            if len(context) > MAX_PROMPT_LEN:
                context = context[:MAX_PROMPT_LEN]
            prompt = f"""你是一个文档理解助手，具备结构化内容分析能力。请基于以下文档内容尽可能准确地回答用户提出的问题。文档内容可能来自 OCR 图片、PDF、Word、Excel 等。

文档内容如下：
{context}

用户问题：{query}
"""
        else:
            history = session_data.get(session_id, {}).get("history", [])
            prompt = f"""你是一个多模态文档理解助手，能够总结上传的文件、图像或纯文本内容。请根据以下用户输入的问题进行回答。

用户问题：{query}
"""

    model_list = [model_name] if model_name != "auto" else ["claude", "deepseek", "openai"]
    result = safe_call_model(prompt, session_id, model_list, history)

    papers = fetch_related_papers(query)
    paper_text = "\n\n📚 相关研究论文：\n" + "\n".join([f"- {p['title']}（{p['url']}）" for p in papers])
    combined_answer = result["answer"] + paper_text

    history += [{"role": "user", "content": query}, {"role": "assistant", "content": combined_answer}]
    if session_id in session_data:
        session_data[session_id]["history"] = history
    else:
        session_data[session_id] = {"history": history, "chunks": [], "index": None}

    return history, f"模型：{result['model']}  耗时：{result['elapsed']:.2f}s  估算费用：${result['cost']:.6f}"

with gr.Blocks() as demo:
    session_id = gr.State(f"session_{int(time.time())}_{np.random.randint(1000,9999)}")

    gr.Markdown("# 📷 图片问答 + 多模型文本对话系统")

    with gr.Row():
        image_upload = gr.File(label="上传文件（支持图片 / PDF / Word / 文本 / Excel）", file_types=[".jpg", ".png", ".jpeg", ".pdf", ".docx", ".txt", ".md", ".csv", ".xlsx"])
        upload_btn = gr.Button("识别并建立索引")

    upload_status = gr.Textbox(label="上传状态", interactive=False)
    gallery = gr.Gallery(label="上传图像")

    model_select = gr.Radio(
        ["auto", "claude", "deepseek", "openai", "text-only"],
        label="选择模型",
        value="auto"
    )

    task_mode = gr.Radio(
        ["default", "summary", "translation", "qa-extract"],
        label="任务类型",
        value="default"
    )

    chatbot = gr.Chatbot(label="聊天记录", height=400, type="messages")
    query_input = gr.Textbox(label="请输入问题")
    query_btn = gr.Button("发送查询")
    status_output = gr.Textbox(label="状态", interactive=False)

    image_uploaded = gr.State(False)

    upload_btn.click(
        fn=handle_file_upload,
        inputs=[image_upload, session_id],
        outputs=[upload_status, gallery, image_uploaded]
    )

    query_btn.click(
        fn=handle_user_query,
        inputs=[query_input, session_id, model_select, chatbot, image_uploaded, task_mode],
        outputs=[chatbot, status_output]
    )

    # 支持回车提交查询
    query_input.submit(
        fn=handle_user_query,
        inputs=[query_input, session_id, model_select, chatbot, image_uploaded, task_mode],
        outputs=[chatbot, status_output]
    )

if __name__ == "__main__":
    demo.launch()