feat(ai): 阶段 2 Ask|Agent 共用 aitools Runtime 真查库

- 侧栏 Ask|Agent 分段;Ask 仍走 /api/ai/chat,Agent 走 /api/ai/agent
- RunAgent 多轮 tool call → 进程内 aitools(与 MCP 同源只读/白名单/审计)
- 回复下步骤时间线;openai-chat tools(DeepSeek);agent-accept 冒烟
This commit is contained in:
logan committed 2026-10-15 11:51:39 +08:00
1 parent b8134c4503
commit cd319ce7bd
14 files changed
+1270 -182

No files matched your search

+252
View File
@@ -0,0 +1,252 @@
package ai
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"opsxdb/server/internal/aitools"
)
const (
defaultAgentMaxRounds = 8
maxToolResultFeedChars = 12000
maxStepArgsChars = 200
)
// AgentReq 侧栏 Agent 请求(与 Ask 的 /api/ai/chat 分路径)。
type AgentReq struct {
Prompt string `json:"prompt"`
ConnectionID string `json:"connectionId,omitempty"`
Driver string `json:"driver,omitempty"`
Database string `json:"database,omitempty"`
Model string `json:"model,omitempty"`
MaxRounds int `json:"maxRounds,omitempty"`
}
// AgentStep 挂在回复下的步骤时间线一项。
type AgentStep struct {
Tool string `json:"tool"`
Args string `json:"args,omitempty"`
OK bool `json:"ok"`
Summary string `json:"summary"`
Error string `json:"error,omitempty"`
Ms int64 `json:"ms"`
}
// AgentRes Agent 最终回答 + 步骤。
type AgentRes struct {
Content string `json:"content"`
SQL string `json:"sql,omitempty"`
Steps []AgentStep `json:"steps"`
Rounds int `json:"rounds"`
}
// RunAgent 多轮:LLM tool_call → aitools.Runtime.Call → 回灌,直到最终文本或轮次耗尽。
// ConfirmDanger 永不自动 true(由 Runtime 保证)。Ask 模式不要走这里。
func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (AgentRes, error) {
if rt == nil {
return AgentRes{}, fmt.Errorf("aitools runtime required")
}
if llm == nil {
return AgentRes{}, fmt.Errorf("tool chat LLM required")
}
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
return AgentRes{}, fmt.Errorf("prompt required")
}
maxRounds := req.MaxRounds
if maxRounds <= 0 {
maxRounds = defaultAgentMaxRounds
}
if maxRounds > 16 {
maxRounds = 16
}
tools := aitools.OpenAITools()
msgs := []Message{
{Role: "system", Content: buildAgentSystem(req)},
{Role: "user", Content: prompt},
}
var steps []AgentStep
rounds := 0
for round := 0; round < maxRounds; round++ {
rounds = round + 1
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
if err != nil {
return AgentRes{Steps: steps, Rounds: rounds}, err
}
if len(turn.ToolCalls) == 0 {
content := turn.Content
if content == "" {
content = "(模型未返回文本,也未调用工具)"
}
return AgentRes{
Content: content,
SQL: extractSQL(content),
Steps: steps,
Rounds: rounds,
}, nil
}
// assistant 消息带 tool_calls
msgs = append(msgs, Message{
Role: "assistant",
Content: turn.Content,
ToolCalls: turn.ToolCalls,
})
for _, tc := range turn.ToolCalls {
step, toolMsg := execOneTool(ctx, rt, tc)
steps = append(steps, step)
msgs = append(msgs, toolMsg)
}
}
return AgentRes{
Content: fmt.Sprintf("已达最大工具轮次(%d)。请根据已执行步骤自行判断,或缩小问题后重试。", maxRounds),
Steps: steps,
Rounds: rounds,
}, nil
}
func buildAgentSystem(req AgentReq) string {
var b strings.Builder
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具只读查询用户已白名单的连接。\n")
b.WriteString("规则:\n")
b.WriteString("1. 需要查库时必须调用工具;不要假装已执行 SQL。\n")
b.WriteString("2. 仅只读:execute_query 只允许 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT;Redis 仅读命令。\n")
b.WriteString("3. 先确认 connection_id(可用 opsxdb_list_connections);用户给出偏好连接时优先用它。\n")
b.WriteString("4. 用中文简明回答;可执行 SQL 放在 ```sql 代码块中。\n")
b.WriteString("5. 工具结果可能截断,注意说明。\n")
if req.ConnectionID != "" {
b.WriteString(fmt.Sprintf("\n用户当前偏好连接 id: %s", req.ConnectionID))
if req.Driver != "" {
b.WriteString(fmt.Sprintf("(%s)", req.Driver))
}
if req.Database != "" {
b.WriteString(fmt.Sprintf(",数据库: %s", req.Database))
}
b.WriteString("。\n")
}
return b.String()
}
func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentStep, Message) {
name := strings.TrimSpace(tc.Function.Name)
argsRaw := strings.TrimSpace(tc.Function.Arguments)
if argsRaw == "" {
argsRaw = "{}"
}
step := AgentStep{
Tool: name,
Args: truncateRunes(argsRaw, maxStepArgsChars),
}
start := time.Now()
var args map[string]any
if err := json.Unmarshal([]byte(argsRaw), &args); err != nil {
step.OK = false
step.Error = "invalid tool arguments JSON: " + err.Error()
step.Summary = step.Error
step.Ms = time.Since(start).Milliseconds()
return step, Message{
Role: "tool",
ToolCallID: tc.ID,
Name: name,
Content: step.Error,
}
}
out, err := rt.Call(ctx, name, args)
step.Ms = time.Since(start).Milliseconds()
if err != nil {
step.OK = false
step.Error = err.Error()
step.Summary = err.Error()
return step, Message{
Role: "tool",
ToolCallID: tc.ID,
Name: name,
Content: err.Error(),
}
}
step.OK = true
step.Summary = summarizeToolResult(name, out)
feed := aitools.ResultJSON(out)
if len(feed) > maxToolResultFeedChars {
feed = feed[:maxToolResultFeedChars] + "\n…(truncated)"
}
return step, Message{
Role: "tool",
ToolCallID: tc.ID,
Name: name,
Content: feed,
}
}
func summarizeToolResult(tool string, out any) string {
m, ok := out.(map[string]any)
if !ok {
return truncateRunes(fmt.Sprint(out), 120)
}
lenAny := func(key string) (int, bool) {
switch v := m[key].(type) {
case []any:
return len(v), true
case []string:
return len(v), true
default:
return 0, false
}
}
switch tool {
case aitools.ToolListConnections:
if n, ok := lenAny("connections"); ok {
return fmt.Sprintf("%d connections", n)
}
case aitools.ToolListDatabases:
if n, ok := lenAny("databases"); ok {
return fmt.Sprintf("%d databases", n)
}
case aitools.ToolListTables:
if n, ok := lenAny("tables"); ok {
return fmt.Sprintf("%d tables", n)
}
case aitools.ToolDescribeTable:
if t, _ := m["table"].(string); t != "" {
if n, ok := lenAny("columns"); ok {
return fmt.Sprintf("%s (%d cols)", t, n)
}
return t
}
case aitools.ToolGetSchemaContext:
if s, ok := m["schema"].(string); ok {
return fmt.Sprintf("schema %d chars", len(s))
}
case aitools.ToolExecuteQuery, aitools.ToolExecuteRedisCommand:
rc := 0
switch v := m["rowCount"].(type) {
case int:
rc = v
case int64:
rc = int(v)
case float64:
rc = int(v)
}
return fmt.Sprintf("rows=%d", rc)
}
return truncateRunes(aitools.ResultJSON(out), 120)
}
func truncateRunes(s string, n int) string {
s = strings.Join(strings.Fields(s), " ")
r := []rune(s)
if len(r) <= n {
return s
}
return string(r[:n]) + "…"
}
+209
View File
@@ -0,0 +1,209 @@
package ai
import (
"context"
"encoding/json"
"strings"
"testing"
"opsxdb/server/internal/aitools"
"opsxdb/server/internal/db"
)
// scriptedLLM 按调用次序返回预设 ToolsTurn,用于证明 Agent tool loop 不依赖真 LLM。
type scriptedLLM struct {
turns []ToolsTurn
i int
seen [][]Message
}
func (s *scriptedLLM) ChatWithTools(_ context.Context, _ string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
cp := append([]Message(nil), messages...)
s.seen = append(s.seen, cp)
if len(tools) < 7 {
return ToolsTurn{}, nil // 测试里仍继续;下面断言
}
if s.i >= len(s.turns) {
return ToolsTurn{Content: "fallback"}, nil
}
t := s.turns[s.i]
s.i++
return t, nil
}
func TestOpenAIToolsMatchRuntime(t *testing.T) {
tools := aitools.OpenAITools()
names := aitools.ToolNames()
if len(tools) != len(names) {
t.Fatalf("tools %d vs names %d", len(tools), len(names))
}
for i, n := range names {
if tools[i].Function.Name != n {
t.Fatalf("tool[%d]=%s want %s", i, tools[i].Function.Name, n)
}
if tools[i].Type != "function" {
t.Fatalf("type want function")
}
}
}
func TestRunAgentToolLoop(t *testing.T) {
rt := aitools.NewRuntime(aitools.Policy{
Mode: aitools.ModeReadOnly,
AllowedConnectionIDs: []string{"c1"},
DefaultLimit: 100,
MaxLimit: 500,
})
rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
return &aitools.SavedConnection{ID: id, Name: "demo", Driver: "mysql", Host: "127.0.0.1", Port: 3306, Database: "demo"}, nil
}
rt.ListTbl = func(_ context.Context, c db.ConnReq) ([]string, error) {
if c.Database != "demo" && c.Database != "" {
return nil, nil
}
return []string{"users", "orders"}, nil
}
rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
if !strings.Contains(strings.ToUpper(c.SQL), "SELECT") {
return db.ExecResult{Error: "unexpected"}
}
return db.ExecResult{
Columns: []string{"id", "name"},
Rows: []map[string]interface{}{{"id": 1, "name": "alice"}, {"id": 2, "name": "bob"}},
RowCount: 2,
TimeMs: 1,
}
}
rt.Audit = func(aitools.AuditEntry) {}
argsList, _ := json.Marshal(map[string]any{"connection_id": "c1", "database": "demo"})
argsQuery, _ := json.Marshal(map[string]any{
"connection_id": "c1",
"database": "demo",
"sql": "SELECT id, name FROM users LIMIT 5",
})
llm := &scriptedLLM{turns: []ToolsTurn{
{
ToolCalls: []ToolCall{{
ID: "call_1", Type: "function",
Function: ToolCallFunc{Name: aitools.ToolListTables, Arguments: string(argsList)},
}},
},
{
ToolCalls: []ToolCall{{
ID: "call_2", Type: "function",
Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: string(argsQuery)},
}},
},
{Content: "demo 库有 users、orders。users 样例两行。\n\n```sql\nSELECT id, name FROM users LIMIT 5;\n```"},
}}
res, err := RunAgent(context.Background(), llm, rt, AgentReq{
Prompt: "列出 demo 库有哪些表,并查 users 几行",
ConnectionID: "c1",
Driver: "mysql",
Database: "demo",
})
if err != nil {
t.Fatalf("RunAgent: %v", err)
}
if len(res.Steps) != 2 {
t.Fatalf("want 2 steps, got %d: %+v", len(res.Steps), res.Steps)
}
if !res.Steps[0].OK || res.Steps[0].Tool != aitools.ToolListTables {
t.Fatalf("step0: %+v", res.Steps[0])
}
if !strings.Contains(res.Steps[0].Summary, "2 tables") {
t.Fatalf("list_tables summary: %s", res.Steps[0].Summary)
}
if !res.Steps[1].OK || res.Steps[1].Tool != aitools.ToolExecuteQuery {
t.Fatalf("step1: %+v", res.Steps[1])
}
if !strings.Contains(res.Steps[1].Summary, "rows=2") {
t.Fatalf("execute_query summary: %s", res.Steps[1].Summary)
}
if !strings.Contains(res.Content, "users") {
t.Fatalf("final content: %s", res.Content)
}
if res.SQL == "" {
t.Fatal("expected extracted SQL")
}
if res.Rounds != 3 {
t.Fatalf("rounds=%d", res.Rounds)
}
// 第二轮 messages 应已含 tool 角色回灌
if len(llm.seen) < 2 {
t.Fatal("expected multiple LLM calls")
}
foundTool := false
for _, m := range llm.seen[1] {
if m.Role == "tool" && m.ToolCallID == "call_1" {
foundTool = true
if !strings.Contains(m.Content, "users") {
t.Fatalf("tool result missing tables: %s", m.Content)
}
}
}
if !foundTool {
t.Fatal("second LLM call should include tool result for call_1")
}
}
func TestRunAgentRejectsWriteViaRuntime(t *testing.T) {
rt := aitools.NewRuntime(aitools.Policy{
Mode: aitools.ModeReadOnly,
AllowedConnectionIDs: []string{"c1"},
DefaultLimit: 100,
MaxLimit: 500,
})
rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
return &aitools.SavedConnection{ID: id, Name: "demo", Driver: "mysql"}, nil
}
rt.Audit = func(aitools.AuditEntry) {}
args, _ := json.Marshal(map[string]any{"connection_id": "c1", "sql": "DROP TABLE users"})
llm := &scriptedLLM{turns: []ToolsTurn{
{ToolCalls: []ToolCall{{
ID: "x", Type: "function",
Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: string(args)},
}}},
{Content: "无法删除表:只读策略拒绝。"},
}}
res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "删掉 users"})
if err != nil {
t.Fatal(err)
}
if res.Steps[0].OK {
t.Fatal("DROP should fail")
}
errMsg := strings.ToLower(res.Steps[0].Error + " " + res.Steps[0].Summary)
if !strings.Contains(errMsg, "read-only") && !strings.Contains(errMsg, "not a read-only") {
t.Fatalf("expected read-only deny, got %+v", res.Steps[0])
}
}
func TestParseToolsTurn(t *testing.T) {
content := "hi"
raw, _ := json.Marshal(map[string]any{
"choices": []map[string]any{{
"finish_reason": "tool_calls",
"message": map[string]any{
"role": "assistant",
"content": content,
"tool_calls": []map[string]any{{
"id": "c1",
"type": "function",
"function": map[string]string{
"name": "opsxdb_list_tables",
"arguments": `{"connection_id":"x"}`,
},
}},
},
}},
})
turn, err := parseToolsTurn(raw, 200)
if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
t.Fatalf("parse: %+v %v", turn, err)
}
}
+25 -10
View File
@@ -2,10 +2,25 @@ package ai
import "context"
// Message 对话消息(OpenAI 兼容协议)。
// Message 对话消息(OpenAI 兼容协议)。Agent 多轮可带 tool_calls / tool 角色。
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
Role string `json:"role"`
Content string `json:"content,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
Name string `json:"name,omitempty"`
}
// ToolCall OpenAI function tool_call。
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function ToolCallFunc `json:"function"`
}
type ToolCallFunc struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
// ChatReq 在 stub 时代字段上演进,前端契约不变;新增 schema/model 可选字段。
@@ -32,14 +47,14 @@ type Provider interface {
Complete(ctx context.Context, req ChatReq) (ChatRes, error)
}
// Default 按当前配置选择 Provider;未配置 API Key 时回落 stub(离线占位)。
// Default 按当前配置选择 Provider;stub 时或远端配置缺 key(免鉴权除外)回落 stub(离线占位)。
func Default() Provider {
cfg := LoadConfig()
switch cfg.Provider {
case "openai":
if cfg.APIKey != "" {
return NewOpenAIProvider(cfg)
}
if cfg.Provider == "" || cfg.Provider == "stub" {
return stubProvider{}
}
return stubProvider{}
if cfg.APIKey == "" && cfg.AuthStyle != AuthNone {
return stubProvider{}
}
return NewRemoteProvider(cfg)
}
+143
View File
@@ -0,0 +1,143 @@
package ai
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"opsxdb/server/internal/aitools"
)
// ToolsTurn 一轮带 tools 的模型回复:要么最终文本,要么一组 tool_calls。
type ToolsTurn struct {
Content string
ToolCalls []ToolCall
RawFinish string
}
// ToolChat 可挂载 tools 的 LLM(Agent 循环依赖此接口,便于单测注入)。
type ToolChat interface {
ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error)
}
// ChatWithTools openai-chat 兼容协议的 tools 调用(DeepSeek / OpenAI / 多数兼容网关)。
// anthropic / responses 本轮不支持 Agent tool calling,返回明确错误。
func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil {
return ToolsTurn{}, err
}
if p.protocol != "" && p.protocol != ProtoChat {
return ToolsTurn{}, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议(当前: %s);请改用 DeepSeek / OpenAI 兼容端点", p.protocol)
}
m := strings.TrimSpace(model)
if m == "" {
m = p.model
}
urlStr := endpoint(p.baseURL, "chat/completions")
payload := map[string]any{
"model": m,
"messages": messages,
"tools": tools,
"tool_choice": "auto",
"max_tokens": remoteMaxTokens,
"temperature": 0.2,
}
body, err := json.Marshal(payload)
if err != nil {
return ToolsTurn{}, err
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body))
if err != nil {
return ToolsTurn{}, err
}
httpReq.Header.Set("Content-Type", "application/json")
switch p.authStyle {
case AuthXAPIKey:
httpReq.Header.Set("x-api-key", p.apiKey)
httpReq.Header.Set("anthropic-version", "2023-06-01")
case AuthNone:
default:
httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
}
resp, err := p.hc.Do(httpReq)
if err != nil {
return ToolsTurn{}, fmt.Errorf("AI 请求失败: %w", err)
}
defer resp.Body.Close()
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
if err != nil {
return ToolsTurn{}, fmt.Errorf("读取 AI 响应失败: %w", err)
}
return parseToolsTurn(raw, resp.StatusCode)
}
func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
var errBody struct {
Error *struct {
Message string `json:"message"`
} `json:"error"`
Message string `json:"message"`
}
_ = json.Unmarshal(raw, &errBody)
if status < 200 || status >= 300 {
msg := "未知错误"
if errBody.Error != nil && errBody.Error.Message != "" {
msg = errBody.Error.Message
} else if errBody.Message != "" {
msg = errBody.Message
}
return ToolsTurn{}, fmt.Errorf("AI 服务 HTTP %d: %s", status, msg)
}
var out struct {
Choices []struct {
FinishReason string `json:"finish_reason"`
Message struct {
Role string `json:"role"`
Content *string `json:"content"`
ToolCalls []ToolCall `json:"tool_calls"`
} `json:"message"`
} `json:"choices"`
Error *struct {
Message string `json:"message"`
} `json:"error"`
}
if err := json.Unmarshal(raw, &out); err != nil {
return ToolsTurn{}, fmt.Errorf("解析 AI 响应失败(HTTP %d)", status)
}
if out.Error != nil {
return ToolsTurn{}, fmt.Errorf("AI 服务返回错误: %s", out.Error.Message)
}
if len(out.Choices) == 0 {
return ToolsTurn{}, fmt.Errorf("AI 服务未返回候选结果")
}
ch := out.Choices[0]
content := ""
if ch.Message.Content != nil {
content = strings.TrimSpace(*ch.Message.Content)
}
return ToolsTurn{
Content: content,
ToolCalls: ch.Message.ToolCalls,
RawFinish: ch.FinishReason,
}, nil
}
// DefaultToolChat 返回可跑 Agent 的远端 Provider;stub / 缺 key 时返回错误说明。
func DefaultToolChat() (ToolChat, error) {
cfg := LoadConfig()
if cfg.Provider == "" || cfg.Provider == "stub" {
return nil, fmt.Errorf("Agent 需要真实 LLM:请在设置中配置 DeepSeek / OpenAI 等(stub 仅支持 Ask 生成)")
}
if cfg.APIKey == "" && cfg.AuthStyle != AuthNone {
return nil, fmt.Errorf("Agent 需要 API Key:请在设置中填写后重试")
}
p := NewRemoteProvider(cfg)
if p.protocol != "" && p.protocol != ProtoChat {
return nil, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议(当前: %s)", p.protocol)
}
return p, nil
}