在线查询 支持AI根据描述生成查询语句 (#2726)

* support openai generate query sql

* fix lint

* use SysConfig

* fix

* support user update prompt

* use sysconfig default_query_template

* set openai default config

common/templates/config.html

@@ -932,6 +932,53 @@ <h4 style="color: darkgrey; display: inline;"><b>OIDC 配置</b></h4>
                                 </div>
                             </div>
                         </div>
+
+                        <h4 style="color: darkgrey; display: inline;"><b>OPENAI 配置</b></h4>
+                        <h6 style="color:red">注1：若无OPENAI_API_KEY配置，则不开启相关功能</h6>
+                        <h6 style="color:red">注2：DEFAULT_CHAT_MODEL 默认配置为gpt-3.5-turbo，DEFAULT_QUERY_TEMPLATE 默认配置为系统定义的模板</h6>
+                        <hr/>
+                        <div class="form-horizontal">
+                            <div class="form-group">
+                                <label for="openai_base_url"
+                                       class="col-sm-4 control-label">OPENAI_BASE_URL</label>
+                                <div class="col-sm-5">
+                                    <input type="text" class="form-control" id="openai_base_url"
+                                           key="openai_base_url"
+                                           value="{{ config.openai_base_url }}"
+                                           placeholder="openai base url" />
+                                </div>
+                            </div>
+                            <div class="form-group">
+                                <label for="openai_api_key"
+                                       class="col-sm-4 control-label">OPENAI_API_KEY</label>
+                                <div class="col-sm-5">
+                                    <input type="text" class="form-control" id="openai_api_key"
+                                           key="openai_api_key"
+                                           value="{{ config.openai_api_key }}"
+                                           placeholder="openai api key" />
+                                </div>
+                            </div>
+                            <div class="form-group">
+                                <label for="default_chat_model"
+                                       class="col-sm-4 control-label">DEFAULT_CHAT_MODEL</label>
+                                <div class="col-sm-5">
+                                    <input type="text" class="form-control" id="default_chat_model"
+                                           key="default_chat_model"
+                                           value="{{ config.default_chat_model }}"
+                                           placeholder="openai default chat model" />
+                                </div>
+                            </div>
+                            <div class="form-group">
+                                <label for="default_query_template"
+                                       class="col-sm-4 control-label">DEFAULT_QUERY_TEMPLATE</label>
+                                <div class="col-sm-5">
+                                    <input type="text" class="form-control" id="default_query_template"
+                                           key="default_query_template"
+                                           value="{{ config.default_query_template }}"
+                                           placeholder="默认生成SQL语句的Django模板, 无模板会导致生成结果失败, 需提供db_type/table_schema/user_input" />
+                                </div>
+                            </div>
+                        </div>
 
                         <h4 style="color: darkgrey"><b>其他配置</b></h4>
                         <hr/>

common/utils/openai.py

@@ -0,0 +1,49 @@
+from openai import OpenAI
+import logging
+from common.config import SysConfig
+from django.template import Context, Template
+
+logger = logging.getLogger("default")
+
+
+class OpenaiClient:
+    def __init__(self):
+        all_config = SysConfig()
+        self.base_url = all_config.get("openai_base_url", "")
+        self.api_key = all_config.get("openai_api_key", "")
+        self.default_chat_model = all_config.get("default_chat_model", "gpt-3.5-turbo")
+        self.default_query_template = all_config.get(
+            "default_query_template",
+            "你是一个熟悉 {{db_type}} 的工程师, 我会给你一些基本信息和要求, 你会生成一个查询语句给我使用, 不要返回任何注释和序号, 仅返回查询语句：{{table_schema}} \n {{user_input}}",
+        )
+        self.client = OpenAI(base_url=self.base_url, api_key=self.api_key)
+
+    def request_chat_completion(self, messages, **kwargs):
+        """chat_completion"""
+        completion = self.client.chat.completions.create(
+            model=self.default_chat_model, messages=messages, **kwargs
+        )
+        return completion
+
+    def generate_sql_by_openai(self, db_type: str, table_schema: str, user_input: str):
+        """根据传入的基本信息生成查询语句"""
+        template = Template(self.default_query_template)
+        current_context = Context(
+            dict(db_type=db_type, table_schema=table_schema, user_input=user_input)
+        )
+        messages = [dict(role="user", content=template.render(current_context))]
+        logger.info(messages)
+        try:
+            res = self.request_chat_completion(messages)
+            return res.choices[0].message.content
+        except Exception as e:
+            raise ValueError(f"请求openai生成查询语句失败: {e}")
+
+
+def check_openai_config():
+    """校验openai必需配置openai_api_key是否存在"""
+    all_config = SysConfig()
+    api_key = all_config.get("openai_api_key")
+    if api_key:
+        return True
+    return False

requirements.txt

@@ -42,3 +42,5 @@ mozilla-django-oidc==3.0.0
 django-auth-dingding==0.0.3
 django-cas-ng==4.3.0
 cassandra-driver
+httpx
+OpenAI

sql/query.py

@@ -12,6 +12,7 @@
 from django.http import HttpResponse
 from common.config import SysConfig
 from common.utils.extend_json_encoder import ExtendJSONEncoder, ExtendJSONEncoderFTime
+from common.utils.openai import OpenaiClient, check_openai_config
 from common.utils.timer import FuncTimer
 from sql.query_privileges import query_priv_check
 from sql.utils.resource_group import user_instances
@@ -313,3 +314,77 @@ def kill_query_conn(instance_id, thread_id):
     instance = Instance.objects.get(pk=instance_id)
     query_engine = get_engine(instance)
     query_engine.kill_connection(thread_id)
+
+
+@permission_required("sql.menu_sqlquery", raise_exception=True)
+def generate_sql(request):
+    """
+    利用AI生成查询SQL, 传入数据基本结构和查询描述
+    :param request:
+    :return:
+    """
+    query_desc = request.POST.get("query_desc")
+    db_type = request.POST.get("db_type")
+    if not query_desc or not db_type:
+        return HttpResponse(
+            json.dumps({"status": 1, "msg": "query_desc or db_type不存在", "data": []}),
+            content_type="application/json",
+        )
+
+    instance_name = request.POST.get("instance_name")
+    try:
+        instance = Instance.objects.get(instance_name=instance_name)
+    except Instance.DoesNotExist:
+        return HttpResponse(
+            json.dumps({"status": 1, "msg": "实例不存在", "data": []}),
+            content_type="application/json",
+        )
+    db_name = request.POST.get("db_name")
+    schema_name = request.POST.get("schema_name")
+    tb_name = request.POST.get("tb_name")
+
+    result = {"status": 0, "msg": "ok", "data": ""}
+    try:
+        query_engine = get_engine(instance=instance)
+        query_result = query_engine.describe_table(
+            db_name, tb_name, schema_name=schema_name
+        )
+        openai_client = OpenaiClient()
+        # 有些不存在表结构, 例如 redis
+        if len(query_result.rows) != 0:
+            result["data"] = openai_client.generate_sql_by_openai(
+                db_type, query_result.rows[0][-1], query_desc
+            )
+        else:
+            result["data"] = openai_client.generate_sql_by_openai(
+                db_type, "", query_desc
+            )
+    except Exception as msg:
+        result["status"] = 1
+        result["msg"] = str(msg)
+    return HttpResponse(json.dumps(result), content_type="application/json")
+
+
+def check_openai(request):
+    """
+    校验openai配置是否存在
+    :param request:
+    :return:
+    """
+    config_validate = check_openai_config()
+    if not config_validate:
+        return HttpResponse(
+            json.dumps(
+                {
+                    "status": 1,
+                    "msg": "openai 缺少配置, 必需配置[openai_base_url, openai_api_key, default_chat_model]",
+                    "data": False,
+                }
+            ),
+            content_type="application/json",
+        )
+
+    return HttpResponse(
+        json.dumps({"status": 0, "msg": "ok", "data": True}),
+        content_type="application/json",
+    )

sql/templates/sqlquery.html

@@ -57,6 +57,11 @@ <h4 class="modal-title text-danger">收藏语句</h4>
                             <option value={{ sql.id }}>{{ sql.alias }}</option>
                         {% endfor %}
                     </select>
+                    <input id="generateDesc" class="form-control" style="display: none" placeholder="AI 查询描述" />
+                    <input id="btn-generatesql" type="button" class="btn btn-info" style="display: none" value="生成SQL"/>
+                    <button type="button" class="btn"  data-toggle="tooltip" title="仅此操作会与 AI 交互, 收集数据库类型、表结构以及输入框信息, 交给 AI 生成 SQL 并置于下面的语句框中" >
+                        <span class="fa fa-question-circle" />
+                    </button>
                 </div>
                 <div class="panel-body">
                     <form id="form-sqlquery" action="/sqlquery/" method="post" class="form-horizontal" role="form">
@@ -495,6 +500,27 @@ <h4 class="modal-title text-danger">收藏语句</h4>
             }
             sessionStorage.removeItem('re_query');
         }
