mirror of
https://github.com/siyuan-note/siyuan.git
synced 2025-05-02 23:42:36 +08:00
112 lines
2.8 KiB
Go
112 lines
2.8 KiB
Go
// SiYuan - Build Your Eternal Digital Garden
|
|
// 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 <https://www.gnu.org/licenses/>.
|
|
|
|
package util
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
gogpt "github.com/sashabaranov/go-gpt3"
|
|
"github.com/siyuan-note/logging"
|
|
)
|
|
|
|
var (
|
|
OpenAIAPIKey = ""
|
|
OpenAIAPITimeout = 15 * time.Second
|
|
OpenAIAPIProxy = ""
|
|
OpenAIAPIMaxTokens = 4096
|
|
)
|
|
|
|
func ChatGPT(msg string) (ret string) {
|
|
if "" == OpenAIAPIKey {
|
|
return
|
|
}
|
|
|
|
config := gogpt.DefaultConfig(OpenAIAPIKey)
|
|
if "" != OpenAIAPIProxy {
|
|
proxyUrl, err := url.Parse(OpenAIAPIProxy)
|
|
if nil != err {
|
|
logging.LogErrorf("OpenAI API proxy error: %v", err)
|
|
} else {
|
|
config.HTTPClient = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyUrl)}}
|
|
}
|
|
}
|
|
|
|
c := gogpt.NewClientWithConfig(config)
|
|
ctx, cancel := context.WithTimeout(context.Background(), OpenAIAPITimeout)
|
|
defer cancel()
|
|
req := gogpt.ChatCompletionRequest{
|
|
Model: gogpt.GPT3Dot5Turbo,
|
|
MaxTokens: OpenAIAPIMaxTokens,
|
|
Messages: []gogpt.ChatCompletionMessage{
|
|
{
|
|
Role: "user",
|
|
Content: msg,
|
|
},
|
|
},
|
|
}
|
|
resp, err := c.CreateChatCompletion(ctx, req)
|
|
if nil != err {
|
|
logging.LogErrorf("create chat completion failed: %s", err)
|
|
return
|
|
}
|
|
|
|
if 0 < len(resp.Choices) {
|
|
ret = resp.Choices[0].Message.Content
|
|
ret = strings.TrimSpace(ret)
|
|
}
|
|
return
|
|
}
|
|
|
|
func initOpenAI() {
|
|
OpenAIAPIKey = os.Getenv("SIYUAN_OPENAI_API_KEY")
|
|
if "" == OpenAIAPIKey {
|
|
return
|
|
}
|
|
|
|
timeout := os.Getenv("SIYUAN_OPENAI_API_TIMEOUT")
|
|
if "" != timeout {
|
|
timeoutInt, err := strconv.Atoi(timeout)
|
|
if nil == err {
|
|
OpenAIAPITimeout = time.Duration(timeoutInt) * time.Second
|
|
}
|
|
}
|
|
|
|
proxy := os.Getenv("SIYUAN_OPENAI_API_PROXY")
|
|
if "" != proxy {
|
|
OpenAIAPIProxy = proxy
|
|
}
|
|
|
|
maxTokens := os.Getenv("SIYUAN_OPENAI_API_MAX_TOKENS")
|
|
if "" != maxTokens {
|
|
maxTokensInt, err := strconv.Atoi(maxTokens)
|
|
if nil == err {
|
|
OpenAIAPIMaxTokens = maxTokensInt
|
|
}
|
|
}
|
|
if 1 > OpenAIAPIMaxTokens {
|
|
OpenAIAPIMaxTokens = 4096
|
|
}
|
|
|
|
logging.LogInfof("OpenAI API enabled [maxTokens=%d, timeout=%ds, proxy=%s]", OpenAIAPIMaxTokens, int(OpenAIAPITimeout.Seconds()), OpenAIAPIProxy)
|
|
}
|