lily_llm_api/api/routers/document_router.py

"""
Document processing and RAG router for Lily LLM API
"""
from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from typing import Optional, List
import logging
import time
import os
import uuid

from ...models.schemas import (
    DocumentUploadResponse, RAGQueryRequest, RAGQueryResponse,
    DocumentProcessResponse, MultimodalRAGResponse
)
from ...services.session_registry import set_user_for_room, set_flag_for_room, set_last_document_for_room

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

@router.post("/document/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """문서 업로드 및 처리"""
    try:
        start_time = time.time()
        
        # 파일 읽기 및 임시 저장 (파일 경로 기반 처리기 호환)
        content = await file.read()
        filename = file.filename
        temp_dir = os.path.join("data", "uploads")
        os.makedirs(temp_dir, exist_ok=True)
        temp_name = f"{int(time.time()*1000)}_{uuid.uuid4().hex}_{filename}"
        temp_path = os.path.join(temp_dir, temp_name)
        with open(temp_path, "wb") as f:
            f.write(content)
        
        # 문서 처리기 사용 (우선 RAG에 저장 포함 경로)
        try:
            from lily_llm_core.rag_processor import rag_processor
            document_id = f"doc_{int(time.time()*1000)}_{uuid.uuid4().hex}"
            result = rag_processor.process_and_store_document(
                user_id=user_id,
                document_id=document_id,
                file_path=temp_path,
            )
            # 업로드 시 방-사용자 매핑 저장 (후속 생성에서 자동 보정)
            try:
                set_user_for_room(room_id, user_id)
            except Exception:
                pass
            
            if result.get("success"):
                processing_time = time.time() - start_time
                # 업로드 직후, 같은 방에서 다음 1회 생성은 이미지 복구를 허용
                try:
                    set_flag_for_room(room_id, "use_rag_images_once", True)
                except Exception:
                    pass
                # 방별 마지막 문서ID 기록 (후속 요청에서 기본 문서 선택)
                try:
                    set_last_document_for_room(room_id, result.get("document_id", document_id))
                except Exception:
                    pass
                return DocumentUploadResponse(
                    success=True,
                    document_id=result.get("document_id", document_id),
                    message="문서 업로드 및 처리 완료",
                    chunks=result.get("chunks", 0),
                    latex_count=result.get("latex_count", 0),
                    auto_response=result.get("auto_response")
                )
            else:
                return DocumentUploadResponse(
                    success=False,
                    document_id="",
                    message="문서 처리 실패",
                    error=result.get("error", "Unknown error")
                )
                
        except ImportError:
            # 폴백: 순수 문서 파서로 처리만 수행
            try:
                from lily_llm_core.document_processor import document_processor
                docs = document_processor.process_document(temp_path)
                processing_time = time.time() - start_time
                return DocumentUploadResponse(
                    success=True,
                    document_id="",
                    message="문서 업로드 및 처리 완료 (벡터 저장 미수행)",
                    chunks=len(docs) if docs else 0,
                    latex_count=0,
                    auto_response=None
                )
            except Exception as e:
                return DocumentUploadResponse(
                    success=False,
                    document_id="",
                    message="문서 처리기 import 실패",
                    error=str(e)
                )
            
    except Exception as e:
        logger.error(f"문서 업로드 실패: {e}")
        return DocumentUploadResponse(
            success=False,
            document_id="",
            message="문서 업로드 중 오류 발생",
            error=str(e)
        )

@router.post("/document/upload/", response_model=DocumentUploadResponse, include_in_schema=False)
async def upload_document_trailing_slash(
    file: UploadFile = File(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    # 슬래시 유무 모두 허용 (FastAPI 307 리디렉션 회피)
    return await upload_document(file=file, user_id=user_id, room_id=room_id)
@router.post("/rag/query", response_model=RAGQueryResponse)
async def rag_query(
    query: str = Form(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default"),
    max_results: int = Form(5),
    include_sources: bool = Form(True)
):
    """RAG 쿼리 처리"""
    try:
        start_time = time.time()
        
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            # RAG 쿼리 실행
            result = rag_processor.query(
                query=query,
                user_id=user_id,
                room_id=room_id,
                max_results=max_results,
                include_sources=include_sources
            )
            
            if result.get("success"):
                processing_time = time.time() - start_time
                return RAGQueryResponse(
                    success=True,
                    response=result.get("response", ""),
                    sources=result.get("sources", []),
                    search_results=len(result.get("sources", [])),
                    processing_time=processing_time
                )
            else:
                return RAGQueryResponse(
                    success=False,
                    response="",
                    sources=[],
                    search_results=0,
                    processing_time=0,
                    error=result.get("error", "RAG 쿼리 실패")
                )
                
        except ImportError:
            return RAGQueryResponse(
                success=False,
                response="",
                sources=[],
                search_results=0,
                processing_time=0,
                error="RAG processor not available"
            )
            
    except Exception as e:
        logger.error(f"RAG 쿼리 실패: {e}")
        return RAGQueryResponse(
            success=False,
            response="",
            sources=[],
            search_results=0,
            processing_time=0,
            error=str(e)
        )

@router.post("/rag/generate", response_model=RAGQueryResponse)
async def rag_generate(
    prompt: str = Form(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default"),
    max_results: int = Form(5)
):
    """RAG 기반 텍스트 생성"""
    try:
        start_time = time.time()
        
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            # RAG 생성 실행
            result = rag_processor.generate_with_context(
                prompt=prompt,
                user_id=user_id,
                room_id=room_id,
                max_results=max_results
            )
            
            if result.get("success"):
                processing_time = time.time() - start_time
                return RAGQueryResponse(
                    success=True,
                    response=result.get("response", ""),
                    sources=result.get("sources", []),
                    search_results=len(result.get("sources", [])),
                    processing_time=processing_time
                )
            else:
                return RAGQueryResponse(
                    success=False,
                    response="",
                    sources=[],
                    search_results=0,
                    processing_time=0,
                    error=result.get("error", "RAG 생성 실패")
                )
                
        except ImportError:
            return RAGQueryResponse(
                success=False,
                response="",
                sources=[],
                search_results=0,
                processing_time=0,
                error="RAG processor not available"
            )
            
    except Exception as e:
        logger.error(f"RAG 생성 실패: {e}")
        return RAGQueryResponse(
            success=False,
            response="",
            sources=[],
            search_results=0,
            processing_time=0,
            error=str(e)
        )

@router.post("/rag/summary")
async def generate_rag_summary(
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """RAG 문서 요약 생성"""
    try:
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            # RAG 요약 생성
            result = rag_processor.generate_summary(
                user_id=user_id,
                room_id=room_id
            )
            
            if result.get("success"):
                return {"status": "success", "summary": result.get("summary", "")}
            else:
                raise HTTPException(status_code=500, detail=result.get("error", "RAG 요약 생성 실패"))
                
        except ImportError:
            raise HTTPException(status_code=500, detail="RAG processor not available")
            
    except Exception as e:
        logger.error(f"RAG 요약 생성 실패: {e}")
        raise HTTPException(status_code=500, detail=f"RAG 요약 생성 실패: {str(e)}")

@router.post("/rag/clear")
async def clear_rag_context(
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """RAG 컨텍스트 정리"""
    try:
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            # RAG 컨텍스트 정리
            success = rag_processor.clear_context(
                user_id=user_id,
                room_id=room_id
            )
            
            if success:
                return {"status": "success", "message": "RAG 컨텍스트 정리 완료"}
            else:
                raise HTTPException(status_code=500, detail="RAG 컨텍스트 정리 실패")
                
        except ImportError:
            raise HTTPException(status_code=500, detail="RAG processor not available")
            
    except Exception as e:
        logger.error(f"RAG 컨텍스트 정리 실패: {e}")
        raise HTTPException(status_code=500, detail=f"RAG 컨텍스트 정리 실패: {str(e)}")

@router.post("/rag/batch-process")
async def batch_process_documents(
    files: List[UploadFile] = File(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """여러 문서 일괄 처리"""
    try:
        start_time = time.time()
        results = []
        
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            for file in files:
                content = await file.read()
                filename = file.filename
                # 임시 저장 후 RAG에 저장 포함 처리
                temp_dir = os.path.join("data", "uploads")
                os.makedirs(temp_dir, exist_ok=True)
                temp_name = f"{int(time.time()*1000)}_{uuid.uuid4().hex}_{filename}"
                temp_path = os.path.join(temp_dir, temp_name)
                with open(temp_path, "wb") as f:
                    f.write(content)
                document_id = f"doc_{int(time.time()*1000)}_{uuid.uuid4().hex}"
                result = rag_processor.process_and_store_document(
                    user_id=user_id,
                    document_id=document_id,
                    file_path=temp_path,
                )
                
                results.append({
                    "filename": filename,
                    "success": result.get("success", False),
                    "document_id": result.get("document_id", document_id),
                    "chunks": result.get("chunks", 0),
                    "error": result.get("error")
                })
            
            processing_time = time.time() - start_time
            return {
                "status": "success",
                "results": results,
                "total_files": len(files),
                "processing_time": processing_time
            }
            
        except ImportError:
            # 폴백: 저장 없이 처리만 수행
            try:
                from lily_llm_core.document_processor import document_processor
                for file in files:
                    content = await file.read()
                    filename = file.filename
                    temp_dir = os.path.join("data", "uploads")
                    os.makedirs(temp_dir, exist_ok=True)
                    temp_name = f"{int(time.time()*1000)}_{uuid.uuid4().hex}_{filename}"
                    temp_path = os.path.join(temp_dir, temp_name)
                    with open(temp_path, "wb") as f:
                        f.write(content)
                    docs = document_processor.process_document(temp_path)
                    results.append({
                        "filename": filename,
                        "success": bool(docs),
                        "document_id": "",
                        "chunks": len(docs) if docs else 0,
                        "error": None if docs else "processing failed"
                    })
                processing_time = time.time() - start_time
                return {
                    "status": "success",
                    "results": results,
                    "total_files": len(files),
                    "processing_time": processing_time
                }
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))
            
    except Exception as e:
        logger.error(f"일괄 문서 처리 실패: {e}")
        raise HTTPException(status_code=500, detail=f"일괄 문서 처리 실패: {str(e)}")

@router.get("/rag/search-history")
async def search_rag_history(
    user_id: str = "anonymous",
    room_id: str = "default",
    query: str = "",
    limit: int = 10
):
    """RAG 검색 히스토리 조회"""
    try:
        try:
            from lily_llm_core.rag_processor import rag_processor
            
            # RAG 검색 히스토리 조회
            history = rag_processor.get_search_history(
                user_id=user_id,
                room_id=room_id,
                query=query,
                limit=limit
            )
            
            return {"status": "success", "history": history}
            
        except ImportError:
            raise HTTPException(status_code=500, detail="RAG processor not available")
            
    except Exception as e:
        logger.error(f"RAG 검색 히스토리 조회 실패: {e}")
        raise HTTPException(status_code=500, detail=f"RAG 검색 히스토리 조회 실패: {str(e)}")

@router.post("/multimodal-rag/upload")
async def upload_multimodal_document(
    file: UploadFile = File(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """멀티모달 문서 업로드"""
    try:
        start_time = time.time()
        
        # 파일 읽기
        content = await file.read()
        filename = file.filename
        
        try:
            from lily_llm_core.hybrid_rag_processor import hybrid_rag_processor
            # 임시 저장 후 파일 경로 기반 처리
            temp_dir = os.path.join("data", "uploads")
            os.makedirs(temp_dir, exist_ok=True)
            temp_name = f"{int(time.time()*1000)}_{uuid.uuid4().hex}_{filename}"
            temp_path = os.path.join(temp_dir, temp_name)
            with open(temp_path, "wb") as f:
                f.write(content)
            result = hybrid_rag_processor.process_document(
                file_path=temp_path,
                user_id=user_id,
                room_id=room_id
            )
            
            if result.get("success"):
                processing_time = time.time() - start_time
                return {
                    "status": "success",
                    "document_id": result.get("document_id", ""),
                    "processing_time": processing_time,
                    "message": "멀티모달 문서 업로드 완료"
                }
            else:
                raise HTTPException(status_code=500, detail=result.get("error", "멀티모달 문서 처리 실패"))
                
        except ImportError:
            raise HTTPException(status_code=500, detail="Hybrid RAG processor not available")
            
    except Exception as e:
        logger.error(f"멀티모달 문서 업로드 실패: {e}")
        raise HTTPException(status_code=500, detail=f"멀티모달 문서 업로드 실패: {str(e)}")

@router.post("/multimodal-rag/generate", response_model=MultimodalRAGResponse)
async def generate_multimodal_rag(
    prompt: str = Form(...),
    user_id: str = Form("anonymous"),
    room_id: str = Form("default")
):
    """멀티모달 RAG 기반 텍스트 생성"""
    try:
        start_time = time.time()
        
        try:
            from lily_llm_core.hybrid_rag_processor import hybrid_rag_processor
            
            # 멀티모달 RAG 생성
            result = hybrid_rag_processor.generate(
                prompt=prompt,
                user_id=user_id,
                room_id=room_id
            )
            
            if result.get("success"):
                processing_time = time.time() - start_time
                return MultimodalRAGResponse(
                    success=True,
                    response=result.get("response", ""),
                    image_processed=result.get("image_processed", False),
                    processing_time=processing_time
                )
            else:
                return MultimodalRAGResponse(
                    success=False,
                    response="",
                    image_processed=False,
                    processing_time=0,
                    error=result.get("error", "멀티모달 RAG 생성 실패")
                )
                
        except ImportError:
            return MultimodalRAGResponse(
                success=False,
                response="",
                image_processed=False,
                processing_time=0,
                error="Hybrid RAG processor not available"
            )
            
    except Exception as e:
        logger.error(f"멀티모달 RAG 생성 실패: {e}")
        return MultimodalRAGResponse(
            success=False,
            response="",
            image_processed=False,
            processing_time=0,
            error=str(e)
        )