+
+        // 获取sysconfig
+        function check_openai() {
+            $.ajax({
+                type: "get",
+                url: "/check/openai/",
+                dataType: "json",
+                data: false,
+                complete: function () {
+                },
+                success: function (data) {
+                    if (data["data"]) {
+                        $("#generateDesc").show()
+                        $("#btn-generatesql").show()
+                    }
+                },
+                error: function (XMLHttpRequest, textStatus, errorThrown) {
+                    alert(errorThrown);
+                }
+            });
+        }
     </script>
     <!-- 执行结果  -->
     <script>
@@ -624,6 +650,32 @@ <h4 class="modal-title text-danger">收藏语句</h4>
             return result;
         }
 
+        //提交AI生成sql语句请求
+        $("#btn-generatesql").click(function () {
+                var check = false
+                var optgroup = $('#instance_name :selected').parent().attr('label')
+                var instance_name = $("#instance_name").val()
+                var db_name = $("#db_name").val()
+                var tb_name = $("#table_name").val()
+                var query_desc = $("#generateDesc").val()
+
+                if (!instance_name) {
+                    alert("请选择实例！")
+                } else if (!db_name) {
+                    alert("请选择数据库！")
+                } else if (optgroup !== 'Redis' && !tb_name){
+                    alert("请选择表结构！")
+                } else if (!query_desc) {
+                    alert("请输入查询描述！")
+                } else {
+                    check = true
+                }
+                if (check) {
+                    generatesql()
+                }
+            }
+        );
+
         //先做表单验证，验证成功再成功提交查询请求
         $("#btn-sqlquery").click(function () {
                 dosqlquery();
@@ -1023,6 +1075,37 @@ <h4 class="modal-title text-danger">收藏语句</h4>
             });
         }
 
