sqlbot_agent/main_service.py

import copy
import json
import logging
import uuid
import time
from functools import wraps
import util.utils
from logging_config import LOGGING_CONFIG
from service.cus_vanna_srevice import CustomVanna, QdrantClient, TTLCacheWrapper
from service.question_feedback_service import save_save_question_async, query_predefined_question_list, \
    update_user_feedBack, query_feedBack_question_list
from service.conversation_service import (save_conversation, update_conversation,
                                          get_qa_by_id, get_latest_question,get_all_conversations_by_user)
from decouple import config
import flask
from util import load_ddl_doc
from flask import Flask, Response, jsonify, request
from graph_chat.gen_sql_chart_agent import SqlAgentState, sql_chart_agent
from graph_chat.gen_data_report_agent import result_report_agent, DateReportAgentState
import traceback

logger = logging.getLogger(__name__)
logging.getLogger('sqlalchemy.engine').setLevel(logging.DEBUG)

def generate_timestamp_id():
    """生成基于时间戳的ID"""
    # 获取当前时间戳（秒级）
    timestamp = int(time.time() * 1000)
    return f"Q{timestamp}"



def connect_database(vn):
    db_type = config('DATA_SOURCE_TYPE', default='sqlite')
    if db_type == 'sqlite':
        vn.connect_to_sqlite(config('SQLITE_DATABASE_URL', default=''))
    elif db_type == 'mysql':
        vn.connect_to_mysql(host=config('MYSQL_DATABASE_HOST', default=''),
                            port=int(config('MYSQL_DATABASE_PORT', default=3306)),
                            user=config('MYSQL_DATABASE_USER', default=''),
                            password=config('MYSQL_DATABASE_PASSWORD', default=''),
                            database=config('MYSQL_DATABASE_DBNAME', default=''))
    elif db_type == 'dameng':
        # 待补充
        vn.connect_to_dameng(
            host=config('DAMENG_DATABASE_HOST', default=''),
            port=config('DAMENG_DATABASE_PORT', default=3306),
            user=config('DAMENG_DATABASE_USER', default=''),
            password=config('DAMENG_DATABASE_PASSWORD', default=''),
        )
    else:
        pass


def load_train_data_ddl(vn: CustomVanna):
    vn.train()


def create_vana():
    logger.info("----------------create vana ---------")
    q_client = QdrantClient(":memory:") if config('QDRANT_TYPE', default='memory') == 'memory' else QdrantClient(
        url=config('QDRANT_DB_HOST', default=''), port=config('QDRANT_DB_PORT', default=6333))
    vn = CustomVanna(
        vector_store_config={"client": q_client},
        llm_config={
            "api_key": config('CHAT_MODEL_API_KEY', default=''),
            "api_base": config('CHAT_MODEL_BASE_URL', default=''),
            "model": config('CHAT_MODEL_NAME', default=''),
            'temperature': config('CHAT_MODEL_TEMPERATURE', default=0.7, cast=float),
            'max_tokens': config('CHAT_MODEL_MAX_TOKEN', default=5000),
        },
    )

    return vn


def init_vn(vn):
    logger.info("--------------init vana-----connect to datasouce db----")
    connect_database(vn)
    if config('IS_FIRST_LOAD', default=False, cast=bool):
        load_ddl_doc.add_ddl(vn)
        load_ddl_doc.add_documentation(vn)
        load_train_data_ddl(vn)
    return vn


from vanna.flask import VannaFlaskApp

vn = create_vana()
app = VannaFlaskApp(vn, chart=False)
app.cache = TTLCacheWrapper(app.cache, ttl=config('TTL_CACHE', cast=int, default=60 * 60))
init_vn(vn)
cache = app.cache


@app.flask_app.route("/yj_sqlbot/api/v0/generate_sql_2", methods=["GET"])
def generate_sql_2():
    """
    Generate SQL from a question
    ---
    parameters:
      - name: user_id
        in: query
      - name: question
        in: query
      - name: question_id
    responses:
      200:
        schema:
          type: object
          properties:
            type:
              type: string
              default: sql
            id:
              type: string
            text:
              type: string
    """
    logger.info("Start to generate sql in main")
    question = flask.request.args.get("question")
    if question is None:
        return jsonify({"type": "error", "error": "No question provided"})

    user_id = request.args.get("user_id")
    cvs_id = request.args.get("cvs_id")
    need_context = bool(request.args.get("need_context"))
    if user_id is None or cvs_id is None:
        return jsonify({"type": "error", "error": "No user_id or cvs_id provided"})
    id = generate_timestamp_id()
    logger.info(f"question_id: {id} user_id: {user_id} cvs_id: {cvs_id} question: {question}")
    save_conversation(id, user_id, cvs_id, question)
    try:
        logger.info(f"Generate sql for {question}")
        data = vn.generate_sql_2(user_id, cvs_id, question, id, need_context)
        logger.info("Generate sql result is {0}".format(data))
        data['id'] = id
        sql = data["resp"]["sql"]
        logger.info("generate sql is : " + sql)
        update_conversation(id, sql)
        save_save_question_async(id, user_id, question, sql)
        data["type"] = "success"
        return jsonify(data)
    except Exception as e:
        logger.error(f"generate sql failed:{e}")
        return jsonify({"type": "error", "error": str(e)})


