RAG/api/main.py

"""
FastAPI main application
"""
import warnings
# Suppress pkg_resources deprecation warning from debugpy extension
# This warning is harmless and comes from VS Code/Cursor debugpy extension
# It appears when debugpy is loaded, which happens during debugging
warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*", category=UserWarning)

from contextlib import asynccontextmanager
from pathlib import Path
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.responses import StreamingResponse, Response
from fastapi.staticfiles import StaticFiles
from fastapi.openapi.docs import get_swagger_ui_html, get_redoc_html
from pydantic import BaseModel, Field
from typing import Optional, List
from loguru import logger
from config import settings
from rag import VectorStoreManager, RAGEngine, FileParser, DocumentProcessor
from sync_service import SyncService
import requests
from datetime import datetime
import uuid
import threading
import logging
from fastapi import Request
from fastapi.responses import PlainTextResponse
import sqlite3
import json
import os
import asyncio
import re
import markdown2
from pathlib import Path


# Global instances
vector_store_manager: Optional[VectorStoreManager] = None
rag_engine: Optional[RAGEngine] = None
sync_service: Optional[SyncService] = None
file_parser: Optional[FileParser] = None
document_processor: Optional[DocumentProcessor] = None
auto_sync_task = None  # Keep reference to prevent garbage collection


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan context manager for startup and shutdown events"""
    global vector_store_manager, rag_engine, sync_service, file_parser, document_processor, auto_sync_task
    
    # Startup
    try:
        logger.info("Initializing RAG services...")
        
        # Initialize vector store
        try:
            vector_store_manager = VectorStoreManager()
            logger.info("✓ VectorStoreManager initialized")
        except Exception as e:
            logger.error(f"Failed to initialize VectorStoreManager: {e}")
            raise
        
        # Initialize RAG engine
        try:
            rag_engine = RAGEngine(vector_store_manager)
            logger.info("✓ RAGEngine initialized")
        except Exception as e:
            logger.error(f"Failed to initialize RAGEngine: {e}")
            raise
        
        # Initialize file parser
        try:
            file_parser = FileParser()
            logger.info("✓ FileParser initialized")
        except Exception as e:
            logger.error(f"Failed to initialize FileParser: {e}")
            raise
        
        # Initialize document processor
        try:
            document_processor = DocumentProcessor()
            logger.info("✓ DocumentProcessor initialized")
        except Exception as e:
            logger.error(f"Failed to initialize DocumentProcessor: {e}")
            raise
        
        logger.info("✓ Core RAG services initialized")
        
        # Initialize sync service in background (non-blocking)
        # This allows API to start immediately even if MySQL connection fails
        async def init_sync_service():
            """Initialize sync service and start sync in background"""
            global sync_service
            try:
                logger.info("Initializing MySQL sync service in background...")
                
                # Run SyncService initialization in thread pool to avoid blocking event loop
                # SyncService.__init__() contains synchronous MySQL connection checks
                loop = asyncio.get_event_loop()
                sync_service = await loop.run_in_executor(None, SyncService)
                
                # Start initial sync in background and wait for it to complete before starting auto sync
                # This ensures API starts immediately without blocking, but auto sync waits for initial sync
                async def run_sync():
                    """Run sync_all in background without blocking"""
                    try:
                        logger.info("Starting initial data sync in background...")
                        await sync_service.sync_all()  # 全量同步 MySQL 数据到向量库
                        logger.info("✓ Initial sync completed")
                        return True
                    except Exception as e:
                        logger.error(f"Error during initial sync: {e}")
                        return False
                
                # Create task for initial sync (don't await - allows API to start immediately)
                initial_sync_task = asyncio.create_task(run_sync())
                
                # Start auto sync if enabled
                # Auto sync service will wait for initial sync to complete before starting incremental syncs
                # 自动同步服务将在初始同步（全量同步）完成后开始增量同步
                if settings.AUTO_SYNC:
                    logger.info(f"Starting auto sync service (interval: {settings.SYNC_INTERVAL}s)...")
                    # Start auto sync in background - it will wait for initial sync to complete
                    async def start_auto_sync_after_init():
                        """Wait for initial sync to complete, then start auto sync with error recovery"""
                        try:
                            # Wait for initial sync to complete (no timeout - wait indefinitely)
                            logger.info("Auto sync service waiting for initial sync to complete (no timeout, will wait until completion)...")
                            try:
                                initial_sync_success = await initial_sync_task
                                if initial_sync_success:
                                    logger.info("✓ Initial sync completed successfully")
                                else:
                                    logger.warning("Initial sync failed, but starting auto sync service anyway...")
                            except Exception as sync_error:
                                logger.error(f"Initial sync encountered an error: {sync_error}", exc_info=True)
                                logger.warning("Starting auto sync service despite initial sync error...")
                            
                            # Start auto sync service with error recovery (it will skip initial sync since we already did it)
                            await sync_service.start_auto_sync_with_recovery(skip_initial_sync=True)
                        except Exception as e:
                            logger.error(f"Error starting auto sync service: {e}", exc_info=True)
                    
                    # Start auto sync service in background (don't await - allows API to start immediately)
                    asyncio.create_task(start_auto_sync_after_init())
                else:
                    # If auto sync is disabled, just run initial sync in background
                    # Don't await it to avoid blocking
                    pass
            except Exception as e:
                logger.error(f"Failed to initialize sync service: {e}")
                logger.warning("API is still available, but MySQL sync is disabled. Please check MySQL connection.")
                # Don't raise - allow API to continue running
        
        # Start sync service initialization in background (non-blocking)
        # This ensures API can respond immediately while sync runs in background
        sync_init_task = asyncio.create_task(init_sync_service())
        auto_sync_task = sync_init_task  # Store reference for shutdown
        
        logger.info("✓ RAG API is now available (sync running in background)")
    except Exception as e:
        logger.error(f"Failed to initialize core services: {e}")
        raise

    yield

    # Shutdown
    if auto_sync_task and not auto_sync_task.done():
        logger.info("Stopping sync services...")
        # Wait a bit for the task to finish gracefully
        try:
            await asyncio.wait_for(auto_sync_task, timeout=10.0)
        except asyncio.TimeoutError:
            logger.warning("Sync task did not stop in time, cancelling...")
            auto_sync_task.cancel()
            try:
                await auto_sync_task
            except asyncio.CancelledError:
                pass
    if sync_service:
        if hasattr(sync_service, 'stop_auto_sync'):
            sync_service.stop_auto_sync()
        if hasattr(sync_service, 'close'):
            sync_service.close()
    logger.info("RAG services shut down")


# Check if static assets are available for offline use
STATIC_DIR = Path(__file__).parent.parent / "static" / "swagger-ui"
SWAGGER_UI_BUNDLE = STATIC_DIR / "swagger-ui-bundle.js"
SWAGGER_UI_CSS = STATIC_DIR / "swagger-ui.css"
REDOC_BUNDLE = STATIC_DIR / "redoc.standalone.js"

# Determine if we should use local assets or CDN
USE_LOCAL_ASSETS = (
    SWAGGER_UI_BUNDLE.exists() and 
    SWAGGER_UI_CSS.exists() and 
    REDOC_BUNDLE.exists()
)

if USE_LOCAL_ASSETS:
    logger.info("Using local Swagger UI assets for offline mode")
else:
    logger.warning(
        "Local Swagger UI assets not found. Swagger UI will use CDN resources. "
        "For offline use, run: python download_swagger_assets.py"
    )

# Initialize FastAPI app
app = FastAPI(
    title=settings.API_TITLE,
    version=settings.API_VERSION,
    description="RAG API for local knowledge base retrieval and generation",
    lifespan=lifespan,
    swagger_ui_parameters={
        "persistAuthorization": True,
    },
)

# Mount static files directory if it exists
if STATIC_DIR.exists():
    app.mount("/static/swagger-ui", StaticFiles(directory=str(STATIC_DIR)), name="swagger-ui-static")


# Add favicon route to prevent 404 errors
@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    """Return empty favicon to prevent 404 errors"""
    # Return a minimal 1x1 transparent PNG
    # This prevents browser from requesting favicon and getting 404
    return Response(
        content=b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xdb\x00\x00\x00\x00IEND\xaeB`\x82',
        media_type="image/png"
    )


# Override Swagger UI to use local assets when available
@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    """Custom Swagger UI that uses local assets in offline mode"""
    if USE_LOCAL_ASSETS:
        # Use local static files
        return get_swagger_ui_html(
            openapi_url=app.openapi_url,
            title=app.title + " - Swagger UI",
            swagger_js_url="/static/swagger-ui/swagger-ui-bundle.js",
            swagger_css_url="/static/swagger-ui/swagger-ui.css",
            swagger_favicon_url="/favicon.ico",
            oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
            init_oauth=app.swagger_ui_init_oauth,
            swagger_ui_parameters=app.swagger_ui_parameters,
        )
    else:
        # Fallback to default (CDN)
        return get_swagger_ui_html(
            openapi_url=app.openapi_url,
            title=app.title + " - Swagger UI",
            swagger_favicon_url="/favicon.ico",
            oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
            init_oauth=app.swagger_ui_init_oauth,
            swagger_ui_parameters=app.swagger_ui_parameters,
        )


# Override ReDoc to use local assets when available
@app.get("/redoc", include_in_schema=False)
async def custom_redoc_html():
    """Custom ReDoc that uses local assets in offline mode"""
    if USE_LOCAL_ASSETS:
        # Use local static files
        return get_redoc_html(
            openapi_url=app.openapi_url,
            title=app.title + " - ReDoc",
            redoc_js_url="/static/swagger-ui/redoc.standalone.js",
            redoc_favicon_url="/favicon.ico",
            with_google_fonts=False,  # Disable Google Fonts for offline use
        )
    else:
        # Fallback to default (CDN)
        return get_redoc_html(
            openapi_url=app.openapi_url,
            title=app.title + " - ReDoc",
            redoc_favicon_url="/favicon.ico",
            with_google_fonts=True,
        )


# Request/Response models
class QueryRequest(BaseModel):
    """Query request model"""
    query: str = Field(..., description="User query string", min_length=1)
    top_k: Optional[int] = Field(None, description="Number of documents to retrieve", ge=1, le=20)
    stream: bool = Field(True, description="Whether to stream the response")


class RetrieveRequest(BaseModel):
    """Retrieve request model - only retrieves documents from ChromaDB, no LLM generation"""
    query: str = Field(..., description="Query string for retrieval", min_length=1)
    top_k: Optional[int] = Field(None, description="Number of documents to retrieve", ge=1, le=20)


class RetrievedDocument(BaseModel):
    """Retrieved document model"""
    content: str = Field(..., description="Document content/text")
    score: Optional[float] = Field(None, description="Similarity score")
    metadata: dict = Field(default_factory=dict, description="Document metadata")
    doc_id: Optional[str] = Field(None, description="Document ID")


class RetrieveResponse(BaseModel):
    """Retrieve response model"""
    query: str = Field(..., description="Original query")
    documents: list[RetrievedDocument] = Field(..., description="Retrieved documents")
    count: int = Field(..., description="Number of documents retrieved")


class SyncRequest(BaseModel):
    """Manual sync request model"""
    full_sync: bool = Field(False, description="Whether to perform full sync")
    force: bool = Field(False, description="Whether to force re-processing of all documents (even if they exist)")


class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    message: str


class DocumentChunk(BaseModel):
    """Document chunk model"""
    id: str = Field(..., description="Chunk ID")
    text: str = Field(..., description="Chunk content")
    metadata: dict = Field(default_factory=dict, description="Chunk metadata")
    chunk_index: Optional[int] = Field(None, description="Chunk index (if document was chunked)")


class DocumentResponse(BaseModel):
    """Document response model"""
    doc_id: str = Field(..., description="Document ID")
    chunks: List[DocumentChunk] = Field(..., description="Document chunks")
    total_chunks: int = Field(..., description="Total number of chunks")
    full_text: str = Field(..., description="Full document text (all chunks combined)")


class UploadResponse(BaseModel):
    """File upload response model"""
    doc_id: str = Field(..., description="Document ID")
    filename: str = Field(..., description="Original filename")
    file_type: str = Field(..., description="File type/extension")
    chunks: int = Field(..., description="Number of chunks created")
    message: str = Field(..., description="Upload status message")


# --- Compatibility layer: endpoints copied from example.py (Flask) ---
# These endpoints provide a backward-compatible interface so existing clients
# that used the Flask-based API can keep working.

# In-memory conversation store and lock (thread-safe)
conversations = {}
conv_lock = threading.Lock()

# Configure logging for this compatibility layer
compat_logger = logging.getLogger("compat_api")
compat_logger.setLevel(logging.INFO)

# Session persistence (SQLite)
DATA_DIR = Path(__file__).parent.parent / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DB_PATH = DATA_DIR / "sessions.db"


def init_session_db():
    conn = sqlite3.connect(DB_PATH)
    try:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                user_login TEXT,
                title TEXT,
                data TEXT,
                update_time TEXT
            )
            """
        )
        # Improve concurrency for writes
        conn.execute("PRAGMA journal_mode=WAL;")
        conn.execute("PRAGMA synchronous=NORMAL;")
        conn.commit()
    finally:
        conn.close()


