lily_llm_api/api/routers/generation_router.py

"""
Generation router for Lily LLM API
"""
from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Form, Depends
import os
from typing import Optional, List
import logging
import time

from ...models.schemas import GenerateResponse, MultimodalGenerateResponse
from ...services.generation_service import generate_sync
from ...services.model_service import is_model_loaded
from ...utils.system_utils import select_model_interactive
from ...services.session_registry import get_user_for_room, set_user_for_room, set_user_for_session, pop_flag_for_room

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

@router.post("/generate", response_model=GenerateResponse)
async def generate(request: Request,
                  prompt: str = Form(...), 
                  image1: UploadFile = File(None),
                  image2: UploadFile = File(None),
                  image3: UploadFile = File(None),
                  image4: UploadFile = File(None),
                  user_id: str = Form("anonymous"),
                  room_id: str = Form("default"),
                  use_context: bool = Form(True),
                  session_id: str = Form(None),
                  use_rag_images: bool = Form(False),
                  use_rag_text: bool = Form(False),
                  document_id: str = Form(None),
                  image_short_side: int = Form(None),
                  # 새 옵션: 생성 토큰 수/입력 길이 제어
                  max_new_tokens: int = Form(None),
                  input_max_length: int = Form(None)):
    
    if not is_model_loaded():
        raise HTTPException(status_code=503, detail="모델이 로드되지 않았습니다.")

    start_time = time.time()

    # 사용자가 비어있거나 anonymous면 룸 기반 최근 사용자 보정
    if not user_id or user_id == "anonymous":
        try:
            recovered_user = get_user_for_room(room_id)
            if recovered_user:
                print(f"🔍 [DEBUG] 룸 기반 사용자 보정: {user_id} -> {recovered_user} (room={room_id})")
                user_id = recovered_user
        except Exception:
            pass

    # 방/사용자 기반 고정 세션 ID로 정규화
    stable_session_id = f"room_{room_id}_{user_id}"
    if not session_id or session_id != stable_session_id:
        original = session_id
        session_id = stable_session_id
        print(f"🔍 [DEBUG] 세션 ID 정규화: {original} -> {session_id}")
    else:
        # 제공된 세션에도 사용자 매핑 저장
        try:
            set_user_for_session(session_id, user_id)
        except Exception:
            pass
    
    if use_context:
        try:
            from lily_llm_core.context_manager import context_manager
            # 중기/장기 메모리 주입: room/user 기반
            try:
                from lily_llm_core.memory_store import memory_store
                rm = memory_store.get_room_memory(room_id)
                um = memory_store.get_user_memory(user_id)
                note_parts = []
                if rm and (rm.get("summary") or rm.get("key_topics")):
                    note_parts.append(f"[Room Memory]\n{rm.get('summary','')}\nTopics: {rm.get('key_topics','')}")
                if um and um.get("notes"):
                    note_parts.append(f"[User Memory]\n{um.get('notes','')}")
                if note_parts:
                    context_manager.set_system_note(session_id, "hier_mem", "\n\n".join(note_parts))
            except Exception as _me:
                print(f"⚠️ [DEBUG] 메모리 주입 실패: {_me}")
            # 컨텍스트 영속화 설정: 필요 시 세션 파일에서 로드
            try:
                persist = os.getenv('LILY_CONTEXT_PERSIST', '0') in ['1', 'true', 'True']
                if persist:
                    base_dir = os.getenv('LILY_CONTEXT_DIR', 'data/context')
                    os.makedirs(base_dir, exist_ok=True)
                    session_file = os.path.join(base_dir, f"{session_id}.json")
                    if hasattr(context_manager, 'load_session_from_file'):
                        context_manager.load_session_from_file(session_id, session_file)
            except Exception as _e:
                print(f"⚠️ [DEBUG] 컨텍스트 로드 실패: {_e}")
            context_manager.add_user_message(prompt, metadata={"session_id": session_id})
            print(f"🔍 [DEBUG] 사용자 메시지 추가됨 (세션: {session_id})")
        except Exception as e:
            logger.warning(f"⚠️ 컨텍스트 관리자 사용 불가: {e}")

    # 이미지 데이터 처리
    image_data_list = []
    for img_file in [image1, image2, image3, image4]:
        if img_file:
            try:
                data = await img_file.read()
                image_data_list.append(data)
            except Exception as e:
                logger.warning(f"이미지 로드 실패: {e}")

    try:
        # UX: 첨부 이미지가 있으면 자동으로 멀티모달 허용
        if image_data_list and len([img for img in image_data_list if img]) > 0:
            use_rag_images = True
            # 이미지가 있으면 텍스트 컨텍스트는 유지하고, RAG 텍스트는 기본 그대로 사용

        # 텍스트-only 강제: 첨부 이미지가 전혀 없고 플래그가 False면 이미지 리스트를 비움
        if (not image_data_list or len([img for img in image_data_list if img]) == 0) and not use_rag_images:
            # 업로드 직후 1회 한정 자동 이미지 허용 플래그가 있으면 소비(pop)하여 True로 전환
            try:
                once_flag = pop_flag_for_room(room_id, "use_rag_images_once")
                if once_flag:
                    use_rag_images = True
                    print(f"🔍 [DEBUG] 룸 {room_id}에 저장된 1회성 이미지 복구 플래그 사용 -> use_rag_images=True")
            except Exception:
                pass
            image_data_list = []
        # generate_sync 함수 호출 (컨텍스트 포함)
        result = generate_sync(
            prompt,
            image_data_list,
            # 생성 길이: 신/구 파라미터 모두 지원 (신>구 우선)
            max_length=max_new_tokens,
            use_context=use_context,
            session_id=session_id,
            user_id=user_id,
            room_id=room_id,
            use_rag_images=use_rag_images,
            use_rag_text=use_rag_text,
            document_id=document_id,
            image_short_side=image_short_side,
            input_max_length=input_max_length
        )
        
        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])
        
        if use_context:
            try:
                from lily_llm_core.context_manager import context_manager
                context_manager.add_assistant_message(result["generated_text"], metadata={"session_id": session_id})
                # 컨텍스트 영속화: 턴 저장 후 파일로 저장
                try:
                    persist = os.getenv('LILY_CONTEXT_PERSIST', '0') in ['1', 'true', 'True']
                    if persist:
                        base_dir = os.getenv('LILY_CONTEXT_DIR', 'data/context')
                        os.makedirs(base_dir, exist_ok=True)
                        session_file = os.path.join(base_dir, f"{session_id}.json")
                        if hasattr(context_manager, 'save_session_to_file'):
                            context_manager.save_session_to_file(session_id, session_file)
                    # 중기(room) 메모리 갱신: 최근 턴 요약을 압축 저장
                    try:
                        from lily_llm_core.memory_store import memory_store
                        # 간단히 최근 요약 컨텍스트를 저장하고 키토픽은 비움(후속 확장 지점)
                        summary_text = ""
                        if hasattr(context_manager, 'get_summary_context'):
                            summary_text = context_manager.get_summary_context(session_id)
                        last_ts = time.time()
                        memory_store.upsert_room_memory(room_id, summary_text, "", last_ts)
                    except Exception as _ms:
                        print(f"⚠️ [DEBUG] 룸 메모리 저장 실패: {_ms}")
                except Exception as _e:
                    print(f"⚠️ [DEBUG] 컨텍스트 저장 실패: {_e}")
            except Exception as e:
                logger.warning(f"⚠️ 컨텍스트 관리자 사용 불가: {e}")

        return GenerateResponse(
            generated_text=result["generated_text"],
            processing_time=result["processing_time"],
            model_name=result["model_name"],
            image_processed=result["image_processed"]
        )
        
    except Exception as e:
        logger.error(f"❌ 생성 중 오류 발생: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"모델 생성 중 오류 발생: {str(e)}")


@router.post("/generate-multimodal", response_model=MultimodalGenerateResponse)
async def generate_multimodal(prompt: str = Form(...),
                             image: UploadFile = File(None),
                             model_id: Optional[str] = Form(None),
                             max_length: Optional[int] = Form(None),
                             temperature: Optional[float] = Form(None),
                             top_p: Optional[float] = Form(None),
                             do_sample: Optional[bool] = Form(None)):
    
    if not is_model_loaded():
        raise HTTPException(status_code=500, detail="모델이 로드되지 않았습니다")

    start_time = time.time()

    # 이미지 데이터 처리
    image_data_list = []
    if image:
        try:
            data = await image.read()
            image_data_list.append(data)
        except Exception as e:
            logger.error(f"이미지 처리 실패: {e}")

    try:
        # generate_sync 함수 호출
        result = generate_sync(prompt, image_data_list, max_length=max_length, 
                             temperature=temperature, top_p=top_p, do_sample=do_sample)
        
        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])
        
        from ...services.model_service import get_current_profile
        current_profile = get_current_profile()
        
        return MultimodalGenerateResponse(
            generated_text=result["generated_text"],
            processing_time=result["processing_time"],
            model_name=current_profile.display_name,
            model_id=model_id or current_profile.get_model_info().get("model_name"),
            image_processed=bool(image_data_list)
        )
    except Exception as e:
        logger.error(f"❌ 멀티모달 생성 오류: {e}")
        raise HTTPException(status_code=500, detail=f"멀티모달 생성 실패: {str(e)}")