# def requires_cache_2(required_keys):
#     def decorator(f):
#         @wraps(f)
#         def decorated(*args, **kwargs):
#             id = request.args.get("id")
#             user_id = request.args.get("user_id")
#             if user_id is None:
#                 user_id = request.json.get("user_id")
#                 if user_id is None:
#                     return jsonify({"type": "error", "error": "No user_id provided"})
#             if id is None:
#                 id = request.json.get("id")
#                 if id is None:
#                     return jsonify({"type": "error", "error": "No id provided"})
#             all_v = cache.items()
#             logger.info(f"all values {all_v}")
#             logger.info(f"user {user_id} id {id}")
#             qa_list = cache.get(id=user_id, field="qa_list")
#             if qa_list is None:
#                 return jsonify({"type": "error", "error": f"No qa_list found"})
#             logger.info(f"qa_list {qa_list}")
#             q_a = list(filter(lambda x: x["id"] == id, qa_list))
#             logger.info(f"q_a {q_a}")
#             for key in required_keys:
#                 if q_a[0][key] is None:
#                     return jsonify({"type": "error", "error": f"No {key} found for id:{id}"})
#             values = {key:q_a[0][key] for key in required_keys}
#             values["id"] = id
#             logger.info("cache values {0}".format(values))
#
#             return f(*args, **values, **kwargs)
#
#         return decorated
#
#     return decorator




# def session_save(func):
#     @wraps(func)
#     def wrapper(*args, **kwargs):
#         id = request.args.get("id")
#         user_id = request.args.get("user_id")
#         logger.info(f"   id: {id},user_id: {user_id}")
#         result = func(*args, **kwargs)
#
#         datas = []
#         session_len = int(config("SESSION_LENGTH", default=2))
#         if cache.exists(id=user_id, field="qa_list"):
#             datas = copy.deepcopy(cache.get(id=user_id, field="qa_list"))
#         logger.info("datas is {0}".format(datas))
#         if len(datas) > session_len and session_len > 0:
#             logger.info(f"开始裁剪-------------------------------------")
#             datas=datas[-session_len:]
#         # 删除id对应的所有缓存值,因为已经run_sql完毕，改用user_id保存为上下文
#         # cache.delete(id=id, field="question")
#         print("datas---------------------{0}".format(datas))
#         cache.set(id=user_id, field="qa_list", value=copy.deepcopy(datas))
#         logger.info(f" user data {cache.get(user_id, field='qa_list')}")
#         return result
#
#     return wrapper



@app.flask_app.route("/yj_sqlbot/api/v0/run_sql_2", methods=["GET"])
# @session_save
# @requires_cache_2(required_keys=["sql"])
def run_sql_2():
    """
    Run SQL
    ---
    parameters:
      - name: user_id
        in: query
        required: true
      - name: id
        in: query|body
        type: string
        required: true
    responses:
      200:
        schema:
          type: object
          properties:
            type:
              type: string
              default: df
            id:
              type: string
            df:
              type: object
            should_generate_chart:
              type: boolean
    """
    logger.info("Start to run sql in main")
    try:
        id = request.args.get("id")
        qa = get_qa_by_id(id)
        sql = qa["sql"]
        logger.info(f"sql is {sql}")
        if not vn.run_sql_is_set:
            return jsonify(
                {
                    "type": "error",
                    "error": "Please connect to a database using vn.connect_to_... in order to run SQL queries.",
                }
            )

        df = vn.run_sql_2(sql=sql)
        result = df.to_dict(orient='records')
        logger.info("df ---------------{0}   {1}".format(result, type(result)))
        return jsonify(
            {
                "type": "success",
                "id": id,
                "df": result,
            }
        )

    except Exception as e:
        logger.error(f"run sql failed:{e}")
        return jsonify({"type": "sql_error", "error": str(e)})