def load_sessions():
    conn = sqlite3.connect(DB_PATH)
    try:
        cur = conn.execute("SELECT id, user_login, title, data, update_time FROM sessions")
        rows = cur.fetchall()
        with conv_lock:
            for rid, user_login, title, data_text, update_time in rows:
                try:
                    data = json.loads(data_text) if data_text else {"messages": []}
                except Exception:
                    data = {"messages": []}
                conversations[rid] = {
                    "title": title or "",
                    "user_login": user_login or "",
                    "messages": data.get("messages", []),
                    "create_time": update_time,
                    "update_time": update_time,
                    "active_stream": False
                }
    finally:
        conn.close()

def serialize_conversation(conv):
    """安全的序列化函数，处理可能包含协程的情况"""
    messages = conv.get("messages", [])
    
    # 清理消息列表，移除不可序列化的对象
    cleaned_messages = []
    for msg in messages:
        if asyncio.iscoroutine(msg):
            # 如果是协程，记录错误或跳过
            compat_logger.warning(f"发现协程对象在消息列表中: {msg}")
            continue
        elif isinstance(msg, dict):
            # 递归清理字典中的值
            cleaned_msg = {}
            for key, value in msg.items():
                if not asyncio.iscoroutine(value):
                    cleaned_msg[key] = value
                else:
                    cleaned_msg[key] = "<coroutine>"
            cleaned_messages.append(cleaned_msg)
        else:
            cleaned_messages.append(msg)
    
    return json.dumps(
        {"messages": cleaned_messages}, 
        ensure_ascii=False, 
        default=str  # 处理其他不可序列化的类型
    )


