Files
poc_system/test_chat.py

72 lines
2.4 KiB
Python

import logging
import sys
import os
# Để import được các module từ thư mục gốc
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from chat.rag_engine import RAGEngine
from core.config import settings
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger("TestChat")
def start_chat():
# Ép buộc dùng localhost cho OpenSearch khi chạy ngoài Docker
if settings.opensearch_host == "opensearch":
settings.opensearch_host = "localhost"
logger.info("=== HỆ THỐNG RAG SHAREPOINT ĐÃ SẴN SÀNG ===")
logger.info(f"Đang sử dụng bộ não: {settings.llm_provider.upper()}")
try:
engine = RAGEngine()
except Exception as e:
logger.error(f"Không thể khởi động RAG Engine: {e}")
return
print("\n" + "="*50)
print("MỜI BẠN NHẬP CÂU HỎI (Gõ 'exit' để thoát)")
print("="*50)
history = []
while True:
# Đọc trực tiếp dữ liệu thô (Bytes) từ bàn phím để tránh mọi lỗi mã hóa của Terminal
print("\nBạn: ", end='', flush=True)
try:
raw_line = sys.stdin.buffer.readline()
if not raw_line:
break
# Tự tay giải mã về UTF-8, bỏ qua các byte lỗi nếu có
query = raw_line.decode('utf-8', errors='ignore').strip()
except Exception as e:
logger.error(f"Lỗi khi đọc dữ liệu nhập: {e}")
break
if not query:
continue
if query.lower() in ['exit', 'quit']:
break
print("\nAI đang suy nghĩ...")
result = engine.chat(query, history=history)
print(f"\nAI: {result['answer']}")
print("\n--- NGUỒN TRÍCH DẪN (CITATIONS) ---")
for i, src in enumerate(result['sources']):
print(f"[{i+1}] {src['file_name']} (Trang {src['page']})")
# Lưu vào lịch sử chat để có ngữ cảnh cho câu hỏi tiếp theo
history.append({"role": "user", "content": query})
history.append({"role": "assistant", "content": result['answer']})
# Giữ lịch sử ngắn gọn (3 cặp câu hỏi - trả lời gần nhất)
if len(history) > 6:
history = history[-6:]
if __name__ == "__main__":
start_chat()