+        function generatesql() {
+            var optgroup = $('#instance_name :selected').parent().attr('label');
+            const data = {
+                db_type: optgroup,
+                instance_name: $("#instance_name").val(),
+                db_name: $("#db_name").val(),
+                schema_name: $("#schema_name").val(),
+                tb_name: $("#table_name").val(),
+                query_desc: $("#generateDesc").val(),
+            }
+            //提交请求
+            $.ajax({
+                type: "post",
+                url: "/query/generate_sql/",
+                dataType: "json",
+                data: data,
+                complete: function () {
+                    $('input[type=button]').removeClass('disabled');
+                    $('input[type=button]').prop('disabled', false);
+                    optgroup_control();
+                },
+                success: function (data) {
+                    editor.setValue(data["data"]);
+                    editor.clearSelection();
+                },
+                error: function (XMLHttpRequest, textStatus, errorThrown) {
+                    alert(errorThrown);
+                }
+            });
+        }
+
         function dosqlquery() {
             if (sqlquery_validate()) {
                 $('input[type=button]').addClass('disabled');
@@ -1325,6 +1408,9 @@ <h4 class="modal-title text-danger">收藏语句</h4>
             } else {
                 editor.setValue("");
             }
+
+            // check openai 配置是否存在以支持AI生成查询语句功能
+            check_openai()
 
             //默认获取查询历史
             get_querylog();