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
1 parent
b8134c4503
commit
cd319ce7bd
14 files changed
+1270
-182
No files matched your search
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user