Files
opsxdb/server/internal/ai/agent.go
T
logan 5416e86379 feat(ai): session harness contract — incomplete, run_id, confirm bind
Add docs/16 for session state machine, incomplete/resume, queue semantics,
and agent-loop/harness/MCP gap review. Wire RunOutcome on max-rounds,
run_id through audit, TrimMessages/Redact hooks, ConfirmBindHash on write
confirm path; session file skeleton without HTTP yet. No push.
2026-10-15 13:58:10 +08:00

476 lines
14 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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"`
}
// 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()
}
rt.RunID = runID
// 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 _, m := range msgs {
if m.Role == "system" {
hasSystem = true
break
}
}
if !hasSystem {
msgs = append([]Message{{Role: "system", Content: buildAgentSystem(req)}}, msgs...)
}
if prompt != "" {
msgs = append(msgs, Message{Role: "user", Content: prompt})
}
msgs = TrimMessages(msgs, h)
} else {
msgs = []Message{
{Role: "system", Content: buildAgentSystem(req)},
{Role: "user", Content: prompt},
}
}
baseMeta := func(steps []AgentStep, rounds int) AgentRes {
return AgentRes{
Steps: steps,
Rounds: rounds,
RunID: runID,
SessionID: req.SessionID,
StepsBudget: maxRounds,
}
}
var steps []AgentStep
var callHistory []string
rounds := 0
for round := 0; round < maxRounds; round++ {
if err := ctx.Err(); err != nil {
res := baseMeta(steps, rounds)
res.Content = "已取消"
res.Outcome = OutcomeCancelled
res.Resumable = len(steps) > 0
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds, RunID: runID, Outcome: OutcomeCancelled})
return res, err
}
rounds = round + 1
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
if err != nil {
res := baseMeta(steps, rounds)
res.Outcome = OutcomeFailed
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds, RunID: runID, 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 := baseMeta(steps, rounds)
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, 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 := baseMeta(steps, rounds)
res.Content = content
res.Outcome = OutcomeAwaitingConfirm
res.Resumable = true
emitEvent(emit, AgentEvent{
Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds,
RunID: runID, 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 := baseMeta(steps, rounds)
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, Outcome: OutcomeIncomplete, IncompleteReason: ReasonMaxSteps,
Resumable: true, StepsBudget: maxRounds,
})
return res, 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. 写操作若被拒绝并提示需要确认,向用户说明须在界面确认后执行;不要声称已写入。\n")
b.WriteString("4. 先确认 connection_id(可用 opsxdb_list_connections);用户给出偏好连接时优先用它。\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]) + "…"
}