（八）现有架构上新增一个通用Excel导出工具

新增 Excel 导出工具完整方案

我会在你现有架构上新增一个通用Excel导出工具，Agent 可以自动把查询到的数据生成 Excel 文件，返回下载链接。支持固化接口数据和 SQL 查询数据导出，前端点击链接就能下载文件。

---

一、先做准备

1. 安装 Excel 生成依赖

pip install openpyxl

2. 新增功能说明

• 新增 export_to_excel 工具：接收结构化数据，自动生成 Excel 文件
• 新增文件下载接口：返回可直接下载的文件流
• Agent 自动识别导出需求：用户说「导出、下载Excel、生成报表」时，自动先查数据再导出
• 自动调整列宽、支持自定义文件名、临时文件自动过期清理

---

二、完整修改后的代码（标注新增部分）

在你之前的 main.py 基础上修改，新增部分都有注释标注，直接替换对应部分即可。

from dotenv import load_dotenv
import os
import ast
import uuid
from contextlib import asynccontextmanager
from typing import List, Dict
from datetime import datetime, timedelta

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_core.prompts import ChatPromptTemplate

# ===================== 【新增】配置临时文件目录 =====================
# 存放生成的Excel文件，定期自动清理
TEMP_DIR = "temp_excel_files"
# 文件过期时间（小时）
EXPIRE_HOURS = 1

# 确保目录存在
os.makedirs(TEMP_DIR, exist_ok=True)

load_dotenv()

# 全局变量
sql_agent = None
agent_executor = None

# ===================== 1. 业务查询函数（固化接口） =====================
def get_fleet_vehicles(company_name: str, fleet_name: str) -> List[Dict]:
    """查询指定分公司、指定车队的车辆清单"""
    return [
        {"车辆自编号": "11901", "车牌号": "沪A12345", "车型": "宇通ZK6125", "状态": "运营中"},
        {"车辆自编号": "11902", "车牌号": "沪A12346", "车型": "宇通ZK6125", "状态": "维修中"},
        {"车辆自编号": "11903", "车牌号": "沪A12347", "车型": "申龙SLK6119", "状态": "运营中"},
    ]

def get_vehicle_repair_history(bus_no: str) -> List[Dict]:
    """查询指定车辆的历史维修记录"""
    return [
        {"工单号": "R202609001", "故障类型": "电机故障", "维修时间": "2026-09-01", "维修内容": "更换轴承"},
        {"工单号": "R202608015", "故障类型": "电路故障", "维修时间": "2026-08-15", "维修内容": "更换线束"},
    ]

def get_fleet_monthly_stats(fleet_name: str, month: str) -> Dict:
    """查询指定车队某月的维修统计"""
    return {
        "车队名称": fleet_name,
        "统计月份": month,
        "总工单数量": 28,
        "已完成工单": 25,
        "故障类型TOP3": "电机系统、底盘系统、电气系统",
        "平均维修工时(小时)": 4.2
    }

# ===================== 2. 业务工具封装（返回纯数据，方便复用导出） =====================
@tool
def query_fleet_vehicles(company_name: str, fleet_name: str) -> str:
    """
    查询指定分公司和车队的车辆清单，返回车辆数据列表。
    参数：
    - company_name: 分公司名称，比如"第一运营分公司"、"一公司"
    - fleet_name: 车队名称，比如"119车队"
    当用户查询某个车队的车辆清单、车辆列表时，优先使用此工具。
    """
    if company_name in ["一公司", "第一分公司"]:
        company_name = "第一运营分公司"
    
    result = get_fleet_vehicles(company_name, fleet_name)
    return str(result)

@tool
def query_vehicle_repair_history(bus_no: str) -> str:
    """
    查询指定车辆的历史维修记录，返回维修记录列表。
    参数：
    - bus_no: 车辆自编号，比如"11901"
    当用户查询某辆车的维修记录、维修历史、修过什么时，优先使用此工具。
    """
    result = get_vehicle_repair_history(bus_no)
    return str(result)

@tool
def query_fleet_monthly_stats(fleet_name: str, month: str) -> str:
    """
    查询指定车队某个月的维修统计数据。
    参数：
    - fleet_name: 车队名称，比如"119车队"
    - month: 月份，格式为YYYY-MM，比如"2026-09"
    当用户查询车队月度统计、维修数据、工单数量时，优先使用此工具。
    """
    result = get_fleet_monthly_stats(fleet_name, month)
    return str(result)

BUSINESS_TOOLS = [query_fleet_vehicles, query_vehicle_repair_history, query_fleet_monthly_stats]

