app.py

"""
HuggingFace Spaces용 Gradio 앱
Financial RAG with Metacognitive Agent
"""

import gradio as gr
import os
import sys
from loguru import logger
import asyncio
from typing import Dict, Tuple

# 로깅 설정
logger.remove()
logger.add(
    sys.stdout,
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
    level="INFO"
)

# 프로젝트 루트를 Python 경로에 추가
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app.metacognitive_agent import MetaCognitiveAgent
from app.rag_pipeline import RAGPipeline
from services.vector_store import VectorStore
from services.embedder import Embedder
from utils.config import settings

# 글로벌 변수
rag_pipeline = None


def setup_vector_db():
    """벡터 DB 자동 설정 (없으면 다운로드 또는 생성)"""
    db_path = settings.chroma_persist_directory

    # 벡터 DB가 이미 존재하고 비어있지 않은지 확인
    if os.path.exists(db_path):
        if os.listdir(db_path):
            logger.info("✅ Vector DB already exists. Skipping setup.")
            return True

    logger.info("📥 Vector DB not found. Setting up...")
    os.makedirs(db_path, exist_ok=True)

    # 옵션 1: GitHub Release에서 다운로드 시도
    try:
        import urllib.request
        import tarfile

        release_url = "https://github.com/csjjin2025/Hallucination_and_Deception_for_financial_RAG/releases/download/v1.0/chroma_db.tar.gz"
        tar_path = "./data/chroma_db.tar.gz"

        logger.info(f"Attempting to download from {release_url}...")
        urllib.request.urlretrieve(release_url, tar_path)

        # 파일 크기 확인
        file_size = os.path.getsize(tar_path)
        if file_size > 1000:
            logger.info(f"📦 Extracting vector DB ({file_size} bytes)...")
            with tarfile.open(tar_path, 'r:gz') as tar:
                tar.extractall(path='./data/')
            os.remove(tar_path)
            logger.info("✅ Vector DB downloaded and extracted!")
            return True
        else:
            logger.warning(f"Downloaded file too small ({file_size} bytes)")
            os.remove(tar_path)
    except Exception as e:
        logger.warning(f"Failed to download from Release: {e}")

    # 옵션 2: 테스트 DB 생성
    try:
        logger.info("⚠️ Creating test DB with sample data...")
        import subprocess
        result = subprocess.run(
            ["python", "scripts/quick_setup_test_db.py"],
            capture_output=True,
            text=True,
            timeout=300
        )
        if result.returncode == 0:
            logger.info("✅ Test DB created successfully!")
            return True
        else:
            logger.error(f"Test DB creation failed: {result.stderr}")
            return False
    except Exception as e:
        logger.error(f"Failed to create test DB: {e}")
        return False


