scripts/test_query.py

"""
RAG 시스템 테스트 스크립트

API 서버가 실행 중일 때 사용
"""

import requests
import json
from typing import Dict


def test_query(
    question: str,
    top_k: int = 5,
    enable_metacognition: bool = True,
    api_url: str = "http://localhost:8000"
) -> Dict:
    """
    질문 테스트

    Args:
        question: 질문
        top_k: 검색할 문서 개수
        enable_metacognition: 메타인지 활성화
        api_url: API URL

    Returns:
        응답 데이터
    """
    print("=" * 80)
    print(f"질문: {question}")
    print("=" * 80)

    # 요청
    response = requests.post(
        f"{api_url}/query",
        json={
            "question": question,
            "top_k": top_k,
            "enable_metacognition": enable_metacognition
        }
    )

    if response.status_code != 200:
        print(f"❌ 오류: {response.status_code}")
        print(response.text)
        return {}

    result = response.json()

    # 결과 출력
    print("\n📝 답변:")
    print("-" * 80)
    print(result["answer"])
    print("-" * 80)

    print(f"\n📚 참고 문헌: {len(result['sources'])}개")
    for i, source in enumerate(result['sources'][:3], 1):
        print(f"\n[{i}] {source['source_filename']}")
        print(f"    유사도: {source['similarity']:.3f}")
        print(f"    내용: {source['text'][:100]}...")

    if result.get('metacognition'):
        print(f"\n🧠 메타인지 정보:")
        print(f"    반복 횟수: {result['metacognition']['iterations']}")
        print(f"    사고 과정 단계: {len(result['metacognition']['thinking_history'])}")

    print("\n" + "=" * 80)
    return result


def test_health(api_url: str = "http://localhost:8000"):
    """헬스 체크"""
    print("🏥 헬스 체크 중...")
    response = requests.get(f"{api_url}/health")

    if response.status_code == 200:
        data = response.json()
        print("✅ 서버 정상")
        print(f"   Vector Store: {data['vector_store']['total_documents']}개 문서")
        print(f"   Embedding: {data['embedding_model']['type']} ({data['embedding_model']['dimension']}차원)")
    else:
        print(f"❌ 서버 오류: {response.status_code}")


if __name__ == "__main__":
    # 헬스 체크
    test_health()

    print("\n")

    # 샘플 질문들
    questions = [
        "금융위기의 주요 원인은 무엇인가요?",
        "포트폴리오 다각화의 효과는?",
        "중앙은행의 금리 정책이 시장에 미치는 영향은?",
    ]

    for question in questions:
        try:
            test_query(question, top_k=5, enable_metacognition=True)
            print("\n\n")
        except Exception as e:
            print(f"❌ 오류: {str(e)}\n\n")

    # 사용자 입력 질문
    print("\n커스텀 질문을 입력하세요 (Enter를 누르면 종료):")
    while True:
        question = input("\n질문: ").strip()
        if not question:
            break

        try:
            test_query(question, top_k=5, enable_metacognition=True)
        except Exception as e:
            print(f"❌ 오류: {str(e)}")