第12章: Capstone · Enterprise AI Copilot

csv_file = st.file_uploader("📊 CSV", type=["csv"])
if csv_file is not None:
    csv_name = csv_file.name
    # 检查是否是新文件（避免重复处理）
    if st.session_state.uploaded_files.get("csv_name") != csv_name:
        csv_file.seek(0)  # 重置文件指针！
        df = pd.read_csv(csv_file)
        st.session_state.uploaded_files["csv_name"] = csv_name
        st.session_state.uploaded_files["csv_df"] = df

# 处理新消息前清理旧图表
if os.path.exists("chart.png"):
    os.remove("chart.png")

# 只有数据分析时才显示图表
if os.path.exists("chart.png") and "数据分析" in response:
    st.image("chart.png")

def extract_json(text: str) -> str:
    """从 LLM 输出中提取 JSON（容错处理）"""
    t = (text or "").strip()
    t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE)
    t = re.sub(r"\s*```$", "", t)
    m = re.search(r"\{.*\}", t, flags=re.S)
    return m.group(0) if m else t

# 使用时加 try-catch
try:
    obj = json.loads(extract_json(text))
except json.JSONDecodeError:
    obj = {"raw": text}  # 兜底

if "plt" in code or "matplotlib" in code:
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        exec_globals["plt"] = plt
    except ImportError:
        # matplotlib 未安装，移除绘图代码
        code = re.sub(r".*plt\..*\n?", "", code)
        print("注意：matplotlib 未安装，已跳过绘图代码\n")

"""
chatbot.py — Enterprise AI Copilot 聊天机器人

功能：
1. 连续多轮对话（聊天历史管理）
2. 意图识别（自动判断用户想做什么）
3. 文件上传（CSV/PDF/图片）并自动关联到对话
4. 能力路由（RAG/Agent/Vision/普通问答）
5. 全链路 Trace

启动方式：
    streamlit run chatbot.py
"""

import io
import json
import os
import re
import tempfile
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional

import pandas as pd
import streamlit as st

from langchain_community.chat_models import ChatTongyi
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pydantic import BaseModel, Field, ValidationError


# ═══════════════════════════════════════════════════════════════════════════
# 配置
# ═══════════════════════════════════════════════════════════════════════════

class Config:

    # 主要对话模型 - 用于RAG问答、数据分析等需要高质量输出的任务
    LLM_MODEL = "qwen-max"
    
    # 意图识别模型 - 使用轻量级模型降低成本，提高响应速度
    INTENT_MODEL = "qwen-turbo"
    
    # 视觉识别模型 - 支持图片理解和OCR
    VISION_MODEL = "qwen-vl-max"
    
    # 向量嵌入模型 - 用于文档语义检索
    EMBEDDING_MODEL = "text-embedding-v2"
    
    # 意图识别置信度阈值 - 低于此值的识别结果可能需要人工确认
    CONFIDENCE_THRESHOLD = 0.65
    
    # 文档分块配置 - 影响检索精度和token成本
    CHUNK_SIZE = 800        # 每个分块的最大字符数
    CHUNK_OVERLAP = 120     # 分块间的重叠字符数，避免语义断裂
    
    # 检索配置 - 控制返回的相关文档数量
    RETRIEVAL_K = 4         # 检索返回的最相关文档数量


# ═══════════════════════════════════════════════════════════════════════════
# 意图识别
# ═══════════════════════════════════════════════════════════════════════════

# 支持的意图类型 - 使用Literal类型确保类型安全
# rag_qa: 文档检索问答
# data_agent: 数据分析代理
# vision_extract: 图片识别提取
# general: 普通对话问答
IntentName = Literal["rag_qa", "data_agent", "vision_extract", "general"]


class IntentSlots(BaseModel):
    """
    意图槽位类 - 存储从用户输入中提取的具体信息
    
    属性说明：
    - query: 用户的核心查询内容
    - file_path: 相关文件路径（如果有）
    - image_path: 相关图片路径（如果有）
    """
    query: Optional[str] = None           # 用户的核心问题或请求
    file_path: Optional[str] = None        # 关联的文件路径
    image_path: Optional[str] = None       # 关联的图片路径


