feat: add sql agent (#152)

4 files changed, 64 insertions, 1 deletions

diff --git a/agents/sql/README.md b/agents/sql/README.md
new file mode 100644
index 0000000..ea0e63c
--- /dev/null
+++ b/agents/sql/README.md
@@ -0,0 +1,7 @@
+# SQL
+
+An AI agent that helps you manage a SQL database.
+
+> The tool script uses [usql](https://github.com/xo/usql) to interact with SQL, it supports all mainstream databases. 
+
+![image](https://github.com/user-attachments/assets/28bc1118-5f87-4571-a1c9-6c8cec4636d5)
diff --git a/agents/sql/index.yaml b/agents/sql/index.yaml
new file mode 100644
index 0000000..4659fc9
--- /dev/null
+++ b/agents/sql/index.yaml
@@ -0,0 +1,13 @@
+name: Sql
+description: An AI agent that helps you manage a SQL database
+version: 0.1.0
+instructions: |
+  You are an AI agent that manages a SQL database.
+
+  Available tools:
+  {{__tools__}}
+variables:
+  - name: dsn
+    description: The database connection url. e.g. pgsql://user:pass@host:port
+conversation_starters:
+  - What you can do?
\ No newline at end of file
diff --git a/agents/sql/tools.sh b/agents/sql/tools.sh
new file mode 100755
index 0000000..dc3062a
--- /dev/null
+++ b/agents/sql/tools.sh
@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+set -e
+
+ROOT_DIR="${LLM_ROOT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
+
+# @meta require-tools usql
+# @env LLM_AGENT_VAR_DSN! The database connection url. e.g. pgsql://user:pass@host:port 
+
+# @cmd Execute a SELECT query
+# @option --query!                  SELECT SQL query to execute
+read_query() {
+    if ! grep -qi '^select' <<<"$argc_query"; then
+        echo "error: only SELECT query is allowed" >&2
+        exit 1
+    fi
+    _run_sql "$argc_query"
+}
+
+# @cmd Execute an SQL query
+# @option --query!                  SQL query to execute
+write_query() {
+    "$ROOT_DIR/utils/guard_operation.sh" "Execute SQL?"
+    _run_sql "$argc_query"
+}
+
+# @cmd List all tables
+list_tables() {
+    _run_sql "\dt+"
+}
+
+# @cmd Get the schema information for a specific table
+# @option --table-name!             Name of the table to describe
+describe_table() {
+    _run_sql "\d $argc_table_name"
+}
+
+_run_sql() {
+    usql "$LLM_AGENT_VAR_DSN" -c "$1" >> "$LLM_OUTPUT"
+}
+
+# See more details at https://github.com/sigoden/argc
+eval "$(argc --argc-eval "$0" "$@")"
diff --git a/tools/execute_sql_code.sh b/tools/execute_sql_code.sh
index 34cd948..eb27ce2 100755
--- a/tools/execute_sql_code.sh
+++ b/tools/execute_sql_code.sh
@@ -6,7 +6,7 @@ set -e
 
 # @meta require-tools usql
 
-# @env USQL_DSN! The database url, e.g. pgsql://user:pass@host/dbname
+# @env USQL_DSN! The database connection url. e.g. pgsql://user:pass@host:port
 # @env LLM_OUTPUT=/dev/stdout The output path
 
 ROOT_DIR="${LLM_ROOT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"