backend/backend_app/api/llm_api/newUpdateFile/new_updateFile.py

from fastapi import APIRouter, Request, UploadFile, File, Form, HTTPException
from typing import List, Dict
from datetime import datetime
import logging
import os
import asyncio
import hashlib
import json
import tempfile
import shutil

from helloAgents.tools.async_executor import AsyncToolExecutor
from helloAgents.tools.builtin.rag_tool import RAGTool
from helloAgents.tools.builtin.memory_tool import MemoryTool
from helloAgents.tools.registry import global_registry

from redis_config import get_redis, QUEUE_RAG_QDRANT, QUEUE_RAG_NEO4J, QUEUE_MEMORY
redis = get_redis()

logger = logging.getLogger(__name__)
new_updateFile_router = APIRouter()

# ------------------------------
# ⚙️ 企业级配置
# ------------------------------
MANIFEST_FILE = ".file_manifest.json"
CHUNK_SIZE = 8192
MAX_PARALLEL_FILES = 4

# ------------------------------
# 🛠️ 核心工具函数
# ------------------------------

def load_manifest(ns_path: str) -> Dict[str, str]:
    manifest_path = os.path.join(ns_path, MANIFEST_FILE)
    if os.path.exists(manifest_path):
        try:
            with open(manifest_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            logger.warning(f"读取 manifest 失败: {manifest_path}")
            return {}
    return {}

def save_manifest(ns_path: str, manifest: Dict[str, str]):
    manifest_path = os.path.join(ns_path, MANIFEST_FILE)
    temp_path = manifest_path + ".tmp"
    with open(temp_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)
    os.replace(temp_path, manifest_path)

def calculate_file_hash_stream(file_path: str) -> str:
    sha256 = hashlib.sha256()
    with open(file_path, "rb") as f:
        while chunk := f.read(CHUNK_SIZE):
            sha256.update(chunk)
    return sha256.hexdigest()

def calculate_upload_hash_stream(file_obj) -> str:
    sha256 = hashlib.sha256()
    file_obj.file.seek(0)
    while chunk := file_obj.file.read(CHUNK_SIZE):
        sha256.update(chunk)
    file_obj.file.seek(0)
    return sha256.hexdigest()

async def save_file_with_version_and_deduplicate(
    upload_file: UploadFile,
    user_id: str,
    save_dir: str = "./knowledge_base"
) -> dict:
    original_name = upload_file.filename.strip()
    name, ext = os.path.splitext(original_name)
    ns_path = os.path.join(save_dir, user_id)

    os.makedirs(ns_path, exist_ok=True)
    manifest = load_manifest(ns_path)

    try:
        new_file_hash = calculate_upload_hash_stream(upload_file)

        for existing_name, existing_hash in manifest.items():
            if existing_hash == new_file_hash:
                return {
                    "success": True,
                    "is_duplicate": True,
                    "filename": original_name,
                    "file_path": os.path.join(ns_path, existing_name),
                    "message": f"内容已存在（同名：{existing_name}），自动去重"
                }

        target_name = original_name
        version = 1
        while target_name in manifest:
            version += 1
            target_name = f"{name}_v{version}{ext}"
        
        final_path = os.path.join(ns_path, target_name)

        with tempfile.NamedTemporaryFile(dir=ns_path, delete=False) as tmp_file:
            shutil.copyfileobj(upload_file.file, tmp_file)
            temp_file_path = tmp_file.name
        
        os.replace(temp_file_path, final_path)

        manifest[target_name] = new_file_hash
        save_manifest(ns_path, manifest)

        return {
            "success": True,
            "is_duplicate": False,
            "filename": original_name,
            "file_path": final_path,
            "hash": new_file_hash,
            "message": f"已保存（版本 v{version})" if version > 1 else "已保存"
        }

    except Exception as e:
        logger.error(f"保存失败 {upload_file.filename}: {str(e)}")
        return {
            "success": False,
            "filename": upload_file.filename,
            "error": str(e)
        }

# ------------------------------
# 主处理逻辑（已补全记忆）
# ------------------------------
async def process_uploaded_files(files: List[UploadFile], user_id: str) -> dict:
    memory_tool: MemoryTool = global_registry.get_tool("memory")
    rag_tool: RAGTool = global_registry.get_tool("rag")

    if not files:
        raise HTTPException(status_code=400, detail="请上传至少一个文件")
    if not user_id.strip():
        raise HTTPException(status_code=400, detail="命名空间不能为空")

    semaphore = asyncio.Semaphore(MAX_PARALLEL_FILES)
    
    async def save_with_limit(f):
        async with semaphore:
            return await save_file_with_version_and_deduplicate(f, user_id)

    save_tasks = [save_with_limit(f) for f in files]
    save_results = await asyncio.gather(*save_tasks)

    saved_files = []
    save_errors = []
    duplicate_files = []

    for res in save_results:
        if not res["success"]:
            save_errors.append(res)
        elif res.get("is_duplicate"):
            duplicate_files.append(res)
        else:
            saved_files.append(res)

    # ============================
    # ✅ 企业级精简版：只存 1 条总结记忆（包含所有文件路径+状态）
    # ============================
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    total = len(files)
    success_cnt = len(saved_files)
    dup_cnt = len(duplicate_files)
    fail_cnt = len(save_errors)

    # 构建清晰的文件清单
    success_files = "\n".join([f"- {f['filename']} | 路径：{f['file_path']}" for f in saved_files])
    dup_files_list = "\n".join([f"- {d['filename']}" for d in duplicate_files])
    err_files_list = "\n".join([f"- {e['filename']} | 原因：{e['error']}" for e in save_errors])

    # ============================
    # ✅ 记忆存储结束
    # ============================
    summary_content = f"""【文件上传总结】{now} | 用户：{user_id}
        总文件数：{total}
        ✅ 上传成功：{success_cnt} 个
        {success_files if success_files else '无'}

        ⚠️ 重复文件：{dup_cnt} 个
        {dup_files_list if dup_files_list else '无'}

        ❌ 上传失败：{fail_cnt} 个
        {err_files_list if err_files_list else '无'}
        """
    summary_tasks = {
        "action": "add",
        "user_id": user_id,
        "memory_type": "episodic",
        "content": summary_content.strip(),
        "importance": 0.8,
        "session_id": None
    }
    redis.rpush(QUEUE_MEMORY, json.dumps(summary_tasks))

    if not saved_files:
        return {
            "success": len(save_errors) == 0,
            "msg": f"去重 {len(duplicate_files)} 个 | 失败 {len(save_errors)} 个",
            "duplicate_count": len(duplicate_files),
            "results": [],
            "errors": save_errors,
            "duplicates": duplicate_files
        }

    # RAG qdrant 入库
    task = {
        "action": "add_document",
        "file_path": saved_files,
        "user_id": user_id
    }
    redis.rpush(QUEUE_RAG_QDRANT, json.dumps(task))

    # RAG neo4j 入库
    task_neo4j = {
        "action": "add_neo4j_document",
        "file_path": saved_files,
        "user_id": user_id
    }
    redis.rpush(QUEUE_RAG_NEO4J, json.dumps(task_neo4j))


    return {
        "success": True,
        "msg": f"成功 {len(saved_files)} | 去重 {len(duplicate_files)}",
        "data": {
            "total": len(files),
            "saved": len(saved_files),
            "success": saved_files,
            "duplicate_files": duplicate_files,
            "namespace": user_id
        }
    }

# ------------------------------
# 路由
# ------------------------------
@new_updateFile_router.post("/new_update_file")
async def new_update_file(
    request: Request,
    files: List[UploadFile] = File(..., description="支持 txt, md, pdf, docx, doc, json"),
    namespace: str = Form(..., description="命名空间")
) -> dict:
    return await process_uploaded_files(files, namespace)