def save_session(conv_id: str):
    # Upsert session into SQLite
    with conv_lock:
        conv = conversations.get(conv_id)
        if not conv:
            return
        title = conv.get("title", "")
        user_login = conv.get("user_login", "")
        data_text = serialize_conversation(conv)
        update_time = conv.get("update_time", datetime.now().isoformat())

    conn = sqlite3.connect(DB_PATH)
    try:
        conn.execute(
            "REPLACE INTO sessions (id, user_login, title, data, update_time) VALUES (?, ?, ?, ?, ?)",
            (conv_id, user_login, title, data_text, update_time)
        )
        conn.commit()
    finally:
        conn.close()


# Initialize DB and load sessions at import
init_session_db()
load_sessions()


def build_prompt_from_history(conv_id: str, max_messages: int = 10) -> str:
    """Build a single prompt string that includes recent conversation history."""
    
    def filter_coroutines_and_join(parts):
        """
        过滤掉协程对象和协程字符串，只保留真正的内容
        """
        filtered_parts = []
        
        for item in parts:
            # 如果是协程对象，直接跳过
            if asyncio.iscoroutine(item):
                continue
            
            # 如果是字符串，检查是否包含协程描述
            if isinstance(item, str):
                # 检查是否包含协程标记
                if contains_coroutine_marker(item):
                    continue
                
                # 保留非协程相关的字符串
                filtered_parts.append(item)
            # 其他非字符串类型（如果不是协程）
            else:
                filtered_parts.append(str(item))
        
        return "\n".join(filtered_parts)

    def contains_coroutine_marker(text):
        """检查文本是否包含协程标记"""
        coroutine_patterns = [
            r'<coroutine object',
            r'<coroutine>',
            r'coroutine object at 0x',
            r'coroutine at 0x'
        ]
        
        for pattern in coroutine_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                return True
        
        return False
    
    with conv_lock:
        conv = conversations.get(conv_id)
        if not conv:
            return None
        msgs = conv.get("messages", [])[-max_messages:]

    if len(msgs) <= 1: # only new query content
        return None
    
    parts = []
    for m in msgs[:-1]:  # Exclude the last message (current user query)
        role = m.get("role", "user")
        content = m.get("content", "")
        if role == "user":
            parts.append(f"User: {content}")
        else:
            parts.append(f"Assistant: {content}")
    if len(parts) ==  0:
        return None
    else:
        return filter_coroutines_and_join(parts) # 防止coroutine对象被放入messages中