class IntentResult(BaseModel):
    """
    意图识别结果类 - 存储完整的意图识别信息
    
    属性说明：
    - intent: 识别出的意图列表（支持多意图）
    - slots: 提取的槽位信息
    - confidence: 识别置信度（0.0-1.0）
    """
    intent: List[IntentName] = Field(default_factory=list)  # 识别出的意图列表
    slots: IntentSlots = Field(default_factory=IntentSlots)  # 提取的槽位
    confidence: float = Field(ge=0.0, le=1.0)              # 置信度，限制在0-1之间


INTENT_PROMPT = ChatPromptTemplate.from_messages([
    ("system", 
     "你是企业智能助手【意图识别器】。你必须只输出严格 JSON（不要 Markdown，不要解释）。\n"
     "要求：\n"
     "- 支持多意图（intent 是数组）。\n"
     "- intent 只能取：rag_qa / data_agent / vision_extract / general。\n"
     "- slots 至少包含：query / file_path / image_path。缺失填 '?'。\n"
     "- confidence 为 0~1 的小数。\n"
     "\n"
     "意图判断规则：\n"
     "- 如果用户提到【数据分析】【统计】【画图】【CSV】【表格】【分析数据】等，用 data_agent\n"
     "- 如果用户提到【图片】【发票】【截图】【识别】【照片】【看图】等，用 vision_extract\n"
     "- 如果用户提到【文档】【PDF】【手册】【规定】【制度】【知识库】【查询文档】等，用 rag_qa\n"
     "- 其他情况用 general\n"),
    
    ("human", 
     "用户输入：用条形图展示各类别的数量分布\n"
     "只输出 JSON：\n"
     '{{"intent":["data_agent"],"slots":{{"query":"用条形图展示各类别的数量分布","file_path":"?","image_path":"?"}},"confidence":0.86}}'),
    
    ("human", 
     "用户输入：这张发票截图里金额是多少？\n"
     "只输出 JSON：\n"
     '{{"intent":["vision_extract"],"slots":{{"query":"金额是多少","file_path":"?","image_path":"?"}},"confidence":0.82}}'),
    
    ("human", 
     "用户输入：员工手册里病假最多几天？\n"
     "只输出 JSON：\n"
     '{{"intent":["rag_qa"],"slots":{{"query":"病假最多几天","file_path":"?","image_path":"?"}},"confidence":0.88}}'),
    
    ("human", "用户输入：{input}\n只输出 JSON："),
])


def extract_json(text: str) -> str:

    t = (text or "").strip()  # 去除首尾空白
    
    # 移除开头的Markdown代码块标记
    t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE)
    
    # 移除结尾的Markdown代码块标记
    t = re.sub(r"\s*```$", "", t)
    
    # 提取第一个完整的JSON对象
    m = re.search(r"\{.*\}", t, flags=re.S)
    return m.group(0) if m else t


def recognize_intent(user_input: str) -> IntentResult:

    # 使用轻量级模型进行意图识别，降低成本
    llm = ChatTongyi(model=Config.INTENT_MODEL, temperature=0)
    
    # 调用LLM进行意图识别
    resp = (INTENT_PROMPT | llm).invoke({"input": user_input})
    
    # 提取并解析JSON结果
    raw = json.loads(extract_json(resp.content))
    
    # 验证并返回结构化结果
    return IntentResult.model_validate(raw)


# ═══════════════════════════════════════════════════════════════════════════
# RAG 能力
# ═══════════════════════════════════════════════════════════════════════════

def process_pdf(pdf_bytes: bytes) -> List[Any]:

    # 创建临时文件保存PDF内容
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        tmp.write(pdf_bytes)
        tmp_path = tmp.name
    
    try:
        # 使用PyMuPDF加载PDF文档
        docs = PyMuPDFLoader(tmp_path).load()
        
        # 创建文档分块器
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=Config.CHUNK_SIZE,        # 每块最大字符数
            chunk_overlap=Config.CHUNK_OVERLAP,  # 块间重叠字符数
            is_separator_regex=True,             # 使用正则表达式分隔符
            keep_separator=False,                # 不保留分隔符在chunk中
            separators=[
                "(?<=。)",   # 在句号之后切分（句号留在当前 chunk 结尾）
                "(?<=！)",   # 在感叹号之后切分
                "(?<=？)",   # 在问号之后切分
                "(?<=，)",   # 在逗号之后切分（次优先级）
                " ",         # 空格（最低优先级）
            ],
        )
        
        # 执行文档分块
        return splitter.split_documents(docs)
        
    finally:
        # 清理临时文件
        os.unlink(tmp_path)


