knowledge-agent/agent.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用 FastAPI + LangGraph 重构的 Agent 服务
提供强大的工作流管理和状态控制
"""

import json
import sys
import os
import time
import threading
import asyncio
import concurrent.futures
import fcntl
import errno
import multiprocessing
from typing import Any, Dict, List, Optional, TypedDict, Annotated
from contextlib import asynccontextmanager

# 设置环境变量以抑制 gRPC fork 警告
os.environ.setdefault('GRPC_POLL_STRATEGY', 'poll')
from utils.mysql_db import MysqlHelper
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn
from agents.clean_agent.agent import execute_agent_with_api, execute
from agents.expand_agent.agent import execute_expand_agent_with_api, _update_expansion_status

# LangGraph 相关导入
try:
    from langgraph.graph import StateGraph, END
    HAS_LANGGRAPH = True
except ImportError:
    HAS_LANGGRAPH = False
    print("警告: LangGraph 未安装")

from utils.logging_config import get_logger
from tools.agent_tools import QueryDataTool, IdentifyTool, UpdateDataTool, StructureTool

# 保证可以导入本项目模块
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 创建 logger
logger = get_logger('Agent')

# 状态定义
class AgentState(TypedDict):
    request_id: str
    items: List[Dict[str, Any]]
    details: List[Dict[str, Any]]
    processed: int
    success: int
    error: Optional[str]
    status: str

class ExpandRequest(BaseModel):
    requestId: str = Field(..., description="请求ID")
    query: str = Field(..., description="查询词")

# 请求模型
class TriggerRequest(BaseModel):
    requestId: str = Field(..., description="请求ID")

# 响应模型
class TriggerResponse(BaseModel):
    requestId: str
    processed: int
    success: int
    details: List[Dict[str, Any]]

class ExtractRequest(BaseModel):
    requestId: str = Field(..., description="请求ID")
    query: str = Field(..., description="查询词")

# 全局变量
identify_tool = None
# 全局线程池
THREAD_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=20)

def update_request_status(request_id: str, status: int):
    """
    更新 knowledge_request 表中的 parsing_status
    
    Args:
        request_id: 请求ID
        status: 状态值 (1: 处理中, 2: 处理完成, 3: 处理失败)
    """
    try:
        from utils.mysql_db import MysqlHelper
        
        sql = "UPDATE knowledge_request SET parsing_status = %s WHERE request_id = %s"
        result = MysqlHelper.update_values(sql, (status, request_id))
        
        if result is not None:
            logger.info(f"更新请求状态成功: requestId={request_id}, status={status}")
        else:
            logger.error(f"更新请求状态失败: requestId={request_id}, status={status}")
            
    except Exception as e:
        logger.error(f"更新请求状态异常: requestId={request_id}, status={status}, error={e}")

def _update_expansion_status(requestId: str, status: int):
    """更新扩展查询状态"""
    try:
        from utils.mysql_db import MysqlHelper
        sql = "UPDATE knowledge_request SET expansion_status = %s WHERE request_id = %s"
        MysqlHelper.update_values(sql, (status, requestId))
        logger.info(f"更新扩展查询状态成功: requestId={requestId}, status={status}")
    except Exception as e:
        logger.error(f"更新扩展查询状态失败: requestId={requestId}, status={status}, error={e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时执行
    logger.info("🚀 启动 Knowledge Agent 服务...")
    
    # 初始化全局工具
    global identify_tool
    identify_tool = IdentifyTool()
    
    # 启动后恢复中断的流程
    # 使用线程池恢复中断流程，避免阻塞启动
    thread = threading.Thread(target=restore_interrupted_processes)
    thread.daemon = True
    thread.start()
    app.state.restore_thread = thread
    
    yield
    
    # 关闭时执行
    logger.info("🛑 关闭 Knowledge Agent 服务...")
    # 关闭线程池
    THREAD_POOL.shutdown(wait=False)
    logger.info("✅ 已关闭线程池")

def restore_interrupted_processes():
    """
    启动后恢复中断的流程
    1. 找到knowledge_request表中parsing_status=1的request_id，去请求 /parse/async
    2. 找到knowledge_request表中extraction_status=1的request_id和query，去请求 /extract
    3. 找到knowledge_request表中expansion_status=1的request_id和query，去请求 /expand
    
    使用文件锁确保只有一个进程执行恢复操作
    """
    # 定义锁文件路径
    lock_file_path = "/tmp/knowledge_agent_restore.lock"
    
    try:
        # 创建或打开锁文件
        lock_file = open(lock_file_path, 'w')
        
        try:
            # 尝试获取文件锁（非阻塞模式）
            fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
            logger.info("🔄 获取恢复锁成功，开始恢复中断的流程...")
            
            # 等待服务完全启动
            time.sleep(3)
            
            # 1. 恢复解析中断的流程
            restore_parsing_processes()
            
            # 2. 恢复提取中断的流程
            restore_extraction_processes()

            # 3. 恢复扩展中断的流程
            restore_expansion_processes()
            
            logger.info("✅ 流程恢复完成")
            
            # 释放锁
            fcntl.flock(lock_file, fcntl.LOCK_UN)
            
        except IOError as e:
            # 如果错误是因为无法获取锁（资源暂时不可用），说明已有其他进程在执行恢复
            if e.errno == errno.EAGAIN:
                logger.info("⏩ 另一个进程正在执行恢复操作，跳过本次恢复")
            else:
                logger.error(f"❌ 获取恢复锁时发生错误: {e}")
        finally:
            # 关闭锁文件
            lock_file.close()
            
    except Exception as e:
        logger.error(f"❌ 流程恢复失败: {e}")
        # 尝试清理锁文件
        try:
            if os.path.exists(lock_file_path):
                os.remove(lock_file_path)
        except:
            pass

def restore_parsing_processes():
    """恢复解析中断的流程"""
    try:
        # 查询parsing_status=1的请求
        sql = "SELECT request_id FROM knowledge_request WHERE parsing_status = 1"
        rows = MysqlHelper.get_values(sql)
        
        if not rows:
            logger.info("📋 没有发现中断的解析流程")
            return
        
        logger.info(f"🔄 发现 {len(rows)} 个中断的解析流程，开始恢复...")
        
        for row in rows:
            request_id = row[0]
            try:
                # 调用 /parse/async 接口，带重试机制
                call_parse_async_with_retry(request_id)
                logger.info(f"✅ 恢复解析流程成功: request_id={request_id}")
            except Exception as e:
                logger.error(f"❌ 恢复解析流程失败: request_id={request_id}, error={e}")
                
    except Exception as e:
        logger.error(f"❌ 恢复解析流程时发生错误: {e}")

def restore_extraction_processes():
    """恢复提取中断的流程"""
    try:
        # 查询extraction_status=1的请求和query
        sql = "SELECT request_id, query FROM knowledge_request WHERE extraction_status = 1"
        rows = MysqlHelper.get_values(sql)
        
        if not rows:
            logger.info("📋 没有发现中断的提取流程")
            return
        
        logger.info(f"🔄 发现 {len(rows)} 个中断的提取流程，开始恢复...")
        
        for row in rows:
            request_id = row[0]
            query = row[1] if len(row) > 1 else ""
            try:
                # 直接调用提取函数，带重试机制（函数内部会处理状态更新）
                call_extract_with_retry(request_id, query)
                logger.info(f"✅ 恢复提取流程成功: request_id={request_id}")
            except Exception as e:
                logger.error(f"❌ 恢复提取流程失败: request_id={request_id}, error={e}")
                
    except Exception as e:
        logger.error(f"❌ 恢复提取流程时发生错误: {e}")

def restore_expansion_processes():
    """恢复扩展中断的流程"""
    try:
        # 查询expansion_status=1的请求和query
        sql = "SELECT request_id, query FROM knowledge_request WHERE expansion_status = 1"
        rows = MysqlHelper.get_values(sql)
        
        if not rows:
            logger.info("📋 没有发现中断的扩展流程")
            return
        
        logger.info(f"🔄 发现 {len(rows)} 个中断的扩展流程，开始恢复...")
        
        for row in rows:
            request_id = row[0]
            query = row[1] if len(row) > 1 else ""
            try:
                # 直接调用扩展函数，带重试机制（函数内部会处理状态更新）
                call_expand_with_retry(request_id, query)
                logger.info(f"✅ 恢复扩展流程成功: request_id={request_id}")
            except Exception as e:
                logger.error(f"❌ 恢复扩展流程失败: request_id={request_id}, error={e}")
                
    except Exception as e:
        logger.error(f"❌ 恢复扩展流程时发生错误: {e}")

def call_parse_async_with_retry(request_id: str, max_retries: int = 3):
    """直接调用解析函数，带重试机制"""
    for attempt in range(max_retries):
        try:
            # 直接调用后台处理函数，使用线程池
            future = THREAD_POOL.submit(process_request_background_sync, request_id)
            result = future.result()
            logger.info(f"直接调用解析函数成功: request_id={request_id}")
            return
                    
        except Exception as e:
            logger.warning(f"直接调用解析函数异常: request_id={request_id}, error={e}, attempt={attempt+1}")
        
        # 如果不是最后一次尝试，等待后重试
        if attempt < max_retries - 1:
            time.sleep(2 ** attempt)  # 指数退避
    
    logger.error(f"直接调用解析函数最终失败: request_id={request_id}, 已重试{max_retries}次")

def call_extract_with_retry(request_id: str, query: str, max_retries: int = 3):
    """直接调用提取函数，带重试机制"""
    for attempt in range(max_retries):
        try:
            # 更新状态为处理中
            update_extract_status(request_id, 1)
            
            # 直接调用提取函数（同步函数，在线程池中执行）
            
            # 在全局线程池中执行同步函数
            future = THREAD_POOL.submit(
                execute_agent_with_api, 
                json.dumps({"query_word": query, "request_id": request_id})
            )
            result = future.result()
            
            # 更新状态为处理完成
            update_extract_status(request_id, 2)
            logger.info(f"直接调用提取函数成功: request_id={request_id}, result={result}")
            return
                    
        except Exception as e:
            logger.warning(f"直接调用提取函数异常: request_id={request_id}, error={e}, attempt={attempt+1}")
            # 更新状态为处理失败
            update_extract_status(request_id, 3)
        
        # 如果不是最后一次尝试，等待后重试
        if attempt < max_retries - 1:
            time.sleep(2 ** attempt)  # 指数退避
    
    logger.error(f"直接调用提取函数最终失败: request_id={request_id}, 已重试{max_retries}次")

def call_expand_with_retry(request_id: str, query: str, max_retries: int = 3):
    """直接调用扩展函数，带重试机制"""
    for attempt in range(max_retries):
        try:
            # 直接调用扩展函数
            
            # 在全局线程池中执行同步函数
            future = THREAD_POOL.submit(execute_expand_agent_with_api, request_id, query)
            result = future.result()
            
            logger.info(f"直接调用扩展函数成功: request_id={request_id}")
            return
                    
        except Exception as e:
            logger.warning(f"直接调用扩展函数异常: request_id={request_id}, error={e}, attempt={attempt+1}")
        
        # 如果不是最后一次尝试，等待后重试
        if attempt < max_retries - 1:
            time.sleep(2 ** attempt)  # 指数退避
    
    logger.error(f"直接调用扩展函数最终失败: request_id={request_id}, 已重试{max_retries}次")

# 这些函数已被删除，因为我们现在直接调用相应的函数而不是通过HTTP请求

# 创建 FastAPI 应用
app = FastAPI(
    title="Knowledge Agent API",
    description="基于 LangGraph 的智能内容识别和结构化处理服务",
    version="2.0.0",
    lifespan=lifespan
)

# 并发控制：跟踪正在处理的 requestId，防止重复并发提交
RUNNING_REQUESTS: set = set()
RUNNING_LOCK = asyncio.Lock()

# =========================
# LangGraph 工作流定义
# =========================

def process_single_item(args):
    """处理单个数据项的函数，用于多进程 (模块级，便于pickle)"""
    idx, item, request_id = args
    try:
        crawl_data = item.get('crawl_data') or {}
        content_id = item.get('content_id') or ''
        task_id = item.get('task_id') or ''

        # 先在库中查询是否已经处理过
        check_sql = "SELECT id,status,indentify_data FROM knowledge_parsing_content WHERE request_id = %s AND content_id = %s"
        check_result = MysqlHelper.get_values(check_sql, (request_id, content_id))
        result_status = 0
        result_id = 0
        result_indentify_data = {}
        if check_result:
            id, status, indentify_data = check_result[0]
            logger.info(f"查询到待结构化处理的条目，id: {id}, status: {status}, indentify_data: {str(indentify_data)[:100]}")
            result_status = status
            result_id = id
            result_indentify_data = indentify_data
            if status == 5:
                return {
                    "index": idx,
                    "dbInserted": True,
                    "identifyError": None,
                    "status": 2,
                    "success": True
                }

        # 0 未识别  3识别失败，需要重新进行识别
        if result_status == 0 or result_status == 3:
            # Step 1: 识别
            identify_result = identify_tool.run(
                crawl_data if isinstance(crawl_data, dict) else {}
            )
            
            # Step 2: 结构化并入库
            affected = UpdateDataTool.store_indentify_result(
                request_id, 
                {
                    "content_id": content_id,
                    "task_id": task_id
                }, 
                identify_result
            )
        else:
            # result_indentify_data是JSON字符串，需要解析为对象
            identify_result = json.loads(result_indentify_data) if isinstance(result_indentify_data, str) else result_indentify_data
            affected = result_id
        
        # 使用StructureTool进行内容结构化处理
        structure_tool = StructureTool()
        structure_result = structure_tool.process_content_structure(identify_result)
        
        # 存储结构化解析结果
        parsing_affected = UpdateDataTool.store_parsing_result(
            request_id,
            {
                "id": affected,
                "content_id": content_id,
                "task_id": task_id
            },
            structure_result
        )
        logger.info(f"调试信息: affected={affected}, content_id={content_id}, result_status={result_status}")
        ok = affected is not None and affected > 0 and parsing_affected is not None and parsing_affected > 0
        if ok:
            success = True
        else:
            success = True
            logger.error(f"处理第 {idx} 项时出错: {identify_result.get('error') or structure_result.get('error')}")
        
        # 记录处理详情
        detail = {
            "index": idx,
            "dbInserted": ok,
            "identifyError": identify_result.get('error') or structure_result.get('error'),
            "status": 2 if ok else 3,
            "success": success
        }
        
        logger.info(f"处理进度: {idx} - {'成功' if ok else '失败'}")
        return detail
        
    except Exception as e:
        logger.error(f"处理第 {idx} 项时出错: {e}")
        return {
            "index": idx,
            "dbInserted": False,
            "identifyError": str(e),
            "status": 3,
            "success": False
        }


def create_langgraph_workflow():
    """创建 LangGraph 工作流"""
    if not HAS_LANGGRAPH:
        return None
    
    # 工作流节点定义
    
    def fetch_data(state: AgentState) -> AgentState:
        """获取待处理数据"""
        try:
            request_id = state["request_id"]
            logger.info(f"开始获取数据: requestId={request_id}")
            
            # 更新状态为处理中
            update_request_status(request_id, 1)
            
            items = QueryDataTool.fetch_crawl_data_list(request_id)
            state["items"] = items
            state["processed"] = len(items)
            state["status"] = "data_fetched"
            
            logger.info(f"数据获取完成: requestId={request_id}, 数量={len(items)}")
            return state
            
        except Exception as e:
            logger.error(f"获取数据失败: {e}")
            state["error"] = str(e)
            state["status"] = "error"
            return state
    
    

    def process_items_batch(state: AgentState) -> AgentState:
        """批量处理所有数据项 - 使用多进程并行处理"""
        try:
            items = state["items"]
            if not items:
                state["status"] = "completed"
                return state
            
            # 准备多进程参数
            process_args = [(idx, item, state["request_id"]) for idx, item in enumerate(items, start=1)]
            
            # 使用3个进程并行处理，添加多进程保护
            if __name__ == '__main__' or multiprocessing.current_process().name == 'MainProcess':
                # 设置多进程启动方法为 'spawn' 以避免 gRPC fork 问题
                original_start_method = multiprocessing.get_start_method()
                try:
                    multiprocessing.set_start_method('spawn', force=True)
                except RuntimeError:
                    pass  # 如果已经设置过，忽略错误
                
                with multiprocessing.Pool(processes=3) as pool:
                    logger.info(f"开始多进程处理: 数量={len(process_args)}")
                    results = pool.map(process_single_item, process_args)
                
                # 恢复原始启动方法
                try:
                    multiprocessing.set_start_method(original_start_method, force=True)
                except RuntimeError:
                    pass
            else:
                # 如果不在主进程中，回退到串行处理
                logger.warning("不在主进程中，回退到串行处理")
                results = [process_single_item(args) for args in process_args]
            
            # 统计结果
            success_count = sum(1 for result in results if result.get("success", False))
            details = [result for result in results]
            
            state["success"] = success_count
            state["details"] = details
            state["status"] = "completed"
            
            logger.info(f"多进程处理完成: 成功 {success_count}/{len(items)} 项")
            return state
            
        except Exception as e:
            logger.error(f"批量处理失败: {e}")
            state["error"] = str(e)
            state["status"] = "error"
            return state
    
    def should_continue(state: AgentState) -> str:
        """判断是否继续处理"""
        if state.get("error"):
            # 处理失败，更新状态为3
            update_request_status(state["request_id"], 3)
            return "end"
        
        # 所有数据处理完毕，更新状态为2
        update_request_status(state["request_id"], 2)
        return "end"
    
    # 构建工作流图
    workflow = StateGraph(AgentState)
    
    # 添加节点
    workflow.add_node("fetch_data", fetch_data)
    workflow.add_node("process_items_batch", process_items_batch)
    
    # 设置入口点
    workflow.set_entry_point("fetch_data")
    
    # 添加边
    workflow.add_edge("fetch_data", "process_items_batch")
    workflow.add_edge("process_items_batch", END)
    
    # 编译工作流
    return workflow.compile()

# 全局工作流实例
WORKFLOW = create_langgraph_workflow() if HAS_LANGGRAPH else None

# =========================
# FastAPI 接口定义
# =========================

@app.get("/")
async def root():
    """根路径，返回服务信息"""
    return {
        "service": "Knowledge Agent API",
        "version": "2.0.0",
        "status": "running",
        "langgraph_enabled": HAS_LANGGRAPH,
        "endpoints": {
            "parse": "/parse",
            "parse/async": "/parse/async",
            "health": "/health",
            "docs": "/docs"
        }
    }

@app.get("/health")
async def health_check():
    """健康检查接口"""
    return {
        "status": "healthy", 
        "timestamp": time.time(),
        "langgraph_enabled": HAS_LANGGRAPH
    }

@app.post("/parse", response_model=TriggerResponse)
async def parse_processing(request: TriggerRequest, background_tasks: BackgroundTasks):
    """
    解析内容处理
    
    - **requestId**: 请求ID，用于标识处理任务
    """
    try:
        logger.info(f"收到解析请求: requestId={request.requestId}")
        
        if WORKFLOW and HAS_LANGGRAPH:
            # 使用 LangGraph 工作流
            logger.info("使用 LangGraph 工作流处理")
            
            # 初始化状态
            initial_state = AgentState(
                request_id=request.requestId,
                items=[],
                details=[],
                processed=0,
                success=0,
                error=None,
                status="started"
            )
            
            # 执行工作流
            final_state = WORKFLOW.invoke(
                initial_state,
                config={
                    "configurable": {"thread_id": f"thread_{request.requestId}"},
                    "recursion_limit": 100  # 增加递归限制
                }
            )
            
            # 构建响应
            result = TriggerResponse(
                requestId=request.requestId,
                processed=final_state.get("processed", 0),
                success=final_state.get("success", 0),
                details=final_state.get("details", [])
            )
            return result
        
    except Exception as e:
        logger.error(f"处理请求失败: {e}")
        # 处理失败，更新状态为3
        update_request_status(request.requestId, 3)
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

@app.post("/parse/async", status_code=200)
async def parse_processing_async(request: TriggerRequest, background_tasks: BackgroundTasks):
    """
    异步解析内容处理（后台任务）
    
    - **requestId**: 请求ID，用于标识处理任务
    
    行为：立即返回 200，并在后台继续处理任务。
    若同一个 requestId 已有任务进行中，则立即返回失败（status=3）。
    """
    try:
        logger.info(f"收到异步解析请求: requestId={request.requestId}")
        
        # 并发防抖：同一 requestId 只允许一个在运行
        async with RUNNING_LOCK:
            if request.requestId in RUNNING_REQUESTS:
                return {
                    "requestId": request.requestId,
                    "status": 3,
                    "message": "已有任务进行中，稍后再试",
                    "langgraph_enabled": HAS_LANGGRAPH
                }
            RUNNING_REQUESTS.add(request.requestId)
        
        def _background_wrapper_sync(rid: str):
            try:
                process_request_background_sync(rid)
            finally:
                # 使用线程安全的方式移除请求ID
                with threading.Lock():
                    RUNNING_REQUESTS.discard(rid)
        
        # 使用全局线程池提交后台任务
        THREAD_POOL.submit(_background_wrapper_sync, request.requestId)
        
        # 立即返回（不阻塞）
        return {
            "requestId": request.requestId,
            "status": 1,
            "message": "任务已进入队列并在后台处理",
            "langgraph_enabled": HAS_LANGGRAPH
        }
        
    except Exception as e:
        logger.error(f"提交异步任务失败: {e}")
        raise HTTPException(status_code=500, detail=f"提交任务失败: {str(e)}")

def process_request_background_sync(request_id: str):
    """后台处理请求（同步版本）"""
    try:
        logger.info(f"开始后台处理: requestId={request_id}")
        
        if WORKFLOW and HAS_LANGGRAPH:
            # 使用 LangGraph 工作流
            # 更新状态为处理中
            update_request_status(request_id, 1)
            
            initial_state = AgentState(
                request_id=request_id,
                items=[],
                details=[],
                processed=0,
                success=0,
                error=None,
                status="started"
            )
            
            final_state = WORKFLOW.invoke(
                initial_state,
                config={
                    "configurable": {"thread_id": f"thread_{request_id}"},
                    "recursion_limit": 100  # 增加递归限制
                }
            )
            # 所有数据处理完毕，更新状态为2
            update_request_status(request_id, 2)
            logger.info(f"LangGraph 后台处理完成: requestId={request_id}, processed={final_state.get('processed', 0)}, success={final_state.get('success', 0)}")
        
    except Exception as e:
        logger.error(f"后台处理失败: requestId={request_id}, error={e}")
        # 处理失败，更新状态为3
        update_request_status(request_id, 3)

async def process_request_background(request_id: str):
    """后台处理请求（异步版本，为了兼容性保留）"""
    try:
        # 在线程池中执行同步版本
        loop = asyncio.get_event_loop()
        with concurrent.futures.ThreadPoolExecutor() as executor:
            await loop.run_in_executor(executor, process_request_background_sync, request_id)
    except Exception as e:
        logger.error(f"后台处理失败: requestId={request_id}, error={e}")
        # 处理失败，更新状态为3
        update_request_status(request_id, 3)


extraction_requests: set = set()

@app.post("/extract")
async def extract(request: ExtractRequest):
    """
    执行提取处理（异步方式）
    
    Args:
        request: 包含请求ID和查询词的请求体
        
    Returns:
        dict: 包含执行状态的字典
    """
    try:
        requestId = request.requestId
        query = request.query
        logger.info(f"收到提取请求: requestId={requestId}, query={query}")
        # 并发防抖：同一 requestId 只允许一个在运行
        if requestId in extraction_requests:
            return {"status": 1, "requestId": requestId, "message": "请求已在处理中"}
        extraction_requests.add(requestId)
        
        # 更新状态为处理中
        update_extract_status(requestId, 1)
        
        # 创建线程池任务执行Agent
        def _execute_extract_sync():
            try:
                # result = execute_agent_with_api(json.dumps({"query_word": query, "request_id": requestId}))
                result = execute(query, requestId)
                # 更新状态为处理完成
                update_extract_status(requestId, 2)
                logger.info(f"异步提取任务完成: requestId={requestId}")
                return result
            except Exception as e:
                logger.error(f"异步提取任务失败: requestId={requestId}, error={e}")
                # 更新状态为处理失败
                update_extract_status(requestId, 3)
            finally:
                # 移除请求ID
                extraction_requests.discard(requestId)
        
        # 使用全局线程池提交任务
        THREAD_POOL.submit(_execute_extract_sync)
        
        # 立即返回状态
        return {"status": 1, "requestId": requestId, "message": "提取任务已启动并在后台处理"}
    except Exception as e:
        logger.error(f"启动提取任务失败: requestId={requestId}, error={e}")
        # 发生异常，更新状态为处理失败
        update_extract_status(requestId, 3)
        # 从运行集合中移除
        extraction_requests.discard(requestId)
        raise HTTPException(status_code=500, detail=f"启动提取任务失败: {str(e)}")

@app.post("/expand")
async def expand(request: ExpandRequest):
    """
    执行扩展查询处理（异步方式）
    
    Args:
        request: 包含请求ID和查询词的请求体
        
    Returns:
        dict: 包含执行状态的字典
    """
    try:
        requestId = request.requestId
        query = request.query
        logger.info(f"收到扩展查询请求: requestId={requestId}, query={query}")
        
        # 并发防抖：同一 requestId 只允许一个在运行
        expansion_requests = getattr(app.state, 'expansion_requests', set())
        # 使用线程锁而不是asyncio锁
        with threading.Lock():
            if requestId in expansion_requests:
                return {"status": 1, "requestId": requestId, "message": "扩展查询已在处理中"}
            # 如果集合不存在，创建它
            if not hasattr(app.state, 'expansion_requests'):
                app.state.expansion_requests = set()
            app.state.expansion_requests.add(requestId)
        
        # 立即更新状态为处理中
        _update_expansion_status(requestId, 1)
        
        # 创建线程池任务执行扩展Agent
        def _execute_expand_sync():
            try:
                # 直接调用同步函数，使用线程池
                result = execute_expand_agent_with_api(requestId, query)
                # 更新状态为处理完成
                _update_expansion_status(requestId, 2)
                logger.info(f"异步扩展查询任务完成: requestId={requestId}")
                return result
            except Exception as e:
                logger.error(f"异步扩展查询任务失败: requestId={requestId}, error={e}")
                # 更新状态为处理失败
                _update_expansion_status(requestId, 3)
            finally:
                # 无论成功失败，都从运行集合中移除
                with threading.Lock():
                    if hasattr(app.state, 'expansion_requests'):
                        app.state.expansion_requests.discard(requestId)
        
        # 使用全局线程池提交任务
        THREAD_POOL.submit(_execute_expand_sync)
        
        # 立即返回状态
        return {"status": 1, "requestId": requestId, "message": "扩展查询任务已启动并在后台处理"}
    except Exception as e:
        logger.error(f"启动扩展查询任务失败: requestId={requestId}, error={e}")
        # 发生异常，更新状态为处理失败
        _update_expansion_status(requestId, 3)
        # 从运行集合中移除
        async with RUNNING_LOCK:
            if hasattr(app.state, 'expansion_requests'):
                app.state.expansion_requests.discard(requestId)
        raise HTTPException(status_code=500, detail=f"启动扩展查询任务失败: {str(e)}")

    except Exception as e:
        # 发生异常，更新状态为处理失败
        _update_expansion_status(request.requestId, 3)
        # 从运行集合中移除
        async with RUNNING_LOCK:
            if hasattr(app.state, 'expansion_requests'):
                app.state.expansion_requests.discard(request.requestId)
        raise HTTPException(status_code=500, detail=f"启动扩展查询任务失败: {str(e)}")
def update_extract_status(request_id: str, status: int):
    try:
        from utils.mysql_db import MysqlHelper
        
        sql = "UPDATE knowledge_request SET extraction_status = %s WHERE request_id = %s"
        result = MysqlHelper.update_values(sql, (status, request_id))
        
        if result is not None:
            logger.info(f"更新请求状态成功: requestId={request_id}, status={status}")
        else:
            logger.error(f"更新请求状态失败: requestId={request_id}, status={status}")
            
    except Exception as e:
        logger.error(f"更新请求状态异常: requestId={request_id}, status={status}, error={e}")

if __name__ == "__main__":
    # 从环境变量获取配置
    import os
    reload_enabled = os.getenv("RELOAD_ENABLED", "false").lower() == "true"
    log_level = os.getenv("LOG_LEVEL", "info")
    
    # 启动服务
    uvicorn.run(
        "agent:app",
        host="0.0.0.0",
        port=8080,
        reload=reload_enabled,  # 通过环境变量控制
        log_level=log_level
    )