Skip to content

Commit 34fd092

Browse files
author
junjie.miao
committed
add output format supported markdown and json
1 parent c41a628 commit 34fd092

File tree

5 files changed

+68
-16
lines changed

5 files changed

+68
-16
lines changed

‎db_query/tools/sql_query.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ def _invoke(
4545
if statement.get_type() != 'SELECT':
4646
raise ValueError("Query SQL can only be a single SELECT statement")
4747

48+
output_format = tool_parameters.get("output_format", "markdown").lower()
49+
4850
try:
4951
db = DbUtil(db_type=db_type,
5052
username=db_username, password=db_password,
@@ -57,9 +59,13 @@ def _invoke(
5759

5860
try:
5961
records = db.run_query(query_sql)
60-
text = tabulate.tabulate(records, headers="keys", tablefmt="github")
61-
yield self.create_text_message(text)
6262
except Exception as e:
6363
message = "SQL query execution exception."
6464
logging.exception(message)
6565
raise Exception(message + " {}".format(e))
66+
67+
if output_format == "json":
68+
yield self.create_json_message({"records": records})
69+
else:
70+
text = tabulate.tabulate(records, headers="keys", tablefmt="github")
71+
yield self.create_text_message(text)

‎db_query/tools/sql_query.yaml‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ parameters:
4949
en_US: Database Host
5050
zh_Hans: 数据库地址
5151
human_description:
52-
en_US: database hostname or IP address (Original string, not URL-encoded string).
52+
en_US: Database hostname or IP address (Original string, not URL-encoded string).
5353
zh_Hans: 数据库的主机名或IP地址（原始字符串，非URL编码字符串）。
5454
form: llm
5555
- name: db_port
@@ -59,7 +59,7 @@ parameters:
5959
en_US: Port
6060
zh_Hans: 端口
6161
human_description:
62-
en_US: database port.
62+
en_US: Database port.
6363
zh_Hans: 数据库的端口。
6464
form: llm
6565
- name: db_username
@@ -69,7 +69,7 @@ parameters:
6969
en_US: Username
7070
zh_Hans: 用户名
7171
human_description:
72-
en_US: database username (Original string, not URL-encoded string).
72+
en_US: Database username (Original string, not URL-encoded string).
7373
zh_Hans: 数据库的用户名（原始字符串，非URL编码字符串）。
7474
form: llm
7575
- name: db_password
@@ -79,7 +79,7 @@ parameters:
7979
en_US: Password
8080
zh_Hans: 密码
8181
human_description:
82-
en_US: database password (Original string, not URL-encoded string).
82+
en_US: Database password (Original string, not URL-encoded string).
8383
zh_Hans: 数据库的密码（原始字符串，非URL编码字符串）。
8484
form: llm
8585
- name: db_name
@@ -89,7 +89,7 @@ parameters:
8989
en_US: Database name
9090
zh_Hans: 库名
9191
human_description:
92-
en_US: database name (Original string, not URL-encoded string).
92+
en_US: Database name (Original string, not URL-encoded string).
9393
zh_Hans: 数据库的名称（原始字符串，非URL编码字符串）。
9494
form: llm
9595
- name: db_properties
@@ -112,4 +112,24 @@ parameters:
112112
en_US: 'SQL query statement, for example: select * from tbl_name'
113113
zh_Hans: SQL查询语句，例如：select * from tbl_name
114114
llm_description: 'SQL query statement, for example: select * from tbl_name'
115-
form: llm
115+
form: llm
116+
- name: output_format
117+
type: select
118+
required: true
119+
label:
120+
en_US: Output format
121+
zh_Hans: 输出格式
122+
human_description:
123+
en_US: Used for selecting the output format, markdown or json.
124+
zh_Hans: 用于选择输出格式，markdown或json。
125+
options:
126+
- value: markdown
127+
label:
128+
en_US: MARKDOWN
129+
zh_Hans: MARKDOWN
130+
- value: json
131+
label:
132+
en_US: JSON
133+
zh_Hans: JSON
134+
default: markdown
135+
form: form

‎db_query_pre_auth/provider/db_query.yaml‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ credentials_for_provider:
4848
en_US: Database Host
4949
zh_Hans: 数据库地址
5050
human_description:
51-
en_US: database hostname or IP address (Original string, not URL-encoded string).
51+
en_US: Database hostname or IP address (Original string, not URL-encoded string).
5252
zh_Hans: 数据库的主机名或IP地址（原始字符串，非URL编码字符串）。
5353
db_port:
5454
type: text-input
@@ -57,7 +57,7 @@ credentials_for_provider:
5757
en_US: Port
5858
zh_Hans: 端口
5959
human_description:
60-
en_US: database port (Original string, not URL-encoded string).
60+
en_US: Database port (Original string, not URL-encoded string).
6161
zh_Hans: 数据库的端口（原始字符串，非URL编码字符串）。
6262
db_username:
6363
type: text-input
@@ -66,7 +66,7 @@ credentials_for_provider:
6666
en_US: Username
6767
zh_Hans: 用户名
6868
human_description:
69-
en_US: database username (Original string, not URL-encoded string).
69+
en_US: Database username (Original string, not URL-encoded string).
7070
zh_Hans: 数据库的用户名（原始字符串，非URL编码字符串）。
7171
db_password:
7272
type: secret-input
@@ -75,7 +75,7 @@ credentials_for_provider:
7575
en_US: Password
7676
zh_Hans: 密码
7777
human_description:
78-
en_US: database password (Original string, not URL-encoded string).
78+
en_US: Database password (Original string, not URL-encoded string).
7979
zh_Hans: 数据库的密码（原始字符串，非URL编码字符串）。
8080
db_name:
8181
type: text-input
@@ -84,7 +84,7 @@ credentials_for_provider:
8484
en_US: Database name
8585
zh_Hans: 库名
8686
human_description:
87-
en_US: database name.
87+
en_US: Database name.
8888
zh_Hans: 数据库的名称。
8989
db_properties:
9090
type: text-input

‎db_query_pre_auth/tools/sql_query.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ def _invoke(
4444
if statement.get_type() != 'SELECT':
4545
raise ValueError("Query SQL can only be a single SELECT statement")
4646

47+
output_format = tool_parameters.get("output_format", "markdown").lower()
48+
4749
try:
4850
db = DbUtil(db_type=db_type,
4951
username=db_username, password=db_password,
@@ -56,9 +58,13 @@ def _invoke(
5658

5759
try:
5860
records = db.run_query(query_sql)
59-
text = tabulate.tabulate(records, headers="keys", tablefmt="github")
60-
yield self.create_text_message(text)
6161
except Exception as e:
6262
message = "SQL query execution exception."
6363
logging.exception(message)
6464
raise Exception(message + " {}".format(e))
65+
66+
if output_format == "json":
67+
yield self.create_json_message({"records": records})
68+
else:
69+
text = tabulate.tabulate(records, headers="keys", tablefmt="github")
70+
yield self.create_text_message(text)

‎db_query_pre_auth/tools/sql_query.yaml‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,24 @@ parameters:
2323
en_US: 'SQL query statement, for example: select * from tbl_name'
2424
zh_Hans: SQL查询语句，例如：select * from tbl_name
2525
llm_description: 'SQL query statement, for example: select * from tbl_name'
26-
form: llm
26+
form: llm
27+
- name: output_format
28+
type: select
29+
required: true
30+
label:
31+
en_US: Output format
32+
zh_Hans: 输出格式
33+
human_description:
34+
en_US: Used for selecting the output format, markdown or json.
35+
zh_Hans: 用于选择输出格式，markdown或json。
36+
options:
37+
- value: markdown
38+
label:
39+
en_US: MARKDOWN
40+
zh_Hans: MARKDOWN
41+
- value: json
42+
label:
43+
en_US: JSON
44+
zh_Hans: JSON
45+
default: markdown
46+
form: form

0 commit comments

Comments
 (0)