Andreas/siyuan
1
0
Fork 0
mirror of https://github.com/siyuan-note/siyuan.git synced 2025-12-24 02:20:13 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

🎨 Support custom AI request User-Agent header https://github.com/siyuan-note/siyuan/issues/10351

Browse source
This commit is contained in:
Daniel 2024-02-15 10:59:37 +08:00
parent 3df54adcc5
commit 7e87c8c8ad
No known key found for this signature in database
GPG key ID: 86211BA83DF03017
11 changed files with 59 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

12
kernel/conf/ai.go
View file

@ -17,6 +17,7 @@
package conf
import (
"github.com/siyuan-note/siyuan/kernel/util"
"os"
"strconv"
@ -34,13 +35,15 @@ type OpenAI struct {
APIModel string `json:"apiModel"`
APIMaxTokens int `json:"apiMaxTokens"`
APIBaseURL string `json:"apiBaseURL"`
APIUserAgent string `json:"apiUserAgent"`
}
func NewAI() *AI {
openAI := &OpenAI{
APITimeout: 30,
APIModel: openai.GPT3Dot5Turbo,
APIBaseURL: "https://api.openai.com/v1",
APITimeout: 30,
APIModel: openai.GPT3Dot5Turbo,
APIBaseURL: "https://api.openai.com/v1",
APIUserAgent: util.UserAgent,
}
openAI.APIKey = os.Getenv("SIYUAN_OPENAI_API_KEY")
@ -67,5 +70,8 @@ func NewAI() *AI {
openAI.APIBaseURL = baseURL
}
if userAgent := os.Getenv("SIYUAN_OPENAI_API_USER_AGENT"); "" != userAgent {
openAI.APIUserAgent = userAgent
}
return &AI{OpenAI: openAI}
}

2
kernel/model/ai.go
View file

@ -92,7 +92,7 @@ func chatGPTContinueWrite(msg string, contextMsgs []string, cloud bool) (ret str
if cloud {
gpt = &CloudGPT{}
} else {
gpt = &OpenAIGPT{c: util.NewOpenAIClient(Conf.AI.OpenAI.APIKey, Conf.AI.OpenAI.APIProxy, Conf.AI.OpenAI.APIBaseURL)}
gpt = &OpenAIGPT{c: util.NewOpenAIClient(Conf.AI.OpenAI.APIKey, Conf.AI.OpenAI.APIProxy, Conf.AI.OpenAI.APIBaseURL, Conf.AI.OpenAI.APIUserAgent)}
}
buf := &bytes.Buffer{}

11
kernel/model/conf.go
View file

@ -405,15 +405,24 @@ func InitConf() {
if "" == Conf.AI.OpenAI.APIModel {
Conf.AI.OpenAI.APIModel = openai.GPT3Dot5Turbo
}
if "" == Conf.AI.OpenAI.APIUserAgent {
Conf.AI.OpenAI.APIUserAgent = util.UserAgent
}
if "" != Conf.AI.OpenAI.APIKey {
logging.LogInfof("OpenAI API enabled\n"+
" userAgent=%s\n"+
" baseURL=%s\n"+
" timeout=%ds\n"+
" proxy=%s\n"+
" model=%s\n"+
" maxTokens=%d",
Conf.AI.OpenAI.APIBaseURL, Conf.AI.OpenAI.APITimeout, Conf.AI.OpenAI.APIProxy, Conf.AI.OpenAI.APIModel, Conf.AI.OpenAI.APIMaxTokens)
Conf.AI.OpenAI.APIUserAgent,
Conf.AI.OpenAI.APIBaseURL,
Conf.AI.OpenAI.APITimeout,
Conf.AI.OpenAI.APIProxy,
Conf.AI.OpenAI.APIModel,
Conf.AI.OpenAI.APIMaxTokens)
}
Conf.ReadOnly = util.ReadOnly

21
kernel/util/openai.go
View file

@ -75,17 +75,32 @@ func ChatGPT(msg string, contextMsgs []string, c *openai.Client, model string, m
return
}
func NewOpenAIClient(apiKey, apiProxy, apiBaseURL string) *openai.Client {
func NewOpenAIClient(apiKey, apiProxy, apiBaseURL, apiUserAgent string) *openai.Client {
config := openai.DefaultConfig(apiKey)
transport := &http.Transport{}
if "" != apiProxy {
proxyUrl, err := url.Parse(apiProxy)
if nil != err {
logging.LogErrorf("OpenAI API proxy failed: %v", err)
} else {
config.HTTPClient = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyUrl)}}
transport.Proxy = http.ProxyURL(proxyUrl)
}
}
config.HTTPClient = &http.Client{Transport: newAddHeaderTransport(transport, apiUserAgent)}
config.BaseURL = apiBaseURL
return openai.NewClientWithConfig(config)
}
type AddHeaderTransport struct {
RoundTripper http.RoundTripper
UserAgent string
}
func (adt *AddHeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
req.Header.Add("User-Agent", adt.UserAgent)
return adt.RoundTripper.RoundTrip(req)
}
func newAddHeaderTransport(transport *http.Transport, userAgent string) *AddHeaderTransport {
return &AddHeaderTransport{RoundTripper: transport, UserAgent: userAgent}
}