up

agents/SQL_query/sql_query.go

@@ -0,0 +1,763 @@
+package sqlquery
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+	"unicode/utf8"
+
+	_ "github.com/go-sql-driver/mysql"
+	"gopkg.in/yaml.v3"
+	_ "modernc.org/sqlite"
+)
+
+const (
+	defaultActivationPrompt = `判断用户问题是否需要查询业务数据库。
+仅当用户询问数据库表、记录、字段、时间、状态、内容、统计、最近/最早/某时间范围内的数据时返回 activate=true。
+普通知识问答、代码问题、闲聊、联网搜索问题返回 activate=false。
+只返回 JSON: {"activate": true/false, "reason": "..."}`
+	defaultDatabaseName      = "default"
+	defaultSQLiteDSN         = "file:data/app.db?mode=ro"
+	defaultTimeout           = 10
+	defaultMaxRows           = 50
+	defaultMaxCellBytes      = 4096
+	defaultSchemaCacheSecond = 300
+)
+
+type Config struct {
+	Enabled            bool             `yaml:"enabled" json:"enabled"`
+	ActivationPrompt   string           `yaml:"activation_prompt" json:"activation_prompt"`
+	DefaultDatabase    string           `yaml:"default_database" json:"default_database"`
+	SchemaCacheSeconds int              `yaml:"schema_cache_seconds" json:"schema_cache_seconds"`
+	Databases          []DatabaseConfig `yaml:"databases" json:"databases"`
+}
+
+type DatabaseConfig struct {
+	Name         string       `yaml:"name" json:"name"`
+	Active       bool         `yaml:"active" json:"active"`
+	Driver       string       `yaml:"driver" json:"driver"`
+	DSN          string       `yaml:"dsn" json:"-"`
+	Timeout      int          `yaml:"timeout" json:"timeout"`
+	MaxRows      int          `yaml:"max_rows" json:"max_rows"`
+	MaxCellBytes int          `yaml:"max_cell_bytes" json:"max_cell_bytes"`
+	Schema       SchemaConfig `yaml:"schema" json:"schema"`
+}
+
+type SchemaConfig struct {
+	IncludeTables []string `yaml:"include_tables" json:"include_tables"`
+	ExcludeTables []string `yaml:"exclude_tables" json:"exclude_tables"`
+}
+
+type State struct {
+	cfg       *Config
+	dbs       map[string]*database
+	order     []string
+	cacheMu   sync.Mutex
+	cacheText string
+	cacheAt   time.Time
+}
+
+type database struct {
+	cfg DatabaseConfig
+	db  *sql.DB
+}
+
+type QueryResult struct {
+	Database  string     `json:"database"`
+	SQL       string     `json:"sql"`
+	Columns   []string   `json:"columns"`
+	Rows      [][]string `json:"rows"`
+	Truncated bool       `json:"truncated"`
+	MaxRows   int        `json:"max_rows"`
+}
+
+func LoadConfig(path string) (*Config, error) {
+	if _, err := os.Stat(path); err != nil {
+		if !os.IsNotExist(err) {
+			return nil, fmt.Errorf("检查 SQL 查询插件配置失败: %w", err)
+		}
+		cfg := defaultConfig()
+		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
+			return nil, fmt.Errorf("创建 SQL 查询插件目录失败: %w", err)
+		}
+		data, err := yaml.Marshal(&cfg)
+		if err != nil {
+			return nil, fmt.Errorf("生成 SQL 查询插件配置失败: %w", err)
+		}
+		if err := os.WriteFile(path, data, 0644); err != nil {
+			return nil, fmt.Errorf("写入 SQL 查询插件配置失败: %w", err)
+		}
+	}
+
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("读取 SQL 查询插件配置失败: %w", err)
+	}
+	var cfg Config
+	if err := yaml.Unmarshal(data, &cfg); err != nil {
+		return nil, fmt.Errorf("解析 SQL 查询插件配置失败: %w", err)
+	}
+	if err := normalizeConfig(&cfg); err != nil {
+		return nil, err
+	}
+	return &cfg, nil
+}
+
+func NewState(cfg *Config) (*State, error) {
+	state := &State{cfg: cfg, dbs: map[string]*database{}}
+	if cfg == nil || !cfg.Enabled {
+		return state, nil
+	}
+
+	for _, item := range cfg.Databases {
+		db, err := sql.Open(driverName(item.Driver), item.DSN)
+		if err != nil {
+			return nil, fmt.Errorf("打开数据库 %s 失败: %w", item.Name, err)
+		}
+		if item.Driver == "sqlite" {
+			db.SetMaxOpenConns(1)
+			db.SetMaxIdleConns(1)
+		} else {
+			db.SetMaxOpenConns(5)
+			db.SetMaxIdleConns(2)
+		}
+		db.SetConnMaxLifetime(30 * time.Minute)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(item.Timeout)*time.Second)
+		if err := db.PingContext(ctx); err != nil {
+			cancel()
+			db.Close()
+			return nil, fmt.Errorf("连接数据库 %s 失败: %w", item.Name, err)
+		}
+		if item.Driver == "sqlite" {
+			if _, err := db.ExecContext(ctx, "PRAGMA query_only = ON"); err != nil {
+				cancel()
+				db.Close()
+				return nil, fmt.Errorf("启用 SQLite 只读模式失败: %w", err)
+			}
+		}
+		cancel()
+
+		state.dbs[item.Name] = &database{cfg: item, db: db}
+		state.order = append(state.order, item.Name)
+	}
+	return state, nil
+}
+
+func (s *State) Close() error {
+	if s == nil {
+		return nil
+	}
+	var errs []string
+	for _, db := range s.dbs {
+		if err := db.db.Close(); err != nil {
+			errs = append(errs, err.Error())
+		}
+	}
+	if len(errs) > 0 {
+		return errors.New(strings.Join(errs, "; "))
+	}
+	return nil
+}
+
+func (s *State) Enabled() bool {
+	return s != nil && s.cfg != nil && s.cfg.Enabled && len(s.dbs) > 0
+}
+
+func (s *State) ActivationPrompt() string {
+	if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.ActivationPrompt) == "" {
+		return defaultActivationPrompt
+	}
+	return strings.TrimSpace(s.cfg.ActivationPrompt)
+}
+
+func (s *State) DefaultDatabase() string {
+	if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.DefaultDatabase) == "" {
+		return defaultDatabaseName
+	}
+	return strings.TrimSpace(s.cfg.DefaultDatabase)
+}
+
+func (s *State) SchemaContext(ctx context.Context) (string, error) {
+	if !s.Enabled() {
+		return "", errors.New("SQL 查询插件未启用")
+	}
+
+	s.cacheMu.Lock()
+	if s.cacheText != "" && time.Since(s.cacheAt) < time.Duration(s.cfg.SchemaCacheSeconds)*time.Second {
+		text := s.cacheText
+		s.cacheMu.Unlock()
+		return text, nil
+	}
+	s.cacheMu.Unlock()
+
+	var b strings.Builder
+	fmt.Fprintf(&b, "可查询数据库列表（只能生成 SELECT/WITH 查询）：\n")
+	for _, name := range s.order {
+		handle := s.dbs[name]
+		fmt.Fprintf(&b, "\n数据库 %s，类型 %s，单次最多返回 %d 行：\n", handle.cfg.Name, handle.cfg.Driver, handle.cfg.MaxRows)
+		schema, err := handle.schemaContext(ctx)
+		if err != nil {
+			return "", err
+		}
+		b.WriteString(schema)
+	}
+	text := b.String()
+
+	s.cacheMu.Lock()
+	s.cacheText = text
+	s.cacheAt = time.Now()
+	s.cacheMu.Unlock()
+	return text, nil
+}
+
+func (s *State) ExecuteReadOnly(ctx context.Context, databaseName string, query string) (*QueryResult, error) {
+	if !s.Enabled() {
+		return nil, errors.New("SQL 查询插件未启用")
+	}
+	if err := ValidateReadOnlySQL(query); err != nil {
+		return nil, err
+	}
+
+	handle := s.databaseByName(databaseName)
+	if handle == nil {
+		return nil, fmt.Errorf("数据库配置不存在: %s", databaseName)
+	}
+	if err := handle.rejectExcludedTables(query); err != nil {
+		return nil, err
+	}
+
+	timeout := time.Duration(handle.cfg.Timeout) * time.Second
+	queryCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	if handle.cfg.Driver == "mysql" {
+		tx, err := handle.db.BeginTx(queryCtx, &sql.TxOptions{ReadOnly: true})
+		if err != nil {
+			return nil, fmt.Errorf("开启只读事务失败: %w", err)
+		}
+		defer tx.Rollback()
+		rows, err := tx.QueryContext(queryCtx, query)
+		if err != nil {
+			return nil, fmt.Errorf("执行 SQL 查询失败: %w", err)
+		}
+		result, err := scanRows(rows, handle.cfg, query)
+		if err != nil {
+			return nil, err
+		}
+		if err := tx.Commit(); err != nil {
+			return nil, fmt.Errorf("提交只读事务失败: %w", err)
+		}
+		return result, nil
+	}
+
+	rows, err := handle.db.QueryContext(queryCtx, query)
+	if err != nil {
+		return nil, fmt.Errorf("执行 SQL 查询失败: %w", err)
+	}
+	return scanRows(rows, handle.cfg, query)
+}
+
+func BuildResultContext(userQuery string, generatedSQL string, result *QueryResult) string {
+	var b strings.Builder
+	fmt.Fprintf(&b, "用户问题需要查询本地数据库。请仅根据以下 SQL 查询结果回答；不要编造结果中不存在的记录。\n")
+	fmt.Fprintf(&b, "用户问题：%s\n", userQuery)
+	fmt.Fprintf(&b, "数据库：%s\n", result.Database)
+	fmt.Fprintf(&b, "已执行只读 SQL：%s\n", generatedSQL)
+	if len(result.Columns) == 0 {
+		b.WriteString("查询没有返回列。\n")
+		return b.String()
+	}
+	if len(result.Rows) == 0 {
+		b.WriteString("查询结果：没有匹配记录。\n")
+		return b.String()
+	}
+	b.WriteString("查询结果：\n")
+	b.WriteString(markdownTable(result.Columns, result.Rows))
+	if result.Truncated {
+		fmt.Fprintf(&b, "\n结果已按配置截断，只展示前 %d 行。\n", result.MaxRows)
+	}
+	return b.String()
+}
+
+func BuildErrorContext(userQuery string, err error) string {
+	return fmt.Sprintf("用户问题可能需要查询本地数据库，但 SQL 查询插件执行失败：%s\n用户问题：%s\n请向用户说明无法完成数据库查询，不要编造数据库记录。", err.Error(), userQuery)
+}
+
+func ValidateReadOnlySQL(query string) error {
+	trimmed := strings.TrimSpace(query)
+	if trimmed == "" {
+		return errors.New("SQL 不能为空")
+	}
+	if strings.Contains(trimmed, "--") || strings.Contains(trimmed, "/*") || strings.Contains(trimmed, "*/") {
+		return errors.New("SQL 不允许包含注释")
+	}
+	body := strings.TrimSuffix(trimmed, ";")
+	if strings.Contains(body, ";") {
+		return errors.New("SQL 只允许单条语句")
+	}
+	upper := strings.ToUpper(body)
+	first := firstToken(upper)
+	if first != "SELECT" && first != "WITH" {
+		return fmt.Errorf("SQL 只允许 SELECT/WITH 查询，当前为 %s", first)
+	}
+
+	stripped := stripSingleQuotedStrings(upper)
+	forbidden := []string{
+		"INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "CREATE", "TRUNCATE", "REPLACE", "MERGE",
+		"GRANT", "REVOKE", "VACUUM", "ANALYZE", "ATTACH", "DETACH", "LOAD", "CALL", "EXEC",
+		"SET", "USE", "LOCK", "UNLOCK", "BEGIN", "COMMIT", "ROLLBACK",
+	}
+	for _, word := range forbidden {
+		if hasSQLWord(stripped, word) {
+			return fmt.Errorf("SQL 包含禁止关键词: %s", word)
+		}
+	}
+	risky := []string{"INTO OUTFILE", "INTO DUMPFILE", "LOAD_FILE", "SLEEP", "BENCHMARK", "LOAD_EXTENSION"}
+	for _, phrase := range risky {
+		if strings.Contains(stripped, phrase) {
+			return fmt.Errorf("SQL 包含禁止函数或语法: %s", phrase)
+		}
+	}
+	return nil
+}
+
+func defaultConfig() Config {
+	return Config{
+		Enabled:            false,
+		ActivationPrompt:   defaultActivationPrompt,
+		DefaultDatabase:    defaultDatabaseName,
+		SchemaCacheSeconds: defaultSchemaCacheSecond,
+		Databases: []DatabaseConfig{{
+			Name:         defaultDatabaseName,
+			Active:       true,
+			Driver:       "sqlite",
+			DSN:          defaultSQLiteDSN,
+			Timeout:      defaultTimeout,
+			MaxRows:      defaultMaxRows,
+			MaxCellBytes: defaultMaxCellBytes,
+		}},
+	}
+}
+
+func normalizeConfig(cfg *Config) error {
+	if strings.TrimSpace(cfg.ActivationPrompt) == "" {
+		cfg.ActivationPrompt = defaultActivationPrompt
+	}
+	if strings.TrimSpace(cfg.DefaultDatabase) == "" {
+		cfg.DefaultDatabase = defaultDatabaseName
+	}
+	if cfg.SchemaCacheSeconds <= 0 {
+		cfg.SchemaCacheSeconds = defaultSchemaCacheSecond
+	}
+	if len(cfg.Databases) == 0 {
+		cfg.Databases = defaultConfig().Databases
+	}
+
+	seen := map[string]bool{}
+	activeIndex := -1
+	for i := range cfg.Databases {
+		item := &cfg.Databases[i]
+		item.Name = strings.TrimSpace(item.Name)
+		if item.Name == "" {
+			item.Name = fmt.Sprintf("database-%d", i+1)
+		}
+		if seen[item.Name] {
+			return fmt.Errorf("SQL 查询插件数据库名称重复: %s", item.Name)
+		}
+		seen[item.Name] = true
+
+		item.Driver = strings.ToLower(strings.TrimSpace(item.Driver))
+		if item.Driver == "" {
+			item.Driver = "sqlite"
+		}
+		if item.Driver != "sqlite" && item.Driver != "mysql" {
+			return fmt.Errorf("SQL 查询插件暂不支持数据库类型: %s", item.Driver)
+		}
+		if strings.TrimSpace(item.DSN) == "" {
+			return fmt.Errorf("数据库 %s 缺少 dsn", item.Name)
+		}
+		if item.Timeout <= 0 {
+			item.Timeout = defaultTimeout
+		}
+		if item.MaxRows <= 0 {
+			item.MaxRows = defaultMaxRows
+		}
+		if item.MaxCellBytes <= 0 {
+			item.MaxCellBytes = defaultMaxCellBytes
+		}
+		item.Schema.IncludeTables = cleanList(item.Schema.IncludeTables)
+		item.Schema.ExcludeTables = cleanList(item.Schema.ExcludeTables)
+		if item.Active {
+			if activeIndex == -1 {
+				activeIndex = i
+			} else {
+				item.Active = false
+			}
+		}
+	}
+	if activeIndex == -1 {
+		cfg.Databases[0].Active = true
+	}
+	if !seen[cfg.DefaultDatabase] {
+		for _, item := range cfg.Databases {
+			if item.Active {
+				cfg.DefaultDatabase = item.Name
+				break
+			}
+		}
+	}
+	return nil
+}
+
+func driverName(driver string) string {
+	if driver == "sqlite" {
+		return "sqlite"
+	}
+	return driver
+}
+
+func (s *State) databaseByName(name string) *database {
+	name = strings.TrimSpace(name)
+	if name == "" {
+		name = s.DefaultDatabase()
+	}
+	if db := s.dbs[name]; db != nil {
+		return db
+	}
+	return s.dbs[s.DefaultDatabase()]
+}
+
+func (d *database) schemaContext(ctx context.Context) (string, error) {
+	timeout := time.Duration(d.cfg.Timeout) * time.Second
+	schemaCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+	if d.cfg.Driver == "mysql" {
+		return d.mysqlSchemaContext(schemaCtx)
+	}
+	return d.sqliteSchemaContext(schemaCtx)
+}
+
+func (d *database) sqliteSchemaContext(ctx context.Context) (string, error) {
+	rows, err := d.db.QueryContext(ctx, `SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name`)
+	if err != nil {
+		return "", fmt.Errorf("读取 SQLite 表列表失败: %w", err)
+	}
+	defer rows.Close()
+
+	var tables []string
+	for rows.Next() {
+		var name string
+		if err := rows.Scan(&name); err != nil {
+			return "", err
+		}
+		if d.tableAllowed(name) {
+			tables = append(tables, name)
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return "", err
+	}
+
+	var b strings.Builder
+	for _, table := range tables {
+		fmt.Fprintf(&b, "- 表 %s\n", table)
+		colRows, err := d.db.QueryContext(ctx, "PRAGMA table_info("+quoteSQLiteString(table)+")")
+		if err != nil {
+			return "", fmt.Errorf("读取 SQLite 表 %s 字段失败: %w", table, err)
+		}
+		for colRows.Next() {
+			var cid int
+			var name, typ string
+			var notNull int
+			var defaultValue any
+			var pk int
+			if err := colRows.Scan(&cid, &name, &typ, &notNull, &defaultValue, &pk); err != nil {
+				colRows.Close()
+				return "", err
+			}
+			extra := ""
+			if pk > 0 {
+				extra = " primary_key"
+			}
+			fmt.Fprintf(&b, "  - %s %s%s\n", name, typ, extra)
+		}
+		if err := colRows.Close(); err != nil {
+			return "", err
+		}
+	}
+	if len(tables) == 0 {
+		b.WriteString("- 没有可查询表，或表被 include/exclude 规则过滤。\n")
+	}
+	return b.String(), nil
+}
+
+func (d *database) mysqlSchemaContext(ctx context.Context) (string, error) {
+	rows, err := d.db.QueryContext(ctx, `SELECT table_name FROM information_schema.tables WHERE table_schema = DATABASE() AND table_type = 'BASE TABLE' ORDER BY table_name`)
+	if err != nil {
+		return "", fmt.Errorf("读取 MySQL 表列表失败: %w", err)
+	}
+	defer rows.Close()
+
+	var tables []string
+	for rows.Next() {
+		var name string
+		if err := rows.Scan(&name); err != nil {
+			return "", err
+		}
+		if d.tableAllowed(name) {
+			tables = append(tables, name)
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return "", err
+	}
+
+	var b strings.Builder
+	for _, table := range tables {
+		fmt.Fprintf(&b, "- 表 %s\n", table)
+		colRows, err := d.db.QueryContext(ctx, `SELECT column_name, data_type, is_nullable, column_key FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = ? ORDER BY ordinal_position`, table)
+		if err != nil {
+			return "", fmt.Errorf("读取 MySQL 表 %s 字段失败: %w", table, err)
+		}
+		for colRows.Next() {
+			var name, typ, nullable, key string
+			if err := colRows.Scan(&name, &typ, &nullable, &key); err != nil {
+				colRows.Close()
+				return "", err
+			}
+			extra := ""
+			if key == "PRI" {
+				extra = " primary_key"
+			}
+			if nullable == "NO" {
+				extra += " not_null"
+			}
+			fmt.Fprintf(&b, "  - %s %s%s\n", name, typ, extra)
+		}
+		if err := colRows.Close(); err != nil {
+			return "", err
+		}
+	}
+	if len(tables) == 0 {
+		b.WriteString("- 没有可查询表，或表被 include/exclude 规则过滤。\n")
+	}
+	return b.String(), nil
+}
+
+func (d *database) tableAllowed(table string) bool {
+	name := strings.ToLower(strings.TrimSpace(table))
+	include := lowerSet(d.cfg.Schema.IncludeTables)
+	if len(include) > 0 && !include[name] {
+		return false
+	}
+	exclude := lowerSet(d.cfg.Schema.ExcludeTables)
+	return !exclude[name]
+}
+
+func (d *database) rejectExcludedTables(query string) error {
+	cleaned := strings.ToLower(stripSingleQuotedStrings(query))
+	include := lowerSet(d.cfg.Schema.IncludeTables)
+	if len(include) > 0 {
+		matched := false
+		for table := range include {
+			if hasSQLWord(cleaned, table) {
+				matched = true
+				break
+			}
+		}
+		if !matched {
+			return errors.New("SQL 未访问 include_tables 中允许的表")
+		}
+	}
+	exclude := lowerSet(d.cfg.Schema.ExcludeTables)
+	for table := range exclude {
+		if hasSQLWord(cleaned, table) {
+			return fmt.Errorf("SQL 访问了被排除的表: %s", table)
+		}
+	}
+	return nil
+}
+
+func scanRows(rows *sql.Rows, cfg DatabaseConfig, query string) (*QueryResult, error) {
+	defer rows.Close()
+	columns, err := rows.Columns()
+	if err != nil {
+		return nil, fmt.Errorf("读取查询列失败: %w", err)
+	}
+	result := &QueryResult{
+		Database: cfg.Name,
+		SQL:      query,
+		Columns:  columns,
+		MaxRows:  cfg.MaxRows,
+	}
+	for rows.Next() {
+		values := make([]any, len(columns))
+		ptrs := make([]any, len(columns))
+		for i := range values {
+			ptrs[i] = &values[i]
+		}
+		if err := rows.Scan(ptrs...); err != nil {
+			return nil, fmt.Errorf("读取查询结果失败: %w", err)
+		}
+		if len(result.Rows) >= cfg.MaxRows {
+			result.Truncated = true
+			break
+		}
+		row := make([]string, len(columns))
+		for i, value := range values {
+			row[i] = formatCell(value, cfg.MaxCellBytes)
+		}
+		result.Rows = append(result.Rows, row)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("读取查询结果失败: %w", err)
+	}
+	return result, nil
+}
+
+func formatCell(value any, maxBytes int) string {
+	if value == nil {
+		return "NULL"
+	}
+	switch v := value.(type) {
+	case time.Time:
+		return v.Format(time.RFC3339)
+	case []byte:
+		if !utf8.Valid(v) {
+			return fmt.Sprintf("<binary %d bytes>", len(v))
+		}
+		return truncateString(string(v), maxBytes)
+	case string:
+		return truncateString(v, maxBytes)
+	default:
+		return truncateString(fmt.Sprint(v), maxBytes)
+	}
+}
+
+func truncateString(s string, maxBytes int) string {
+	if maxBytes <= 0 || len(s) <= maxBytes {
+		return s
+	}
+	cut := s[:maxBytes]
+	for !utf8.ValidString(cut) && len(cut) > 0 {
+		cut = cut[:len(cut)-1]
+	}
+	return cut + "...<truncated>"
+}
+
+func markdownTable(columns []string, rows [][]string) string {
+	var b strings.Builder
+	b.WriteString("| ")
+	for i, col := range columns {
+		if i > 0 {
+			b.WriteString(" | ")
+		}
+		b.WriteString(escapeMarkdownCell(col))
+	}
+	b.WriteString(" |\n| ")
+	for i := range columns {
+		if i > 0 {
+			b.WriteString(" | ")
+		}
+		b.WriteString("---")
+	}
+	b.WriteString(" |\n")
+	for _, row := range rows {
+		b.WriteString("| ")
+		for i, cell := range row {
+			if i > 0 {
+				b.WriteString(" | ")
+			}
+			b.WriteString(escapeMarkdownCell(cell))
+		}
+		b.WriteString(" |\n")
+	}
+	return b.String()
+}
+
+func escapeMarkdownCell(s string) string {
+	s = strings.ReplaceAll(s, "|", "\\|")
+	s = strings.ReplaceAll(s, "\r\n", " ")
+	s = strings.ReplaceAll(s, "\n", " ")
+	s = strings.ReplaceAll(s, "\r", " ")
+	return s
+}
+
+func firstToken(s string) string {
+	fields := strings.Fields(s)
+	if len(fields) == 0 {
+		return ""
+	}
+	return strings.Trim(fields[0], "();")
+}
+
+func stripSingleQuotedStrings(s string) string {
+	var b strings.Builder
+	inString := false
+	for i := 0; i < len(s); i++ {
+		ch := s[i]
+		if ch == '\'' {
+			if inString && i+1 < len(s) && s[i+1] == '\'' {
+				i++
+				continue
+			}
+			inString = !inString
+			b.WriteByte(' ')
+			continue
+		}
+		if inString {
+			b.WriteByte(' ')
+		} else {
+			b.WriteByte(ch)
+		}
+	}
+	return b.String()
+}
+
+func hasSQLWord(s string, word string) bool {
+	pattern := `(?i)(^|[^a-zA-Z0-9_])` + regexp.QuoteMeta(word) + `([^a-zA-Z0-9_]|$)`
+	return regexp.MustCompile(pattern).FindStringIndex(s) != nil
+}
+
+func quoteSQLiteString(s string) string {
+	return "'" + strings.ReplaceAll(s, "'", "''") + "'"
+}
+
+func lowerSet(items []string) map[string]bool {
+	set := map[string]bool{}
+	for _, item := range items {
+		item = strings.ToLower(strings.TrimSpace(item))
+		if item != "" {
+			set[item] = true
+		}
+	}
+	return set
+}
+
+func cleanList(items []string) []string {
+	seen := map[string]bool{}
+	var cleaned []string
+	for _, item := range items {
+		item = strings.TrimSpace(item)
+		if item == "" {
+			continue
+		}
+		key := strings.ToLower(item)
+		if seen[key] {
+			continue
+		}
+		seen[key] = true
+		cleaned = append(cleaned, item)
+	}
+	sort.Strings(cleaned)
+	return cleaned
+}

agents/SQL_query/sql_query_test.go

@@ -0,0 +1,36 @@
+package sqlquery
+
+import "testing"
+
+func TestValidateReadOnlySQLAllowsSelectAndWith(t *testing.T) {
+	queries := []string{
+		"SELECT * FROM events LIMIT 10",
+		"select id, created_at from events where content = 'delete keyword in text' limit 5;",
+		"WITH recent AS (SELECT * FROM events LIMIT 10) SELECT * FROM recent",
+	}
+	for _, query := range queries {
+		if err := ValidateReadOnlySQL(query); err != nil {
+			t.Fatalf("ValidateReadOnlySQL(%q) returned error: %v", query, err)
+		}
+	}
+}
+
+func TestValidateReadOnlySQLRejectsUnsafeStatements(t *testing.T) {
+	queries := []string{
+		"",
+		"DELETE FROM events",
+		"UPDATE events SET content='x'",
+		"DROP TABLE events",
+		"SELECT * FROM events; DELETE FROM events",
+		"SELECT * INTO OUTFILE '/tmp/x' FROM events",
+		"SELECT SLEEP(10)",
+		"ATTACH DATABASE 'x' AS y",
+		"VACUUM",
+		"SELECT * FROM events -- comment",
+	}
+	for _, query := range queries {
+		if err := ValidateReadOnlySQL(query); err == nil {
+			t.Fatalf("ValidateReadOnlySQL(%q) returned nil, want error", query)
+		}
+	}
+}

go.mod

@@ -4,15 +4,19 @@ go 1.25.0
 
 require (
 	github.com/gin-gonic/gin v1.12.0
+	github.com/go-sql-driver/mysql v1.10.0
 	github.com/volcengine/volcengine-go-sdk v1.2.28
 	gopkg.in/yaml.v3 v3.0.1
+	modernc.org/sqlite v1.52.0
 )
 
 require (
+	filippo.io/edwards25519 v1.2.0 // indirect
 	github.com/bytedance/gopkg v0.1.3 // indirect
 	github.com/bytedance/sonic v1.15.0 // indirect
 	github.com/bytedance/sonic/loader v0.5.0 // indirect
 	github.com/cloudwego/base64x v0.1.6 // indirect
+	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.12 // indirect
 	github.com/gin-contrib/sse v1.1.0 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
@@ -20,7 +24,7 @@ require (
 	github.com/go-playground/validator/v10 v10.30.1 // indirect
 	github.com/goccy/go-json v0.10.5 // indirect
 	github.com/goccy/go-yaml v1.19.2 // indirect
-	github.com/google/uuid v1.3.0 // indirect
+	github.com/google/uuid v1.6.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/cpuid/v2 v2.3.0 // indirect
@@ -28,9 +32,11 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/ncruces/go-strftime v1.0.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.4 // indirect
 	github.com/quic-go/qpack v0.6.0 // indirect
 	github.com/quic-go/quic-go v0.59.0 // indirect
+	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.3.1 // indirect
 	github.com/volcengine/volc-sdk-golang v1.0.23 // indirect
@@ -38,8 +44,11 @@ require (
 	golang.org/x/arch v0.22.0 // indirect
 	golang.org/x/crypto v0.48.0 // indirect
 	golang.org/x/net v0.51.0 // indirect
-	golang.org/x/sys v0.41.0 // indirect
+	golang.org/x/sys v0.42.0 // indirect
 	golang.org/x/text v0.34.0 // indirect
 	google.golang.org/protobuf v1.36.10 // indirect
 	gopkg.in/yaml.v2 v2.2.8 // indirect
+	modernc.org/libc v1.72.3 // indirect
+	modernc.org/mathutil v1.7.1 // indirect
+	modernc.org/memory v1.11.0 // indirect
 )

go.sum

@@ -1,4 +1,6 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
+filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo=
+filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY=
 github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=
@@ -14,6 +16,8 @@ github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gE
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/gabriel-vasile/mimetype v1.4.12 h1:e9hWvmLYvtp846tLHam2o++qitpguFiYCKbn0w9jyqw=
@@ -30,6 +34,8 @@ github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJn
 github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
 github.com/go-playground/validator/v10 v10.30.1 h1:f3zDSN/zOma+w6+1Wswgd9fLkdwy06ntQJp0BBvFG0w=
 github.com/go-playground/validator/v10 v10.30.1/go.mod h1:oSuBIQzuJxL//3MelwSLD5hc2Tu889bF0Idm9Dg26cM=
+github.com/go-sql-driver/mysql v1.10.0 h1:Q+1LV8DkHJvSYAdR83XzuhDaTykuDx0l6fkXxoWCWfw=
+github.com/go-sql-driver/mysql v1.10.0/go.mod h1:M+cqaI7+xxXGG9swrdeUIoPG3Y3KCkF0pZej+SK+nWk=
 github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
 github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
 github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM=
@@ -55,8 +61,13 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
-github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
+github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
@@ -81,6 +92,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
+github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
+github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
 github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
 github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -91,6 +104,8 @@ github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8=
 github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII=
 github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw=
 github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
 github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -126,6 +141,8 @@ golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
+golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
+golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -136,11 +153,13 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
+golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
+golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
-golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
 golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
@@ -148,6 +167,8 @@ golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
+golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
+golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
@@ -182,3 +203,31 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
+modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY=
+modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
+modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ=
+modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A=
+modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
+modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
+modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
+modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
+modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo=
+modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
+modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
+modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
+modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU=
+modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs=
+modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
+modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
+modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
+modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
+modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
+modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
+modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
+modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
+modernc.org/sqlite v1.52.0 h1:p4dhYh2tXZCiyaqHwRVJDjIGKWyXayiQpThxgDzJaxo=
+modernc.org/sqlite v1.52.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM=
+modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
+modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
+modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
+modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=

main.go

@@ -19,6 +19,8 @@ import (
 	"sync"
 	"time"
 
+	sqlquery "aichat/agents/SQL_query"
+
 	"github.com/gin-gonic/gin"
 	ark "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
 	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
@@ -624,6 +626,7 @@ var (
 	cfg         *Config
 	aiState     *OpenAIState
 	searchState *SearchState
+	sqlState    *sqlquery.State
 	store       *ConvStore
 )
 
@@ -737,8 +740,16 @@ func chatHandler(c *gin.Context) {
 	}
 
 	chatMessages := req.Messages
+	if sqlState != nil && sqlState.Enabled() {
+		withSQL, err := enrichMessagesWithSQL(c.Request.Context(), profile, chatMessages)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "SQL 查询插件调用失败:", err)
+		} else {
+			chatMessages = withSQL
+		}
+	}
 	if req.WebSearch {
-		withSearch, err := enrichMessagesWithSearch(c.Request.Context(), req.Messages)
+		withSearch, err := enrichMessagesWithSearch(c.Request.Context(), chatMessages)
 		if err != nil {
 			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
 			return
@@ -882,6 +893,146 @@ func latestUserQuery(messages []ChatMessage) string {
 	return ""
 }
 
+type sqlActivationDecision struct {
+	Activate bool   `json:"activate"`
+	Reason   string `json:"reason"`
+}
+
+type sqlGenerationResult struct {
+	Database string `json:"database"`
+	SQL      string `json:"sql"`
+	Reason   string `json:"reason"`
+}
+
+func enrichMessagesWithSQL(ctx context.Context, profile *OpenAIProfile, messages []ChatMessage) ([]ChatMessage, error) {
+	query := latestUserQuery(messages)
+	if query == "" {
+		return messages, nil
+	}
+
+	activate, reason, err := classifySQLActivation(ctx, profile, messages)
+	if err != nil {
+		return messages, err
+	}
+	if !activate {
+		return messages, nil
+	}
+
+	schemaContext, err := sqlState.SchemaContext(ctx)
+	if err != nil {
+		return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil
+	}
+
+	generated, err := generateSQLForUserQuery(ctx, profile, query, schemaContext)
+	if err != nil {
+		return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil
+	}
+	generated.Database = strings.TrimSpace(generated.Database)
+	generated.SQL = strings.TrimSpace(generated.SQL)
+	if generated.SQL == "" {
+		err := fmt.Errorf("模型未生成可执行 SQL: %s", generated.Reason)
+		return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil
+	}
+
+	result, err := sqlState.ExecuteReadOnly(ctx, generated.Database, generated.SQL)
+	if err != nil {
+		return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil
+	}
+	contextText := sqlquery.BuildResultContext(query, generated.SQL, result)
+	if strings.TrimSpace(reason) != "" {
+		contextText += "\n激活原因：" + reason
+	}
+	return prependSQLContext(messages, contextText), nil
+}
+
+func prependSQLContext(messages []ChatMessage, content string) []ChatMessage {
+	withSQL := make([]ChatMessage, 0, len(messages)+1)
+	withSQL = append(withSQL, ChatMessage{Role: "system", Content: content, Hidden: true})
+	withSQL = append(withSQL, messages...)
+	return withSQL
+}
+
+func classifySQLActivation(ctx context.Context, profile *OpenAIProfile, messages []ChatMessage) (bool, string, error) {
+	query := latestUserQuery(messages)
+	prompt := fmt.Sprintf("%s\n\n最新用户问题：%s", sqlState.ActivationPrompt(), query)
+	text, err := completeText(ctx, profile, []ChatMessage{{Role: "system", Content: prompt}}, 512)
+	if err != nil {
+		return false, "", err
+	}
+	var decision sqlActivationDecision
+	if err := json.Unmarshal([]byte(extractJSONObject(text)), &decision); err != nil {
+		return false, "", fmt.Errorf("解析 SQL 查询激活结果失败: %w", err)
+	}
+	return decision.Activate, decision.Reason, nil
+}
+
+func generateSQLForUserQuery(ctx context.Context, profile *OpenAIProfile, userQuery string, schemaContext string) (*sqlGenerationResult, error) {
+	prompt := fmt.Sprintf(`你是只读 SQL 生成器。请根据用户问题和数据库 schema 生成一条只读 SQL。
+要求：
+- 只能返回 JSON，不要使用 Markdown。
+- JSON 格式：{"database":"数据库名称","sql":"SELECT ... LIMIT N","reason":"生成原因"}
+- 只能生成 SELECT 或 WITH 查询，禁止 INSERT/UPDATE/DELETE/DROP/ALTER/CREATE 等任何修改语句。
+- 必须只使用 schema 中出现的数据库、表和字段。
+- 必须添加 LIMIT，且 LIMIT 不超过插件配置的 max_rows。
+- 如果无法根据 schema 回答，返回 {"database":"","sql":"","reason":"无法根据已知表结构生成查询"}。
+
+%s
+
+用户问题：%s`, schemaContext, userQuery)
+	text, err := completeText(ctx, profile, []ChatMessage{{Role: "system", Content: prompt}}, 1024)
+	if err != nil {
+		return nil, err
+	}
+	var generated sqlGenerationResult
+	if err := json.Unmarshal([]byte(extractJSONObject(text)), &generated); err != nil {
+		return nil, fmt.Errorf("解析 SQL 生成结果失败: %w", err)
+	}
+	return &generated, nil
+}
+
+func completeText(ctx context.Context, profile *OpenAIProfile, chatMessages []ChatMessage, maxTokens int) (string, error) {
+	messages, err := buildArkMessages(chatMessages)
+	if err != nil {
+		return "", err
+	}
+	timeout := time.Duration(profile.Config.Timeout) * time.Second
+	completionCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+	stream, err := profile.Client.CreateChatCompletionStream(completionCtx, model.CreateChatCompletionRequest{
+		Model:     profile.Config.Model,
+		Messages:  messages,
+		MaxTokens: intPtr(maxTokens),
+	}.WithStream(true))
+	if err != nil {
+		return "", err
+	}
+	defer stream.Close()
+
+	var b strings.Builder
+	for {
+		resp, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			return b.String(), nil
+		}
+		if err != nil {
+			return "", err
+		}
+		if len(resp.Choices) > 0 {
+			b.WriteString(resp.Choices[0].Delta.Content)
+		}
+	}
+}
+
+func extractJSONObject(text string) string {
+	text = strings.TrimSpace(text)
+	start := strings.Index(text, "{")
+	end := strings.LastIndex(text, "}")
+	if start >= 0 && end > start {
+		return text[start : end+1]
+	}
+	return text
+}
+
 func webSearch(ctx context.Context, config SearchConfig, query string) ([]searchResult, error) {
 	switch strings.ToLower(config.Provider) {
 	case "duckduckgo", "ddg":
@@ -1363,6 +1514,17 @@ func main() {
 		fmt.Fprintln(os.Stderr, "搜索配置初始化失败:", err)
 		os.Exit(1)
 	}
+	sqlConfig, err := sqlquery.LoadConfig("agents/SQL_query/config.yaml")
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "SQL 查询插件配置加载失败:", err)
+		os.Exit(1)
+	}
+	sqlState, err = sqlquery.NewState(sqlConfig)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "SQL 查询插件初始化失败:", err)
+		os.Exit(1)
+	}
+	defer sqlState.Close()
 	store = NewConvStore("conversations")
 
 	// Gin 路由