api.py

"""
AI선배 My Co-Pilot - FastAPI Backend
=====================================
RESTful API 서버

엔드포인트:
- POST /onboard: 사용자 온보딩 (연구 주제 설정)
- POST /inject_fake_history: 시연용 과거 데이터 주입
- POST /chat: LLM 대화 처리
- GET /stats: 시스템 통계

Author: AI선배 Team (공모전 MVP)
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
import uvicorn

# RAG 엔진 임포트
from rag_core import RAGEngine, inject_demo_data_2024

# ===========================================
# FastAPI 앱 초기화
# ===========================================
app = FastAPI(
    title="AI선배 My Co-Pilot API",
    description="대학원생의 영구 연구 기억 서비스 - 30년 뒤에도 기억하는 Private RAG",
    version="1.0.0",
    docs_url="/docs",  # Swagger UI
    redoc_url="/redoc"  # ReDoc
)

# CORS 설정 (Streamlit과의 통신 허용)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 개발용, 프로덕션에서는 제한 필요
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ===========================================
# 전역 RAG 엔진 인스턴스
# ===========================================
# Singleton 패턴: 앱 전체에서 하나의 엔진 공유
rag_engine: Optional[RAGEngine] = None


def get_rag_engine() -> RAGEngine:
    """
    RAG 엔진 싱글톤 접근자
    
    처음 호출 시 엔진을 초기화하고, 이후에는 기존 인스턴스를 반환합니다.
    """
    global rag_engine
    if rag_engine is None:
        rag_engine = RAGEngine()
    return rag_engine


# ===========================================
# Pydantic 모델 정의 (Request/Response)
# ===========================================
class OnboardRequest(BaseModel):
    """사용자 온보딩 요청"""
    topic: str = Field(..., description="연구 주제", example="Private LLM")
    keywords: List[str] = Field(..., description="관심 키워드", example=["Security", "Efficiency"])


class OnboardResponse(BaseModel):
    """온보딩 응답"""
    success: bool
    greeting: str
    message: str


class ChatRequest(BaseModel):
    """채팅 요청"""
    query: str = Field(..., description="사용자 질문", example="2024년에 했던 실험 결과가 뭐였지?")
    model_name: str = Field(
        default="gpt-4o-mini", 
        description="사용할 LLM 모델",
        example="claude-3.5-sonnet"
    )
    include_memories: bool = Field(
        default=True, 
        description="기억 검색 포함 여부"
    )


class MemoryReference(BaseModel):
    """참조된 기억 정보"""
    text: str
    date: str
    category: str
    score: float


class ChatResponse(BaseModel):
    """채팅 응답"""
    answer: str
    memories: List[MemoryReference]
    model: str
    context_used: bool


class InjectRequest(BaseModel):
    """기억 주입 요청"""
    text: str = Field(..., description="저장할 기억 텍스트")
    date: str = Field(..., description="기억 날짜", example="2024-03-15")
    category: str = Field(default="experiment", description="카테고리")


class InjectResponse(BaseModel):
    """기억 주입 응답"""
    success: bool
    message: str


class StatsResponse(BaseModel):
    """시스템 통계 응답"""
    total_memories: int
    collection_name: str
    status: str


# ===========================================
# API 엔드포인트
# ===========================================

@app.get("/", tags=["Health"])
async def root():
    """
    헬스체크 엔드포인트
    """
    return {
        "service": "AI선배 My Co-Pilot",
        "status": "running",
        "version": "1.0.0",
        "description": "대학원생의 영구 연구 기억 서비스"
    }


@app.post("/onboard", response_model=OnboardResponse, tags=["Onboarding"])
async def onboard_user(request: OnboardRequest):
    """
    [Scenario A] 입학 첫날 - 사용자 온보딩
    
    연구 주제와 관심 키워드를 입력받아 AI를 개인화합니다.
    시스템 프롬프트가 재설정되고, 맞춤형 첫인사가 반환됩니다.
    
    **시연 포인트:**
    - 연구 주제와 키워드가 AI에 각인됨
    - 개인화된 연구 파트너로 변신
    """
    try:
        engine = get_rag_engine()
        
        # 온보딩 수행
        greeting = engine.onboard_user(request.topic, request.keywords)
        
        return OnboardResponse(
            success=True,
            greeting=greeting,
            message=f"'{request.topic}' 주제로 온보딩 완료! 키워드: {', '.join(request.keywords)}"
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"온보딩 실패: {str(e)}")


@app.post("/inject_fake_history", response_model=InjectResponse, tags=["Demo"])
async def inject_fake_history():
    """
    [Scenario B] 시연용 과거 데이터 주입
    
    2024년의 연구 기록을 한번에 주입합니다.
    심사위원 시연에서 "3년 뒤 회상" 시나리오를 구현합니다.
    
    **포함되는 기록:**
    - 2024-03-15: Ablation Study 실험 결과
    - 2024-02-20: 프로젝트 킥오프 미팅
    - 2024-04-10: Quantization 실험
    - 2024-05-01: 논문 아이디어
    
    **시연 멘트:**
    "이 버튼을 누르면 1년 전의 연구 기록이 AI의 기억에 주입됩니다."
    """
    try:
        engine = get_rag_engine()
        
        # 데모 데이터 주입
        success = inject_demo_data_2024(engine)
        
        if success:
            return InjectResponse(
                success=True,
                message="✅ 2024년 데이터 주입 완료! 이제 과거 기억을 물어보세요."
            )
        else:
            return InjectResponse(
                success=False,
                message="⚠️ 일부 데이터 주입에 실패했습니다."
            )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"데이터 주입 실패: {str(e)}")


@app.post("/inject", response_model=InjectResponse, tags=["Memory"])
async def inject_single_memory(request: InjectRequest):
    """
    단일 기억 주입
    
    특정 텍스트를 특정 날짜의 기억으로 저장합니다.
    연구 노트, 실험 결과, 아이디어 등을 기록할 때 사용합니다.
    """
    try:
        engine = get_rag_engine()
        
        success = engine.inject_memory(request.text, request.date, request.category)
        
        if success:
            return InjectResponse(
                success=True,
                message=f"✅ 기억 저장 완료 (날짜: {request.date})"
            )
        else:
            return InjectResponse(
                success=False,
                message="❌ 기억 저장 실패"
            )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"기억 저장 실패: {str(e)}")


@app.post("/chat", response_model=ChatResponse, tags=["Chat"])
async def chat(request: ChatRequest):
    """
    [핵심 기능] RAG 기반 채팅
    
    사용자의 질문에 대해:
    1. 벡터 DB에서 관련 기억 검색
    2. Flashrank로 검색 결과 재정렬
    3. OpenRouter를 통해 선택한 LLM 모델로 응답 생성
    
    **지원 모델:**
    - claude-3.5-sonnet
    - gpt-4o, gpt-4o-mini
    - gemini-flash, gemini-pro
    - llama-3.1-70b
    
    **시연 질문 예시:**
    "2024년 3월에 했던 Ablation Study 결과가 뭐였지?"
    """
    try:
        engine = get_rag_engine()
        
        # RAG 응답 생성
        result = engine.get_answer(
            query=request.query,
            model_name=request.model_name,
            include_memories=request.include_memories
        )
        
        # 메모리 정보 변환
        memories = [
            MemoryReference(
                text=m["text"],
                date=m["date"],
                category=m["category"],
                score=m.get("score", 0.0)
            )
            for m in result["memories"]
        ]
        
        return ChatResponse(
            answer=result["answer"],
            memories=memories,
            model=result["model"],
            context_used=result["context_used"]
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"채팅 처리 실패: {str(e)}")


@app.get("/stats", response_model=StatsResponse, tags=["System"])
async def get_stats():
    """
    시스템 통계 조회
    
    저장된 기억의 수, 컬렉션 상태 등을 반환합니다.
    """
    try:
        engine = get_rag_engine()
        stats = engine.get_stats()
        
        return StatsResponse(
            total_memories=stats.get("total_memories", 0),
            collection_name=stats.get("collection_name", "unknown"),
            status=stats.get("status", "unknown")
        )
    
    except Exception as e:
        return StatsResponse(
            total_memories=0,
            collection_name="error",
            status=f"error: {str(e)}"
        )


@app.get("/models", tags=["System"])
async def get_available_models():
    """
    사용 가능한 LLM 모델 목록
    
    OpenRouter를 통해 사용할 수 있는 모델들을 반환합니다.
    """
    engine = get_rag_engine()
    return {
        "models": list(engine.SUPPORTED_MODELS.keys()),
        "default": "gpt-4o-mini",
        "provider": "OpenRouter"
    }


# ===========================================
# 서버 실행
# ===========================================
if __name__ == "__main__":
    print("=" * 50)
    print("AI선배 My Co-Pilot - FastAPI Server")
    print("=" * 50)
    print("Swagger UI: http://localhost:8000/docs")
    print("ReDoc: http://localhost:8000/redoc")
    print("=" * 50)
    
    uvicorn.run(
        "api:app",
        host="0.0.0.0",
        port=8000,
        reload=True  # 개발 모드: 코드 변경 시 자동 재시작
    )