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
+1305
-178
No files matched your search
@@ -26,6 +26,7 @@ func main() {
|
||||
fmt.Fprintf(os.Stderr, "opsxdb-mcp: load policy: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
rt.Source = "mcp"
|
||||
|
||||
s := server.NewMCPServer(
|
||||
"opsxdb",
|
||||
|
||||
+141
-15
@@ -14,6 +14,7 @@ const (
|
||||
defaultAgentMaxRounds = 8
|
||||
maxToolResultFeedChars = 12000
|
||||
maxStepArgsChars = 200
|
||||
maxGridRows = 200
|
||||
)
|
||||
|
||||
// AgentReq 侧栏 Agent 请求(与 Ask 的 /api/ai/chat 分路径)。
|
||||
@@ -24,6 +25,7 @@ type AgentReq struct {
|
||||
Database string `json:"database,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
MaxRounds int `json:"maxRounds,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
}
|
||||
|
||||
// AgentStep 挂在回复下的步骤时间线一项。
|
||||
@@ -34,6 +36,10 @@ type AgentStep struct {
|
||||
Summary string `json:"summary"`
|
||||
Error string `json:"error,omitempty"`
|
||||
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 最终回答 + 步骤。
|
||||
@@ -44,9 +50,18 @@ type AgentRes struct {
|
||||
Rounds int `json:"rounds"`
|
||||
}
|
||||
|
||||
// RunAgent 多轮:LLM tool_call → aitools.Runtime.Call → 回灌,直到最终文本或轮次耗尽。
|
||||
// ConfirmDanger 永不自动 true(由 Runtime 保证)。Ask 模式不要走这里。
|
||||
// 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")
|
||||
}
|
||||
@@ -65,6 +80,9 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
||||
maxRounds = 16
|
||||
}
|
||||
|
||||
// Agent 路径禁止模型自行带确认令牌
|
||||
rt.RejectCallerConfirm = true
|
||||
|
||||
tools := aitools.OpenAITools()
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: buildAgentSystem(req)},
|
||||
@@ -74,9 +92,15 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
||||
var steps []AgentStep
|
||||
rounds := 0
|
||||
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
|
||||
turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
|
||||
if err != nil {
|
||||
emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds})
|
||||
return AgentRes{Steps: steps, Rounds: rounds}, err
|
||||
}
|
||||
if len(turn.ToolCalls) == 0 {
|
||||
@@ -84,15 +108,19 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
||||
if content == "" {
|
||||
content = "(模型未返回文本,也未调用工具)"
|
||||
}
|
||||
return AgentRes{
|
||||
for _, part := range chunkDelta(content, 36) {
|
||||
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||||
}
|
||||
res := AgentRes{
|
||||
Content: content,
|
||||
SQL: extractSQL(content),
|
||||
Steps: steps,
|
||||
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{
|
||||
Role: "assistant",
|
||||
Content: turn.Content,
|
||||
@@ -100,28 +128,41 @@ func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentR
|
||||
})
|
||||
|
||||
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)
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
return AgentRes{
|
||||
Content: fmt.Sprintf("已达最大工具轮次(%d)。请根据已执行步骤自行判断,或缩小问题后重试。", maxRounds),
|
||||
Steps: steps,
|
||||
Rounds: rounds,
|
||||
}, nil
|
||||
content := fmt.Sprintf("已达最大工具轮次(%d)。请根据已执行步骤自行判断,或缩小问题后重试。", maxRounds)
|
||||
for _, part := range chunkDelta(content, 36) {
|
||||
emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
|
||||
}
|
||||
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 {
|
||||
var b strings.Builder
|
||||
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具只读查询用户已白名单的连接。\n")
|
||||
b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具查询用户已白名单的连接。\n")
|
||||
b.WriteString("规则:\n")
|
||||
b.WriteString("1. 需要查库时必须调用工具;不要假装已执行 SQL。\n")
|
||||
b.WriteString("2. 仅只读:execute_query 只允许 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT;Redis 仅读命令。\n")
|
||||
b.WriteString("3. 先确认 connection_id(可用 opsxdb_list_connections);用户给出偏好连接时优先用它。\n")
|
||||
b.WriteString("4. 用中文简明回答;可执行 SQL 放在 ```sql 代码块中。\n")
|
||||
b.WriteString("5. 工具结果可能截断,注意说明。\n")
|
||||
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 != "" {
|
||||
@@ -160,6 +201,9 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
|
||||
Content: step.Error,
|
||||
}
|
||||
}
|
||||
// 永不让模型自行带确认令牌
|
||||
delete(args, "confirm_danger")
|
||||
delete(args, "confirmDanger")
|
||||
|
||||
out, err := rt.Call(ctx, name, args)
|
||||
step.Ms = time.Since(start).Milliseconds()
|
||||
@@ -174,6 +218,37 @@ func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentSt
|
||||
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.Summary = summarizeToolResult(name, 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 {
|
||||
m, ok := out.(map[string]any)
|
||||
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" {
|
||||
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"
|
||||
)
|
||||
|
||||
// Config AI 配置;provider: stub|openai(OpenAI 兼容协议)。
|
||||
// Config AI 配置;provider 为预设 id(stub|openai|deepseek|anthropic|ollama|custom…),
|
||||
// protocol/authStyle/allowInsecureHTTP 决定远端调用行为(见 remote.go)。
|
||||
type Config struct {
|
||||
Provider string `json:"provider"`
|
||||
BaseURL string `json:"baseURL"`
|
||||
APIKey string `json:"apiKey,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Protocol string `json:"protocol,omitempty"`
|
||||
AuthStyle string `json:"authStyle,omitempty"`
|
||||
AllowInsecureHTTP bool `json:"allowInsecureHttp,omitempty"`
|
||||
}
|
||||
|
||||
func configPath() (string, error) {
|
||||
@@ -47,6 +51,13 @@ func LoadConfig() Config {
|
||||
if 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 openai-chat 兼容协议的 tools 调用(DeepSeek / OpenAI / 多数兼容网关)。
|
||||
// anthropic / responses 本轮不支持 Agent tool calling,返回明确错误。
|
||||
// ChatWithTools 按协议分派:openai-chat / anthropic messages / openai-responses。
|
||||
func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
|
||||
if err := validateEndpoint(p.baseURL, p.allowHTTP); err != nil {
|
||||
return ToolsTurn{}, err
|
||||
}
|
||||
if p.protocol != "" && p.protocol != ProtoChat {
|
||||
return ToolsTurn{}, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议(当前: %s);请改用 DeepSeek / OpenAI 兼容端点", p.protocol)
|
||||
}
|
||||
m := strings.TrimSpace(model)
|
||||
if m == "" {
|
||||
m = p.model
|
||||
}
|
||||
urlStr := endpoint(p.baseURL, "chat/completions")
|
||||
payload := map[string]any{
|
||||
"model": m,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"tool_choice": "auto",
|
||||
"max_tokens": remoteMaxTokens,
|
||||
"temperature": 0.2,
|
||||
switch p.protocol {
|
||||
case ProtoAnthropic:
|
||||
return p.chatToolsAnthropic(ctx, m, messages, tools)
|
||||
case ProtoResponses:
|
||||
return p.chatToolsResponses(ctx, m, messages, tools)
|
||||
default:
|
||||
return p.chatToolsOpenAI(ctx, m, messages, tools)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *RemoteProvider) doJSON(ctx context.Context, urlStr string, payload any) ([]byte, int, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return ToolsTurn{}, err
|
||||
return nil, 0, err
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, urlStr, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return ToolsTurn{}, err
|
||||
return nil, 0, err
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
switch p.authStyle {
|
||||
@@ -65,17 +63,174 @@ func (p *RemoteProvider) ChatWithTools(ctx context.Context, model string, messag
|
||||
}
|
||||
resp, err := p.hc.Do(httpReq)
|
||||
if err != nil {
|
||||
return ToolsTurn{}, fmt.Errorf("AI 请求失败: %w", err)
|
||||
return nil, 0, fmt.Errorf("AI 请求失败: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
if err != nil {
|
||||
return ToolsTurn{}, fmt.Errorf("读取 AI 响应失败: %w", err)
|
||||
return 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 {
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
@@ -90,7 +245,14 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
|
||||
} else if 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 {
|
||||
Choices []struct {
|
||||
@@ -119,25 +281,108 @@ func parseToolsTurn(raw []byte, status int) (ToolsTurn, error) {
|
||||
if ch.Message.Content != nil {
|
||||
content = strings.TrimSpace(*ch.Message.Content)
|
||||
}
|
||||
return ToolsTurn{
|
||||
Content: content,
|
||||
ToolCalls: ch.Message.ToolCalls,
|
||||
RawFinish: ch.FinishReason,
|
||||
}, nil
|
||||
return ToolsTurn{Content: content, ToolCalls: ch.Message.ToolCalls, RawFinish: ch.FinishReason}, nil
|
||||
}
|
||||
|
||||
func parseToolsTurnAnthropic(raw []byte, status int) (ToolsTurn, error) {
|
||||
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 时返回错误说明。
|
||||
func DefaultToolChat() (ToolChat, error) {
|
||||
cfg := LoadConfig()
|
||||
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 {
|
||||
return nil, fmt.Errorf("Agent 需要 API Key:请在设置中填写后重试")
|
||||
}
|
||||
p := NewRemoteProvider(cfg)
|
||||
if p.protocol != "" && p.protocol != ProtoChat {
|
||||
return nil, fmt.Errorf("Agent 模式目前仅支持 openai-chat 协议(当前: %s)", p.protocol)
|
||||
}
|
||||
return p, nil
|
||||
return NewRemoteProvider(cfg), 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"`
|
||||
ConnectionID string `json:"connection_id,omitempty"`
|
||||
Mode string `json:"mode"`
|
||||
Source string `json:"source,omitempty"` // mcp | agent
|
||||
OK bool `json:"ok"`
|
||||
Summary string `json:"summary,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
|
||||
@@ -81,10 +81,14 @@ func LoadPolicy() (Policy, error) {
|
||||
}
|
||||
|
||||
func (p Policy) IsReadOnly() bool {
|
||||
// P0:即使配置写了 read_write,执行层仍按只读拦截;此 helper 表示「策略意图」。
|
||||
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 {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
|
||||
@@ -46,6 +46,13 @@ type Runtime struct {
|
||||
ListTbl func(ctx context.Context, c db.ConnReq) ([]string, error)
|
||||
ListCol func(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error)
|
||||
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 {
|
||||
@@ -67,6 +74,9 @@ func (r *Runtime) audit(e AuditEntry) {
|
||||
e.Mode = ModeReadOnly
|
||||
}
|
||||
}
|
||||
if e.Source == "" {
|
||||
e.Source = r.Source
|
||||
}
|
||||
if r.Audit != nil {
|
||||
r.Audit(e)
|
||||
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) {
|
||||
id := argString(args, "connection_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 == "" {
|
||||
return fail("sql required")
|
||||
}
|
||||
// P0 纯只读:写/DDL 一律拒绝;不做 write mode / confirm token。
|
||||
if ok, reason := IsReadOnlySQL(sqlText); !ok {
|
||||
return fail(reason)
|
||||
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||||
if r.RejectCallerConfirm {
|
||||
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)
|
||||
@@ -371,7 +422,7 @@ func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, e
|
||||
}
|
||||
req.SQL = sqlText
|
||||
req.Limit = limit
|
||||
req.ConfirmDanger = false // 永不带确认放行
|
||||
req.ConfirmDanger = confirmed && !readOK
|
||||
|
||||
res := r.execSQL(ctx, req)
|
||||
if res.NeedConfirm {
|
||||
@@ -407,11 +458,27 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e
|
||||
if command == "" {
|
||||
return fail("command required")
|
||||
}
|
||||
cmd, err := CheckRedisReadOnly(command)
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||||
if r.RejectCallerConfirm {
|
||||
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)
|
||||
if err != nil {
|
||||
@@ -426,7 +493,7 @@ func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, e
|
||||
req.Database = database
|
||||
}
|
||||
req.SQL = command
|
||||
req.ConfirmDanger = false
|
||||
req.ConfirmDanger = confirmed && rerr != nil
|
||||
|
||||
res := r.execSQL(ctx, req)
|
||||
if res.NeedConfirm {
|
||||
|
||||
@@ -227,3 +227,75 @@ func TestToolNamesStable(t *testing.T) {
|
||||
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,
|
||||
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{
|
||||
"connection_id": connID,
|
||||
"sql": map[string]any{"type": "string", "description": "Read-only SQL"},
|
||||
"sql": map[string]any{"type": "string", "description": "SQL statement"},
|
||||
"database": database,
|
||||
"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"},
|
||||
},
|
||||
{
|
||||
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{
|
||||
"connection_id": connID,
|
||||
"command": map[string]any{"type": "string", "description": "Single Redis command line"},
|
||||
"database": database,
|
||||
"confirm_danger": map[string]any{"type": "boolean", "description": "Required for writes when mode=read_write"},
|
||||
},
|
||||
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。
|
||||
// stream=true 或 Accept: text/event-stream 时走 SSE(步骤边推 + 正文分块);否则 JSON fallback。
|
||||
api.POST("/ai/agent", func(c *gin.Context) {
|
||||
var req ai.AgentReq
|
||||
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()})
|
||||
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)
|
||||
if err != nil {
|
||||
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)
|
||||
})
|
||||
|
||||
// 用户在气泡确认条点「执行」后:单次写操作(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) {
|
||||
var req ai.ChatReq
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
|
||||
Reference in new issue
Block a user