def initialize_rag_system():
    """RAG 시스템 초기화"""
    global rag_pipeline

    try:
        logger.info("=" * 80)
        logger.info("🚀 Financial RAG 시스템 초기화 중...")
        logger.info("=" * 80)

        # 0. Vector DB 자동 설정
        logger.info("0️⃣ Vector DB 설정 확인 중...")
        if not setup_vector_db():
            logger.error("❌ Vector DB 설정 실패")
            return False

        # 1. Vector Store 초기화
        logger.info("1️⃣ Vector Store 로딩 중...")
        vector_store = VectorStore(
            persist_directory=settings.chroma_persist_directory,
            collection_name=settings.collection_name
        )
        doc_count = vector_store.collection.count()
        logger.info(f"✅ Vector Store 로딩 완료 ({doc_count}개 문서)")

        # 2. Embedder 초기화
        logger.info("2️⃣ Embedder 초기화 중...")
        embedder = Embedder(
            model_type=settings.embedding_model,
            model_name=settings.embedding_model_name,
            openai_api_key=settings.openai_api_key,
            cohere_api_key=settings.cohere_api_key
        )
        logger.info(f"✅ Embedder 초기화 완료 ({embedder.get_embedding_dimension()}차원)")

        # 3. Metacognitive Agent 초기화
        logger.info("3️⃣ Metacognitive Agent 초기화 중...")
        agent = MetaCognitiveAgent(api_key=settings.anthropic_api_key)
        logger.info(f"✅ Agent 초기화 완료 ({agent.model})")

        # 4. RAG Pipeline 생성
        logger.info("4️⃣ RAG Pipeline 생성 중...")
        rag_pipeline = RAGPipeline(
            vector_store=vector_store,
            embedder=embedder,
            metacognitive_agent=agent
        )
        logger.info("✅ RAG Pipeline 생성 완료")

        logger.info("=" * 80)
        logger.info("✨ 시스템 준비 완료!")
        logger.info(f"📚 Vector DB: {doc_count}개 문서")
        logger.info(f"🤖 Model: {agent.model}")
        logger.info("=" * 80)

        return True

    except Exception as e:
        logger.error(f"❌ 초기화 실패: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False


def format_sources(sources: list) -> str:
    """출처 문서 포맷팅"""
    if not sources:
        return "출처 문서가 없습니다."

    formatted = "### 📚 참고 문서\n\n"
    for idx, source in enumerate(sources[:3], 1):  # 상위 3개만 표시
        similarity = source.get('similarity', 0) * 100
        filename = source.get('source_filename', 'unknown')
        text = source.get('text', '')[:300]  # 앞 300자만

        formatted += f"**{idx}. {filename}** (유사도: {similarity:.1f}%)\n"
        formatted += f"> {text}...\n\n"

    return formatted


def format_metacognition(metacognition: Dict) -> str:
    """메타인지 과정 포맷팅"""
    if not metacognition:
        return ""

    history = metacognition.get('thinking_history', [])
    iterations = metacognition.get('iterations', 0)

    formatted = f"\n\n### 🧠 메타인지 과정 ({iterations}회 반복)\n\n"

    for idx, step in enumerate(history, 1):
        stage = step.get('stage', 'unknown')
        content = step.get('content', '')

        stage_emoji = {
            'planning': '📋',
            'monitoring': '👁️',
            'evaluation': '✅',
            'revision': '🔄'
        }.get(stage, '💭')

        formatted += f"**{stage_emoji} {stage.capitalize()}**\n"
        formatted += f"{content}\n\n"

    return formatted


async def process_query_async(question: str, top_k: int, enable_metacognition: bool) -> Tuple[str, str]:
    """비동기 쿼리 처리"""
    if not rag_pipeline:
        return "❌ 시스템이 초기화되지 않았습니다.", ""

    if not question.strip():
        return "❌ 질문을 입력해주세요.", ""

    try:
        logger.info(f"📝 질문: {question}")

        # RAG 파이프라인으로 쿼리 처리
        result = await rag_pipeline.query(
            question=question,
            top_k=top_k,
            enable_metacognition=enable_metacognition
        )

        # 답변 포맷팅
        answer = result.get('answer', '답변을 생성할 수 없습니다.')
        sources = result.get('sources', [])
        metacognition = result.get('metacognition', None)

        # 출력 구성
        main_output = f"## 💬 답변\n\n{answer}\n\n"
        main_output += format_sources(sources)

        # 메타인지 과정 (별도 탭)
        meta_output = format_metacognition(metacognition) if metacognition else "메타인지가 비활성화되었습니다."

        logger.info("✅ 답변 생성 완료")
        return main_output, meta_output

    except Exception as e:
        error_msg = f"❌ 오류 발생: {str(e)}"
        logger.error(error_msg)
        import traceback
        logger.error(traceback.format_exc())
        return error_msg, ""


def process_query(question: str, top_k: int, enable_metacognition: bool) -> Tuple[str, str]:
    """Gradio용 동기 래퍼"""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(process_query_async(question, top_k, enable_metacognition))
    finally:
        loop.close()


# Gradio 인터페이스 구성
def create_interface():
    """Gradio 인터페이스 생성"""

    with gr.Blocks(theme=gr.themes.Soft(), title="Financial RAG with Metacognitive Agent") as demo:
        gr.Markdown("""
        # 💼 Financial RAG System
        ### 메타인지 에이전트 기반 금융/경제 질의응답 시스템

        이 시스템은 금융/경제 논문을 기반으로 질문에 답변합니다.
        메타인지 기능을 활성화하면 더 깊이 있는 사고 과정을 거쳐 답변을 생성합니다.
        """)

        with gr.Row():
            with gr.Column(scale=2):
                question_input = gr.Textbox(
                    label="💬 질문을 입력하세요",
                    placeholder="예: 포트폴리오 다각화는 무엇인가요?",
                    lines=3
                )

                with gr.Row():
                    top_k_slider = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="🔍 검색할 문서 개수"
                    )

                    metacognition_check = gr.Checkbox(
                        label="🧠 메타인지 활성화",
                        value=True,
                        info="더 깊이 있는 사고 과정 (처리 시간 증가)"
                    )

                submit_btn = gr.Button("🚀 질문하기", variant="primary", size="lg")

                gr.Markdown("""
                ### 💡 사용 팁
                - **메타인지 활성화**: Planning → Monitoring → Evaluation → Revision 과정을 거쳐 신중한 답변 생성
                - **메타인지 비활성화**: 빠른 답변 생성
                - **검색 문서 개수**: 많을수록 더 많은 정보를 참고하지만 처리 시간 증가
                """)

            with gr.Column(scale=3):
                with gr.Tabs():
                    with gr.Tab("📝 답변 및 출처"):
                        answer_output = gr.Markdown(label="답변")

                    with gr.Tab("🧠 메타인지 과정"):
                        metacognition_output = gr.Markdown(label="사고 과정")

        # 이벤트 핸들러
        submit_btn.click(
            fn=process_query,
            inputs=[question_input, top_k_slider, metacognition_check],
            outputs=[answer_output, metacognition_output]
        )

        # Enter 키로도 제출
        question_input.submit(
            fn=process_query,
            inputs=[question_input, top_k_slider, metacognition_check],
            outputs=[answer_output, metacognition_output]
        )

        gr.Markdown("""
        ---
        ### 📌 시스템 정보
        - **모델**: Claude 3.5 Sonnet
        - **임베딩**: sentence-transformers/all-MiniLM-L6-v2
        - **벡터 DB**: ChromaDB
        """)

    return demo


# 메인 실행
if __name__ == "__main__":
    # 시스템 초기화
    logger.info("시스템 초기화 시작...")
    success = initialize_rag_system()

    if not success:
        logger.error("시스템 초기화 실패. 종료합니다.")
        sys.exit(1)

    # Gradio 앱 실행
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )