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


# 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
    import asyncio
    
    # 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()
                        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")


# 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
            async def generate():
                try:
                    async for chunk in rag_engine.query_stream(request.query, request.top_k):
                        yield chunk
                except Exception as e:
                    logger.error(f"Error in stream generation: {e}")
                    yield f"\n\nError: {str(e)}"
            
            return StreamingResponse(
                generate(),
                media_type="text/plain",
                headers={
                    "X-Accel-Buffering": "no",
                    "Cache-Control": "no-cache"
                }
            )
        else:
            # Return complete response (run in thread pool for better concurrency)
            import asyncio
            response = await asyncio.to_thread(rag_engine.query, request.query, request.top_k)
            return {"response": response, "query": request.query}
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")


@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:
        import asyncio
        
        # 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:
        import json
        import asyncio
        
        # 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:
            import asyncio
            loop = asyncio.get_event_loop()
            documents = await loop.run_in_executor(
                None,
                file_parser.parse_file_content,
                content,
                file.filename,
                doc_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
            import asyncio
            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)}")