# API endpoints
@app.get("/", response_model=HealthResponse)
async def root():
    """Root endpoint"""
    return HealthResponse(
        status="ok",
        message="RAG API is running"
    )


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    try:
        if vector_store_manager is None or rag_engine is None:
            raise HTTPException(status_code=503, detail="Services not initialized")
        
        return HealthResponse(
            status="healthy",
            message="All services are running"
        )
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"Service unhealthy: {str(e)}")


@app.post("/query")
async def query(request: QueryRequest):
    """
    Query the RAG system (retrieves documents and generates LLM response)
    
    Args:
        request: Query request with query string and options
        
    Returns:
        Streaming response or complete response
    """
    if rag_engine is None:
        raise HTTPException(status_code=503, detail="RAG engine not initialized")
    
    try:
        if request.stream:
            # Stream response                  
            
            return StreamingResponse(
                rag_engine.query_stream(request.query, None, request.top_k),
                media_type="text/event-stream",
                headers={
                    "X-Accel-Buffering": "no",
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive"
                }
            )
        else:
            # Return complete response (run in thread pool for better concurrency)
            response = await rag_engine.query(request.query, None, request.top_k)
            return response
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")


# ----------------------
# Backward-compatible Flask-style endpoints (from example.py)
# ----------------------


@app.get('/conversations')
def list_conversations():
    with conv_lock:
        conv_list = [
            {
                "id": conv_id,
                "title": conv["title"],
                "update_time": conv["update_time"],
                "active_stream": conv.get("active_stream", False)
            }
            for conv_id, conv in conversations.items()
        ]

    conv_list.sort(key=lambda x: x["update_time"], reverse=True)
    return conv_list