@app.flask_app.route("/yj_sqlbot/api/v0/verify", methods=["GET"])
def verify_user():
    try:
        id = request.args.get("user_id")
        users = config('ALLOWED_USERS', default='')
        users = users.split(',')
        logger.info(f"allowed users {users}")
        for user in users:
            if user == id:
                return jsonify({"type": "success", "verify": True})
            else:
                return jsonify({"type": "success", "verify": False})
    except Exception as e:
        logger.error(f"verify user failed:{e}")
        return jsonify({"type": "error", "error": str(e)})


@app.flask_app.route("/yj_sqlbot/api/v0/query_present_question", methods=["GET"])
def query_present_question():
    try:
        data = query_predefined_question_list()
        return jsonify({"type": "success", "data": data})
    except Exception as e:
        logger.error(f"查询预制问题失败 failed:{e}")
        return jsonify({"type": "error", "error": f'查询预制问题失败:{str(e)}'})


@app.flask_app.route("/yj_sqlbot/api/v0/query_feedback_question", methods=["POST"])
def query_feedback_question():
    id_list = request.json.get("id_list", [])
    try:
        data = query_feedBack_question_list(id_list)
        return jsonify({"type": "success", "data": data})
    except Exception as e:
        logger.error(f"查询用户反馈问题失败 failed:{e}")
        return jsonify({"type": "error", "error": f'查询用户反馈问题失败:{str(e)}'})


@app.flask_app.route("/yj_sqlbot/api/v0/question_feed_back", methods=["PUT"])
def update_question_feed_back():
    id = request.json.get("id")
    user_feedback = request.json.get("user_feedback")
    if not id or not user_feedback:
        return jsonify({"type": "error", "error": "id 或者用户反馈为空"})
    try:
        update_user_feedBack(id, '', user_feedback)
        return jsonify({"type": "success"})
    except Exception as e:
        logger.error(f"查询预制问题失败 failed:{e}")
        return jsonify({"type": "error", "error": f'查询预制问题失败:{str(e)}'})


@app.flask_app.route("/yj_sqlbot/api/v0/gen_graph_question", methods=["GET"])
def gen_graph_question():
    try:
        user_id = request.args.get("user_id")
        logger.info(f"Into gen_graph_question => user {user_id}")
        cvs_id = request.args.get("cvs_id")
        config = {"configurable": {"thread_id": str(uuid.uuid4())}}
        question = flask.request.args.get("question")
        logger.info(f"Start to get question context")
        question_context = get_latest_question(cvs_id, user_id, limit_count=2)
        history = []
        i = 0
        if not question_context:
            question_context = []
        for q in question_context:
            is_latest=False
            if i==0:
                is_latest=True
            history.append({"role": "user", 'content': q, 'order': i,'is_latest':is_latest})
            i +=1
        logger.info(f"question context is {history}")
        initial_state: SqlAgentState = {
            "user_id": user_id,
            "user_question": question,
            "history": history,
            "sql_retry_count": 0,
            "chart_retry_count": 0
        }
        id = str(uuid.uuid4())
        logger.info(f"Start to save conversation info (id,user_id,cvs_id,question)")
        save_conversation(id, user_id, cvs_id, question)
        logger.info(f"Enter the graph node=>gen_sql, gen_chart")
        result = sql_chart_agent.invoke(initial_state, config=config)
        logger.info(f"End====>gen_sql, gen_chart")
        new_question = result.get('rewritten_user_question', question)
        # save_conversation(id, user_id, cvs_id, new_question)
        if new_question:
            logger.info(f"new_question is {new_question}")
            update_conversation(id=id, meta={'new_question': new_question})
        logger.info("gen_sql_result => {0}".format(result.get("gen_sql_result", {})))
        data = {
            'id': id,
            'sql': result.get("gen_sql_result", {}),
            'chart': result.get("gen_chart_result", {}),
            'gen_sql_error': result.get("gen_sql_error", None),
            'gen_chart_error': result.get("gen_chart_error", None),
        }
        sql = data.get('sql', {}).get('sql', '')
        state = data.get('sql', {}).get('success', False)
        if not state:
            logger.info("SQL generation failed.save error info to table")
            error_msg = data.get('sql', {}).get('message', '')
            update_conversation(id=id, answer={'type_error':'error','error':error_msg})
        chart_cfg = data.get('chart', {})
        update_conversation(id=id, sql=sql, chart_cfg=chart_cfg)
        save_save_question_async(id, user_id, new_question, sql)
        return jsonify(data)
    except Exception as e:
        traceback.print_exc()
        logger.error(f"查询预制问题失败 failed:{e}")
        return jsonify({"type": "error", "error": f'查询预制问题失败:{str(e)}'})


