Files
chatgpt-dingtalk/pkg/chatgpt/chatgpt.go
2023-09-07 23:01:14 +08:00

92 lines
2.3 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package chatgpt
import (
"context"
"net/http"
"net/url"
"time"
openai "github.com/sashabaranov/go-openai"
"github.com/eryajf/chatgpt-dingtalk/public"
)
type ChatGPT struct {
client *openai.Client
ctx context.Context
userId string
maxQuestionLen int
maxText int
maxAnswerLen int
timeOut time.Duration // 超时时间, 0表示不超时
doneChan chan struct{}
cancel func()
ChatContext *ChatContext
}
func New(userId string) *ChatGPT {
var ctx context.Context
var cancel func()
ctx, cancel = context.WithTimeout(context.Background(), 600*time.Second)
timeOutChan := make(chan struct{}, 1)
go func() {
<-ctx.Done()
timeOutChan <- struct{}{} // 发送超时信号或是提示结束用于聊天机器人场景配合GetTimeOutChan() 使用
}()
config := openai.DefaultConfig(public.Config.ApiKey)
if public.Config.AzureOn {
config = openai.DefaultAzureConfig(
public.Config.AzureOpenAIToken,
"https://"+public.Config.AzureResourceName+".openai.azure.com",
)
config.APIVersion = public.Config.AzureApiVersion
config.AzureModelMapperFunc = func(model string) string {
return public.Config.AzureDeploymentName
}
} else {
if public.Config.HttpProxy != "" {
config.HTTPClient.Transport = &http.Transport{
// 设置代理
Proxy: func(req *http.Request) (*url.URL, error) {
return url.Parse(public.Config.HttpProxy)
}}
}
if public.Config.BaseURL != "" {
config.BaseURL = public.Config.BaseURL + "/v1"
}
}
return &ChatGPT{
client: openai.NewClientWithConfig(config),
ctx: ctx,
userId: userId,
maxQuestionLen: public.Config.MaxQuestionLen, // 最大问题长度
maxAnswerLen: public.Config.MaxAnswerLen, // 最大答案长度
maxText: public.Config.MaxText, // 最大文本 = 问题 + 回答, 接口限制
timeOut: public.Config.SessionTimeout,
doneChan: timeOutChan,
cancel: func() {
cancel()
},
ChatContext: NewContext(),
}
}
func (c *ChatGPT) Close() {
c.cancel()
}
func (c *ChatGPT) GetDoneChan() chan struct{} {
return c.doneChan
}
func (c *ChatGPT) SetMaxQuestionLen(maxQuestionLen int) int {
if maxQuestionLen > c.maxText-c.maxAnswerLen {
maxQuestionLen = c.maxText - c.maxAnswerLen
}
c.maxQuestionLen = maxQuestionLen
return c.maxQuestionLen
}