opsxdb/server/internal/ai/ai_test.go

package ai

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
)

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }

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")
	}
}

func TestAskChatsWithoutClaimingDatabaseExecution(t *testing.T) {
	request := ChatReq{Action: "ask", Prompt: "这个库里有哪些表？", Driver: "mysql", Database: "demo"}
	msgs := BuildMessages(request)
	if !strings.Contains(msgs[0].Content, "demo(mysql)") || !strings.Contains(msgs[0].Content, "不要声称已经查询") {
		t.Fatalf("Ask should carry context without pretending to query: %s", msgs[0].Content)
	}
	if strings.Contains(msgs[0].Content, "只输出可执行 SQL") {
		t.Fatal("Ask must not force a SQL response")
	}
	generate := BuildMessages(ChatReq{Action: "generate", Prompt: "查 users", Driver: "mysql", Database: "demo"})
	if !strings.Contains(generate[0].Content, "生成一条 mysql 方言的 SQL") {
		t.Fatalf("SQL generation lost its dedicated instruction: %s", generate[0].Content)
	}
	if !strings.Contains(generate[0].Content, "澄清") || strings.Contains(generate[0].Content, "只输出可执行 SQL") {
		t.Fatalf("generate must allow clarification and not force executable-only: %s", generate[0].Content)
	}
	if !strings.Contains(generate[0].Content, "禁止编造") || !strings.Contains(generate[0].Content, "UPDATE_TIME") {
		t.Fatalf("system rules missing invent/UPDATE_TIME guards: %s", generate[0].Content)
	}
	if !strings.Contains(generate[0].Content, "换表") || !strings.Contains(generate[0].Content, "created_at") {
		t.Fatalf("follow-up no-invent column rules missing: %s", generate[0].Content)
	}
	offline := chatStub(request)
	if offline.SQL != "" || !strings.Contains(offline.Content, "离线体验无法回答") {
		t.Fatalf("offline Ask must not pretend to generate or query: %+v", offline)
	}
}

func TestAskRemoteRequestIsChatAndReturnsNoUnrequestedSQL(t *testing.T) {
	p := NewRemoteProvider(Config{Provider: "deepseek", Protocol: ProtoChat, BaseURL: "https://api.example.test/v1", APIKey: "test-secret", Model: "test-model"})
	p.hc = &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
		if req.URL.Path != "/v1/chat/completions" {
			t.Fatalf("unexpected endpoint: %s", req.URL)
		}
		if req.Header.Get("Authorization") != "Bearer test-secret" {
			t.Fatal("missing auth")
		}
		var payload struct {
			Messages []Message `json:"messages"`
			Tools    []any     `json:"tools"`
		}
		if err := json.NewDecoder(req.Body).Decode(&payload); err != nil {
			t.Fatal(err)
		}
		if len(payload.Messages) != 2 || len(payload.Tools) != 0 || !strings.Contains(payload.Messages[0].Content, "demo(mysql)") || strings.Contains(payload.Messages[0].Content, "只输出可执行 SQL") {
			t.Fatalf("Ask should be contextual chat without tools: %+v", payload)
		}
		return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"当前表结构上下文不足，无法确认表名。"}}]}`)), Header: make(http.Header)}, nil
	})}
	res, err := p.Complete(context.Background(), ChatReq{Action: "ask", Prompt: "有哪些表？", Driver: "mysql", Database: "demo"})
	if err != nil || res.SQL != "" || !strings.Contains(res.Content, "无法确认表名") {
		t.Fatalf("Ask must return a plain answer without invented SQL: %+v err=%v", res, err)
	}
}