def rag_answer(question: str, chunks: List[Any]) -> Dict[str, Any]:

    # 初始化向量嵌入模型
    embeddings = DashScopeEmbeddings(model=Config.EMBEDDING_MODEL)
    
    # 创建向量数据库并添加文档
    # 注意：这里每次都重新创建，生产环境应该使用持久化存储
    store = Chroma.from_documents(chunks, embeddings)
    
    # 创建检索器，返回最相关的K个文档块
    retriever = store.as_retriever(search_kwargs={"k": Config.RETRIEVAL_K})
    
    # 检索相关文档
    results = retriever.invoke(question)
    
    # 构造上下文文本
    context_parts = []
    for i, doc in enumerate(results):
        page_num = doc.metadata.get("page", doc.metadata.get("page_number", "未知"))
        context_parts.append(f"【第 {page_num} 页】\n{doc.page_content}")
    context = "\n\n".join(context_parts)
    
    # 构造引用信息（页码和内容片段）
    citations = [
        {
            "page": doc.metadata.get("page", doc.metadata.get("page_number", "?")),  # 文档页码（支持多种metadata key）
            "content": doc.page_content[:150] + "..."  # 内容片段（前150字符）
        }
        for doc in results
    ]
    
    # 构造引用信息字符串（用于LLM理解如何标注来源）
    citation_guide = "\n".join([f"- 第 {c['page']} 页: {c['content'][:100]}..." for c in citations])
    
    # 构造RAG问答prompt - 明确要求LLM在答案中标注具体页码
    llm = ChatTongyi(model=Config.LLM_MODEL, temperature=0)
    prompt = f"""根据以下文档内容回答问题。

文档内容（已标注页码）：
{context}

问题：{question}

引用来源参考：
{citation_guide}

要求：
1. 只根据文档内容回答，不要编造
2. 如果文档中没有相关信息，请说"文档中未找到相关信息"
3. 在答案中标注来源页码，格式为"（来源页码：第X页）"或"（第X页）"
4. 对于每个关键信息，都要注明具体来源页码
"""
    
    # 调用LLM生成答案
    resp = llm.invoke(prompt)
    answer = resp.content if isinstance(resp.content, str) else resp.content[0].get("text", "")
    
    return {"answer": answer, "citations": citations}


# ═══════════════════════════════════════════════════════════════════════════
# 数据分析 Agent 能力
# ═══════════════════════════════════════════════════════════════════════════

def agent_answer(df: pd.DataFrame, question: str) -> Dict[str, Any]:
    """数据分析 Agent"""
    import sys
    import numpy as np
    
    df_meta = {
        "shape": df.shape,
        "columns": list(df.columns),
        "dtypes": df.dtypes.astype(str).to_dict(),
        "sample": df.sample(n=min(5, len(df)), random_state=42).to_string(index=False),
    }
    
    llm = ChatTongyi(model=Config.LLM_MODEL, temperature=0)
    prompt = f"""你是数据分析专家。根据以下 DataFrame 元信息和用户问题，生成 Pandas 代码。

DataFrame 元信息：
- shape: {df_meta['shape']}
- columns: {df_meta['columns']}
- dtypes: {df_meta['dtypes']}
- sample:
{df_meta['sample']}

用户问题：{question}

要求：
1. 只输出 Python 代码，不要解释
2. 代码中用 df 变量代表 DataFrame
3. 最后用 print() 输出结果
4. 如果需要画图，用 matplotlib 并保存到 chart.png
"""
    resp = llm.invoke(prompt)
    code = resp.content.strip()
    code = re.sub(r"^```(?:python)?\s*", "", code, flags=re.IGNORECASE)
    code = re.sub(r"\s*```$", "", code)
    
    # 执行代码
    old_stdout = sys.stdout
    sys.stdout = io.StringIO()
    
    try:
        exec_globals = {"df": df, "pd": pd, "np": np}
        
        # 尝试导入 matplotlib（如果代码中用到了）
        if "plt" in code or "matplotlib" in code:
            try:
                import matplotlib
                matplotlib.use("Agg")
                import matplotlib.pyplot as plt
                exec_globals["plt"] = plt
                exec_globals["matplotlib"] = matplotlib
            except ImportError:
                # matplotlib 未安装，从代码中移除绘图相关代码
                code = re.sub(r".*plt\..*\n?", "", code)
                code = re.sub(r".*matplotlib.*\n?", "", code)
                print("注意：matplotlib 未安装，已跳过绘图代码\n")
        
        exec(code, exec_globals)
        output = sys.stdout.getvalue()
    except Exception as e:
        output = f"执行出错: {e}"
    finally:
        sys.stdout = old_stdout
    
    chart_path = "chart.png" if os.path.exists("chart.png") else None
    return {"answer": output, "code": code, "chart": chart_path}


