logan
2c0bfafd8d
- 补入库：server/internal/db/dump.go、redis 分页/扫描预算测试、opsxdb-mcp main_test、 packages/opsxdb-mcp package-lock、docs 审计与对比文档等 37 个新文件 - 含结构扁平化进行中的 80 处修改与 155 处旧结构删除（server/ src/ src-tauri/ 为当前有效源） - 构建产物（bin/target/dist）与签名私钥仍按规则排除
495 lines
15 KiB
Go
495 lines
15 KiB
Go
package ai
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
|
||
"opsxdb/server/internal/aitools"
|
||
)
|
||
|
||
const (
|
||
defaultAgentMaxRounds = 8
|
||
maxToolResultFeedChars = 12000
|
||
maxStepArgsChars = 200
|
||
maxGridRows = 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"`
|
||
Stream bool `json:"stream,omitempty"`
|
||
SessionID string `json:"sessionId,omitempty"`
|
||
Continue bool `json:"continue,omitempty"` // 真·续跑：注入 Prior
|
||
Prior []Message `json:"prior,omitempty"` // 含 tool 轨迹的历史（不含本轮 user）
|
||
RunID string `json:"runId,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"`
|
||
NeedConfirm bool `json:"needConfirm,omitempty"`
|
||
ConfirmSQL string `json:"confirmSql,omitempty"`
|
||
ConfirmBind string `json:"confirmBind,omitempty"`
|
||
Pending bool `json:"pending,omitempty"`
|
||
Result *AgentQueryResult `json:"result,omitempty"`
|
||
}
|
||
|
||
// AgentRes Agent 最终回答 + 步骤 + 运行结果契约。
|
||
type AgentRes struct {
|
||
Content string `json:"content"`
|
||
SQL string `json:"sql,omitempty"`
|
||
Steps []AgentStep `json:"steps"`
|
||
Rounds int `json:"rounds"`
|
||
RunID string `json:"runId,omitempty"`
|
||
SessionID string `json:"sessionId,omitempty"`
|
||
Outcome string `json:"runOutcome,omitempty"`
|
||
IncompleteReason string `json:"incompleteReason,omitempty"`
|
||
Resumable bool `json:"resumable,omitempty"`
|
||
StepsBudget int `json:"stepsBudget,omitempty"`
|
||
// Trajectory 含 system/user/assistant.tool_calls/tool 的完整对话（供落盘续跑；JSON 默认省略大字段可选）
|
||
Trajectory []Message `json:"trajectory,omitempty"`
|
||
}
|
||
|
||
// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌。
|
||
// ConfirmDanger 永不由模型自动带上（execOneTool 会剥离 confirm_danger）。
|
||
func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (AgentRes, error) {
|
||
return runAgent(ctx, llm, rt, req, nil)
|
||
}
|
||
|
||
// RunAgentStream 同 RunAgent，经 emit 边跑边推（步骤 / 正文块 / done）。
|
||
func RunAgentStream(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (AgentRes, error) {
|
||
return runAgent(ctx, llm, rt, req, emit)
|
||
}
|
||
|
||
func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (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 == "" && !(req.Continue && len(req.Prior) > 0) {
|
||
return AgentRes{}, fmt.Errorf("prompt required")
|
||
}
|
||
h := DefaultHarness()
|
||
maxRounds := h.ClampRounds(req.MaxRounds)
|
||
|
||
runID := strings.TrimSpace(req.RunID)
|
||
if runID == "" {
|
||
runID = NewRunID()
|
||
}
|
||
sessionID := strings.TrimSpace(req.SessionID)
|
||
if sessionID == "" {
|
||
sessionID = NewSessionID()
|
||
req.SessionID = sessionID
|
||
}
|
||
rt.RunID = runID
|
||
if req.ConnectionID != "" {
|
||
rt.BoundConnectionID = req.ConnectionID
|
||
rt.BoundDatabase = req.Database
|
||
}
|
||
// Agent 路径禁止模型自行带确认令牌
|
||
rt.RejectCallerConfirm = true
|
||
|
||
tools := aitools.OpenAITools()
|
||
var msgs []Message
|
||
if req.Continue && len(req.Prior) > 0 {
|
||
msgs = append([]Message{}, req.Prior...)
|
||
hasSystem := false
|
||
for i, m := range msgs {
|
||
if m.Role == "system" {
|
||
msgs[i].Content = buildAgentSystem(req, rt.Policy)
|
||
hasSystem = true
|
||
}
|
||
}
|
||
if !hasSystem {
|
||
msgs = append([]Message{{Role: "system", Content: buildAgentSystem(req, rt.Policy)}}, msgs...)
|
||
}
|
||
if prompt != "" {
|
||
msgs = append(msgs, Message{Role: "user", Content: prompt})
|
||
}
|
||
msgs = TrimMessages(msgs, h)
|
||
} else {
|
||
msgs = []Message{
|
||
{Role: "system", Content: buildAgentSystem(req, rt.Policy)},
|
||
{Role: "user", Content: prompt},
|
||
}
|
||
}
|
||
|
||
finish := func(steps []AgentStep, rounds int, msgs []Message) AgentRes {
|
||
return AgentRes{
|
||
Steps: steps,
|
||
Rounds: rounds,
|
||
RunID: runID,
|
||
SessionID: sessionID,
|
||
StepsBudget: maxRounds,
|
||
Trajectory: append([]Message{}, msgs...),
|
||
}
|
||
}
|
||
|
||
var steps []AgentStep
|
||
var callHistory []string
|
||
rounds := 0
|
||
for round := 0; round < maxRounds; round++ {
|
||
if err := ctx.Err(); err != nil {
|
||
res := finish(steps, rounds, msgs)
|
||
res.Content = "已取消"
|
||
res.Outcome = OutcomeCancelled
|
||
res.Resumable = len(steps) > 0
|
||
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds, RunID: runID, SessionID: sessionID, Outcome: OutcomeCancelled})
|
||
return res, err
|
||
}
|
||
rounds = round + 1
|
||
if tip := CheckpointHint(rounds); tip != "" {
|
||
msgs = append(msgs, Message{Role: "system", Content: tip})
|
||
}
|
||
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
|
||
if err != nil {
|
||
res := finish(steps, rounds, msgs)
|
||
res.Outcome = OutcomeFailed
|
||
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds, RunID: runID, SessionID: sessionID, Outcome: OutcomeFailed})
|
||
return res, err
|
||
}
|
||
if len(turn.ToolCalls) == 0 {
|
||
content := turn.Content
|
||
if content == "" {
|
||
content = "（模型未返回文本，也未调用工具）"
|
||
}
|
||
for _, part := range chunkDelta(content, 36) {
|
||
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||
}
|
||
res := finish(steps, rounds, msgs)
|
||
res.Content = content
|
||
res.SQL = extractSQL(content)
|
||
res.Outcome = OutcomeCompleted
|
||
emitEvent(emit, AgentEvent{
|
||
Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds,
|
||
RunID: runID, SessionID: sessionID, Outcome: OutcomeCompleted, StepsBudget: maxRounds,
|
||
})
|
||
return res, nil
|
||
}
|
||
|
||
msgs = append(msgs, Message{
|
||
Role: "assistant",
|
||
Content: turn.Content,
|
||
ToolCalls: turn.ToolCalls,
|
||
})
|
||
|
||
awaiting := false
|
||
for _, tc := range turn.ToolCalls {
|
||
name := strings.TrimSpace(tc.Function.Name)
|
||
argsRaw := strings.TrimSpace(tc.Function.Arguments)
|
||
fp := ToolCallFingerprint(name, argsRaw)
|
||
if DetectRepeatedToolCall(callHistory, fp, h.SameCallLimit) {
|
||
step := AgentStep{
|
||
Tool: name,
|
||
Args: truncateRunes(argsRaw, maxStepArgsChars),
|
||
OK: false,
|
||
Summary: "repeated tool call rejected (same args)",
|
||
Error: "repeated tool call rejected (same args)",
|
||
}
|
||
steps = append(steps, step)
|
||
cp := step
|
||
emitEvent(emit, AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID})
|
||
msgs = append(msgs, Message{
|
||
Role: "tool", ToolCallID: tc.ID, Name: name,
|
||
Content: step.Error,
|
||
})
|
||
callHistory = append(callHistory, fp)
|
||
continue
|
||
}
|
||
callHistory = append(callHistory, fp)
|
||
|
||
pending := AgentStep{Tool: name, Args: truncateRunes(argsRaw, maxStepArgsChars), Pending: true}
|
||
emitEvent(emit, AgentEvent{Type: AgentEventStepStart, Step: &pending, RunID: runID})
|
||
|
||
step, toolMsg := execOneTool(ctx, rt, tc)
|
||
steps = append(steps, step)
|
||
cp := step
|
||
ev := AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID}
|
||
if step.Result != nil {
|
||
ev.Result = step.Result
|
||
}
|
||
emitEvent(emit, ev)
|
||
msgs = append(msgs, toolMsg)
|
||
if step.NeedConfirm {
|
||
awaiting = true
|
||
}
|
||
}
|
||
if awaiting {
|
||
content := "需要确认写操作后才能继续。"
|
||
res := finish(steps, rounds, msgs)
|
||
res.Content = content
|
||
res.Outcome = OutcomeAwaitingConfirm
|
||
res.Resumable = true
|
||
emitEvent(emit, AgentEvent{
|
||
Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds,
|
||
RunID: runID, SessionID: sessionID, Outcome: OutcomeAwaitingConfirm, Resumable: true, StepsBudget: maxRounds,
|
||
})
|
||
return res, nil
|
||
}
|
||
msgs = TrimMessages(msgs, h)
|
||
}
|
||
|
||
content := fmt.Sprintf("已达最大工具轮次（%d）。可加预算继续，或缩小问题后重试。", maxRounds)
|
||
for _, part := range chunkDelta(content, 36) {
|
||
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||
}
|
||
res := finish(steps, rounds, msgs)
|
||
res.Content = content
|
||
res.Outcome = OutcomeIncomplete
|
||
res.IncompleteReason = ReasonMaxSteps
|
||
res.Resumable = true
|
||
emitEvent(emit, AgentEvent{
|
||
Type: AgentEventIncomplete, Content: content, Steps: steps, Rounds: rounds,
|
||
RunID: runID, SessionID: sessionID, Outcome: OutcomeIncomplete, IncompleteReason: ReasonMaxSteps,
|
||
Resumable: true, StepsBudget: maxRounds,
|
||
})
|
||
return res, nil
|
||
}
|
||
|
||
func buildAgentSystem(req AgentReq, policy aitools.Policy) string {
|
||
var b strings.Builder
|
||
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具查询用户已白名单的连接。\n")
|
||
b.WriteString("规则：\n")
|
||
b.WriteString("1. 需要查库时必须调用工具；不要假装已执行 SQL。\n")
|
||
if policy.WritesEnabled() {
|
||
b.WriteString("2. 当前允许申请写操作：有写入意图时调用对应工具生成待确认命令；工具要求确认后必须停下，等待用户在界面确认，不能自行确认。\n")
|
||
} else {
|
||
b.WriteString("2. 当前策略为只读：仅能查询，写操作会被拒绝，也不会出现写操作确认。用户需要写入时，说明需先在设置 → MCP 选择使用默认权限，再提交任务并逐次确认。\n")
|
||
}
|
||
b.WriteString("3. 写操作若被拒绝并提示需要确认，向用户说明须在界面确认后执行；不要声称已写入。\n")
|
||
b.WriteString("4. 只使用当前会话绑定的 connection_id 和数据库；调用工具时传入它们，不要切换到其他连接或库。\n")
|
||
b.WriteString("5. 用中文简明回答；可执行 SQL 放在 ```sql 代码块中。\n")
|
||
b.WriteString("6. 工具结果可能截断，注意说明。\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,
|
||
}
|
||
}
|
||
// 永不让模型自行带确认令牌
|
||
delete(args, "confirm_danger")
|
||
delete(args, "confirmDanger")
|
||
|
||
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(),
|
||
}
|
||
}
|
||
|
||
if m, ok := out.(map[string]any); ok {
|
||
if nc, _ := m["needConfirm"].(bool); nc {
|
||
step.OK = false
|
||
step.NeedConfirm = true
|
||
if sql, _ := m["sql"].(string); sql != "" {
|
||
step.ConfirmSQL = sql
|
||
} else if cmd, _ := m["command"].(string); cmd != "" {
|
||
step.ConfirmSQL = cmd
|
||
}
|
||
if bind, _ := m["confirmBind"].(string); bind != "" {
|
||
step.ConfirmBind = bind
|
||
}
|
||
msg, _ := m["message"].(string)
|
||
if msg == "" {
|
||
msg = "requires user confirmation (confirm_danger)"
|
||
}
|
||
step.Error = msg
|
||
step.Summary = msg
|
||
feed := aitools.ResultJSON(out)
|
||
return step, Message{
|
||
Role: "tool",
|
||
ToolCallID: tc.ID,
|
||
Name: name,
|
||
Content: feed,
|
||
}
|
||
}
|
||
if name == aitools.ToolExecuteQuery || name == aitools.ToolExecuteRedisCommand {
|
||
if qr := extractQueryResult(name, args, m); qr != nil {
|
||
step.Result = qr
|
||
}
|
||
}
|
||
}
|
||
|
||
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 extractQueryResult(tool string, args map[string]any, m map[string]any) *AgentQueryResult {
|
||
cols, _ := m["columns"].([]string)
|
||
if cols == nil {
|
||
if raw, ok := m["columns"].([]any); ok {
|
||
cols = make([]string, 0, len(raw))
|
||
for _, c := range raw {
|
||
cols = append(cols, fmt.Sprint(c))
|
||
}
|
||
}
|
||
}
|
||
var rows []map[string]interface{}
|
||
switch v := m["rows"].(type) {
|
||
case []map[string]interface{}:
|
||
rows = v
|
||
case []any:
|
||
for _, item := range v {
|
||
if row, ok := item.(map[string]interface{}); ok {
|
||
rows = append(rows, row)
|
||
} else if row, ok := item.(map[string]any); ok {
|
||
rows = append(rows, row)
|
||
}
|
||
}
|
||
}
|
||
if len(cols) == 0 && len(rows) == 0 {
|
||
return nil
|
||
}
|
||
if len(rows) > maxGridRows {
|
||
rows = rows[:maxGridRows]
|
||
}
|
||
rc := len(rows)
|
||
switch v := m["rowCount"].(type) {
|
||
case int:
|
||
rc = v
|
||
case int64:
|
||
rc = int(v)
|
||
case float64:
|
||
rc = int(v)
|
||
}
|
||
sqlText, _ := args["sql"].(string)
|
||
if sqlText == "" {
|
||
sqlText, _ = args["command"].(string)
|
||
}
|
||
return &AgentQueryResult{
|
||
Columns: cols,
|
||
Rows: rows,
|
||
RowCount: rc,
|
||
SQL: sqlText,
|
||
Tool: tool,
|
||
}
|
||
}
|
||
|
||
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]) + "…"
|
||
}
|