diff --git a/app/src/config/tabs/ai/aiProviderUi.ts b/app/src/config/tabs/ai/aiProviderUi.ts index 23e08e0b04e..faed82a219b 100644 --- a/app/src/config/tabs/ai/aiProviderUi.ts +++ b/app/src/config/tabs/ai/aiProviderUi.ts @@ -41,6 +41,7 @@ const PROVIDER_PRESETS: IProviderPreset[] = [ {id: "zhipu", name: "Zhipu AI", baseURL: "https://open.bigmodel.cn/api/paas/v4", category: "official", icon: "/stage/images/ai-providers/zhipu.svg"}, {id: "gemini", name: "Gemini", baseURL: "https://generativelanguage.googleapis.com/v1beta/openai", category: "official", icon: "/stage/images/ai-providers/gemini.svg"}, {id: "mistral", name: "Mistral AI", baseURL: "https://api.mistral.ai/v1", category: "official", icon: "/stage/images/ai-providers/mistral.svg"}, + {id: "aimlapi", name: "aimlapi.com", baseURL: "https://api.aimlapi.com/v1", category: "aggregator"}, {id: "siliconflow", name: "SiliconFlow", baseURL: "https://api.siliconflow.cn/v1", category: "aggregator", icon: "/stage/images/ai-providers/siliconflow.svg"}, {id: "openrouter", name: "OpenRouter", baseURL: "https://openrouter.ai/api/v1", category: "aggregator", icon: "/stage/images/ai-providers/openrouter.svg"}, {id: "groq", name: "Groq", baseURL: "https://api.groq.com/openai/v1", category: "aggregator"}, diff --git a/kernel/util/openai.go b/kernel/util/openai.go index 98052e39272..fd453ca4cb3 100644 --- a/kernel/util/openai.go +++ b/kernel/util/openai.go @@ -146,10 +146,68 @@ func ChatGPT(msg string, contextMsgs []string, c *openai.Client, protocol, model func NewOpenAIClient(apiKey, apiBaseURL string) *openai.Client { config := openai.DefaultConfig(apiKey) config.BaseURL = apiBaseURL - config.HTTPClient = httpclient.NewUserAgentClient(nil) + config.HTTPClient = newProviderHTTPDoer() return openai.NewClientWithConfig(config) } +// newProviderHTTPDoer 返回访问 Provider 所用的 HTTP 客户端:在带 User-Agent 的基础客户端外 +// 再包一层调用方标识请求头注入,供 chat、模型清单等所有 Provider 请求复用。 +func newProviderHTTPDoer() openai.HTTPDoer { + return &attributionTransport{base: httpclient.NewUserAgentClient(nil)} +} + +// providerAttributionHeaders 是「服务商 API 主机名 → 调用方标识请求头」的内置清单。 +// 这些请求头只声明「请求来自思源笔记」,不含 API Key、用户内容或任何可识别用户的信息, +// 服务商据此统计接入来源。按主机名分表是为了把请求头限定在该服务商自有端点上, +// 避免随请求发往其他服务商,或发往仅仅转发同一 API 的第三方代理。 +var providerAttributionHeaders = map[string]map[string]string{ + // AI/ML API:HTTP-Referer 与 X-Title 指向思源笔记自身(OpenRouter 系惯例,标识调用方应用), + // X-AIMLAPI-Source 与 X-AIMLAPI-Partner-ID 是该服务商用于统计接入来源的自有请求头。 + "api.aimlapi.com": { + "HTTP-Referer": "https://github.com/siyuan-note/siyuan", + "X-Title": "SiYuan", + "X-AIMLAPI-Source": "agent/siyuan", + "X-AIMLAPI-Partner-ID": "part_7cceWAMI91xwz7G6FrcOEUwN", + }, +} + +// AttributionHeadersForHost 按主机名大小写不敏感查表,返回该服务商需注入的调用方标识请求头; +// 无匹配返回 nil。返回的是副本,调用方改动不会影响内置清单。 +func AttributionHeadersForHost(host string) map[string]string { + headers := providerAttributionHeaders[strings.ToLower(strings.TrimSpace(host))] + if 0 == len(headers) { + return nil + } + return maps.Clone(headers) +} + +// attributionTransport 包装一个 HTTPDoer,为发往内置清单中服务商域名的请求补上调用方标识请求头。 +// 只在请求实际的目标主机命中清单时注入;已存在的同名请求头一律保留,不覆盖上层已设置的值。 +type attributionTransport struct { + base openai.HTTPDoer +} + +func (t *attributionTransport) Do(req *http.Request) (*http.Response, error) { + if nil == req || nil == req.URL { + return t.base.Do(req) + } + + headers := AttributionHeadersForHost(req.URL.Hostname()) + if 0 == len(headers) { + return t.base.Do(req) + } + + if nil == req.Header { + req.Header = http.Header{} + } + for name, value := range headers { + if "" == req.Header.Get(name) { + req.Header.Set(name, value) + } + } + return t.base.Do(req) +} + // builtinExtraBody 是「模型名前缀 → 额外请求参数」的内置适配清单。 // 仅收录参数语义为「纯输出格式开关」的模型(不改变思考行为、无副作用), // 让这类模型从源头把推理内容拆到 reasoning_content 字段,而非以 标签混在 content 中。 @@ -225,7 +283,7 @@ func NewOpenAIClientWithModel(apiKey, apiBaseURL, model string) *openai.Client { } config := openai.DefaultConfig(apiKey) config.BaseURL = apiBaseURL - var transport openai.HTTPDoer = httpclient.NewUserAgentClient(nil) + transport := newProviderHTTPDoer() if len(extra) > 0 { transport = &extraBodyTransport{base: transport, extraBody: extra} } @@ -355,7 +413,7 @@ func ListAvailableModelsWithContext(apiKey, apiBaseURL string, timeout int) (mod if apiKey != "" { req.Header.Set("Authorization", "Bearer "+apiKey) } - resp, err := httpclient.NewUserAgentClient(nil).Do(req) + resp, err := newProviderHTTPDoer().Do(req) if err != nil { logging.LogErrorf("list models [%s] failed: %s", apiBaseURL, err) return diff --git a/kernel/util/openai_attribution_test.go b/kernel/util/openai_attribution_test.go new file mode 100644 index 00000000000..b16e71706e2 --- /dev/null +++ b/kernel/util/openai_attribution_test.go @@ -0,0 +1,191 @@ +// SiYuan - From thought to insight, with agents +// Copyright (c) 2020-present, b3log.org +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +package util + +import ( + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" +) + +// aimlapiPartnerIDPattern 是 AI/ML API 对 X-AIMLAPI-Partner-ID 的取值约束: +// part_ 前缀加 1 到 64 位字母数字,不允许连字符或下划线。 +// 取值不合法时服务端不会报错,只是把请求算作无来源,因此必须由测试兜住。 +var aimlapiPartnerIDPattern = regexp.MustCompile(`^part_[A-Za-z0-9]{1,64}$`) + +// aimlapiSourcePattern 是 X-AIMLAPI-Source 的取值约束:/, +// channel 为 web、agent、mcp 三选一,client 为小写字母数字与连字符。 +var aimlapiSourcePattern = regexp.MustCompile(`^(web|agent|mcp)/[a-z0-9-]{1,32}$`) + +func TestProviderAttributionHeadersWellFormed(t *testing.T) { + // 内置清单里的每一项都必须是可用的取值,取值写错在运行时没有任何报错信号。 + for host, headers := range providerAttributionHeaders { + if host != strings.ToLower(host) { + t.Errorf("attribution host [%s] must be lower case for case-insensitive lookup", host) + } + if "" == headers["HTTP-Referer"] || "" == headers["X-Title"] { + t.Errorf("attribution headers for [%s] must identify the calling application", host) + } + if !strings.Contains(headers["HTTP-Referer"], "siyuan") { + t.Errorf("HTTP-Referer for [%s] must point at SiYuan, got %q", host, headers["HTTP-Referer"]) + } + if "api.aimlapi.com" != host { + continue + } + if partnerID := headers["X-AIMLAPI-Partner-ID"]; !aimlapiPartnerIDPattern.MatchString(partnerID) { + t.Errorf("X-AIMLAPI-Partner-ID %q does not match %s", partnerID, aimlapiPartnerIDPattern) + } + if source := headers["X-AIMLAPI-Source"]; !aimlapiSourcePattern.MatchString(source) { + t.Errorf("X-AIMLAPI-Source %q does not match %s", source, aimlapiSourcePattern) + } + } +} + +func TestAttributionHeadersForHost(t *testing.T) { + cases := []struct { + host string + want bool + }{ + {"api.aimlapi.com", true}, + {"API.AIMLAPI.COM", true}, // 大小写不敏感 + {" api.aimlapi.com ", true}, + {"api.openai.com", false}, + {"api.aimlapi.com.evil.example", false}, // 后缀伪装的域名不得命中 + {"aimlapi.com", false}, + {"localhost", false}, + {"", false}, + } + for _, tc := range cases { + got := AttributionHeadersForHost(tc.host) + if (0 < len(got)) != tc.want { + t.Errorf("AttributionHeadersForHost(%q) matched = %v, want %v", tc.host, 0 < len(got), tc.want) + } + } +} + +func TestAttributionHeadersForHostReturnsCopy(t *testing.T) { + // 返回值必须是副本,调用方改动不得污染内置清单,否则一次改动会影响此后所有请求。 + first := AttributionHeadersForHost("api.aimlapi.com") + first["X-AIMLAPI-Partner-ID"] = "tampered" + delete(first, "X-Title") + + second := AttributionHeadersForHost("api.aimlapi.com") + if "part_7cceWAMI91xwz7G6FrcOEUwN" != second["X-AIMLAPI-Partner-ID"] { + t.Fatalf("shared attribution map was mutated, got %q", second["X-AIMLAPI-Partner-ID"]) + } + if "" == second["X-Title"] { + t.Fatal("shared attribution map lost a key after the caller deleted it") + } +} + +func TestAttributionTransportInjectsForKnownHost(t *testing.T) { + // 目标主机命中清单时,四个请求头都应出现在实际发出的请求上。 + var captured http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{}`)) + })) + defer server.Close() + + // 请求写成真实的 AI/ML API 地址以命中清单,再由 redirectDoer 改投到本地测试服务器。 + transport := &attributionTransport{base: &redirectDoer{base: server.Client(), addr: server.Listener.Addr().String()}} + req, _ := http.NewRequest(http.MethodPost, "https://api.aimlapi.com/v1/chat/completions", strings.NewReader("{}")) + resp, err := transport.Do(req) + if nil != err { + t.Fatalf("Do failed: %v", err) + } + defer resp.Body.Close() + + want := map[string]string{ + "Http-Referer": "https://github.com/siyuan-note/siyuan", + "X-Title": "SiYuan", + "X-Aimlapi-Source": "agent/siyuan", + "X-Aimlapi-Partner-Id": "part_7cceWAMI91xwz7G6FrcOEUwN", + } + for name, value := range want { + if got := captured.Get(name); got != value { + t.Errorf("header %s = %q, want %q", name, got, value) + } + } +} + +func TestAttributionTransportSkipsOtherHost(t *testing.T) { + // 目标主机不在清单中时(含仅仅转发同一 API 的第三方代理),一个请求头都不能带上。 + var captured http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{}`)) + })) + defer server.Close() + + transport := &attributionTransport{base: server.Client()} + req, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", strings.NewReader("{}")) + resp, err := transport.Do(req) + if nil != err { + t.Fatalf("Do failed: %v", err) + } + defer resp.Body.Close() + + for _, name := range []string{"X-Aimlapi-Partner-Id", "X-Aimlapi-Source", "Http-Referer", "X-Title"} { + if got := captured.Get(name); "" != got { + t.Errorf("header %s leaked to a non-matching host: %q", name, got) + } + } +} + +func TestAttributionTransportKeepsExistingHeader(t *testing.T) { + // 上层已经设置过的同名请求头必须原样保留,注入只补空缺、不覆盖。 + req, _ := http.NewRequest(http.MethodPost, "https://api.aimlapi.com/v1/chat/completions", strings.NewReader("{}")) + req.Header.Set("X-Title", "Custom") + + var captured http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r.Header.Clone() + w.Write([]byte(`{}`)) + })) + defer server.Close() + + transport := &attributionTransport{base: &redirectDoer{base: server.Client(), addr: server.Listener.Addr().String()}} + resp, err := transport.Do(req) + if nil != err { + t.Fatalf("Do failed: %v", err) + } + defer resp.Body.Close() + + if "Custom" != captured.Get("X-Title") { + t.Errorf("existing X-Title was overwritten, got %q", captured.Get("X-Title")) + } + if "part_7cceWAMI91xwz7G6FrcOEUwN" != captured.Get("X-Aimlapi-Partner-Id") { + t.Errorf("missing header was not filled in, got %q", captured.Get("X-Aimlapi-Partner-Id")) + } +} + +// redirectDoer 把请求改投到本地测试服务器,用于在不联网的前提下断言真实发出的请求头。 +type redirectDoer struct { + base *http.Client + addr string +} + +func (d *redirectDoer) Do(req *http.Request) (*http.Response, error) { + req.URL.Scheme = "http" + req.URL.Host = d.addr + return d.base.Do(req) +}