# ═══════════════════════════════════════════════════════════════════════════
# Vision 能力
# ═══════════════════════════════════════════════════════════════════════════

def vision_answer(image_path: str, question: str) -> Dict[str, Any]:
    """Vision 图片识别"""
    if not image_path.startswith("file://"):
        image_path = f"file://{os.path.abspath(image_path)}"
    
    vision_llm = ChatTongyi(model=Config.VISION_MODEL, temperature=0)
    prompt = (
        f"你是识图助手。请根据图片内容回答：{question}\n"
        '请严格输出 JSON：{{"summary": "图片内容概述", "entities": ["识别出的实体"], "answer": "针对问题的答案"}}'
    )
    
    msg = HumanMessage(content=[
        {"text": prompt},
        {"image": image_path},
    ])
    
    resp = vision_llm.invoke([msg])
    text = resp.content[0]["text"] if isinstance(resp.content, list) else resp.content
    
    try:
        obj = json.loads(extract_json(text))
        return {"answer": obj.get("answer", obj.get("summary", "")), "contract": obj}
    except json.JSONDecodeError:
        return {"answer": text, "contract": {"raw": text}}


# ═══════════════════════════════════════════════════════════════════════════
# 普通问答
# ═══════════════════════════════════════════════════════════════════════════

def general_answer(question: str, chat_history: List[Dict]) -> str:
    """普通问答（带上下文）"""
    llm = ChatTongyi(model=Config.LLM_MODEL, temperature=0.7)
    
    # 构造带历史的 prompt
    messages = []
    for msg in chat_history[-10:]:  # 只取最近10轮
        if msg["role"] == "user":
            messages.append(HumanMessage(content=msg["content"]))
        else:
            messages.append(AIMessage(content=msg["content"]))
    messages.append(HumanMessage(content=question))
    
    resp = llm.invoke(messages)
    return resp.content if isinstance(resp.content, str) else resp.content[0].get("text", "")


# ═══════════════════════════════════════════════════════════════════════════
# Streamlit 聊天界面
# ═══════════════════════════════════════════════════════════════════════════

