hxf/backend/th_agenter/api/endpoints/smart_chat.py

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import HTTPBearer
from sqlalchemy.orm import Session
from typing import Optional, Dict, Any
from datetime import datetime

from th_agenter.db.database import get_session
from th_agenter.services.auth import AuthService
from th_agenter.services.smart_workflow import SmartWorkflowManager
from th_agenter.services.conversation import ConversationService
from th_agenter.services.conversation_context import conversation_context_service
from utils.util_schemas import BaseResponse
from pydantic import BaseModel
from loguru import logger
from utils.util_exceptions import HxfResponse

router = APIRouter(prefix="/smart-chat", tags=["smart-chat"])
security = HTTPBearer()

# Request/Response Models
class SmartQueryRequest(BaseModel):
    query: str
    conversation_id: Optional[int] = None
    is_new_conversation: bool = False

class SmartQueryResponse(BaseModel):
    success: bool
    message: str
    data: Optional[Dict[str, Any]] = None
    workflow_steps: Optional[list] = None
    conversation_id: Optional[int] = None

class ConversationContextResponse(BaseModel):
    success: bool
    message: str
    data: Optional[Dict[str, Any]] = None

@router.post("/query", response_model=SmartQueryResponse, summary="智能问数查询")
async def smart_query(
    request: SmartQueryRequest,
    current_user = Depends(AuthService.get_current_user),
    session: Session = Depends(get_session)
):
    """
    智能问数查询接口
    支持新对话时自动加载文件列表，智能选择相关Excel文件，生成和执行pandas代码
    """
    session.desc = f"START: 用户 {current_user.username} 智能问数查询"
    conversation_id = None
    
    try:
        # 验证请求参数
        if not request.query or not request.query.strip():
            session.desc = "ERROR: 用户输入为空, 查询内容不能为空"
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="查询内容不能为空"
            )
        
        if len(request.query) > 1000:
            session.desc = "ERROR: 用户输入过长, 查询内容不能超过1000字符"
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="查询内容过长，请控制在1000字符以内"
            )
        
        # 初始化工作流管理器
        workflow_manager = SmartWorkflowManager(session)
        await workflow_manager.initialize()
        
        conversation_service = ConversationService(session)
        
        # 处理对话上下文
        conversation_id = request.conversation_id
        
        # 如果是新对话或没有指定对话ID，创建新对话
        if request.is_new_conversation or not conversation_id:
            try:
                conversation_id = await conversation_context_service.create_conversation(
                    user_id=current_user.id,
                    title=f"智能问数: {request.query[:20]}..."
                )
                request.is_new_conversation = True
                session.desc = f"创建新对话: {conversation_id}"
            except Exception as e:
                session.desc = f"WARNING: 创建对话失败，使用临时会话: {e}"
                conversation_id = None
        else:
            # 验证对话是否存在且属于当前用户
            try:
                context = await conversation_context_service.get_conversation_context(conversation_id)
                if not context or context.get('user_id') != current_user.id:
                    session.desc = f"ERROR: 对话 {conversation_id} 不存在或无权访问"
                    raise HTTPException(
                        status_code=status.HTTP_404_NOT_FOUND,
                        detail="对话不存在或无权访问"
                    )
                session.desc = f"使用现有对话: {conversation_id}"
            except HTTPException:
                session.desc = f"EXCEPTION: 对话 {conversation_id} 不存在或无权访问"
                raise
            except Exception as e:
                session.desc = f"ERROR: 验证对话失败: {e}"
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail="对话验证失败"
                )
        
        # 保存用户消息
        if conversation_id:
            try:
                await conversation_context_service.save_message(
                    conversation_id=conversation_id,
                    role="user",
                    content=request.query
                )
            except Exception as e:
                session.desc = f"WARNING: 保存用户消息失败: {e}"
                # 不阻断流程，继续执行查询
        
        # 执行智能查询工作流
        try:
            result = await workflow_manager.process_smart_query(
                user_query=request.query,
                user_id=current_user.id,
                conversation_id=conversation_id,
                is_new_conversation=request.is_new_conversation
            )
        except Exception as e:
            session.desc = f"ERROR: 智能查询执行失败: {e}"
            # 返回结构化的错误响应
            response = SmartQueryResponse(
                success=False,
                message=f"查询执行失败: {str(e)}",
                data={'error_type': 'query_execution_error'},
                workflow_steps=[{
                    'step': 'error',
                    'status': 'failed',
                    'message': str(e)
                }],
                conversation_id=conversation_id
            )
            return HxfResponse(response)
        
        # 如果查询成功，保存助手回复和更新上下文
        if result['success'] and conversation_id:
            try:
                # 保存助手回复
                await conversation_context_service.save_message(
                    conversation_id=conversation_id,
                    role="assistant",
                    content=result.get('data', {}).get('summary', '查询完成'),
                    metadata={
                        'query_result': result.get('data'),
                        'workflow_steps': result.get('workflow_steps', []),
                        'selected_files': result.get('data', {}).get('used_files', [])
                    }
                )
                
                # 更新对话上下文
                await conversation_context_service.update_conversation_context(
                    conversation_id=conversation_id,
                    query=request.query,
                    selected_files=result.get('data', {}).get('used_files', [])
                )
                
            except Exception as e:
                session.desc = f"EXCEPTION: 保存消息到对话历史失败: {e}"
                # 不影响返回结果，只记录警告
        
        # 返回结果，包含对话ID
        response_data = result.get('data', {})
        if conversation_id:
            response_data['conversation_id'] = conversation_id
        session.desc = f"SUCCESS: 保存助手回复和更新上下文，对话ID: {conversation_id}"
        response = SmartQueryResponse(
            success=result['success'],
            message=result.get('message', '查询完成'),
            data=response_data,
            workflow_steps=result.get('workflow_steps', []),
            conversation_id=conversation_id
        )
        return HxfResponse(response)
        
    except HTTPException as e:
        session.desc = f"EXCEPTION: HTTP异常: {e}"
        raise e
    except Exception as e:
        session.desc = f"ERROR: 智能查询接口异常: {e}"
        # 返回通用错误响应
        response = SmartQueryResponse(
            success=False,
            message="服务器内部错误，请稍后重试",
            data={'error_type': 'internal_server_error'},
            workflow_steps=[{
                'step': 'error',
                'status': 'failed',
                'message': '系统异常'
            }],
            conversation_id=conversation_id
        )
        return HxfResponse(response)

