enterprise-kb-rag/3.3 基于Mysql库的问答系统/integrated_qa_system/mysql_qa/main.py

# 在项目根目录运行：python -m mysql_qa.main
# 也支持直接运行：python mysql_qa/main.py
import sys
import time
from pathlib import Path

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# 导入 MySQL 客户端
from mysql_qa.db.mysql_client import MySQLClient
# 导入 Redis 客户端
from mysql_qa.cache.redis_client import RedisClient
# 导入 BM25 搜索
from mysql_qa.retrieval.bm25_search import BM25Search
# 导入日志
from base import logger


class MySQLQASystem:
    def __init__(self):
        self.logger = logger
        self.mysql_client = MySQLClient()
        self.redis_client = RedisClient()
        self.bm25_search = BM25Search(self.redis_client, self.mysql_client)

    def query(self, query):
        start_time = time.time()
        self.logger.info(f"处理查询: '{query}'")
        answer, _ = self.bm25_search.search(query, threshold=0.85)
        if answer:
            self.logger.info(f"MySQL 答案: {answer}")
        else:
            self.logger.info("SQL 中未找到答案，需要调用 RAG 系统")
            answer = "SQL未找到答案"
        processing_time = time.time() - start_time
        self.logger.info(f"查询处理耗时 {processing_time:.2f}秒")
        return answer


def main():
    mysql_system = MySQLQASystem()
    try:
        print("\n欢迎使用 MySQL 问答系统！")
        print("输入查询进行问答，输入 'exit' 退出。")
        while True:
            query = input("\n输入查询: ").strip()
            if query.lower() == "exit":
                logger.info("退出 MySQL 系统")
                print("再见！")
                break
            answer = mysql_system.query(query)
            print(f"\n答案: {answer}")
    except Exception as exc:
        logger.error(f"系统错误: {exc}")
        print(f"发生错误: {exc}")
    finally:
        mysql_system.mysql_client.close()


if __name__ == "__main__":
    main()