# ===================== 【新增】3. Excel导出工具 =====================
@tool
def export_to_excel(data: str, file_title: str = "数据导出") -> str:
    """
    将查询到的数据导出为Excel文件，返回下载链接。
    参数：
    - data: 要导出的数据，必须是Python列表字典格式的字符串，比如"[{'车辆编号':'11901'}, ...]"
    - file_title: 导出文件的名称，比如"119车队车辆清单"
    当用户要求导出数据、下载Excel、生成报表、导出表格时，使用此工具。
    使用步骤：先调用查询工具获取数据，再调用此工具导出。
    """
    # 清理过期文件
    _clean_expired_files()
    
    # 解析数据字符串为Python对象
    try:
        data_list = ast.literal_eval(data.strip())
    except Exception as e:
        return f"数据解析失败，无法导出：{str(e)}"
    
    # 数据校验
    if not isinstance(data_list, list) or len(data_list) == 0:
        return "没有有效数据可导出"
    
    if not isinstance(data_list[0], dict):
        return "数据格式错误，必须是字典列表格式"
    
    try:
        # 生成唯一文件ID，防止重名
        file_id = str(uuid.uuid4())
        file_name = f"{file_title}_{datetime.now().strftime('%Y%m%d%H%M%S')}.xlsx"
        file_path = os.path.join(TEMP_DIR, f"{file_id}.xlsx")
        
        # 创建Excel工作簿
        import openpyxl
        wb = openpyxl.Workbook()
        ws = wb.active
        # 工作表名（最多31字符）
        ws.title = file_title[:31] if len(file_title) <= 31 else "数据"
        
        # 写入表头
        headers = list(data_list[0].keys())
        ws.append(headers)
        
        # 写入数据行
        for row in data_list:
            ws.append([str(row.get(h, "")) for h in headers])
        
        # 自动调整列宽
        for col in ws.columns:
            max_len = 0
            col_letter = col[0].column_letter
            for cell in col:
                try:
                    if len(str(cell.value)) > max_len:
                        max_len = len(str(cell.value))
                except:
                    pass
            # 最小宽度8，最大宽度50
            ws.column_dimensions[col_letter].width = min(max(max_len + 2, 8), 50)
        
        # 保存文件
        wb.save(file_path)
        
        # 生成下载链接
        download_url = f"[http://localhost:8000/api/download/{file_id}?filename={file_name}](http://localhost:8000/api/download/{file_id}?filename={file_name})"
        return f"Excel文件已生成，点击下载：{download_url}\n文件有效期1小时，请及时保存。"
    
    except Exception as e:
        return f"生成Excel失败：{str(e)}"

def _clean_expired_files():
    """清理过期的临时Excel文件"""
    try:
        cutoff = datetime.now() - timedelta(hours=EXPIRE_HOURS)
        for filename in os.listdir(TEMP_DIR):
            file_path = os.path.join(TEMP_DIR, filename)
            if os.path.isfile(file_path):
                mtime = datetime.fromtimestamp(os.path.getmtime(file_path))
                if mtime < cutoff:
                    os.remove(file_path)
    except:
        pass

# ===================== 4. 服务启动初始化 =====================
@asynccontextmanager
async def lifespan(app: FastAPI):
    global sql_agent, agent_executor

    print("正在初始化数据库连接...")
    db = SQLDatabase.from_uri(
        os.getenv("MYSQL_URI"),
        include_tables=["bus_vehicle", "bus_fleet", "repair_order"],
        sample_rows_in_table_info=2,
        view_support=False
    )

    print("正在初始化大模型...")
    llm = ChatOpenAI(
        model=os.getenv("MODEL_NAME"),
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url=os.getenv("OPENAI_BASE_URL"),
        temperature=0
    )

    # SQL查询工具
    print("正在创建SQL查询Agent...")
    sql_agent = create_sql_agent(
        llm=llm,
        db=db,
        agent_type="openai-tools",
        verbose=False,
        max_iterations=10,
        handle_parsing_errors=True
    )

    @tool
    def sql_database_query(query: str) -> str:
        """
        通过自然语言灵活查询公交维修数据库，支持筛选、统计、关联查询。
        当没有现成的业务工具可以满足需求时，使用此工具。
        参数：
        - query: 用户的自然语言查询问题
        """
        result = sql_agent.invoke({"input": query})
        return result["output"]

    # 合并所有工具：业务工具 + SQL工具 + Excel导出工具
    all_tools = BUSINESS_TOOLS + [sql_database_query, export_to_excel]

    # 创建总调度Agent
    print("正在创建总调度Agent...")
    
    system_prompt = """
    你是公交机务数据查询助手，可以调用多个工具回答用户问题。

    【工具优先级规则】
    1. 优先使用业务工具（query_fleet_vehicles、query_vehicle_repair_history、query_fleet_monthly_stats）
    2. 业务工具满足不了时，使用 sql_database_query 灵活查询
    3. 用户要求导出、下载Excel、生成报表时，先调用对应查询工具获取数据，再调用 export_to_excel 导出文件

    【导出工具使用说明】
    - export_to_excel 的 data 参数必须传入查询工具返回的纯数据字符串
    - file_title 参数根据查询内容命名，比如"119车队车辆清单"
    - 导出完成后，把下载链接告诉用户

    【回答要求】
    1. 用自然语言回答，先给结论，再列数据
    2. 不要暴露你调用了什么工具，也不要提到SQL、接口等技术术语
    3. 如果参数不明确，直接询问用户补充信息，不要猜测
    """

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("user", "{input}"),
        ("agent_scratchpad", "{agent_scratchpad}"),
    ])

    agent = create_openai_tools_agent(llm, all_tools, prompt)
    agent_executor = AgentExecutor(
        agent=agent,
        tools=all_tools,
        verbose=True,
        max_iterations=10,
        handle_parsing_errors=True
    )

    print("服务启动完成！")
    yield
    print("服务已关闭")

