feat(ai): Agent SSE + Anthropic/responses tools + P1 确认写
- Agent 支持 SSE(步骤边推、正文分块),JSON 非流式 fallback - tool calling 覆盖 openai-chat / anthropic / responses(单测解析回灌) - P1:mcp.json mode=read_write 才可写;需 confirm_danger;Agent 剥离模型自带令牌 - GET /api/ai/mcp-status、POST /api/ai/agent/confirm;审计 source=mcp|agent
This commit is contained in:
1 parent
1d6e89449a
commit
2336c9433c
17 files changed
+1324
-197
No files matched your search
@@ -26,6 +26,7 @@ func main() {
|
|||||||
fmt.Fprintf(os.Stderr, "opsxdb-mcp: load policy: %v\n", err)
|
fmt.Fprintf(os.Stderr, "opsxdb-mcp: load policy: %v\n", err)
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
|
rt.Source = "mcp"
|
||||||
|
|
||||||
s := server.NewMCPServer(
|
s := server.NewMCPServer(
|
||||||
"opsxdb",
|
"opsxdb",
|
||||||
|
|||||||
+150
-24
@@ -11,9 +11,10 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
defaultAgentMaxRounds = 8
|
defaultAgentMaxRounds = 8
|
||||||
maxToolResultFeedChars = 12000
|
maxToolResultFeedChars = 12000
|
||||||
maxStepArgsChars = 200
|
maxStepArgsChars = 200
|
||||||
|
maxGridRows = 200
|
||||||
)
|
)
|
||||||
|
|
||||||
// AgentReq 侧栏 Agent 请求(与 Ask 的 /api/ai/chat 分路径)。
|
// AgentReq 侧栏 Agent 请求(与 Ask 的 /api/ai/chat 分路径)。
|
||||||
@@ -24,16 +25,21 @@ type AgentReq struct {
|
|||||||
Database string `json:"database,omitempty"`
|
Database string `json:"database,omitempty"`
|
||||||
Model string `json:"model,omitempty"`
|
Model string `json:"model,omitempty"`
|
||||||
MaxRounds int `json:"maxRounds,omitempty"`
|
MaxRounds int `json:"maxRounds,omitempty"`
|
||||||
|
Stream bool `json:"stream,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// AgentStep 挂在回复下的步骤时间线一项。
|
// AgentStep 挂在回复下的步骤时间线一项。
|
||||||
type AgentStep struct {
|
type AgentStep struct {
|
||||||
Tool string `json:"tool"`
|
Tool string `json:"tool"`
|
||||||
Args string `json:"args,omitempty"`
|
Args string `json:"args,omitempty"`
|
||||||
OK bool `json:"ok"`
|
OK bool `json:"ok"`
|
||||||
Summary string `json:"summary"`
|
Summary string `json:"summary"`
|
||||||
Error string `json:"error,omitempty"`
|
Error string `json:"error,omitempty"`
|
||||||
Ms int64 `json:"ms"`
|
Ms int64 `json:"ms"`
|
||||||
|
NeedConfirm bool `json:"needConfirm,omitempty"`
|
||||||
|
ConfirmSQL string `json:"confirmSql,omitempty"`
|
||||||
|
Pending bool `json:"pending,omitempty"`
|
||||||
|
Result *AgentQueryResult `json:"result,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// AgentRes Agent 最终回答 + 步骤。
|
// AgentRes Agent 最终回答 + 步骤。
|
||||||
@@ -44,9 +50,18 @@ type AgentRes struct {
|
|||||||
Rounds int `json:"rounds"`
|
Rounds int `json:"rounds"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// RunAgent 多轮:LLM tool_call → aitools.Runtime.Call → 回灌,直到最终文本或轮次耗尽。
|
// RunAgent 多轮:LLM tool_call → aitools.Runtime.Call → 回灌。
|
||||||
// ConfirmDanger 永不自动 true(由 Runtime 保证)。Ask 模式不要走这里。
|
// ConfirmDanger 永不由模型自动带上(execOneTool 会剥离 confirm_danger)。
|
||||||
func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (AgentRes, error) {
|
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 {
|
if rt == nil {
|
||||||
return AgentRes{}, fmt.Errorf("aitools runtime required")
|
return AgentRes{}, fmt.Errorf("aitools runtime required")
|
||||||
}
|
}
|
||||||
@@ -65,6 +80,9 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
|||||||
maxRounds = 16
|
maxRounds = 16
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Agent 路径禁止模型自行带确认令牌
|
||||||
|
rt.RejectCallerConfirm = true
|
||||||
|
|
||||||
tools := aitools.OpenAITools()
|
tools := aitools.OpenAITools()
|
||||||
msgs := []Message{
|
msgs := []Message{
|
||||||
{Role: "system", Content: buildAgentSystem(req)},
|
{Role: "system", Content: buildAgentSystem(req)},
|
||||||
@@ -74,9 +92,15 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
|||||||
var steps []AgentStep
|
var steps []AgentStep
|
||||||
rounds := 0
|
rounds := 0
|
||||||
for round := 0; round < maxRounds; round++ {
|
for round := 0; round < maxRounds; round++ {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
res := AgentRes{Steps: steps, Rounds: rounds, Content: "已取消"}
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds})
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
rounds = round + 1
|
rounds = round + 1
|
||||||
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
|
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds})
|
||||||
return AgentRes{Steps: steps, Rounds: rounds}, err
|
return AgentRes{Steps: steps, Rounds: rounds}, err
|
||||||
}
|
}
|
||||||
if len(turn.ToolCalls) == 0 {
|
if len(turn.ToolCalls) == 0 {
|
||||||
@@ -84,15 +108,19 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
|||||||
if content == "" {
|
if content == "" {
|
||||||
content = "(模型未返回文本,也未调用工具)"
|
content = "(模型未返回文本,也未调用工具)"
|
||||||
}
|
}
|
||||||
return AgentRes{
|
for _, part := range chunkDelta(content, 36) {
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||||||
|
}
|
||||||
|
res := AgentRes{
|
||||||
Content: content,
|
Content: content,
|
||||||
SQL: extractSQL(content),
|
SQL: extractSQL(content),
|
||||||
Steps: steps,
|
Steps: steps,
|
||||||
Rounds: rounds,
|
Rounds: rounds,
|
||||||
}, nil
|
}
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds})
|
||||||
|
return res, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// assistant 消息带 tool_calls
|
|
||||||
msgs = append(msgs, Message{
|
msgs = append(msgs, Message{
|
||||||
Role: "assistant",
|
Role: "assistant",
|
||||||
Content: turn.Content,
|
Content: turn.Content,
|
||||||
@@ -100,28 +128,41 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
|||||||
})
|
})
|
||||||
|
|
||||||
for _, tc := range turn.ToolCalls {
|
for _, tc := range turn.ToolCalls {
|
||||||
|
name := strings.TrimSpace(tc.Function.Name)
|
||||||
|
pending := AgentStep{Tool: name, Args: truncateRunes(strings.TrimSpace(tc.Function.Arguments), maxStepArgsChars), Pending: true}
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventStepStart, Step: &pending})
|
||||||
|
|
||||||
step, toolMsg := execOneTool(ctx, rt, tc)
|
step, toolMsg := execOneTool(ctx, rt, tc)
|
||||||
steps = append(steps, step)
|
steps = append(steps, step)
|
||||||
|
cp := step
|
||||||
|
ev := AgentEvent{Type: AgentEventStep, Step: &cp}
|
||||||
|
if step.Result != nil {
|
||||||
|
ev.Result = step.Result
|
||||||
|
}
|
||||||
|
emitEvent(emit, ev)
|
||||||
msgs = append(msgs, toolMsg)
|
msgs = append(msgs, toolMsg)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return AgentRes{
|
content := fmt.Sprintf("已达最大工具轮次(%d)。请根据已执行步骤自行判断,或缩小问题后重试。", maxRounds)
|
||||||
Content: fmt.Sprintf("已达最大工具轮次(%d)。请根据已执行步骤自行判断,或缩小问题后重试。", maxRounds),
|
for _, part := range chunkDelta(content, 36) {
|
||||||
Steps: steps,
|
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||||||
Rounds: rounds,
|
}
|
||||||
}, nil
|
res := AgentRes{Content: content, Steps: steps, Rounds: rounds}
|
||||||
|
emitEvent(emit, AgentEvent{Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds})
|
||||||
|
return res, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func buildAgentSystem(req AgentReq) string {
|
func buildAgentSystem(req AgentReq) string {
|
||||||
var b strings.Builder
|
var b strings.Builder
|
||||||
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具只读查询用户已白名单的连接。\n")
|
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具查询用户已白名单的连接。\n")
|
||||||
b.WriteString("规则:\n")
|
b.WriteString("规则:\n")
|
||||||
b.WriteString("1. 需要查库时必须调用工具;不要假装已执行 SQL。\n")
|
b.WriteString("1. 需要查库时必须调用工具;不要假装已执行 SQL。\n")
|
||||||
b.WriteString("2. 仅只读:execute_query 只允许 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT;Redis 仅读命令。\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("3. 写操作若被拒绝并提示需要确认,向用户说明须在界面确认后执行;不要声称已写入。\n")
|
||||||
b.WriteString("4. 用中文简明回答;可执行 SQL 放在 ```sql 代码块中。\n")
|
b.WriteString("4. 先确认 connection_id(可用 opsxdb_list_connections);用户给出偏好连接时优先用它。\n")
|
||||||
b.WriteString("5. 工具结果可能截断,注意说明。\n")
|
b.WriteString("5. 用中文简明回答;可执行 SQL 放在 ```sql 代码块中。\n")
|
||||||
|
b.WriteString("6. 工具结果可能截断,注意说明。\n")
|
||||||
if req.ConnectionID != "" {
|
if req.ConnectionID != "" {
|
||||||
b.WriteString(fmt.Sprintf("\n用户当前偏好连接 id: %s", req.ConnectionID))
|
b.WriteString(fmt.Sprintf("\n用户当前偏好连接 id: %s", req.ConnectionID))
|
||||||
if req.Driver != "" {
|
if req.Driver != "" {
|
||||||
@@ -160,6 +201,9 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
|
|||||||
Content: step.Error,
|
Content: step.Error,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// 永不让模型自行带确认令牌
|
||||||
|
delete(args, "confirm_danger")
|
||||||
|
delete(args, "confirmDanger")
|
||||||
|
|
||||||
out, err := rt.Call(ctx, name, args)
|
out, err := rt.Call(ctx, name, args)
|
||||||
step.Ms = time.Since(start).Milliseconds()
|
step.Ms = time.Since(start).Milliseconds()
|
||||||
@@ -174,6 +218,37 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
|
|||||||
Content: err.Error(),
|
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
|
||||||
|
}
|
||||||
|
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.OK = true
|
||||||
step.Summary = summarizeToolResult(name, out)
|
step.Summary = summarizeToolResult(name, out)
|
||||||
feed := aitools.ResultJSON(out)
|
feed := aitools.ResultJSON(out)
|
||||||
@@ -188,6 +263,57 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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 {
|
func summarizeToolResult(tool string, out any) string {
|
||||||
m, ok := out.(map[string]any)
|
m, ok := out.(map[string]any)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AgentEventType SSE / 回调事件类型。
|
||||||
|
const (
|
||||||
|
AgentEventStepStart = "step_start"
|
||||||
|
AgentEventStep = "step"
|
||||||
|
AgentEventDelta = "delta"
|
||||||
|
AgentEventDone = "done"
|
||||||
|
AgentEventError = "error"
|
||||||
|
)
|
||||||
|
|
||||||
|
// AgentQueryResult 可供「在结果区打开」的行列结果。
|
||||||
|
type AgentQueryResult struct {
|
||||||
|
Columns []string `json:"columns"`
|
||||||
|
Rows []map[string]interface{} `json:"rows"`
|
||||||
|
RowCount int `json:"rowCount"`
|
||||||
|
SQL string `json:"sql,omitempty"`
|
||||||
|
Tool string `json:"tool,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// AgentEvent 流式推送的一项。
|
||||||
|
type AgentEvent struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Step *AgentStep `json:"step,omitempty"`
|
||||||
|
Delta string `json:"delta,omitempty"`
|
||||||
|
Content string `json:"content,omitempty"`
|
||||||
|
SQL string `json:"sql,omitempty"`
|
||||||
|
Steps []AgentStep `json:"steps,omitempty"`
|
||||||
|
Rounds int `json:"rounds,omitempty"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
Result *AgentQueryResult `json:"result,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// AgentEmitter 可选;nil 时仅返回最终 AgentRes(非流式 fallback)。
|
||||||
|
type AgentEmitter func(AgentEvent)
|
||||||
|
|
||||||
|
func emitEvent(emit AgentEmitter, ev AgentEvent) {
|
||||||
|
if emit != nil {
|
||||||
|
emit(ev)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func chunkDelta(s string, size int) []string {
|
||||||
|
if s == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if size <= 0 {
|
||||||
|
size = 32
|
||||||
|
}
|
||||||
|
r := []rune(s)
|
||||||
|
var out []string
|
||||||
|
for i := 0; i < len(r); i += size {
|
||||||
|
j := i + size
|
||||||
|
if j > len(r) {
|
||||||
|
j = len(r)
|
||||||
|
}
|
||||||
|
out = append(out, string(r[i:j]))
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteAgentSSE 写一条 SSE 事件并 Flush。
|
||||||
|
func WriteAgentSSE(w http.ResponseWriter, flusher http.Flusher, ev AgentEvent) error {
|
||||||
|
b, err := json.Marshal(ev)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, b); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
flusher.Flush()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -202,7 +202,7 @@ func TestParseToolsTurn(t *testing.T) {
|
|||||||
},
|
},
|
||||||
}},
|
}},
|
||||||
})
|
})
|
||||||
turn, err := parseToolsTurn(raw, 200)
|
turn, err := parseToolsTurnChat(raw, 200)
|
||||||
if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
|
if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
|
||||||
t.Fatalf("parse: %+v %v", turn, err)
|
t.Fatalf("parse: %+v %v", turn, err)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,100 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestValidateEndpoint(t *testing.T) {
|
||||||
|
if err := validateEndpoint("https://api.deepseek.com/v1", false); err != nil {
|
||||||
|
t.Fatalf("https 应直接放行: %v", err)
|
||||||
|
}
|
||||||
|
if err := validateEndpoint("http://127.0.0.1:11434/v1", false); err == nil {
|
||||||
|
t.Fatal("未开启 allowInsecure 时 http 应拒绝")
|
||||||
|
}
|
||||||
|
if err := validateEndpoint("http://127.0.0.1:11434/v1", true); err != nil {
|
||||||
|
t.Fatalf("allowInsecure 下本地 http 应放行: %v", err)
|
||||||
|
}
|
||||||
|
if err := validateEndpoint("ftp://example.com", true); err == nil {
|
||||||
|
t.Fatal("非 HTTP/HTTPS scheme 应拒绝")
|
||||||
|
}
|
||||||
|
if err := validateEndpoint("not a url", false); err == nil {
|
||||||
|
t.Fatal("非法 URL 应拒绝")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEndpointJoin(t *testing.T) {
|
||||||
|
if got := endpoint("https://api.example.com/v1/", "chat/completions"); got != "https://api.example.com/v1/chat/completions" {
|
||||||
|
t.Fatalf("拼接错误: %s", got)
|
||||||
|
}
|
||||||
|
if got := endpoint("https://api.example.com/v1/responses", "responses"); got != "https://api.example.com/v1/responses" {
|
||||||
|
t.Fatalf("已含 suffix 不应重复: %s", got)
|
||||||
|
}
|
||||||
|
if got := anthropicEndpoint("https://api.anthropic.com"); got != "https://api.anthropic.com/v1/messages" {
|
||||||
|
t.Fatalf("anthropic 无 /v1 应补: %s", got)
|
||||||
|
}
|
||||||
|
if got := anthropicEndpoint("https://api.anthropic.com/v1"); got != "https://api.anthropic.com/v1/messages" {
|
||||||
|
t.Fatalf("anthropic 已含 /v1 不应重复: %s", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewRemoteProviderInfersProtocolAndAuth(t *testing.T) {
|
||||||
|
legacy := NewRemoteProvider(Config{Provider: "openai", BaseURL: "https://api.openai.com/v1"})
|
||||||
|
if legacy.protocol != ProtoChat || legacy.authStyle != AuthBearer {
|
||||||
|
t.Fatalf("旧配置应推断 openai-chat/bearer, got %s/%s", legacy.protocol, legacy.authStyle)
|
||||||
|
}
|
||||||
|
anthropic := NewRemoteProvider(Config{Provider: "anthropic", BaseURL: "https://api.anthropic.com"})
|
||||||
|
if anthropic.protocol != ProtoAnthropic || anthropic.authStyle != AuthXAPIKey {
|
||||||
|
t.Fatalf("anthropic 预设应推断原生协议/x-api-key, got %s/%s", anthropic.protocol, anthropic.authStyle)
|
||||||
|
}
|
||||||
|
local := NewRemoteProvider(Config{Provider: "ollama", Protocol: ProtoChat, AuthStyle: AuthNone, BaseURL: "http://127.0.0.1:11434/v1", AllowInsecureHTTP: true})
|
||||||
|
if local.authStyle != AuthNone {
|
||||||
|
t.Fatalf("本地免鉴权应保留 none, got %s", local.authStyle)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseContentByProtocol(t *testing.T) {
|
||||||
|
chatRaw, _ := json.Marshal(map[string]any{
|
||||||
|
"choices": []map[string]any{{"message": map[string]string{"role": "assistant", "content": " ```sql\nSELECT 1;\n``` "}}},
|
||||||
|
})
|
||||||
|
got, err := parseContent(ProtoChat, chatRaw, http.StatusOK)
|
||||||
|
if err != nil || strings.TrimSpace(got) != "```sql\nSELECT 1;\n```" {
|
||||||
|
t.Fatalf("chat 解析失败: %q, %v", got, err)
|
||||||
|
}
|
||||||
|
if s := extractSQL(got); s != "SELECT 1;" {
|
||||||
|
t.Fatalf("sql 代码块提取错误: %q", s)
|
||||||
|
}
|
||||||
|
|
||||||
|
anthropicRaw, _ := json.Marshal(map[string]any{
|
||||||
|
"content": []map[string]any{{"type": "text", "text": "你好"}, {"type": "tool_use", "text": "忽略"}},
|
||||||
|
})
|
||||||
|
got, err = parseContent(ProtoAnthropic, anthropicRaw, http.StatusOK)
|
||||||
|
if err != nil || got != "你好" {
|
||||||
|
t.Fatalf("anthropic 解析失败: %q, %v", got, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
responsesRaw, _ := json.Marshal(map[string]any{
|
||||||
|
"output": []map[string]any{{"content": []map[string]string{{"text": "SELECT "}, {"text": "2;"}}}},
|
||||||
|
})
|
||||||
|
got, err = parseContent(ProtoResponses, responsesRaw, http.StatusOK)
|
||||||
|
if err != nil || got != "SELECT 2;" {
|
||||||
|
t.Fatalf("responses 解析失败: %q, %v", got, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
errRaw, _ := json.Marshal(map[string]any{"error": map[string]string{"message": "invalid api key"}})
|
||||||
|
if _, err := parseContent(ProtoChat, errRaw, http.StatusUnauthorized); err == nil || !strings.Contains(err.Error(), "invalid api key") {
|
||||||
|
t.Fatalf("非 2xx 应提取服务端错误信息: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildMessagesCoversActions(t *testing.T) {
|
||||||
|
msgs := BuildMessages(ChatReq{Action: "optimize", SQL: "SELECT * FROM t", Driver: "postgresql", Database: "db"})
|
||||||
|
if len(msgs) != 2 || msgs[0].Role != "system" || msgs[1].Role != "user" {
|
||||||
|
t.Fatalf("应返回 [system, user], got %+v", msgs)
|
||||||
|
}
|
||||||
|
if !strings.Contains(msgs[1].Content, "SELECT * FROM t") {
|
||||||
|
t.Fatal("user 消息应包含 SQL")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -7,12 +7,16 @@ import (
|
|||||||
"strings"
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Config AI 配置;provider: stub|openai(OpenAI 兼容协议)。
|
// Config AI 配置;provider 为预设 id(stub|openai|deepseek|anthropic|ollama|custom…),
|
||||||
|
// protocol/authStyle/allowInsecureHTTP 决定远端调用行为(见 remote.go)。
|
||||||
type Config struct {
|
type Config struct {
|
||||||
Provider string `json:"provider"`
|
Provider string `json:"provider"`
|
||||||
BaseURL string `json:"baseURL"`
|
BaseURL string `json:"baseURL"`
|
||||||
APIKey string `json:"apiKey,omitempty"`
|
APIKey string `json:"apiKey,omitempty"`
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
|
Protocol string `json:"protocol,omitempty"`
|
||||||
|
AuthStyle string `json:"authStyle,omitempty"`
|
||||||
|
AllowInsecureHTTP bool `json:"allowInsecureHttp,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func configPath() (string, error) {
|
func configPath() (string, error) {
|
||||||
@@ -47,6 +51,13 @@ func LoadConfig() Config {
|
|||||||
if f.Model != "" {
|
if f.Model != "" {
|
||||||
cfg.Model = f.Model
|
cfg.Model = f.Model
|
||||||
}
|
}
|
||||||
|
if f.Protocol != "" {
|
||||||
|
cfg.Protocol = f.Protocol
|
||||||
|
}
|
||||||
|
if f.AuthStyle != "" {
|
||||||
|
cfg.AuthStyle = f.AuthStyle
|
||||||
|
}
|
||||||
|
cfg.AllowInsecureHTTP = f.AllowInsecureHTTP
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,117 +0,0 @@
|
|||||||
package ai
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bytes"
|
|
||||||
"context"
|
|
||||||
"encoding/json"
|
|
||||||
"fmt"
|
|
||||||
"io"
|
|
||||||
"net/http"
|
|
||||||
"regexp"
|
|
||||||
"strings"
|
|
||||||
"time"
|
|
||||||
)
|
|
||||||
|
|
||||||
const (
|
|
||||||
defaultOpenAIBaseURL = "https://api.openai.com/v1"
|
|
||||||
defaultOpenAIModel = "gpt-4o-mini"
|
|
||||||
)
|
|
||||||
|
|
||||||
// OpenAIProvider OpenAI 兼容协议实现(DeepSeek/通义/本地 vLLM 同协议可用)。
|
|
||||||
type OpenAIProvider struct {
|
|
||||||
baseURL string
|
|
||||||
apiKey string
|
|
||||||
model string
|
|
||||||
hc *http.Client
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewOpenAIProvider(cfg Config) *OpenAIProvider {
|
|
||||||
base := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/")
|
|
||||||
if base == "" {
|
|
||||||
base = defaultOpenAIBaseURL
|
|
||||||
}
|
|
||||||
model := strings.TrimSpace(cfg.Model)
|
|
||||||
if model == "" {
|
|
||||||
model = defaultOpenAIModel
|
|
||||||
}
|
|
||||||
return &OpenAIProvider{
|
|
||||||
baseURL: base,
|
|
||||||
apiKey: cfg.APIKey,
|
|
||||||
model: model,
|
|
||||||
hc: &http.Client{Timeout: 60 * time.Second},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *OpenAIProvider) Name() string { return "openai" }
|
|
||||||
|
|
||||||
type chatCompletionReq struct {
|
|
||||||
Model string `json:"model"`
|
|
||||||
Messages []Message `json:"messages"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type chatCompletionRes struct {
|
|
||||||
Choices []struct {
|
|
||||||
Message Message `json:"message"`
|
|
||||||
} `json:"choices"`
|
|
||||||
Error *struct {
|
|
||||||
Message string `json:"message"`
|
|
||||||
} `json:"error"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// Complete 单次非流式调用;ctx 全透传,HTTP 层 60s 超时兜底。
|
|
||||||
func (p *OpenAIProvider) Complete(ctx context.Context, req ChatReq) (ChatRes, error) {
|
|
||||||
model := p.model
|
|
||||||
if strings.TrimSpace(req.Model) != "" {
|
|
||||||
model = strings.TrimSpace(req.Model)
|
|
||||||
}
|
|
||||||
payload := chatCompletionReq{
|
|
||||||
Model: model,
|
|
||||||
Messages: BuildMessages(req),
|
|
||||||
}
|
|
||||||
body, err := json.Marshal(payload)
|
|
||||||
if err != nil {
|
|
||||||
return ChatRes{}, err
|
|
||||||
}
|
|
||||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body))
|
|
||||||
if err != nil {
|
|
||||||
return ChatRes{}, err
|
|
||||||
}
|
|
||||||
httpReq.Header.Set("Content-Type", "application/json")
|
|
||||||
httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
|
|
||||||
|
|
||||||
resp, err := p.hc.Do(httpReq)
|
|
||||||
if err != nil {
|
|
||||||
return ChatRes{}, fmt.Errorf("AI 请求失败: %w", err)
|
|
||||||
}
|
|
||||||
defer resp.Body.Close()
|
|
||||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
|
||||||
if err != nil {
|
|
||||||
return ChatRes{}, fmt.Errorf("读取 AI 响应失败: %w", err)
|
|
||||||
}
|
|
||||||
var out chatCompletionRes
|
|
||||||
if err := json.Unmarshal(raw, &out); err != nil {
|
|
||||||
return ChatRes{}, fmt.Errorf("解析 AI 响应失败(HTTP %d)", resp.StatusCode)
|
|
||||||
}
|
|
||||||
if out.Error != nil {
|
|
||||||
return ChatRes{}, fmt.Errorf("AI 服务返回错误: %s", out.Error.Message)
|
|
||||||
}
|
|
||||||
if resp.StatusCode != http.StatusOK {
|
|
||||||
return ChatRes{}, fmt.Errorf("AI 服务 HTTP %d", resp.StatusCode)
|
|
||||||
}
|
|
||||||
if len(out.Choices) == 0 {
|
|
||||||
return ChatRes{}, fmt.Errorf("AI 服务未返回候选结果")
|
|
||||||
}
|
|
||||||
content := strings.TrimSpace(out.Choices[0].Message.Content)
|
|
||||||
return ChatRes{Content: content, SQL: extractSQL(content)}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
var sqlFence = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```")
|
|
||||||
|
|
||||||
// extractSQL 从模型回复中提取 ```sql 代码块内容。
|
|
||||||
func extractSQL(content string) string {
|
|
||||||
m := sqlFence.FindStringSubmatch(content)
|
|
||||||
if len(m) < 2 {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
return strings.TrimSpace(m[1])
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,292 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultOpenAIBaseURL = "https://api.openai.com/v1"
|
||||||
|
defaultOpenAIModel = "gpt-4o-mini"
|
||||||
|
remoteMaxTokens = 2048
|
||||||
|
)
|
||||||
|
|
||||||
|
// 协议:openai-chat 覆盖绝大多数兼容服务;openai-responses 为 OpenAI 官方;anthropic 为 Claude。
|
||||||
|
const (
|
||||||
|
ProtoChat = "openai-chat"
|
||||||
|
ProtoResponses = "openai-responses"
|
||||||
|
ProtoAnthropic = "anthropic"
|
||||||
|
)
|
||||||
|
|
||||||
|
// 鉴权风格:bearer = Authorization: Bearer;x-api-key = Anthropic 风格头;none = 本地服务免鉴权。
|
||||||
|
const (
|
||||||
|
AuthBearer = "bearer"
|
||||||
|
AuthXAPIKey = "x-api-key"
|
||||||
|
AuthNone = "none"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RemoteProvider 远端 HTTP 模型实现,按协议分派 chat/responses/anthropic(对齐 opsxshell AiService)。
|
||||||
|
type RemoteProvider struct {
|
||||||
|
protocol string
|
||||||
|
authStyle string
|
||||||
|
baseURL string
|
||||||
|
apiKey string
|
||||||
|
model string
|
||||||
|
allowHTTP bool
|
||||||
|
hc *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewRemoteProvider 按配置构造;protocol 缺省按 provider 推断(anthropic → 原生协议)。
|
||||||
|
func NewRemoteProvider(cfg Config) *RemoteProvider {
|
||||||
|
protocol := strings.TrimSpace(cfg.Protocol)
|
||||||
|
if protocol == "" {
|
||||||
|
protocol = ProtoChat
|
||||||
|
if strings.EqualFold(cfg.Provider, ProtoAnthropic) {
|
||||||
|
protocol = ProtoAnthropic
|
||||||
|
}
|
||||||
|
}
|
||||||
|
authStyle := strings.TrimSpace(cfg.AuthStyle)
|
||||||
|
if authStyle == "" {
|
||||||
|
authStyle = AuthBearer
|
||||||
|
if protocol == ProtoAnthropic {
|
||||||
|
authStyle = AuthXAPIKey
|
||||||
|
}
|
||||||
|
}
|
||||||
|
base := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/")
|
||||||
|
if base == "" {
|
||||||
|
base = defaultOpenAIBaseURL
|
||||||
|
}
|
||||||
|
model := strings.TrimSpace(cfg.Model)
|
||||||
|
if model == "" {
|
||||||
|
model = defaultOpenAIModel
|
||||||
|
}
|
||||||
|
return &RemoteProvider{
|
||||||
|
protocol: protocol,
|
||||||
|
authStyle: authStyle,
|
||||||
|
baseURL: base,
|
||||||
|
apiKey: cfg.APIKey,
|
||||||
|
model: model,
|
||||||
|
allowHTTP: cfg.AllowInsecureHTTP,
|
||||||
|
hc: &http.Client{Timeout: 120 * time.Second},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *RemoteProvider) Name() string { return p.protocol }
|
||||||
|
|
||||||
|
// validateEndpoint 默认只放行 https;本地/内网服务(Ollama、vLLM)需显式 AllowInsecureHTTP。
|
||||||
|
func validateEndpoint(baseURL string, allowInsecure bool) error {
|
||||||
|
u, err := url.Parse(strings.TrimSpace(baseURL))
|
||||||
|
if err != nil || u.Host == "" {
|
||||||
|
return fmt.Errorf("Base URL 格式无效")
|
||||||
|
}
|
||||||
|
switch u.Scheme {
|
||||||
|
case "https":
|
||||||
|
return nil
|
||||||
|
case "http":
|
||||||
|
if allowInsecure {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("非 HTTPS Base URL 已拒绝;本地/内网服务请在设置中开启「允许 HTTP」")
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("Base URL 仅支持 HTTP/HTTPS")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// endpoint 拼接 basePath + suffix,已含 suffix 时不重复。
|
||||||
|
func endpoint(base, suffix string) string {
|
||||||
|
b := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||||
|
s := strings.TrimLeft(suffix, "/")
|
||||||
|
if strings.HasSuffix(b, "/"+s) {
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
return b + "/" + s
|
||||||
|
}
|
||||||
|
|
||||||
|
// anthropicEndpoint base 未含 /v1 时补 /v1/messages,否则直接 /messages。
|
||||||
|
func anthropicEndpoint(base string) string {
|
||||||
|
b := strings.TrimRight(strings.TrimSpace(base), "/")
|
||||||
|
if strings.HasSuffix(b, "/v1") {
|
||||||
|
return endpoint(b, "messages")
|
||||||
|
}
|
||||||
|
return endpoint(b, "v1/messages")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Complete 单次非流式调用;req.Model 可临时覆盖配置模型。
|
||||||
|
func (p *RemoteProvider) Complete(ctx context.Context, req ChatReq) (ChatRes, error) {
|
||||||
|
if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil {
|
||||||
|
return ChatRes{}, err
|
||||||
|
}
|
||||||
|
model := p.model
|
||||||
|
if strings.TrimSpace(req.Model) != "" {
|
||||||
|
model = strings.TrimSpace(req.Model)
|
||||||
|
}
|
||||||
|
msgs := BuildMessages(req)
|
||||||
|
var (
|
||||||
|
urlStr string
|
||||||
|
payload any
|
||||||
|
)
|
||||||
|
switch p.protocol {
|
||||||
|
case ProtoAnthropic:
|
||||||
|
urlStr = anthropicEndpoint(p.baseURL)
|
||||||
|
payload = map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"system": msgs[0].Content,
|
||||||
|
"messages": []Message{{Role: "user", Content: msgs[1].Content}},
|
||||||
|
"max_tokens": remoteMaxTokens,
|
||||||
|
"temperature": 0.2,
|
||||||
|
}
|
||||||
|
case ProtoResponses:
|
||||||
|
urlStr = endpoint(p.baseURL, "responses")
|
||||||
|
payload = map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"instructions": msgs[0].Content,
|
||||||
|
"input": msgs[1].Content,
|
||||||
|
"max_output_tokens": remoteMaxTokens,
|
||||||
|
}
|
||||||
|
default: // openai-chat 及未知值按兼容协议处理
|
||||||
|
urlStr = endpoint(p.baseURL, "chat/completions")
|
||||||
|
payload = map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"messages": msgs,
|
||||||
|
"max_tokens": remoteMaxTokens,
|
||||||
|
"temperature": 0.2,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
body, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return ChatRes{}, err
|
||||||
|
}
|
||||||
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
return ChatRes{}, 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 ChatRes{}, fmt.Errorf("AI 请求失败: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||||
|
if err != nil {
|
||||||
|
return ChatRes{}, fmt.Errorf("读取 AI 响应失败: %w", err)
|
||||||
|
}
|
||||||
|
content, err := parseContent(p.protocol, raw, resp.StatusCode)
|
||||||
|
if err != nil {
|
||||||
|
return ChatRes{}, err
|
||||||
|
}
|
||||||
|
return ChatRes{Content: content, SQL: extractSQL(content)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// parseContent 按协议解析响应体;非 2xx 时尽量提取服务端错误信息。
|
||||||
|
func parseContent(protocol string, raw []byte, status int) (string, 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 "", fmt.Errorf("AI 服务 HTTP %d: %s", status, msg)
|
||||||
|
}
|
||||||
|
switch protocol {
|
||||||
|
case ProtoAnthropic:
|
||||||
|
var out struct {
|
||||||
|
Content []struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"content"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status)
|
||||||
|
}
|
||||||
|
var sb strings.Builder
|
||||||
|
for _, c := range out.Content {
|
||||||
|
if c.Type == "text" {
|
||||||
|
sb.WriteString(c.Text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if sb.Len() == 0 {
|
||||||
|
return "", fmt.Errorf("AI 服务未返回文本内容")
|
||||||
|
}
|
||||||
|
return sb.String(), nil
|
||||||
|
case ProtoResponses:
|
||||||
|
var out struct {
|
||||||
|
OutputText string `json:"output_text"`
|
||||||
|
Output []struct {
|
||||||
|
Content []struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"content"`
|
||||||
|
} `json:"output"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status)
|
||||||
|
}
|
||||||
|
if out.OutputText != "" {
|
||||||
|
return out.OutputText, nil
|
||||||
|
}
|
||||||
|
var sb strings.Builder
|
||||||
|
for _, o := range out.Output {
|
||||||
|
for _, c := range o.Content {
|
||||||
|
sb.WriteString(c.Text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if sb.Len() == 0 {
|
||||||
|
return "", fmt.Errorf("AI 服务未返回文本内容")
|
||||||
|
}
|
||||||
|
return sb.String(), nil
|
||||||
|
default:
|
||||||
|
var out struct {
|
||||||
|
Choices []struct {
|
||||||
|
Message Message `json:"message"`
|
||||||
|
} `json:"choices"`
|
||||||
|
Error *struct {
|
||||||
|
Message string `json:"message"`
|
||||||
|
} `json:"error"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return "", fmt.Errorf("解析 AI 响应失败(HTTP %d)", status)
|
||||||
|
}
|
||||||
|
if out.Error != nil {
|
||||||
|
return "", fmt.Errorf("AI 服务返回错误: %s", out.Error.Message)
|
||||||
|
}
|
||||||
|
if len(out.Choices) == 0 {
|
||||||
|
return "", fmt.Errorf("AI 服务未返回候选结果")
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(out.Choices[0].Message.Content), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var sqlFence = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```")
|
||||||
|
|
||||||
|
// extractSQL 从模型回复中提取 ```sql 代码块内容。
|
||||||
|
func extractSQL(content string) string {
|
||||||
|
m := sqlFence.FindStringSubmatch(content)
|
||||||
|
if len(m) < 2 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(m[1])
|
||||||
|
}
|
||||||
@@ -24,35 +24,33 @@ type ToolChat interface {
|
|||||||
ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error)
|
ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ChatWithTools openai-chat 兼容协议的 tools 调用(DeepSeek / OpenAI / 多数兼容网关)。
|
// ChatWithTools 按协议分派:openai-chat / anthropic messages / openai-responses。
|
||||||
// anthropic / responses 本轮不支持 Agent tool calling,返回明确错误。
|
|
||||||
func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
|
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 {
|
if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil {
|
||||||
return ToolsTurn{}, err
|
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)
|
m := strings.TrimSpace(model)
|
||||||
if m == "" {
|
if m == "" {
|
||||||
m = p.model
|
m = p.model
|
||||||
}
|
}
|
||||||
urlStr := endpoint(p.baseURL, "chat/completions")
|
switch p.protocol {
|
||||||
payload := map[string]any{
|
case ProtoAnthropic:
|
||||||
"model": m,
|
return p.chatToolsAnthropic(ctx, m, messages, tools)
|
||||||
"messages": messages,
|
case ProtoResponses:
|
||||||
"tools": tools,
|
return p.chatToolsResponses(ctx, m, messages, tools)
|
||||||
"tool_choice": "auto",
|
default:
|
||||||
"max_tokens": remoteMaxTokens,
|
return p.chatToolsOpenAI(ctx, m, messages, tools)
|
||||||
"temperature": 0.2,
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *RemoteProvider) doJSON(ctx context.Context, urlStr string, payload any) ([]byte, int, error) {
|
||||||
body, err := json.Marshal(payload)
|
body, err := json.Marshal(payload)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return ToolsTurn{}, err
|
return nil, 0, err
|
||||||
}
|
}
|
||||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body))
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return ToolsTurn{}, err
|
return nil, 0, err
|
||||||
}
|
}
|
||||||
httpReq.Header.Set("Content-Type", "application/json")
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
switch p.authStyle {
|
switch p.authStyle {
|
||||||
@@ -65,17 +63,174 @@ func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messag
|
|||||||
}
|
}
|
||||||
resp, err := p.hc.Do(httpReq)
|
resp, err := p.hc.Do(httpReq)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return ToolsTurn{}, fmt.Errorf("AI 请求失败: %w", err)
|
return nil, 0, fmt.Errorf("AI 请求失败: %w", err)
|
||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return ToolsTurn{}, fmt.Errorf("读取 AI 响应失败: %w", err)
|
return nil, resp.StatusCode, fmt.Errorf("读取 AI 响应失败: %w", err)
|
||||||
}
|
}
|
||||||
return parseToolsTurn(raw, resp.StatusCode)
|
return raw, resp.StatusCode, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
|
func (p *RemoteProvider) chatToolsOpenAI(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
|
||||||
|
urlStr := endpoint(p.baseURL, "chat/completions")
|
||||||
|
payload := map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"messages": messages,
|
||||||
|
"tools": tools,
|
||||||
|
"tool_choice": "auto",
|
||||||
|
"max_tokens": remoteMaxTokens,
|
||||||
|
"temperature": 0.2,
|
||||||
|
}
|
||||||
|
raw, status, err := p.doJSON(ctx, urlStr, payload)
|
||||||
|
if err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
|
}
|
||||||
|
return parseToolsTurnChat(raw, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *RemoteProvider) chatToolsAnthropic(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
|
||||||
|
urlStr := anthropicEndpoint(p.baseURL)
|
||||||
|
system, anthMsgs := splitAnthropicMessages(messages)
|
||||||
|
payload := map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"max_tokens": remoteMaxTokens,
|
||||||
|
"temperature": 0.2,
|
||||||
|
"tools": openaiToolsToAnthropic(tools),
|
||||||
|
"messages": anthMsgs,
|
||||||
|
}
|
||||||
|
if system != "" {
|
||||||
|
payload["system"] = system
|
||||||
|
}
|
||||||
|
raw, status, err := p.doJSON(ctx, urlStr, payload)
|
||||||
|
if err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
|
}
|
||||||
|
return parseToolsTurnAnthropic(raw, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *RemoteProvider) chatToolsResponses(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
|
||||||
|
urlStr := endpoint(p.baseURL, "responses")
|
||||||
|
instructions, input := splitResponsesInput(messages)
|
||||||
|
payload := map[string]any{
|
||||||
|
"model": model,
|
||||||
|
"instructions": instructions,
|
||||||
|
"input": input,
|
||||||
|
"tools": openaiToolsToResponses(tools),
|
||||||
|
"max_output_tokens": remoteMaxTokens,
|
||||||
|
}
|
||||||
|
raw, status, err := p.doJSON(ctx, urlStr, payload)
|
||||||
|
if err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
|
}
|
||||||
|
return parseToolsTurnResponses(raw, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitAnthropicMessages(messages []Message) (system string, out []map[string]any) {
|
||||||
|
for _, m := range messages {
|
||||||
|
switch m.Role {
|
||||||
|
case "system":
|
||||||
|
if system == "" {
|
||||||
|
system = m.Content
|
||||||
|
} else {
|
||||||
|
system += "\n" + m.Content
|
||||||
|
}
|
||||||
|
case "assistant":
|
||||||
|
content := []map[string]any{}
|
||||||
|
if strings.TrimSpace(m.Content) != "" {
|
||||||
|
content = append(content, map[string]any{"type": "text", "text": m.Content})
|
||||||
|
}
|
||||||
|
for _, tc := range m.ToolCalls {
|
||||||
|
var input any = map[string]any{}
|
||||||
|
_ = json.Unmarshal([]byte(tc.Function.Arguments), &input)
|
||||||
|
content = append(content, map[string]any{
|
||||||
|
"type": "tool_use",
|
||||||
|
"id": tc.ID,
|
||||||
|
"name": tc.Function.Name,
|
||||||
|
"input": input,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
out = append(out, map[string]any{"role": "assistant", "content": content})
|
||||||
|
case "tool":
|
||||||
|
out = append(out, map[string]any{
|
||||||
|
"role": "user",
|
||||||
|
"content": []map[string]any{{
|
||||||
|
"type": "tool_result",
|
||||||
|
"tool_use_id": m.ToolCallID,
|
||||||
|
"content": m.Content,
|
||||||
|
}},
|
||||||
|
})
|
||||||
|
default: // user
|
||||||
|
out = append(out, map[string]any{"role": "user", "content": m.Content})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return system, out
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitResponsesInput(messages []Message) (instructions string, input []map[string]any) {
|
||||||
|
for _, m := range messages {
|
||||||
|
switch m.Role {
|
||||||
|
case "system":
|
||||||
|
if instructions == "" {
|
||||||
|
instructions = m.Content
|
||||||
|
} else {
|
||||||
|
instructions += "\n" + m.Content
|
||||||
|
}
|
||||||
|
case "assistant":
|
||||||
|
item := map[string]any{"role": "assistant", "content": m.Content}
|
||||||
|
if len(m.ToolCalls) > 0 {
|
||||||
|
calls := make([]map[string]any, 0, len(m.ToolCalls))
|
||||||
|
for _, tc := range m.ToolCalls {
|
||||||
|
calls = append(calls, map[string]any{
|
||||||
|
"id": tc.ID,
|
||||||
|
"type": "function_call",
|
||||||
|
"name": tc.Function.Name,
|
||||||
|
"arguments": tc.Function.Arguments,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
item["tool_calls"] = calls
|
||||||
|
}
|
||||||
|
input = append(input, item)
|
||||||
|
case "tool":
|
||||||
|
input = append(input, map[string]any{
|
||||||
|
"type": "function_call_output",
|
||||||
|
"call_id": m.ToolCallID,
|
||||||
|
"output": m.Content,
|
||||||
|
})
|
||||||
|
default:
|
||||||
|
input = append(input, map[string]any{"role": "user", "content": m.Content})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return instructions, input
|
||||||
|
}
|
||||||
|
|
||||||
|
func openaiToolsToAnthropic(tools []aitools.OpenAITool) []map[string]any {
|
||||||
|
out := make([]map[string]any, 0, len(tools))
|
||||||
|
for _, t := range tools {
|
||||||
|
out = append(out, map[string]any{
|
||||||
|
"name": t.Function.Name,
|
||||||
|
"description": t.Function.Description,
|
||||||
|
"input_schema": t.Function.Parameters,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func openaiToolsToResponses(tools []aitools.OpenAITool) []map[string]any {
|
||||||
|
out := make([]map[string]any, 0, len(tools))
|
||||||
|
for _, t := range tools {
|
||||||
|
out = append(out, map[string]any{
|
||||||
|
"type": "function",
|
||||||
|
"name": t.Function.Name,
|
||||||
|
"description": t.Function.Description,
|
||||||
|
"parameters": t.Function.Parameters,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func httpErrMessage(raw []byte, status int) error {
|
||||||
var errBody struct {
|
var errBody struct {
|
||||||
Error *struct {
|
Error *struct {
|
||||||
Message string `json:"message"`
|
Message string `json:"message"`
|
||||||
@@ -90,7 +245,14 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
|
|||||||
} else if errBody.Message != "" {
|
} else if errBody.Message != "" {
|
||||||
msg = errBody.Message
|
msg = errBody.Message
|
||||||
}
|
}
|
||||||
return ToolsTurn{}, fmt.Errorf("AI 服务 HTTP %d: %s", status, msg)
|
return fmt.Errorf("AI 服务 HTTP %d: %s", status, msg)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseToolsTurnChat(raw []byte, status int) (ToolsTurn, error) {
|
||||||
|
if err := httpErrMessage(raw, status); err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
}
|
}
|
||||||
var out struct {
|
var out struct {
|
||||||
Choices []struct {
|
Choices []struct {
|
||||||
@@ -119,25 +281,108 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
|
|||||||
if ch.Message.Content != nil {
|
if ch.Message.Content != nil {
|
||||||
content = strings.TrimSpace(*ch.Message.Content)
|
content = strings.TrimSpace(*ch.Message.Content)
|
||||||
}
|
}
|
||||||
return ToolsTurn{
|
return ToolsTurn{Content: content, ToolCalls: ch.Message.ToolCalls, RawFinish: ch.FinishReason}, nil
|
||||||
Content: content,
|
}
|
||||||
ToolCalls: ch.Message.ToolCalls,
|
|
||||||
RawFinish: ch.FinishReason,
|
func parseToolsTurnAnthropic(raw []byte, status int) (ToolsTurn, error) {
|
||||||
}, nil
|
if err := httpErrMessage(raw, status); err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
|
}
|
||||||
|
var out struct {
|
||||||
|
Content []struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text string `json:"text"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Input json.RawMessage `json:"input"`
|
||||||
|
} `json:"content"`
|
||||||
|
StopReason string `json:"stop_reason"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return ToolsTurn{}, fmt.Errorf("解析 Anthropic 响应失败(HTTP %d)", status)
|
||||||
|
}
|
||||||
|
var content strings.Builder
|
||||||
|
var calls []ToolCall
|
||||||
|
for _, c := range out.Content {
|
||||||
|
switch c.Type {
|
||||||
|
case "text":
|
||||||
|
content.WriteString(c.Text)
|
||||||
|
case "tool_use":
|
||||||
|
args := "{}"
|
||||||
|
if len(c.Input) > 0 {
|
||||||
|
args = string(c.Input)
|
||||||
|
}
|
||||||
|
calls = append(calls, ToolCall{
|
||||||
|
ID: c.ID,
|
||||||
|
Type: "function",
|
||||||
|
Function: ToolCallFunc{Name: c.Name, Arguments: args},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ToolsTurn{Content: strings.TrimSpace(content.String()), ToolCalls: calls, RawFinish: out.StopReason}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseToolsTurnResponses(raw []byte, status int) (ToolsTurn, error) {
|
||||||
|
if err := httpErrMessage(raw, status); err != nil {
|
||||||
|
return ToolsTurn{}, err
|
||||||
|
}
|
||||||
|
var out struct {
|
||||||
|
OutputText string `json:"output_text"`
|
||||||
|
Output []struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
CallID string `json:"call_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
Content []struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"content"`
|
||||||
|
} `json:"output"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return ToolsTurn{}, fmt.Errorf("解析 responses 响应失败(HTTP %d)", status)
|
||||||
|
}
|
||||||
|
var content strings.Builder
|
||||||
|
if out.OutputText != "" {
|
||||||
|
content.WriteString(out.OutputText)
|
||||||
|
}
|
||||||
|
var calls []ToolCall
|
||||||
|
for _, o := range out.Output {
|
||||||
|
switch o.Type {
|
||||||
|
case "message":
|
||||||
|
for _, c := range o.Content {
|
||||||
|
if c.Type == "output_text" || c.Type == "text" || c.Text != "" {
|
||||||
|
content.WriteString(c.Text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case "function_call":
|
||||||
|
id := o.CallID
|
||||||
|
if id == "" {
|
||||||
|
id = o.ID
|
||||||
|
}
|
||||||
|
args := o.Arguments
|
||||||
|
if args == "" {
|
||||||
|
args = "{}"
|
||||||
|
}
|
||||||
|
calls = append(calls, ToolCall{
|
||||||
|
ID: id,
|
||||||
|
Type: "function",
|
||||||
|
Function: ToolCallFunc{Name: o.Name, Arguments: args},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ToolsTurn{Content: strings.TrimSpace(content.String()), ToolCalls: calls}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// DefaultToolChat 返回可跑 Agent 的远端 Provider;stub / 缺 key 时返回错误说明。
|
// DefaultToolChat 返回可跑 Agent 的远端 Provider;stub / 缺 key 时返回错误说明。
|
||||||
func DefaultToolChat() (ToolChat, error) {
|
func DefaultToolChat() (ToolChat, error) {
|
||||||
cfg := LoadConfig()
|
cfg := LoadConfig()
|
||||||
if cfg.Provider == "" || cfg.Provider == "stub" {
|
if cfg.Provider == "" || cfg.Provider == "stub" {
|
||||||
return nil, fmt.Errorf("Agent 需要真实 LLM:请在设置中配置 DeepSeek / OpenAI 等(stub 仅支持 Ask 生成)")
|
return nil, fmt.Errorf("Agent 需要真实 LLM:请在设置中配置 DeepSeek / OpenAI / Anthropic 等(stub 仅支持 Ask 生成)")
|
||||||
}
|
}
|
||||||
if cfg.APIKey == "" && cfg.AuthStyle != AuthNone {
|
if cfg.APIKey == "" && cfg.AuthStyle != AuthNone {
|
||||||
return nil, fmt.Errorf("Agent 需要 API Key:请在设置中填写后重试")
|
return nil, fmt.Errorf("Agent 需要 API Key:请在设置中填写后重试")
|
||||||
}
|
}
|
||||||
p := NewRemoteProvider(cfg)
|
return NewRemoteProvider(cfg), nil
|
||||||
if p.protocol != "" && p.protocol != ProtoChat {
|
|
||||||
return nil, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议(当前: %s)", p.protocol)
|
|
||||||
}
|
|
||||||
return p, nil
|
|
||||||
}
|
}
|
||||||
@@ -0,0 +1,72 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"opsxdb/server/internal/aitools"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseToolsTurnAnthropic(t *testing.T) {
|
||||||
|
raw, _ := json.Marshal(map[string]any{
|
||||||
|
"stop_reason": "tool_use",
|
||||||
|
"content": []map[string]any{
|
||||||
|
{"type": "text", "text": "先列表"},
|
||||||
|
{"type": "tool_use", "id": "tu_1", "name": "opsxdb_list_tables", "input": map[string]any{"connection_id": "c1"}},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
turn, err := parseToolsTurnAnthropic(raw, 200)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if turn.Content != "先列表" {
|
||||||
|
t.Fatalf("content=%q", turn.Content)
|
||||||
|
}
|
||||||
|
if len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
|
||||||
|
t.Fatalf("calls=%+v", turn.ToolCalls)
|
||||||
|
}
|
||||||
|
if !strings.Contains(turn.ToolCalls[0].Function.Arguments, "c1") {
|
||||||
|
t.Fatalf("args=%s", turn.ToolCalls[0].Function.Arguments)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseToolsTurnResponses(t *testing.T) {
|
||||||
|
raw, _ := json.Marshal(map[string]any{
|
||||||
|
"output": []map[string]any{
|
||||||
|
{"type": "function_call", "call_id": "call_9", "name": "opsxdb_execute_query", "arguments": `{"connection_id":"c1","sql":"SELECT 1"}`},
|
||||||
|
{"type": "message", "content": []map[string]any{{"type": "output_text", "text": "done"}}},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
turn, err := parseToolsTurnResponses(raw, 200)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(turn.ToolCalls) != 1 || turn.ToolCalls[0].ID != "call_9" {
|
||||||
|
t.Fatalf("calls=%+v", turn.ToolCalls)
|
||||||
|
}
|
||||||
|
if turn.Content != "done" {
|
||||||
|
t.Fatalf("content=%q", turn.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpenAIToolsToAnthropic(t *testing.T) {
|
||||||
|
tools := aitools.OpenAITools()
|
||||||
|
anth := openaiToolsToAnthropic(tools)
|
||||||
|
if len(anth) != len(tools) {
|
||||||
|
t.Fatalf("%d vs %d", len(anth), len(tools))
|
||||||
|
}
|
||||||
|
if anth[0]["name"] != tools[0].Function.Name {
|
||||||
|
t.Fatal(anth[0]["name"])
|
||||||
|
}
|
||||||
|
if anth[0]["input_schema"] == nil {
|
||||||
|
t.Fatal("missing input_schema")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestChunkDelta(t *testing.T) {
|
||||||
|
parts := chunkDelta("abcdefghij", 3)
|
||||||
|
if len(parts) != 4 {
|
||||||
|
t.Fatalf("%v", parts)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -14,6 +14,7 @@ type AuditEntry struct {
|
|||||||
Tool string `json:"tool"`
|
Tool string `json:"tool"`
|
||||||
ConnectionID string `json:"connection_id,omitempty"`
|
ConnectionID string `json:"connection_id,omitempty"`
|
||||||
Mode string `json:"mode"`
|
Mode string `json:"mode"`
|
||||||
|
Source string `json:"source,omitempty"` // mcp | agent
|
||||||
OK bool `json:"ok"`
|
OK bool `json:"ok"`
|
||||||
Summary string `json:"summary,omitempty"`
|
Summary string `json:"summary,omitempty"`
|
||||||
Error string `json:"error,omitempty"`
|
Error string `json:"error,omitempty"`
|
||||||
|
|||||||
@@ -81,10 +81,14 @@ func LoadPolicy() (Policy, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (p Policy) IsReadOnly() bool {
|
func (p Policy) IsReadOnly() bool {
|
||||||
// P0:即使配置写了 read_write,执行层仍按只读拦截;此 helper 表示「策略意图」。
|
|
||||||
return p.Mode != ModeReadWrite
|
return p.Mode != ModeReadWrite
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// WritesEnabled P1:仅当 mcp.json mode=read_write 时允许写路径(仍须确认令牌)。
|
||||||
|
func (p Policy) WritesEnabled() bool {
|
||||||
|
return strings.EqualFold(strings.TrimSpace(p.Mode), ModeReadWrite)
|
||||||
|
}
|
||||||
|
|
||||||
func (p Policy) AllowConnection(id string) error {
|
func (p Policy) AllowConnection(id string) error {
|
||||||
id = strings.TrimSpace(id)
|
id = strings.TrimSpace(id)
|
||||||
if id == "" {
|
if id == "" {
|
||||||
|
|||||||
@@ -46,6 +46,13 @@ type Runtime struct {
|
|||||||
ListTbl func(ctx context.Context, c db.ConnReq) ([]string, error)
|
ListTbl func(ctx context.Context, c db.ConnReq) ([]string, error)
|
||||||
ListCol func(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error)
|
ListCol func(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error)
|
||||||
Audit func(AuditEntry)
|
Audit func(AuditEntry)
|
||||||
|
// RejectCallerConfirm 为 true 时(Agent 路径),忽略调用方 confirm_danger,
|
||||||
|
// 写操作只返回 needConfirm,由用户在 UI 显式确认后再走 /api/ai/agent/confirm。
|
||||||
|
RejectCallerConfirm bool
|
||||||
|
// AllowConfirmedWrite 为 true 时,已校验的确认令牌可放行写(仅 confirm 端点设置)。
|
||||||
|
AllowConfirmedWrite bool
|
||||||
|
// Source 写入审计:mcp | agent。
|
||||||
|
Source string
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewRuntime(policy Policy) *Runtime {
|
func NewRuntime(policy Policy) *Runtime {
|
||||||
@@ -67,6 +74,9 @@ func (r *Runtime) audit(e AuditEntry) {
|
|||||||
e.Mode = ModeReadOnly
|
e.Mode = ModeReadOnly
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if e.Source == "" {
|
||||||
|
e.Source = r.Source
|
||||||
|
}
|
||||||
if r.Audit != nil {
|
if r.Audit != nil {
|
||||||
r.Audit(e)
|
r.Audit(e)
|
||||||
return
|
return
|
||||||
@@ -213,6 +223,29 @@ func argInt(args map[string]any, key string, def int) int {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func argBool(args map[string]any, key string) bool {
|
||||||
|
if args == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
v, ok := args[key]
|
||||||
|
if !ok || v == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
switch t := v.(type) {
|
||||||
|
case bool:
|
||||||
|
return t
|
||||||
|
case string:
|
||||||
|
return strings.EqualFold(t, "true") || t == "1"
|
||||||
|
case float64:
|
||||||
|
return t != 0
|
||||||
|
case int:
|
||||||
|
return t != 0
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
func (r *Runtime) listDatabases(ctx context.Context, args map[string]any) (any, error) {
|
func (r *Runtime) listDatabases(ctx context.Context, args map[string]any) (any, error) {
|
||||||
id := argString(args, "connection_id")
|
id := argString(args, "connection_id")
|
||||||
c, err := r.resolveAllowed(id)
|
c, err := r.resolveAllowed(id)
|
||||||
@@ -352,9 +385,27 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e
|
|||||||
if sqlText == "" {
|
if sqlText == "" {
|
||||||
return fail("sql required")
|
return fail("sql required")
|
||||||
}
|
}
|
||||||
// P0 纯只读:写/DDL 一律拒绝;不做 write mode / confirm token。
|
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||||||
if ok, reason := IsReadOnlySQL(sqlText); !ok {
|
if r.RejectCallerConfirm {
|
||||||
return fail(reason)
|
callerConfirm = false
|
||||||
|
}
|
||||||
|
confirmed := r.AllowConfirmedWrite || callerConfirm
|
||||||
|
|
||||||
|
readOK, readReason := IsReadOnlySQL(sqlText)
|
||||||
|
if !readOK {
|
||||||
|
if !r.Policy.WritesEnabled() {
|
||||||
|
return fail(readReason)
|
||||||
|
}
|
||||||
|
// P1:可写已开启,但无确认令牌则返回 needConfirm(永不静默写)。
|
||||||
|
if !confirmed {
|
||||||
|
r.audit(AuditEntry{Tool: ToolExecuteQuery, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(sqlText, 80)})
|
||||||
|
return map[string]any{
|
||||||
|
"needConfirm": true,
|
||||||
|
"sql": sqlText,
|
||||||
|
"message": "write requires explicit user confirmation",
|
||||||
|
"reason": readReason,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
c, err := r.resolveAllowed(id)
|
c, err := r.resolveAllowed(id)
|
||||||
@@ -371,7 +422,7 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e
|
|||||||
}
|
}
|
||||||
req.SQL = sqlText
|
req.SQL = sqlText
|
||||||
req.Limit = limit
|
req.Limit = limit
|
||||||
req.ConfirmDanger = false // 永不带确认放行
|
req.ConfirmDanger = confirmed && !readOK
|
||||||
|
|
||||||
res := r.execSQL(ctx, req)
|
res := r.execSQL(ctx, req)
|
||||||
if res.NeedConfirm {
|
if res.NeedConfirm {
|
||||||
@@ -407,11 +458,27 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e
|
|||||||
if command == "" {
|
if command == "" {
|
||||||
return fail("command required")
|
return fail("command required")
|
||||||
}
|
}
|
||||||
cmd, err := CheckRedisReadOnly(command)
|
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||||||
if err != nil {
|
if r.RejectCallerConfirm {
|
||||||
return fail(err.Error())
|
callerConfirm = false
|
||||||
|
}
|
||||||
|
confirmed := r.AllowConfirmedWrite || callerConfirm
|
||||||
|
|
||||||
|
cmd, rerr := CheckRedisReadOnly(command)
|
||||||
|
if rerr != nil {
|
||||||
|
if redisDeniedAlways[cmd] || !r.Policy.WritesEnabled() {
|
||||||
|
return fail(rerr.Error())
|
||||||
|
}
|
||||||
|
if !confirmed {
|
||||||
|
r.audit(AuditEntry{Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(command, 80)})
|
||||||
|
return map[string]any{
|
||||||
|
"needConfirm": true,
|
||||||
|
"command": command,
|
||||||
|
"message": "redis write requires explicit user confirmation",
|
||||||
|
"reason": rerr.Error(),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
}
|
}
|
||||||
_ = cmd
|
|
||||||
|
|
||||||
c, err := r.resolveAllowed(id)
|
c, err := r.resolveAllowed(id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -426,7 +493,7 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e
|
|||||||
req.Database = database
|
req.Database = database
|
||||||
}
|
}
|
||||||
req.SQL = command
|
req.SQL = command
|
||||||
req.ConfirmDanger = false
|
req.ConfirmDanger = confirmed && rerr != nil
|
||||||
|
|
||||||
res := r.execSQL(ctx, req)
|
res := r.execSQL(ctx, req)
|
||||||
if res.NeedConfirm {
|
if res.NeedConfirm {
|
||||||
|
|||||||
@@ -227,3 +227,75 @@ func TestToolNamesStable(t *testing.T) {
|
|||||||
t.Fatalf("%v", names)
|
t.Fatalf("%v", names)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestExecuteQueryNeedConfirmInReadWrite(t *testing.T) {
|
||||||
|
called := false
|
||||||
|
rt := NewRuntime(Policy{
|
||||||
|
Mode: ModeReadWrite,
|
||||||
|
AllowedConnectionIDs: []string{"c1"},
|
||||||
|
DefaultLimit: 100,
|
||||||
|
MaxLimit: 500,
|
||||||
|
})
|
||||||
|
rt.FindConn = func(id string) (*SavedConnection, error) {
|
||||||
|
return &SavedConnection{ID: id, Name: "d", Driver: "mysql", Host: "127.0.0.1", Port: 3306}, nil
|
||||||
|
}
|
||||||
|
rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
|
||||||
|
called = true
|
||||||
|
return db.ExecResult{RowCount: 1}
|
||||||
|
}
|
||||||
|
rt.Audit = func(AuditEntry) {}
|
||||||
|
|
||||||
|
out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
|
||||||
|
"connection_id": "c1",
|
||||||
|
"sql": "DELETE FROM users WHERE id=1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
m := out.(map[string]any)
|
||||||
|
if m["needConfirm"] != true {
|
||||||
|
t.Fatalf("want needConfirm, got %#v", m)
|
||||||
|
}
|
||||||
|
if called {
|
||||||
|
t.Fatal("must not exec without confirm")
|
||||||
|
}
|
||||||
|
|
||||||
|
out2, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
|
||||||
|
"connection_id": "c1",
|
||||||
|
"sql": "DELETE FROM users WHERE id=1",
|
||||||
|
"confirm_danger": true,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !called {
|
||||||
|
t.Fatal("expected exec after confirm")
|
||||||
|
}
|
||||||
|
_ = out2
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteQueryRejectConfirmWhenAgentFlag(t *testing.T) {
|
||||||
|
called := false
|
||||||
|
rt := NewRuntime(Policy{Mode: ModeReadWrite, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
|
||||||
|
rt.RejectCallerConfirm = true
|
||||||
|
rt.FindConn = func(id string) (*SavedConnection, error) {
|
||||||
|
return &SavedConnection{ID: id, Driver: "mysql", Host: "h", Port: 1}, nil
|
||||||
|
}
|
||||||
|
rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
|
||||||
|
called = true
|
||||||
|
return db.ExecResult{}
|
||||||
|
}
|
||||||
|
rt.Audit = func(AuditEntry) {}
|
||||||
|
out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
|
||||||
|
"connection_id": "c1", "sql": "UPDATE t SET a=1", "confirm_danger": true,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.(map[string]any)["needConfirm"] != true {
|
||||||
|
t.Fatalf("%#v", out)
|
||||||
|
}
|
||||||
|
if called {
|
||||||
|
t.Fatal("agent must not auto-confirm")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -52,22 +52,24 @@ func ToolSpecs() []ToolSpec {
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
Name: ToolExecuteQuery,
|
Name: ToolExecuteQuery,
|
||||||
Description: "Execute a READ-ONLY SQL query (SELECT/SHOW/EXPLAIN/WITH…SELECT). Writes and DDL are rejected in P0. Default limit 100, max 500.",
|
Description: "Execute SQL. Default read-only (SELECT/SHOW/EXPLAIN/WITH…SELECT). Writes require mcp.json mode=read_write AND confirm_danger=true (never auto). Default limit 100, max 500.",
|
||||||
Properties: map[string]any{
|
Properties: map[string]any{
|
||||||
"connection_id": connID,
|
"connection_id": connID,
|
||||||
"sql": map[string]any{"type": "string", "description": "Read-only SQL"},
|
"sql": map[string]any{"type": "string", "description": "SQL statement"},
|
||||||
"database": database,
|
"database": database,
|
||||||
"limit": map[string]any{"type": "integer", "description": "Row limit (default 100, max 500)"},
|
"limit": map[string]any{"type": "integer", "description": "Row limit (default 100, max 500)"},
|
||||||
|
"confirm_danger": map[string]any{"type": "boolean", "description": "Required for writes when mode=read_write; Agent UI sets this only after user confirms"},
|
||||||
},
|
},
|
||||||
Required: []string{"connection_id", "sql"},
|
Required: []string{"connection_id", "sql"},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
Name: ToolExecuteRedisCommand,
|
Name: ToolExecuteRedisCommand,
|
||||||
Description: "Execute a READ-ONLY Redis command (GET/SCAN/TYPE/TTL/…). FLUSH*/KEYS and writes are rejected in P0.",
|
Description: "Execute Redis command. Default read-only whitelist. Writes require mode=read_write AND confirm_danger=true. FLUSH*/KEYS/CONFIG always denied.",
|
||||||
Properties: map[string]any{
|
Properties: map[string]any{
|
||||||
"connection_id": connID,
|
"connection_id": connID,
|
||||||
"command": map[string]any{"type": "string", "description": "Single Redis command line"},
|
"command": map[string]any{"type": "string", "description": "Single Redis command line"},
|
||||||
"database": database,
|
"database": database,
|
||||||
|
"confirm_danger": map[string]any{"type": "boolean", "description": "Required for writes when mode=read_write"},
|
||||||
},
|
},
|
||||||
Required: []string{"connection_id", "command"},
|
Required: []string{"connection_id", "command"},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -0,0 +1,98 @@
|
|||||||
|
package aitools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const mcpActiveWindow = 5 * time.Minute
|
||||||
|
|
||||||
|
// MCPStatus 供侧栏 MCP 状态点 / 设置页轻量指示。
|
||||||
|
type MCPStatus struct {
|
||||||
|
Configured bool `json:"configured"`
|
||||||
|
Mode string `json:"mode"`
|
||||||
|
WhitelistCount int `json:"whitelistCount"`
|
||||||
|
WritesEnabled bool `json:"writesEnabled"`
|
||||||
|
Active bool `json:"active"` // 近 5 分钟有 source=mcp 审计
|
||||||
|
Recent []AuditEntry `json:"recent,omitempty"`
|
||||||
|
PolicyPath string `json:"policyPath,omitempty"`
|
||||||
|
AuditPath string `json:"auditPath,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMCPStatus 读 mcp.json + 尾部审计;不探测外部 Cursor 进程。
|
||||||
|
func GetMCPStatus() MCPStatus {
|
||||||
|
st := MCPStatus{Mode: ModeReadOnly}
|
||||||
|
if path, err := PolicyPath(); err == nil {
|
||||||
|
st.PolicyPath = path
|
||||||
|
if _, err := os.Stat(path); err == nil {
|
||||||
|
st.Configured = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if p, err := LoadPolicy(); err == nil {
|
||||||
|
st.Mode = p.Mode
|
||||||
|
if st.Mode == "" {
|
||||||
|
st.Mode = ModeReadOnly
|
||||||
|
}
|
||||||
|
st.WhitelistCount = len(p.AllowedConnectionIDs)
|
||||||
|
st.WritesEnabled = p.WritesEnabled()
|
||||||
|
st.Configured = st.Configured || st.WhitelistCount > 0
|
||||||
|
}
|
||||||
|
if ap, err := AuditPath(); err == nil {
|
||||||
|
st.AuditPath = ap
|
||||||
|
recent := tailAudit(ap, 20)
|
||||||
|
st.Recent = filterSource(recent, "mcp")
|
||||||
|
if len(st.Recent) == 0 {
|
||||||
|
// 兼容旧审计无 source:仍展示最近条目,但不标 active
|
||||||
|
st.Recent = recent
|
||||||
|
} else {
|
||||||
|
cutoff := time.Now().Add(-mcpActiveWindow)
|
||||||
|
for _, e := range st.Recent {
|
||||||
|
if t, err := time.Parse(time.RFC3339, e.Time); err == nil && t.After(cutoff) {
|
||||||
|
st.Active = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return st
|
||||||
|
}
|
||||||
|
|
||||||
|
func filterSource(in []AuditEntry, source string) []AuditEntry {
|
||||||
|
var out []AuditEntry
|
||||||
|
for _, e := range in {
|
||||||
|
if e.Source == source {
|
||||||
|
out = append(out, e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func tailAudit(path string, n int) []AuditEntry {
|
||||||
|
f, err := os.Open(path)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
defer f.Close()
|
||||||
|
var lines []string
|
||||||
|
sc := bufio.NewScanner(f)
|
||||||
|
sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||||
|
for sc.Scan() {
|
||||||
|
lines = append(lines, sc.Text())
|
||||||
|
if len(lines) > n*4 { // 粗裁
|
||||||
|
lines = lines[len(lines)-n*2:]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(lines) > n {
|
||||||
|
lines = lines[len(lines)-n:]
|
||||||
|
}
|
||||||
|
out := make([]AuditEntry, 0, len(lines))
|
||||||
|
for _, line := range lines {
|
||||||
|
var e AuditEntry
|
||||||
|
if json.Unmarshal([]byte(line), &e) == nil {
|
||||||
|
out = append(out, e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
@@ -475,6 +475,7 @@ func NewRouter(token string) *gin.Engine {
|
|||||||
})
|
})
|
||||||
|
|
||||||
// 侧栏 Agent:多轮 tool call → 同一 aitools.Runtime(进程内,不绕 stdio)。Ask 仍走 /ai/chat。
|
// 侧栏 Agent:多轮 tool call → 同一 aitools.Runtime(进程内,不绕 stdio)。Ask 仍走 /ai/chat。
|
||||||
|
// stream=true 或 Accept: text/event-stream 时走 SSE(步骤边推 + 正文分块);否则 JSON fallback。
|
||||||
api.POST("/ai/agent", func(c *gin.Context) {
|
api.POST("/ai/agent", func(c *gin.Context) {
|
||||||
var req ai.AgentReq
|
var req ai.AgentReq
|
||||||
if err := c.ShouldBindJSON(&req); err != nil {
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
@@ -491,6 +492,30 @@ func NewRouter(token string) *gin.Engine {
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
rt.Source = "agent"
|
||||||
|
rt.RejectCallerConfirm = true
|
||||||
|
|
||||||
|
wantStream := req.Stream || strings.Contains(c.GetHeader("Accept"), "text/event-stream")
|
||||||
|
if wantStream {
|
||||||
|
flusher, ok := c.Writer.(http.Flusher)
|
||||||
|
if !ok {
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming unsupported"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.Header("Content-Type", "text/event-stream")
|
||||||
|
c.Header("Cache-Control", "no-cache")
|
||||||
|
c.Header("Connection", "keep-alive")
|
||||||
|
c.Status(http.StatusOK)
|
||||||
|
flusher.Flush()
|
||||||
|
_, err := ai.RunAgentStream(c.Request.Context(), llm, rt, req, func(ev ai.AgentEvent) {
|
||||||
|
_ = ai.WriteAgentSSE(c.Writer, flusher, ev)
|
||||||
|
})
|
||||||
|
if err != nil && c.Request.Context().Err() == nil {
|
||||||
|
_ = ai.WriteAgentSSE(c.Writer, flusher, ai.AgentEvent{Type: ai.AgentEventError, Error: err.Error()})
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
res, err := ai.RunAgent(c.Request.Context(), llm, rt, req)
|
res, err := ai.RunAgent(c.Request.Context(), llm, rt, req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error(), "steps": res.Steps, "rounds": res.Rounds})
|
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error(), "steps": res.Steps, "rounds": res.Rounds})
|
||||||
@@ -499,6 +524,55 @@ func NewRouter(token string) *gin.Engine {
|
|||||||
c.JSON(http.StatusOK, res)
|
c.JSON(http.StatusOK, res)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
// 用户在气泡确认条点「执行」后:单次写操作(AllowConfirmedWrite),绝不由模型静默触发。
|
||||||
|
api.POST("/ai/agent/confirm", func(c *gin.Context) {
|
||||||
|
var req struct {
|
||||||
|
ConnectionID string `json:"connectionId"`
|
||||||
|
SQL string `json:"sql"`
|
||||||
|
Command string `json:"command"`
|
||||||
|
Database string `json:"database"`
|
||||||
|
Kind string `json:"kind"` // query | redis
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
rt, err := aitools.DefaultRuntime()
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !rt.Policy.WritesEnabled() {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "writes disabled: set mode=read_write in ~/.opsxdb/mcp.json"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
rt.Source = "agent"
|
||||||
|
rt.RejectCallerConfirm = false
|
||||||
|
rt.AllowConfirmedWrite = true
|
||||||
|
kind := strings.ToLower(strings.TrimSpace(req.Kind))
|
||||||
|
args := map[string]any{"connection_id": req.ConnectionID, "database": req.Database, "confirm_danger": true}
|
||||||
|
var out any
|
||||||
|
if kind == "redis" || (kind == "" && req.Command != "") {
|
||||||
|
args["command"] = req.Command
|
||||||
|
if args["command"] == "" || args["command"] == nil {
|
||||||
|
args["command"] = req.SQL
|
||||||
|
}
|
||||||
|
out, err = rt.Call(c.Request.Context(), aitools.ToolExecuteRedisCommand, args)
|
||||||
|
} else {
|
||||||
|
args["sql"] = req.SQL
|
||||||
|
out, err = rt.Call(c.Request.Context(), aitools.ToolExecuteQuery, args)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, out)
|
||||||
|
})
|
||||||
|
|
||||||
|
api.GET("/ai/mcp-status", func(c *gin.Context) {
|
||||||
|
c.JSON(http.StatusOK, aitools.GetMCPStatus())
|
||||||
|
})
|
||||||
|
|
||||||
api.POST("/ai/generate", func(c *gin.Context) {
|
api.POST("/ai/generate", func(c *gin.Context) {
|
||||||
var req ai.ChatReq
|
var req ai.ChatReq
|
||||||
if err := c.ShouldBindJSON(&req); err != nil {
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
|||||||
Reference in new issue
Block a user