@app.get('/conversations/{user_login}')
def list_conversations_by_user(user_login: str):
    with conv_lock:
        conv_list = [
            {
                "id": conv_id,
                "user_login": conv.get("user_login", ""),
                "title": conv["title"],
                "update_time": conv["update_time"],
                "active_stream": conv.get("active_stream", False)
            }
            for conv_id, conv in conversations.items()
            if conv.get("user_login") == user_login
        ]

    conv_list.sort(key=lambda x: x["update_time"], reverse=True)
    return conv_list


@app.post('/conversations')
def create_conversation(payload: dict):
    title = payload.get("title", "新对话")
    user_login = payload.get("user_login")
    if not user_login:
        raise HTTPException(status_code=400, detail="user_login is required")
    conv_id = str(uuid.uuid4())
    with conv_lock:
        conversations[conv_id] = {
            "title": title,
            "user_login": user_login,
            "messages": [],
            "create_time": datetime.now().isoformat(),
            "update_time": datetime.now().isoformat(),
            "active_stream": False
        }
    # Persist
    save_session(conv_id)
    return {"id": conv_id, "title": title}


# @app.get('/conversations/{conv_id}')
# def get_conversation(conv_id: str):
#     with conv_lock:
#         if conv_id not in conversations:
#             raise HTTPException(status_code=404, detail="对话不存在")
#         conv = conversations[conv_id]
#     return {
#         "id": conv_id,
#         "title": conv["title"],
#         "messages": conv["messages"],
#         "create_time": conv["create_time"],
#         "update_time": conv["update_time"],
#         "active_stream": conv.get("active_stream", False)
#     }
    
@app.get('/conversations/{user_login}/{conv_id}')
def get_conversation_by_user(user_login: str, conv_id: str):
    with conv_lock:
        if conv_id not in conversations:
            raise HTTPException(status_code=404, detail="对话不存在")
        conv = conversations[conv_id]
        if conv.get("user_login") != user_login:
            raise HTTPException(status_code=403, detail="无权访问该对话")
    return {
        "id": conv_id,
        "title": conv["title"],
        "messages": conv["messages"],
        "create_time": conv["create_time"],
        "update_time": conv["update_time"],
        "active_stream": conv.get("active_stream", False)
    }


@app.delete('/conversations/{conv_id}')
def delete_conversation(conv_id: str):
    with conv_lock:
        if conv_id not in conversations:
            raise HTTPException(status_code=404, detail="对话不存在")
        del conversations[conv_id]
    return {"message": "删除成功"}

@app.delete('/conversations/{user_login}/{conv_id}')
def delete_conversation_by_user(user_login: str, conv_id: str):
    with conv_lock:
        if conv_id not in conversations:
            raise HTTPException(status_code=404, detail="对话不存在")
        conv = conversations[conv_id]
        if conv.get("user_login") != user_login:
            raise HTTPException(status_code=403, detail="无权删除该对话")
        del conversations[conv_id]
    return {"message": "删除成功"}


@app.post('/conversations/{conv_id}/messages')
async def post_message(conv_id: str, payload: dict):
    if conv_id not in conversations:
        raise HTTPException(status_code=404, detail="对话不存在")

    if not isinstance(payload, dict) or "content" not in payload:
        raise HTTPException(status_code=400, detail="Content required in JSON body")

    user_msg = {"role": "user", "content": payload["content"], "time": datetime.now().isoformat()}
    with conv_lock:
        conv = conversations[conv_id]
        conv["messages"].append(user_msg)
        conv["update_time"] = datetime.now().isoformat()
        if len(conv["messages"]) == 1:
            conv["title"] = user_msg["content"][:20] + ("..." if len(user_msg["content"]) > 20 else "")

    # persist after adding user message
    save_session(conv_id)

    # Use local RAG engine to generate response instead of external API
    if rag_engine is None:
        ai_content = "RAG engine not initialized"
    else:
        try:
            history = build_prompt_from_history(conv_id, max_messages=10)
            ai_content = await rag_engine.query(payload["content"], history)
        except Exception as e:
            compat_logger.error(f"RAG query failed: {e}")
            ai_content = f"RAG错误: {e}"

    ai_msg = {"role": "assistant", "content": ai_content, "time": datetime.now().isoformat()}
    with conv_lock:
        conversations[conv_id]["messages"].append(ai_msg)
        conversations[conv_id]["update_time"] = datetime.now().isoformat()

    # persist after adding assistant message
    save_session(conv_id)

    return ai_msg


