knowledge_search/api/request_context.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
请求上下文管理器
为每个API请求提供独立的临时工作空间
"""

import uuid
import shutil
from pathlib import Path
import logging

logger = logging.getLogger(__name__)


class RequestContextManager:
    """请求级别的临时工作空间管理器"""

    def __init__(self, base_dir: str = "temp_requests"):
        """
        初始化请求上下文管理器

        Args:
            base_dir: 临时目录基础路径（项目内相对路径）
        """
        self.base_dir = Path(base_dir)
        self.request_id = str(uuid.uuid4())[:8]  # 短UUID，便于日志查看
        self.work_dir = None

    async def __aenter__(self):
        """
        创建请求专用工作目录

        Returns:
            self: 返回自身，提供request_id和work_dir访问
        """
        self.work_dir = self.base_dir / f"request_{self.request_id}"
        self.work_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"[{self.request_id}] 创建工作目录: {self.work_dir}")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        清理工作目录

        Args:
            exc_type: 异常类型
            exc_val: 异常值
            exc_tb: 异常追踪
        """
        if self.work_dir and self.work_dir.exists():
            try:
                shutil.rmtree(self.work_dir)
                logger.info(f"[{self.request_id}] 已清理工作目录")
            except Exception as e:
                logger.warning(f"[{self.request_id}] 清理工作目录失败: {e}")