Skip to content

Commit

Permalink
feat: add sql agent
Browse files Browse the repository at this point in the history
  • Loading branch information
sigoden committed Dec 13, 2024
1 parent 411a9e3 commit 07c79d3
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 1 deletion.
5 changes: 5 additions & 0 deletions agents/sql/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
# SQL

An AI agent that helps you manage a SQL database.

![image](https://github.com/user-attachments/assets/28bc1118-5f87-4571-a1c9-6c8cec4636d5)
13 changes: 13 additions & 0 deletions agents/sql/index.yaml
Original file line number Diff line number Diff line change
@@ -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
conversation_starters:
- What you can do?
42 changes: 42 additions & 0 deletions agents/sql/tools.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
#!/usr/bin/env bash

set -e

ROOT_DIR="${LLM_ROOT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"

# @meta require-tools usql

# @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" "$@")"
2 changes: 1 addition & 1 deletion tools/execute_sql_code.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
# @env LLM_OUTPUT=/dev/stdout The output path

ROOT_DIR="${LLM_ROOT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
Expand Down

0 comments on commit 07c79d3

Please sign in to comment.