@router.get("/conversation/{conversation_id}/context", response_model=ConversationContextResponse, summary="获取对话上下文")
async def get_conversation_context(
    conversation_id: int,
    current_user = Depends(AuthService.get_current_user),
    session: Session = Depends(get_session)
):
    """
    获取对话上下文信息，包括已使用的文件和历史查询
    """
    # 获取对话上下文
    session.desc = f"START: 获取对话上下文，对话ID: {conversation_id}"
    context = await conversation_context_service.get_conversation_context(conversation_id)
    
    if not context:
        session.desc = f"ERROR: 对话上下文不存在，对话ID: {conversation_id}"
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="对话上下文不存在"
        )
    
    # 验证用户权限
    if context['user_id'] != current_user.id:
        session.desc = f"ERROR: 无权访问对话上下文，对话ID: {conversation_id}"
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="无权访问此对话"
        )
    
    # 获取对话历史
    history = await conversation_context_service.get_conversation_history(conversation_id)
    context['message_history'] = history
    session.desc = f"SUCCESS: 获取对话上下文成功，对话ID: {conversation_id}"
    response = ConversationContextResponse(
        success=True,
        message="获取对话上下文成功",
        data=context
    )
    return HxfResponse(response)
        

@router.get("/files/status", response_model=ConversationContextResponse, summary="获取用户当前的文件状态和统计信息")
async def get_files_status(
    current_user = Depends(AuthService.get_current_user),
    session: Session = Depends(get_session)
):
    """
    获取用户当前的文件状态和统计信息
    """
    session.desc = f"START: 获取用户文件状态和统计信息，用户ID: {current_user.id}"
    workflow_manager = SmartWorkflowManager()
    await workflow_manager.initialize()
    
    # 获取用户文件列表
    file_list = await workflow_manager.excel_workflow._load_user_file_list(current_user.id)
    
    # 统计信息
    total_files = len(file_list)
    total_rows = sum(f.get('row_count', 0) for f in file_list)
    total_columns = sum(f.get('column_count', 0) for f in file_list)
    
    # 文件类型统计
    file_types = {}
    for file_info in file_list:
        filename = file_info['filename']
        ext = filename.split('.')[-1].lower() if '.' in filename else 'unknown'
        file_types[ext] = file_types.get(ext, 0) + 1
    
    status_data = {
        'total_files': total_files,
        'total_rows': total_rows,
        'total_columns': total_columns,
        'file_types': file_types,
        'files': [{
            'id': f['id'],
            'filename': f['filename'],
            'row_count': f.get('row_count', 0),
            'column_count': f.get('column_count', 0),
            'columns': f.get('columns', []),
            'upload_time': f.get('upload_time')
        } for f in file_list],
        'ready_for_query': total_files > 0
    }
    
    session.desc = f"SUCCESS: 获取用户文件状态和统计信息成功，用户ID: {current_user.id}"
    response = ConversationContextResponse(
        success=True,
        message=f"当前有{total_files}个可用文件" if total_files > 0 else "暂无可用文件，请先上传Excel文件",
        data=status_data
    )
    return HxfResponse(response)

@router.post("/conversation/{conversation_id}/reset", summary="重置对话上下文")
async def reset_conversation_context(
    conversation_id: int,
    current_user = Depends(AuthService.get_current_user),
    session: Session = Depends(get_session)
):
    """
    重置对话上下文，清除历史查询记录但保留文件
    """
    session.desc = f"START: 重置对话上下文，对话ID: {conversation_id}"
    # 验证对话存在和用户权限
    context = await conversation_context_service.get_conversation_context(conversation_id)
    
    if not context:
        session.desc = f"ERROR: 对话上下文不存在，对话ID: {conversation_id}"
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="对话上下文不存在"
        )
    
    if context['user_id'] != current_user.id:
        session.desc = f"ERROR: 无权访问对话上下文，对话ID: {conversation_id}"
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="无权访问此对话"
        )
    
    # 重置对话上下文
    success = await conversation_context_service.reset_conversation_context(conversation_id)
    
    if success:
        session.desc = f"SUCCESS: 重置对话上下文成功，对话ID: {conversation_id}"
        response = ConversationContextResponse(
            success=True,
            message="对话上下文已重置，可以开始新的数据分析会话"
        )
        return HxfResponse(response)
    else:
        session.desc = f"EXCEPTION: 重置对话上下文失败，对话ID: {conversation_id}"
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="重置对话上下文失败"
        )