def init_session_state():
    """初始化 session state"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "uploaded_files" not in st.session_state:
        st.session_state.uploaded_files = {
            "csv_name": None,
            "csv_df": None,
            "pdf_name": None,
            "pdf_chunks": None,
            "image_name": None,
            "image_path": None,
            "image_bytes": None,
        }


def render_sidebar():
    """渲染侧边栏 - 简化为只保留API Key和清空按钮"""
    with st.sidebar:
        st.header("⚙️ 配置")
        
        # API Key
        api_key = st.text_input(
            "DashScope API Key",
            type="password",
            value=os.getenv("DASHSCOPE_API_KEY", ""),
        )
        if api_key:
            os.environ["DASHSCOPE_API_KEY"] = api_key
        
        st.divider()
        
        # 已上传文件状态
        st.header("📁 已上传文件")
        files = st.session_state.uploaded_files
        
        has_files = False
        if files.get("csv_df") is not None:
            st.success(f"📊 CSV: {files.get('csv_name', 'unknown')}")
            has_files = True
        if files.get("pdf_chunks") is not None:
            st.success(f"📄 PDF: {files.get('pdf_name', 'unknown')}")
            has_files = True
        if files.get("image_path") is not None:
            st.success(f"🖼️ 图片: {files.get('image_name', 'unknown')}")
            has_files = True
        
        if not has_files:
            st.info("暂无文件，请在底部聊天区上传")
        
        st.divider()
        
        # 清空按钮
        col1, col2 = st.columns(2)
        with col1:
            if st.button("🗑️ 清空对话", use_container_width=True):
                st.session_state.messages = []
                st.rerun()
        with col2:
            if st.button("📤 清空文件", use_container_width=True):
                st.session_state.uploaded_files = {
                    "csv_name": None, "csv_df": None,
                    "pdf_name": None, "pdf_chunks": None,
                    "image_name": None, "image_path": None, "image_bytes": None,
                }
                st.rerun()


def get_intent_badge(intent: str) -> str:
    """获取意图徽章"""
    badges = {
        "rag_qa": "📄 RAG",
        "data_agent": "📊 数据分析",
        "vision_extract": "🖼️ Vision",
        "general": "💬 普通问答",
    }
    return badges.get(intent, intent)


def process_user_message(user_input: str) -> str:
    """处理用户消息，返回 AI 回复"""
    files = st.session_state.uploaded_files
    
    # 1. 意图识别
    try:
        intent_result = recognize_intent(user_input)
    except Exception as e:
        return f"❌ 意图识别失败: {e}"
    
    # 显示意图识别结果
    intent_str = ", ".join([get_intent_badge(i) for i in intent_result.intent])
    confidence_str = f"{intent_result.confidence:.0%}"
    
    response_parts = [f"🎯 **意图识别**: {intent_str} (置信度: {confidence_str})\n"]
    
    # 2. 根据意图路由
    for intent in intent_result.intent:
        if intent == "rag_qa":
            if files["pdf_chunks"] is None:
                response_parts.append("⚠️ **RAG 问答**需要先上传 PDF 文件（在左侧边栏上传）\n")
            else:
                with st.spinner("正在检索文档并生成答案..."):
                    result = rag_answer(user_input, files["pdf_chunks"])
                response_parts.append(f"📄 **RAG 回答**:\n{result['answer']}\n")
                if result["citations"]:
                    citations_str = "\n".join([f"  - 第 {c['page']} 页: {c['content'][:50]}..." for c in result["citations"][:3]])
                    response_parts.append(f"\n📚 **引用来源**:\n{citations_str}\n")
        
        elif intent == "data_agent":
            if files["csv_df"] is None:
                response_parts.append("⚠️ **数据分析**需要先上传 CSV 文件（在左侧边栏上传）\n")
            else:
                with st.spinner("正在生成并执行代码..."):
                    result = agent_answer(files["csv_df"], user_input)
                response_parts.append(f"📊 **数据分析结果**:\n```\n{result['answer']}\n```\n")
                response_parts.append(f"\n📝 **生成的代码**:\n```python\n{result['code']}\n```\n")
                if result["chart"]:
                    response_parts.append("\n📈 **图表已生成** (chart.png)\n")
        
        elif intent == "vision_extract":
            if files["image_path"] is None:
                response_parts.append("⚠️ **图片识别**需要先上传图片（在左侧边栏上传）\n")
            else:
                with st.spinner("正在识别图片..."):
                    result = vision_answer(files["image_path"], user_input)
                response_parts.append(f"🖼️ **图片识别结果**:\n{result['answer']}\n")
                if "contract" in result and result["contract"]:
                    response_parts.append(f"\n📋 **结构化数据**:\n```json\n{json.dumps(result['contract'], ensure_ascii=False, indent=2)}\n```\n")
        
        else:  # general
            with st.spinner("正在思考..."):
                answer = general_answer(user_input, st.session_state.messages)
            response_parts.append(f"💬 **回答**:\n{answer}\n")
    
    return "\n".join(response_parts)


def render_chat_input_area():
    """渲染聊天输入区域"""
    
    # 初始化
    for k in ["csv", "pdf", "img"]:
        if f"show_{k}" not in st.session_state:
            st.session_state[f"show_{k}"] = False
    
    # 显示已上传文件
    f = st.session_state.uploaded_files
    files = []
    if f.get("csv_df") is not None: files.append("📊CSV")
    if f.get("pdf_chunks") is not None: files.append("📄PDF") 
    if f.get("image_path") is not None: files.append("🖼️图片")
    
    if files:
        st.caption(f"已上传: {' | '.join(files)}")
    
    # 统一的输入区域容器
    with st.container(border=True):
        # 文件上传按钮行
        c1, c2, c3, c4 = st.columns(4)
        
        with c1:
            if st.button("📁 CSV", use_container_width=True, type="secondary"):
                st.session_state.show_csv = not st.session_state.get("show_csv", False)
                for k in ["pdf", "img"]: st.session_state[f"show_{k}"] = False
                st.rerun()
        
        with c2:
            if st.button("📄 PDF", use_container_width=True, type="secondary"):
                st.session_state.show_pdf = not st.session_state.get("show_pdf", False)
                for k in ["csv", "img"]: st.session_state[f"show_{k}"] = False
                st.rerun()
        
        with c3:
            if st.button("🖼️ 图片", use_container_width=True, type="secondary"):
                st.session_state.show_img = not st.session_state.get("show_img", False)
                for k in ["csv", "pdf"]: st.session_state[f"show_{k}"] = False
                st.rerun()
        
        with c4:
            if st.button("🗑️ 清空", use_container_width=True, type="secondary"):
                st.session_state.uploaded_files = {k: None for k in st.session_state.uploaded_files}
                for k in ["csv", "pdf", "img"]: st.session_state[f"show_{k}"] = False
                st.rerun()
        
        # 文件上传器（在按钮下方显示）
        if st.session_state.get("show_csv"):
            file = st.file_uploader("选择CSV文件", type=["csv"], key="up_csv")
            if file:
                st.session_state.uploaded_files["csv_df"] = pd.read_csv(file)
                st.session_state.uploaded_files["csv_name"] = file.name
                st.session_state.show_csv = False
                st.rerun()
        
        if st.session_state.get("show_pdf"):
            file = st.file_uploader("选择PDF文件", type=["pdf"], key="up_pdf")
            if file:
                with st.spinner("处理PDF..."):
                    st.session_state.uploaded_files["pdf_chunks"] = process_pdf(file.read())
                    st.session_state.uploaded_files["pdf_name"] = file.name
                st.session_state.show_pdf = False
                st.rerun()
        
        if st.session_state.get("show_img"):
            file = st.file_uploader("选择图片", type=["png","jpg","jpeg"], key="up_img")
            if file:
                img_bytes = file.read()
                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                    tmp.write(img_bytes)
                    st.session_state.uploaded_files["image_path"] = tmp.name
                st.session_state.uploaded_files["image_name"] = file.name
                st.session_state.uploaded_files["image_bytes"] = img_bytes
                st.session_state.show_img = False
                st.rerun()
    
    # 聊天输入框（在容器外，底部）
    return st.chat_input("发送消息...")


def main():

    st.set_page_config(
        page_title="Enterprise AI Copilot",
        page_icon="🤖",
        layout="wide",
    )
    
    init_session_state()
    
    render_sidebar()
    
    st.title("🤖 Enterprise AI Copilot")
    st.caption("智能对话助手 | 支持 RAG 问答 · 数据分析 · 图片识别 · 多轮对话")
    
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            # 如果有图表，显示图表
            if message.get("chart") and os.path.exists(message["chart"]):
                st.image(message["chart"])
    
    prompt = render_chat_input_area()
    
    if prompt:
        # 检查 API Key
        if not os.getenv("DASHSCOPE_API_KEY"):
            st.error("⚠️ 请先在侧边栏输入 DashScope API Key")
            st.stop()
        
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        if os.path.exists("chart.png"):
            os.remove("chart.png")
        
        with st.chat_message("assistant"):
            response = process_user_message(prompt)
            st.markdown(response)
            
            chart_path = None
            if os.path.exists("chart.png") and "数据分析" in response:
                st.image("chart.png")
                chart_path = "chart.png"
        
        st.session_state.messages.append({
            "role": "assistant",
            "content": response,
            "chart": chart_path,
        })

if __name__ == "__main__":
    main()