from fastapi import FastAPI

app = FastAPI(
    title="Excel Data Analysis API",
    description="基于Pandas DataFrame的數(shù)據(jù)分析API",
    version="1.0.0"
)

@app.on_event("startup")
async def startup_event():
    """應(yīng)用啟動(dòng)時(shí)初始化模型和代理"""
    try:
        initialize_agent()
        print("模型和代理初始化成功")
    except Exception as e:
        print(f"初始化失敗: {e}")
        raise

def initialize_agent():
    global llm, pandas_agent
    
    # 讀取Excel文件中的所有工作表
    excel_file = pd.ExcelFile(path)
    all_sheets = {}
    for sheet_name in excel_file.sheet_names:
        all_sheets[sheet_name] = pd.read_excel(excel_file, sheet_name=sheet_name)
    
    # 合并所有工作表
    df = pd.concat(all_sheets.values(), ignore_index=True)

pandas_agent = create_pandas_dataframe_agent(
    llm,  # 通義千問模型實(shí)例
    df,   # 合并后的DataFrame
    verbose=True,  # 啟用詳細(xì)日志
    agent_type="zero-shot-react-description",  # 代理類型
    allow_dangerous_code=True,  # 允許執(zhí)行代碼
    max_iterations=5  # 最大迭代次數(shù)
)

class StreamingCallbackHandler(BaseCallbackHandler):
    """自定義回調(diào)處理器，用于實(shí)現(xiàn)真正的流式輸出"""

    def __init__(self):
        self.tokens = []
        self.finished = False

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """當(dāng)LLM生成新token時(shí)調(diào)用"""
        self.tokens.append(token)

    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
        """當(dāng)LLM生成結(jié)束時(shí)調(diào)用"""
        self.finished = True

@app.post("/chat", response_model=ChatResponse)
async def chat_with_data(request: ChatRequest):
    # ... 省略其他代碼 ...
    
    if request.stream:
        return StreamingResponse(
            true_stream_response(full_input, session_history, request.session_id),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            }
        )

conversation_history = {}

# 在聊天接口中維護(hù)會(huì)話歷史
session_history = conversation_history.get(request.session_id, [])
session_history.append(request.message)
conversation_history[request.session_id] = session_history

@app.delete("/clear_history/{session_id}")
async def clear_history(session_id: str):
    """清除指定會(huì)話的歷史記錄"""
    if session_id in conversation_history:
        del conversation_history[session_id]
    return {"message": f"會(huì)話 {session_id} 的歷史記錄已清除"}

class ChatRequest(BaseModel):
    session_id: str  # 會(huì)話ID，用于區(qū)分不同用戶的對話歷史
    message: str  # 用戶消息
    history_length: Optional[int] = 5  # 歷史消息長度，默認(rèn)為5
    stream: Optional[bool] = False  # 是否使用流式響應(yīng)

class ChatResponse(BaseModel):
    session_id: str
    response: str
    success: bool
    error: Optional[str] = None

uvicorn main:app --host 0.0.0.0 --port 9113 --workers 4 --timeout-keep-alive 300

# main.py
import pandas as pd
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_experimental.agents import create_pandas_dataframe_agent
import os
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from typing import List, Optional
from fastapi.responses import StreamingResponse
import asyncio
import json
import time
from langchain.callbacks.base import BaseCallbackHandler
from typing import Any, Dict, List
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

# 初始化 FastAPI 應(yīng)用
app = FastAPI(title="Excel Data Analysis API", description="基于 Pandas DataFrame 的數(shù)據(jù)分析 API")

# 全局變量存儲(chǔ)模型和代理
llm = None
pandas_agent = None
conversation_history = {}


class StreamingCallbackHandler(BaseCallbackHandler):
    """自定義回調(diào)處理器，用于實(shí)現(xiàn)真正的流式輸出"""

    def __init__(self):
        self.tokens = []
        self.finished = False

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """當(dāng)LLM生成新token時(shí)調(diào)用"""
        self.tokens.append(token)

    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
        """當(dāng)LLM生成結(jié)束時(shí)調(diào)用"""
        self.finished = True

    def get_tokens(self):
        """獲取已生成的tokens"""
        return self.tokens


class ChatRequest(BaseModel):
    session_id: str  # 會(huì)話ID，用于區(qū)分不同用戶的對話歷史
    message: str  # 用戶消息
    history_length: Optional[int] = 5  # 歷史消息長度，默認(rèn)為5
    stream: Optional[bool] = False  # 是否使用流式響應(yīng)


class ChatResponse(BaseModel):
    session_id: str
    response: str
    success: bool
    error: Optional[str] = None


def initialize_agent():
    """初始化模型和數(shù)據(jù)代理"""
    global llm, pandas_agent

    path = r'./data/1.xlsx'

    # 檢查文件是否存在
    if not os.path.exists(path):
        raise FileNotFoundError(f"文件未找到：{path}")

    # 讀取 Excel 中的所有工作表
    excel_file = pd.ExcelFile(path)
    all_sheets = {}
    for sheet_name in excel_file.sheet_names:
        all_sheets[sheet_name] = pd.read_excel(excel_file, sheet_name=sheet_name)

    # 合并所有工作表到一個(gè)DataFrame中
    df = pd.concat(all_sheets.values(), ignore_index=True)

    os.environ["DASHSCOPE_API_KEY"] = 'sk-50254f6d4df1ab3c9baf30093c4e'

    llm = ChatTongyi(
        model="qwen-max-latest",
        temperature=0.4,
        streaming=True  # 啟用流式輸出
    )

    # 創(chuàng)建excel_agent
    pandas_agent = create_pandas_dataframe_agent(
        llm,
        df,
        verbose=True,
        agent_type="zero-shot-react-description",
        allow_dangerous_code=True,
        max_iterations=5
    )