# ===================== 5. FastAPI应用 =====================
app = FastAPI(
    title="公交机务智能查询系统",
    description="固化接口+灵活SQL+Excel导出双模式智能查询",
    version="2.1.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ===================== 6. 对外接口 =====================
class QueryRequest(BaseModel):
    question: str

class QueryResponse(BaseModel):
    code: int = 200
    message: str = "success"
    data: str = ""

@app.get("/health", summary="健康检查")
async def health_check():
    return {"code": 200, "message": "服务运行正常"}

@app.post("/api/query", summary="自然语言智能查询", response_model=QueryResponse)
async def query_data(request: QueryRequest):
    if not agent_executor:
        raise HTTPException(status_code=500, detail="服务未初始化完成")

    question = request.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="问题不能为空")

    try:
        result = agent_executor.invoke({"input": question})
        return QueryResponse(
            code=200,
            message="success",
            data=result["output"]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败：{str(e)}")

# ===================== 【新增】7. Excel文件下载接口 =====================
@app.get("/api/download/{file_id}", summary="下载Excel文件")
async def download_excel(file_id: str, filename: str = "数据导出.xlsx"):
    """
    下载生成的Excel文件
    - file_id: 文件唯一ID
    - filename: 下载时显示的文件名
    """
    file_path = os.path.join(TEMP_DIR, f"{file_id}.xlsx")
    
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="文件不存在或已过期")
    
    return FileResponse(
        path=file_path,
        filename=filename,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )

---

三、使用测试

启动服务后，访问 [http://localhost:8000/docs](http://localhost:8000/docs)，测试以下场景：

场景1：直接查询（走固化接口）

输入：一公司119车队的车辆清单

效果：正常返回车辆列表

场景2：直接导出Excel

输入：把一公司119车队的车辆清单导出成Excel

效果：Agent 自动调用 query_fleet_vehicles 获取数据 → 调用 export_to_excel 生成文件 → 返回下载链接

点击链接直接下载 Excel 文件。

场景3：先查询再导出

输入：

1. 查11901号车的维修记录
2. 把这些记录导出成Excel
   效果：Agent 记住上下文，用上一轮的数据直接导出。

场景4：SQL查询结果导出

输入：近30天哪个车队工单最多，把结果导出成Excel

效果：先调用 SQL 工具查询统计结果，再自动导出 Excel。

---

四、关键说明

1. 文件有效期

• 生成的 Excel 文件默认保留 1 小时，过期自动清理
• 可以修改 EXPIRE_HOURS 调整有效期

2. 生产环境优化

• 把 localhost:8000 换成你的正式域名或IP
• 可以加下载鉴权，比如校验用户登录状态
• 文件量大的话，存到对象存储（OSS/MinIO） instead 本地磁盘

3. 扩展支持更多数据导出

只要是「列表字典」格式的数据，都可以用这个导出工具。新增业务接口后，不用改导出工具，直接就能用。

4. 常见问题

• 导出失败提示数据格式错误：检查查询工具返回的是不是纯列表字典字符串，不要加多余的文字前缀
• 下载链接打不开：确认服务地址和端口正确，文件没过期
• 中文文件名乱码：FileResponse 会自动处理，大部分浏览器正常显示