@app.flask_app.route("/yj_sqlbot/api/v0/run_sql_3", methods=["GET"])
def run_sql_3():
    id = request.args.get("id")
    qa = get_qa_by_id(id)
    if not qa:
        return jsonify({"type": "error", "error": f'获取对话失败:{id}'})
    logger.info(f"Start to run_sql_3 => {qa}")
    sql = qa["sql"]
    if not sql:
        error_info = qa["answer"]
        if error_info:
            return jsonify({"type": "error", "error": json.loads(error_info).get("error","")})
        return jsonify({"type": "error", "error": f'sql 生成失败，请联系管理员'})
    question = qa["question"]
    logger.info(f"in main sql {sql} question {question}")
    logger.info("Start to run sql in main")
    try:
        user_id = request.args.get("user_id")
        cvs_id = request.args.get("cvs_id")
        if not vn.run_sql_is_set:
            return jsonify(
                {
                    "type": "error",
                    "error": "Please connect to a database using "
                             "vn.connect_to_... in order to run SQL queries.",
                }
            )
        question_context = get_latest_question(cvs_id, user_id, limit_count=2)
        history = []
        i = 0
        for q in question_context:
            is_latest=False
            if i==0:
                is_latest=True
            history.append({"role": "user", 'content': q,
                            'order': i,'is_latest':is_latest})
            i +=1
        initial_state: DateReportAgentState = {
            "id": id,
            "user_id": user_id,
            "sql": sql,
            "question": question,
            "retry_count": 0,
            "history": history,
        }
        config = {"configurable": {"thread_id": str(uuid.uuid4())}}
        rr = result_report_agent.invoke(initial_state, config)
        data = rr.get('data', {})
        summary = rr.get('summary', '')
        run_sql_error = rr.get('run_sql_error', '')
        logger.debug(f"data type  is {type(data)} data is {data} summary is {summary}")
        if data and not run_sql_error:
            update_conversation(id=id, answer={'data': data,'summary':summary})
        elif run_sql_error:
            update_conversation(id=id,  answer={'type_error':'error','error':"sql执行失败"})
        logger.info(f"run_Sql finish run_sql_error => {run_sql_error}")
        return jsonify(
            {
                'data': data,
                'summary': summary,
                'run_sql_error': "sql执行失败" if run_sql_error else None
            }
        )
    except Exception as e:
        logger.error(f"run sql failed:{e}")
        return jsonify({"type": "sql_error", "error": str(e)})


@app.flask_app.route("/yj_sqlbot/api/v0/get_history", methods=["GET"])
def get_history():
    logger.info(f"in main get history {request.args}")
    try:
        user_id = request.args.get("user_id")
        page_num = int(request.args.get("page_num"))
        page_size = int(request.args.get("page_size"))
        if page_num < 1:
            page_num = 1
        if page_size < 1 or page_size > 100:
            page_size = 10
        offset = (page_num-1)*page_size
        result = get_all_conversations_by_user(user_id, offset, page_size)
        data = result.get('data', {})
        total_count = result.get('total_count', 0)
        total_pages = (total_count + page_size - 1) // page_size
        logger.info(f"offset {offset} pagesize {page_size} total_count {total_count} total_pages {total_pages}")
        info = {
            "total_pages": total_pages,
            "total_count": total_count,
        }

        history = []
        ids = [d.id for d in data]
        feedback_questions = query_feedBack_question_list(ids)
        map_list = {q["id"]:q["user_praise"] for q in feedback_questions}
        logger.info(f"map_list {map_list}")
        for item in data:
            cvs = {}
            cvs['id'] = item.id
            cvs['question'] = item.question
            # answer_info = item.answer
            # info_type = answer_info.get('type','')
            # if info_type and info_type=="error":
            #     cvs[info_type] = answer_info.get("error","信息获取失败")
            cvs['answer'] = json.loads(item.answer) if item.answer else None
            cvs['chart_cfg'] = json.loads(item.chart_cfg) if item.chart_cfg else None
            cvs['user_praise'] = map_list[item.id]
            history.append(cvs)
        info["history"] = history
        return jsonify(info)
    except Exception as e:
        logger.error(f"get history failed:{e}")
        return jsonify({"type": "get history error", "error": str(e)})




if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8084, debug=False)