@app.post('/conversations/{conv_id}/stream')
async def stream_message(conv_id: str, payload: dict):
    if conv_id not in conversations:
        raise HTTPException(status_code=404, detail="对话不存在")

    with conv_lock:
        if conversations[conv_id].get("active_stream", False):
            raise HTTPException(status_code=429, detail="该对话已有活动流，请先中止当前流")
        conversations[conv_id]["active_stream"] = True

    if not isinstance(payload, dict) or "content" not in payload:
        with conv_lock:
            conversations[conv_id]["active_stream"] = False
        raise HTTPException(status_code=400, detail="Content required in JSON body")
    
    if not isinstance(payload, dict) or "user_login" not in payload:
        with conv_lock:
            conversations[conv_id]["active_stream"] = False
        raise HTTPException(status_code=400, detail="user_login required in JSON body")
    
    if conversations[conv_id].get("user_login") != payload["user_login"]:
        with conv_lock:
            conversations[conv_id]["active_stream"] = False
        raise HTTPException(status_code=403, detail="用户无权操作该对话")

    user_msg = {"role": "user", "content": payload["content"], "time": datetime.now().isoformat()}
    with conv_lock:
        conversations[conv_id]["messages"].append(user_msg)
        conversations[conv_id]["update_time"] = datetime.now().isoformat()
        if len(conversations[conv_id]["messages"]) == 1:
            conversations[conv_id]["title"] = user_msg["content"][:20] + ("..." if len(user_msg["content"]) > 20 else "")

    # persist after adding user message
    save_session(conv_id)

    # DEEPSEEK_API = getattr(settings, 'DEEPSEEK_API', os.getenv('DEEPSEEK_API', None))
    # MODEL_ID = getattr(settings, 'DEEPSEEK_MODEL_ID', os.getenv('MODEL_ID', None))

    # Use local RAG engine streaming API
    async def generate_async():
        full_response = ""
        try:
            if rag_engine is None:
                yield "RAG engine not initialized"
                return
            
            chunk_counter = 0
            check_interval = 5  # 每生成5个chunk检查一次

            # Use rag_engine.query_stream which yields chunks asynchronously
            try:
                history = build_prompt_from_history(conv_id, max_messages=10)
                async for chunk in rag_engine.query_stream(payload["content"], history):
                    # 定期检查active_stream状态
                    chunk_counter += 1
                    if chunk_counter % check_interval == 0:
                        with conv_lock:
                            if conv_id not in conversations or not conversations[conv_id].get("active_stream", True):
                                compat_logger.info(f"流被中止: {conv_id}")
                                return  # 直接返回，中止生成
                    if chunk:
                        # normalize chunk to str
                        text_chunk = str(chunk.decode('utf-8')) if isinstance(chunk, bytes) else str(chunk)
                        # markdown_text = markdown2.markdown(text_chunk)
                        full_response += text_chunk
                        yield text_chunk.encode('utf-8') if isinstance(chunk, str) else text_chunk
            except Exception as e:
                compat_logger.error(f"RAG stream error: {e}", exc_info=True)
                yield f"流式响应错误: {e}"
                return

            # 最终检查是否被中止
            with conv_lock:
                if conv_id not in conversations or not conversations[conv_id].get("active_stream", True):
                    compat_logger.info(f"流在完成前被中止: {conv_id}")
                    return
            
            # Append final assistant message to conversation
            ai_msg = {"role": "assistant", "content": full_response, "time": datetime.now().isoformat()}
            with conv_lock:
                if conv_id in conversations:
                    conversations[conv_id]["messages"].append(ai_msg)
                    conversations[conv_id]["active_stream"] = False

            # persist after adding assistant message
            save_session(conv_id)

        finally:
            with conv_lock:
                if conv_id in conversations:
                    # 只在正常完成时才设为False，如果被中止，abort接口已经设为False了
                    if conversations[conv_id].get("active_stream", True):
                        conversations[conv_id]["active_stream"] = False

    # return StreamingResponse(generate_async(), media_type='text/plain; charset=utf-8')
    return StreamingResponse(
                generate_async(),
                media_type="text/event-stream",
                headers={
                    "X-Accel-Buffering": "no",
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive"
                }
            )


@app.post('/conversations/{conv_id}/abort')
def abort_stream(conv_id: str):
    with conv_lock:
        if conv_id not in conversations:
            raise HTTPException(status_code=404, detail="对话不存在")
        if not conversations[conv_id].get("active_stream", False):
            raise HTTPException(status_code=400, detail="该对话没有活动流")
        conversations[conv_id]["active_stream"] = False
    return {"message": "流式响应已中止"}


