knowledge-agent/agent.py

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

import json
import sys
import os
import time
from typing import Any, Dict, List, Optional, TypedDict, Annotated
from contextlib import asynccontextmanager
import asyncio
from utils.mysql_db import MysqlHelper

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

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
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

# 创建 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
    current_index: int
    current_item: Optional[Dict[str, Any]]
    identify_result: Optional[Dict[str, Any]]
    error: Optional[str]
    status: str

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

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

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

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

# 全局变量
identify_tool = None

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):
    """应用生命周期管理"""
    # 启动时初始化
    global identify_tool
    identify_tool = IdentifyTool()
    logger.info("Agent 服务启动完成")
    
    yield
    
    # 关闭时清理
    logger.info("Agent 服务正在关闭")

# 创建 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 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_item(state: AgentState) -> AgentState:
        """处理单个数据项"""
        try:
            items = state["items"]
            current_index = state.get("current_index", 0)
            
            if current_index >= len(items):
                state["status"] = "completed"
                return state
            
            item = items[current_index]
            state["current_item"] = item
            state["content_id"] = item.get('content_id') or ''
            state["task_id"] = item.get('task_id') or ''
            state["current_index"] = current_index + 1
            
            # 处理当前项
            crawl_data = item.get('crawl_data') or {}
            
            # Step 1: 识别
            identify_result = identify_tool.run(
                crawl_data if isinstance(crawl_data, dict) else {}
            )
            state["identify_result"] = identify_result
            
            # Step 2: 结构化并入库
            affected = UpdateDataTool.store_indentify_result(
                state["request_id"], 
                {
                    "content_id": state["content_id"],
                    "task_id": state["task_id"]
                }, 
                identify_result
            )
            # 使用StructureTool进行内容结构化处理
            structure_tool = StructureTool()
            structure_result = structure_tool.process_content_structure(identify_result)
            
            # 存储结构化解析结果
            parsing_affected = UpdateDataTool.store_parsing_result(
                state["request_id"],
                {
                    "content_id": state["content_id"],
                    "task_id": state["task_id"]
                },
                structure_result
            )
            
            ok = affected is not None and affected > 0 and parsing_affected is not None and parsing_affected > 0
            if ok:
                state["success"] += 1
            
            # 记录处理详情
            detail = {
                "index": current_index + 1,
                "dbInserted": ok,
                "identifyError": identify_result.get('error'),
                "status": 2 if ok else 3
            }
            state["details"].append(detail)
            
            state["status"] = "item_processed"
            logger.info(f"处理进度: {current_index + 1}/{len(items)} - {'成功' if ok else '失败'}")
            
            return state
            
        except Exception as e:
            logger.error(f"处理第 {current_index + 1} 项时出错: {e}")
            detail = {
                "index": current_index + 1,
                "dbInserted": False,
                "identifyError": str(e),
                "status": 3
            }
            state["details"].append(detail)
            state["status"] = "item_error"
            return state
    
    def should_continue(state: AgentState) -> str:
        """判断是否继续处理"""
        if state.get("error"):
            # 处理失败，更新状态为3
            update_request_status(state["request_id"], 3)
            return "end"
        
        current_index = state.get("current_index", 0)
        items = state.get("items", [])
        if current_index >= len(items):
            # 所有数据处理完毕，更新状态为2
            update_request_status(state["request_id"], 2)
            return "end"
        
        return "continue"
    
    # 构建工作流图
    workflow = StateGraph(AgentState)
    
    # 添加节点
    workflow.add_node("fetch_data", fetch_data)
    workflow.add_node("process_item", process_item)
    
    # 设置入口点
    workflow.set_entry_point("fetch_data")
    
    # 添加边
    workflow.add_edge("fetch_data", "process_item")
    workflow.add_conditional_edges(
        "process_item",
        should_continue,
        {
            "continue": "process_item",
            "end": 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,
                current_index=0,
                current_item=None,
                identify_result=None,
                error=None,
                status="started"
            )
            
            # 执行工作流
            final_state = WORKFLOW.invoke(
                initial_state,
                config={"configurable": {"thread_id": f"thread_{request.requestId}"}}
            )
            
            # 构建响应
            result = TriggerResponse(
                requestId=request.requestId,
                processed=final_state.get("processed", 0),
                success=final_state.get("success", 0),
                details=final_state.get("details", [])
            )
            
        else:
            # 回退到传统模式
            logger.info("使用传统模式处理")
            
            # 更新状态为处理中
            update_request_status(request.requestId, 1)
            
            # 获取待处理数据
            items = QueryDataTool.fetch_crawl_data_list(request.requestId)
            print(f"传统模式---items: {items}")
            if not items:
                # 无数据需要处理，更新状态为完成
                update_request_status(request.requestId, 2)
                return TriggerResponse(
                    requestId=request.requestId,
                    processed=0,
                    success=0,
                    details=[]
                )

            # 处理数据
            success_count = 0
            details: List[Dict[str, Any]] = []
            
            for idx, item in enumerate(items, start=1):
                try:
                    crawl_data = item.get('crawl_data') or {}
                    
                    # Step 1: 识别
                    identify_result = identify_tool.run(
                        crawl_data if isinstance(crawl_data, dict) else {}
                    )
                    
                    # Step 2: 结构化并入库
                    affected = UpdateDataTool.store_indentify_result(
                        request.requestId, 
                        {
                            content_id: item.get('content_id') or '',
                            task_id: item.get('task_id') or ''
                        }, 
                        identify_result
                    )
                    
                    # 使用StructureTool进行内容结构化处理
                    structure_tool = StructureTool()
                    structure_result = structure_tool.process_content_structure(identify_result)
                    
                    # 存储结构化解析结果
                    parsing_affected = UpdateDataTool.store_parsing_result(
                        request.requestId,
                        {
                            content_id: item.get('content_id') or '',
                            task_id: item.get('task_id') or ''
                        },
                        structure_result
                    )
                    
                    ok = affected is not None and affected > 0
                    if ok:
                        success_count += 1
                    
                    details.append({
                        "index": idx,
                        "dbInserted": ok,
                        "identifyError": identify_result.get('error'),
                        "status": 2 if ok else 3
                    })
                    
                except Exception as e:
                    logger.error(f"处理第 {idx} 项时出错: {e}")
                    details.append({
                        "index": idx,
                        "dbInserted": False,
                        "identifyError": str(e),
                        "status": 3
                    })

            result = TriggerResponse(
                requestId=request.requestId,
                processed=len(items),
                success=success_count,
                details=details
            )
            
            # 更新状态为处理完成
            update_request_status(request.requestId, 2)
        
        logger.info(f"处理完成: requestId={request.requestId}, processed={result.processed}, success={result.success}")
        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)
        
        async def _background_wrapper(rid: str):
            try:
                await process_request_background(rid)
            finally:
                async with RUNNING_LOCK:
                    RUNNING_REQUESTS.discard(rid)
        
        # 直接使用 asyncio 创建后台任务（不阻塞当前请求返回）
        asyncio.create_task(_background_wrapper(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)}")

