提示词优化，日志添加

main_service.py

@@ -1,5 +1,8 @@
 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
@@ -7,6 +10,8 @@ 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':
@@ -17,6 +22,8 @@ def connect_database(vn):
                             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
@@ -81,22 +88,78 @@ def generate_sql_2():
             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"})
-    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)
+    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)})
 
-    return jsonify(data)
 
+@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__':