@app.post("/retrieve", response_model=RetrieveResponse)
async def retrieve(request: RetrieveRequest):
    """
    Retrieve documents from ChromaDB based on query (no LLM generation)
    
    This endpoint only performs vector similarity search and returns the retrieved
    documents with their content, metadata, and similarity scores.
    
    Args:
        request: Retrieve request with query string and top_k option
        
    Returns:
        Retrieved documents with content, metadata, and scores
    """
    if vector_store_manager is None:
        raise HTTPException(status_code=503, detail="Vector store not initialized")
    
    try:
               
        # Get retriever (works even if collection is empty, will return empty results)
        retriever = vector_store_manager.get_retriever(top_k=request.top_k or settings.TOP_K)
        
        # Retrieve documents (run in thread pool since retriever.retrieve() is synchronous)
        def retrieve_docs():
            nodes = retriever.retrieve(request.query)
            return nodes
        
        nodes = await asyncio.to_thread(retrieve_docs)
        
        # Convert nodes to response format
        documents = []
        for node in nodes:
            # Extract similarity score if available
            # LlamaIndex retriever returns NodeWithScore objects
            score = None
            if hasattr(node, 'score'):
                score = node.score
            elif hasattr(node, 'node') and hasattr(node.node, 'score'):
                score = node.node.score
            
            # Get the actual node object (NodeWithScore.node or the node itself)
            actual_node = node.node if hasattr(node, 'node') else node
            
            # Get content
            content = actual_node.text if hasattr(actual_node, 'text') else str(actual_node)
            
            # Get metadata
            metadata = {}
            if hasattr(actual_node, 'metadata'):
                metadata = actual_node.metadata.copy() if actual_node.metadata else {}
            
            # Get doc_id
            doc_id = metadata.get('doc_id') or metadata.get('original_doc_id')
            if not doc_id and hasattr(actual_node, 'node_id'):
                doc_id = actual_node.node_id
            
            documents.append(RetrievedDocument(
                content=content,
                score=score,
                metadata=metadata,
                doc_id=doc_id
            ))
        
        return RetrieveResponse(
            query=request.query,
            documents=documents,
            count=len(documents)
        )
    except Exception as e:
        logger.error(f"Error retrieving documents: {e}")
        raise HTTPException(status_code=500, detail=f"Error retrieving documents: {str(e)}")


@app.post("/sync")
async def manual_sync(request: SyncRequest):
    """
    Manually trigger synchronization
    
    Args:
        request: Sync request with options
        
    Returns:
        Sync status
    """
    if sync_service is None:
        raise HTTPException(status_code=503, detail="Sync service not initialized")
    
    try:
        if request.full_sync:
            await sync_service.sync_all(force=request.force)
            message = "Full sync completed" + (" (forced re-processing)" if request.force else "")
            return {"status": "success", "message": message}
        else:
            await sync_service.sync_incremental()
            return {"status": "success", "message": "Incremental sync completed"}
    except Exception as e:
        logger.error(f"Error during sync: {e}")
        raise HTTPException(status_code=500, detail=f"Sync failed: {str(e)}")


@app.get("/stats")
async def get_stats():
    """
    Get system statistics
    
    Returns:
        System statistics
    """
    if vector_store_manager is None:
        raise HTTPException(status_code=503, detail="Vector store not initialized")
    
    try:
        collection = vector_store_manager.collection
        count = collection.count()
        
        return {
            "vector_store": {
                "collection_name": settings.CHROMA_COLLECTION_NAME,
                "document_count": count
            },
            "config": {
                "model": settings.OLLAMA_MODEL,
                "top_k": settings.TOP_K,
                "chunk_size": settings.CHUNK_SIZE
            }
        }
    except Exception as e:
        logger.error(f"Error getting stats: {e}")
        raise HTTPException(status_code=500, detail=f"Error getting stats: {str(e)}")