@app.on_event("startup")
async def startup_event():
    """應(yīng)用啟動(dòng)時(shí)初始化模型和代理"""
    try:
        initialize_agent()
        print("模型和代理初始化成功")
    except Exception as e:
        print(f"初始化失敗: {e}")
        raise


@app.post("/chat", response_model=ChatResponse)
async def chat_with_data(request: ChatRequest):
    """與數(shù)據(jù)進(jìn)行對話（支持流式和非流式響應(yīng)）"""
    global pandas_agent, conversation_history

    if not pandas_agent:
        raise HTTPException(status_code=500, detail="模型未初始化")

    try:
        # 獲取或創(chuàng)建會(huì)話歷史
        session_history = conversation_history.get(request.session_id, [])

        # 構(gòu)建歷史文本
        history_text = "\n".join(session_history[-request.history_length:])

        # 構(gòu)造帶上下文的輸入
        if history_text:
            full_input = f"聊天歷史：{history_text}，當(dāng)前問題：{request.message}，請根據(jù)歷史和當(dāng)前問題，不要截?cái)噍敵?，展示全部?nèi)容，不要總結(jié)，嚴(yán)格按照查詢內(nèi)容輸出,不要多余輸出，如果無法確定指代對象，請?jiān)儐栍脩舫吻?，并且不要編造?.strip()
        else:
            full_input = request.message

        # 如果請求流式響應(yīng)
        if request.stream:
            return StreamingResponse(
                true_stream_response(full_input, session_history, request.session_id),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                }
            )

        # 非流式響應(yīng)
        response = pandas_agent.invoke({"input": full_input})
        output = response['output'] if isinstance(response, dict) else str(response)

        # 更新會(huì)話歷史
        session_history.append(request.message)
        conversation_history[request.session_id] = session_history

        return ChatResponse(
            session_id=request.session_id,
            response=output,
            success=True
        )

    except Exception as e:
        if request.stream:
            # 對于流式請求，返回流式錯(cuò)誤響應(yīng)
            async def error_stream():
                yield f"data: {json.dumps({'error': str(e)})}\n\n"

            return StreamingResponse(error_stream(), media_type="text/event-stream")
        else:
            return ChatResponse(
                session_id=request.session_id,
                response="",
                success=False,
                error=str(e)
            )


async def true_stream_response(input_text: str, session_history: list, session_id: str):
    """真正的流式響應(yīng)生成器"""
    global pandas_agent, conversation_history

    try:
        # 使用回調(diào)處理器捕獲流式輸出
        callback_handler = StreamingCallbackHandler()

        # 在后臺(tái)運(yùn)行代理處理
        async def run_agent():
            try:
                # 調(diào)用代理，傳入回調(diào)處理器
                response = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: pandas_agent.invoke(
                        {"input": input_text},
                        {"callbacks": [callback_handler]}
                    )
                )

                # 確保結(jié)束標(biāo)記被設(shè)置
                callback_handler.finished = True

            except Exception as e:
                print(f"代理執(zhí)行錯(cuò)誤: {e}")
                callback_handler.finished = True

        # 啟動(dòng)代理任務(wù)
        agent_task = asyncio.create_task(run_agent())

        # 流式輸出循環(huán)
        last_token_count = 0
        start_time = time.time()
        max_wait_time = 60  # 最大等待時(shí)間60秒

        while not callback_handler.finished and (time.time() - start_time) < max_wait_time:
            current_tokens = callback_handler.get_tokens()

            # 如果有新token，發(fā)送給客戶端
            if len(current_tokens) > last_token_count:
                for i in range(last_token_count, len(current_tokens)):
                    token_data = {
                        "token": current_tokens[i],
                        "type": "token"
                    }
                    yield f"data: {json.dumps(token_data)}\n\n"

                last_token_count = len(current_tokens)

            # 短暫等待后繼續(xù)檢查
            await asyncio.sleep(0.1)

        # 如果超時(shí)，發(fā)送錯(cuò)誤信息
        if (time.time() - start_time) >= max_wait_time:
            error_data = {"error": "請求超時(shí)", "type": "error"}
            yield f"data: {json.dumps(error_data)}\n\n"
        else:
            # 發(fā)送完成標(biāo)記
            done_data = {"done": True, "type": "done"}
            yield f"data: {json.dumps(done_data)}\n\n"

            # 更新會(huì)話歷史
            current_message = input_text.split("當(dāng)前問題：")[1].split("，")[
                0] if "當(dāng)前問題：" in input_text else input_text
            session_history.append(current_message)
            conversation_history[session_id] = session_history

        # 等待代理任務(wù)完成（如果還未完成）
        if not agent_task.done():
            agent_task.cancel()

    except Exception as e:
        error_data = {"error": f"流式處理錯(cuò)誤: {str(e)}", "type": "error"}
        yield f"data: {json.dumps(error_data)}\n\n"


@app.delete("/clear_history/{session_id}")
async def clear_history(session_id: str):
    """清除指定會(huì)話的歷史記錄"""
    if session_id in conversation_history:
        del conversation_history[session_id]
    return {"message": f"會(huì)話 {session_id} 的歷史記錄已清除"}


@app.get("/health")
async def health_check():
    """健康檢查接口"""
    return {
        "status": "healthy",
        "model_loaded": pandas_agent is not None,
        "active_sessions": len(conversation_history)
    }


@app.get("/stream_test")
async def stream_test():
    """測試流式接口"""

    async def generate_test_data():
        for i in range(10):
            yield f"data: 測試消息 {i}\n\n"
            await asyncio.sleep(1)

    return StreamingResponse(generate_test_data(), media_type="text/event-stream")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=9113)