sqlbot_agent/main_service.py

from email.policy import default
import dmPython
import logging

from logging_config import LOGGING_CONFIG
from service.cus_vanna_srevice import CustomVanna, QdrantClient
from decouple import config
import flask
from util import load_ddl_doc

from flask import Flask, Response, jsonify, request, send_from_directory


logger = logging.getLogger(__name__)
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=''),
                            dbname=config('MYSQL_DATABASE_DBNAME', default=''))
    elif db_type == 'dameng':
        vn.connect_to_dameng( )
    elif db_type == 'postgresql':
        # 待补充
        pass
    else:
        pass


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


def create_vana():
    print("----------------create---------")
    vn = CustomVanna(
        vector_store_config={"client": QdrantClient(":memory:")},
        llm_config={
            "api_key": config('CHAT_MODEL_API_KEY', default=''),
            "api_base": config('CHAT_MODEL_BASE_URL', default=''),
            "model": config('CHAT_MODEL_NAME', default=''),
        },
    )
    return vn


def init_vn(vn):
    print("--------------init vn-----connect----")
    connect_database(vn)
    load_ddl_doc.add_ddl(vn)
    load_ddl_doc.add_documentation(vn)
    if config('IS_FIRST_LOAD', default=False, cast=bool):
        load_train_data_ddl(vn)
    return vn


from vanna.flask import VannaFlaskApp
vn = create_vana()
app = VannaFlaskApp(vn,chart=False)
init_vn(vn)
cache = app.cache
@app.flask_app.route("/api/v0/generate_sql_2", methods=["GET"])
def generate_sql_2():
    """
    Generate SQL from a question
    ---
    parameters:
      - name: user
        in: query
      - name: question
        in: query
        type: string
        required: true
    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"})
    try:
        id = cache.generate_id(question=question)
        data = vn.generate_sql_2(question=question)
        data['id'] = id
        sql = data["resp"]["sql"]
        print("sql:", sql)
        cache.set(id=id, field="question", value=question)
        cache.set(id=id, field="sql", value=sql)
        print("data---------------------------", data)
        return jsonify(data)
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})


@app.flask_app.route("/api/v0/run_sql_2", methods=["GET"])
@app.requires_cache(["sql"])
def run_sql_2(id: str, sql: str):
    """
    Run SQL
    ---
    parameters:
      - name: user
        in: query
      - 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:
        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(sql=sql)
        logger.info("")
        app.cache.set(id=id, field="df", value=df)
        x = df.head(10).to_dict(orient='records')
        logger.info("df ---------------{0}   {1}".format(x,type(x)))
        return jsonify(
            {
                "type": "df",
                "id": id,
                "df": df.head(10).to_dict(orient='records'),
            }
        )

    except Exception as e:
        return jsonify({"type": "sql_error", "error": str(e)})

if __name__ == '__main__':

    app.run(host='0.0.0.0', port=8084, debug=False)