@app.post("/documents/upload", response_model=UploadResponse)
async def upload_document(
    file: UploadFile = File(..., description="File to upload"),
    doc_id: Optional[str] = Form(None, description="Custom document ID (optional, will use filename if not provided)"),
    metadata: Optional[str] = Form(None, description="Additional metadata as JSON string (optional)")
):
    """
    Upload and parse a document file
    
    Supports various file formats:
    - Text files: .txt, .md, .markdown
    - PDF files: .pdf
    - Word documents: .docx, .doc
    - HTML files: .html, .htm
    - CSV files: .csv
    - JSON files: .json
    
    File size limit: Configurable via MAX_UPLOAD_SIZE_MB in .env (default: 10MB)
    
    Args:
        file: File to upload (max size configured in MAX_UPLOAD_SIZE_MB)
        doc_id: Optional custom document ID
        metadata: Optional JSON string with additional metadata
        
    Returns:
        Upload response with document ID and chunk count
        
    Raises:
        HTTPException: 400 if file format is unsupported
        HTTPException: 413 if file size exceeds 10MB
        HTTPException: 400 if document parsing fails
    """
    if file_parser is None or vector_store_manager is None or document_processor is None:
        raise HTTPException(status_code=503, detail="Services not initialized")
    
    try:
        
        # Check if file format is supported
        if not file_parser.is_supported(file.filename):
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported file format. Supported formats: {', '.join(file_parser.SUPPORTED_EXTENSIONS)}"
            )
        
        # Read file content
        content = await file.read()
        
        # Check file size limit (from config, default: 10MB)
        MAX_FILE_SIZE = settings.MAX_UPLOAD_SIZE_MB * 1024 * 1024  # Convert MB to bytes
        file_size = len(content)
        if file_size > MAX_FILE_SIZE:
            raise HTTPException(
                status_code=413,
                detail=f"File size ({file_size / (1024 * 1024):.2f}MB) exceeds maximum allowed size ({settings.MAX_UPLOAD_SIZE_MB}MB)"
            )
        
        # Parse metadata if provided
        extra_metadata = {}
        if metadata:
            try:
                extra_metadata = json.loads(metadata)
            except json.JSONDecodeError:
                raise HTTPException(status_code=400, detail="Invalid metadata JSON format")
        
        # Use custom doc_id or generate from filename
        if not doc_id:
            doc_id = Path(file.filename).stem
        
        # Parse file (run in thread pool to avoid blocking event loop for large files)
        logger.info(f"Parsing uploaded file: {file.filename} (doc_id: {doc_id}, size: {file_size / 1024:.2f}KB)")
        try:
            # 获取事件循环，用于在异步环境中运行同步函数
            loop = asyncio.get_event_loop()
            # 使用线程池执行器运行文件解析函数（同步函数），避免阻塞事件循环
            documents = await loop.run_in_executor(
                None,  # 使用默认线程池
                file_parser.parse_file_content,  # 要执行的同步解析函数
                content,  # 文件内容（字节流）
                file.filename,  # 文件名
                doc_id,  # 文档ID
                extra_metadata# 额外元数据
            )
        except Exception as parse_error:
            logger.error(f"Error parsing file {file.filename}: {parse_error}", exc_info=True)
            raise HTTPException(
                status_code=400,
                detail=f"Failed to parse document: {str(parse_error)}"
            )
        
        if not documents:
            raise HTTPException(status_code=400, detail="Failed to parse document or document is empty")
        
        logger.info(f"Successfully parsed {file.filename}: {len(documents)} document(s)")
        
        # Chunk documents if needed
        logger.info(f"Chunking {len(documents)} document(s)...")
        chunked_documents = document_processor.chunk_documents(documents)
        logger.info(f"Chunked into {len(chunked_documents)} chunk(s)")
        
        # Add to vector store (this may take time for large documents due to embedding generation)
        logger.info(f"Adding {len(chunked_documents)} chunks to vector store (this may take time for large files)...")
        try:
            # Run in thread pool to avoid blocking event loop during embedding generation
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(
                None,
                vector_store_manager.add_documents,
                chunked_documents,
                False  # skip_existing
            )
            logger.info(f"Successfully added {len(chunked_documents)} chunks to vector store")
        except Exception as add_error:
            logger.error(f"Error adding documents to vector store: {add_error}", exc_info=True)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to add documents to vector store: {str(add_error)}"
            )
        
        file_type = Path(file.filename).suffix.lstrip('.')
        
        return UploadResponse(
            doc_id=doc_id,
            filename=file.filename,
            file_type=file_type,
            chunks=len(chunked_documents),
            message=f"Document uploaded and processed successfully"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error uploading document: {e}")
        raise HTTPException(status_code=500, detail=f"Error uploading document: {str(e)}")


@app.get("/documents/{doc_id}", response_model=DocumentResponse)
async def get_document_by_id(doc_id: str):
    """
    Get all content for a document by its ID
    
    This endpoint retrieves all chunks/content that belong to a document,
    even if the document was split into multiple chunks during processing.
    
    Args:
        doc_id: Document ID to retrieve
        
    Returns:
        Document response with all chunks and full text
    """
    if vector_store_manager is None:
        raise HTTPException(status_code=503, detail="Vector store not initialized")
    
    try:
        # Get all chunks for this document
        chunks = vector_store_manager.get_document_by_id(doc_id)
        
        if not chunks:
            raise HTTPException(status_code=404, detail=f"Document with ID '{doc_id}' not found")
        
        # Convert to response format
        document_chunks = []
        full_text_parts = []
        
        for chunk in chunks:
            document_chunks.append(DocumentChunk(
                id=chunk['id'],
                text=chunk['text'],
                metadata=chunk.get('metadata', {}),
                chunk_index=chunk.get('chunk_index')
            ))
            full_text_parts.append(chunk['text'])
        
        # Combine all chunks into full text
        full_text = "\n\n".join(full_text_parts)
        
        return DocumentResponse(
            doc_id=doc_id,
            chunks=document_chunks,
            total_chunks=len(document_chunks),
            full_text=full_text
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error retrieving document {doc_id}: {e}")
        raise HTTPException(status_code=500, detail=f"Error retrieving document: {str(e)}")