async def process_request_background(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,
                current_index=0,
                current_item=None,
                identify_result=None,
                error=None,
                status="started"
            )
            
            final_state = WORKFLOW.invoke(
                initial_state,
                config={"configurable": {"thread_id": f"thread_{request_id}"}}
            )
            logger.info(f"LangGraph 后台处理完成: requestId={request_id}, processed={final_state.get('processed', 0)}, success={final_state.get('success', 0)}")
            
        else:
            # 传统模式
            # 更新状态为处理中
            update_request_status(request_id, 1)
            
            items = QueryDataTool.fetch_crawl_data_list(request_id)
            print(f"传统模式process_request_background---items: {items}")
            if not items:
                logger.info(f"后台处理完成: requestId={request_id}, 无数据需要处理")
                # 无数据需要处理，更新状态为完成
                update_request_status(request_id, 2)
                return

            success_count = 0
            for idx, item in enumerate(items, start=1):
                try:
                    crawl_data = item.get('crawl_data') or {}
                    content_id = item.get('content_id') or ''
                    
                    identify_result = identify_tool.run(
                        crawl_data if isinstance(crawl_data, dict) else {}
                    )
                    
                    affected = UpdateDataTool.store_indentify_result(
                        request_id, 
                        {
                            content_id: item.get('content_id') or '',
                            task_id: item.get('task_id') or ''
                        },
                        identify_result
                    )
                    
                    # 使用StructureTool进行内容结构化处理
                    structure_tool = StructureTool()
                    structure_result = structure_tool.process_content_structure(identify_result)
                    
                    # 存储结构化解析结果
                    parsing_affected = UpdateDataTool.store_parsing_result(
                        request_id,
                        {
                            content_id: item.get('content_id') or '',
                            task_id: item.get('task_id') or ''
                        },
                        structure_result
                    )
                    
                    if affected is not None and affected > 0:
                        success_count += 1
                    
                    logger.info(f"后台处理进度: {idx}/{len(items)} - {'成功' if affected else '失败'} - 结构化{'成功' if parsing_affected else '失败'}")
                    
                except Exception as e:
                    logger.error(f"后台处理第 {idx} 项时出错: {e}")

            logger.info(f"传统模式后台处理完成: requestId={request_id}, processed={len(items)}, success={success_count}")
            # 更新状态为处理完成
            update_request_status(request_id, 2)
        
    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):
    try:
        requestId = request.requestId
        query = request.query
        # 检查请求是否已经在处理中
        async with RUNNING_LOCK:
            if requestId in extraction_requests:
                return {"status": 1, "request_id": requestId, "message": "请求已在处理中"}
            extraction_requests.add(requestId)
        
        try:
            # 更新状态为处理中
            update_extract_status(requestId, 1)
            # 执行Agent
            result = execute_agent_with_api(json.dumps({"query_word": query, "request_id": requestId}))
        finally:
            # 无论成功失败，都从运行集合中移除
            async with RUNNING_LOCK:
                extraction_requests.discard(requestId)
        return {"status": "success", "result": result}
    except Exception as e:
        # 发生异常，更新状态为处理失败
        update_extract_status(requestId, 3)
        raise HTTPException(status_code=500, detail=f"执行Agent时出错: {str(e)}")

@app.post("/expand")
async def expand(request: ExpandRequest, background_tasks: BackgroundTasks):
    """
    执行扩展查询处理
    
    Args:
        request: 包含请求ID的请求体
        background_tasks: FastAPI 后台任务
        
    Returns:
        dict: 包含执行状态的字典
    """
    try:
        # 立即更新状态为处理中
        _update_expansion_status(request.requestId, 1)
        
        # 添加后台任务
        background_tasks.add_task(execute_expand_agent_with_api, request.requestId, request.query)
        
        # 立即返回状态
        return {"status": 1, "requestId": request.requestId, "message": "扩展查询处理已启动"}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"执行Agent时出错: {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__":
    # 启动服务
    uvicorn.run(
        "agent:app",
        host="0.0.0.0",
        port=8080,
        reload=True,  # 开发模式，自动重载
        log_level="info"
    )