clovax-tax-chatbot / simple_test.py
bissal's picture
large modification
3a15ede
# simple_test.py - ๊ฐ„๋‹จํ•œ RAG ํ…Œ์ŠคํŠธ ์Šคํฌ๋ฆฝํŠธ
import sys
import os
# ์ธ์ฝ”๋”ฉ ์„ค์ •
if sys.platform.startswith('win'):
os.system('chcp 65001 > nul')
def test_rag_system():
"""RAG ์‹œ์Šคํ…œ ๊ฐ„๋‹จ ํ…Œ์ŠคํŠธ"""
print("RAG ์‹œ์Šคํ…œ ํ…Œ์ŠคํŠธ ์‹œ์ž‘")
print("=" * 50)
try:
from rag_system import HFSpacesTaxRAG
# RAG ์‹œ์Šคํ…œ ์ดˆ๊ธฐํ™”
print("RAG ์‹œ์Šคํ…œ ์ดˆ๊ธฐํ™” ์ค‘...")
rag = HFSpacesTaxRAG()
if not rag.vector_db or not rag.documents:
print("๋ฒกํ„ฐ DB๊ฐ€ ์—†์Šต๋‹ˆ๋‹ค. ์‹œ์Šคํ…œ์„ ์žฌ๊ตฌ์ถ•ํ•ฉ๋‹ˆ๋‹ค...")
success = rag.build_system()
if not success:
print("RAG ์‹œ์Šคํ…œ ๊ตฌ์ถ• ์‹คํŒจ")
return False
# ๊ฐ„๋‹จํ•œ ํ…Œ์ŠคํŠธ ์ฟผ๋ฆฌ๋“ค
test_queries = [
"์ทจ๋“์„ธ์œจ์ด ์–ผ๋งˆ์ธ๊ฐ€์š”?",
"1์„ธ๋Œ€ 1์ฃผํƒ์ž ๊ฐ๋ฉด ํ˜œํƒ์€?",
"์ทจ๋“์„ธ ์‹ ๊ณ  ๊ธฐํ•œ์€ ์–ธ์ œ์ธ๊ฐ€์š”?"
]
print("\nํ…Œ์ŠคํŠธ ์ฟผ๋ฆฌ ์‹คํ–‰:")
print("-" * 30)
for i, query in enumerate(test_queries, 1):
print(f"\n{i}. {query}")
try:
result = rag.query(query)
if result and result.get('answer'):
print(f"์‘๋‹ต: {result['answer'][:100]}...")
print(f"์†Œ์Šค ๋ฌธ์„œ: {result.get('source_count', 0)}๊ฐœ")
else:
print("์‘๋‹ต ์ƒ์„ฑ ์‹คํŒจ")
except Exception as e:
print(f"์ฟผ๋ฆฌ ์ฒ˜๋ฆฌ ์˜ค๋ฅ˜: {e}")
print("\n" + "=" * 50)
print("ํ…Œ์ŠคํŠธ ์™„๋ฃŒ")
return True
except Exception as e:
print(f"RAG ์‹œ์Šคํ…œ ์˜ค๋ฅ˜: {e}")
return False
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] == '--rebuild':
print("RAG ์‹œ์Šคํ…œ์„ ์žฌ๊ตฌ์ถ•ํ•ฉ๋‹ˆ๋‹ค...")
# ๊ธฐ์กด ๋ฒกํ„ฐ DB ํŒŒ์ผ ์‚ญ์ œ
for file in ['vector_db.faiss', 'documents.pkl', 'metadata.json']:
if os.path.exists(file):
os.remove(file)
print(f"๊ธฐ์กด ํŒŒ์ผ ์‚ญ์ œ: {file}")
success = test_rag_system()
sys.exit(0 if success else 1)