105 lines
2.4 KiB
Go
105 lines
2.4 KiB
Go
package kefuWework
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/patrickmn/go-cache"
|
||
gogpt "github.com/sashabaranov/go-openai"
|
||
)
|
||
|
||
// openai key
|
||
var apiKey = "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"
|
||
|
||
// 这是一个可以自定义的 id，用默认值不会有问题
|
||
var userId = "orgId"
|
||
|
||
// 企业微信 token 缓存，请求频次过高可能有一些额外的问题
|
||
var conversationCache = cache.New(5*time.Minute, 5*time.Minute)
|
||
|
||
type ChatGPT struct {
|
||
client *gogpt.Client
|
||
ctx context.Context
|
||
userId string
|
||
}
|
||
|
||
func Chat(c *gin.Context) {
|
||
question := c.Query("question")
|
||
conversationId := c.Query("conversationId")
|
||
apiKey := c.Query("api_key")
|
||
ret, err := AskOnConversation(apiKey, question, conversationId, weworkConversationSize)
|
||
if err != nil {
|
||
c.JSON(500, err.Error())
|
||
return
|
||
}
|
||
c.JSON(200, ret)
|
||
}
|
||
|
||
func AskOnConversation(apiKey, question, conversationId string, size int) (string, error) {
|
||
var messages = []gogpt.ChatCompletionMessage{}
|
||
key := fmt.Sprintf("cache:conversation:%s", conversationId)
|
||
data, found := conversationCache.Get(key)
|
||
if found {
|
||
messages = data.([]gogpt.ChatCompletionMessage)
|
||
}
|
||
messages = append(messages, gogpt.ChatCompletionMessage{
|
||
Role: "system",
|
||
Content: question,
|
||
})
|
||
fmt.Println(messages)
|
||
pivot := size
|
||
if pivot > len(messages) {
|
||
pivot = len(messages)
|
||
}
|
||
messages = messages[len(messages)-pivot:]
|
||
conversationCache.Set(key, messages, 12*time.Hour)
|
||
k, userId := apiKey, userId
|
||
chat := NewGPT(k, userId)
|
||
defer chat.Close()
|
||
answer, err := chat.Chat(messages)
|
||
if err != nil {
|
||
fmt.Print(err.Error())
|
||
}
|
||
return answer, err
|
||
}
|
||
|
||
func (c *ChatGPT) Chat(messages []gogpt.ChatCompletionMessage) (answer string, err error) {
|
||
var msg = gogpt.ChatCompletionMessage{}
|
||
msg.Role = "system"
|
||
req := gogpt.ChatCompletionRequest{
|
||
Model: gogpt.GPT3Dot5Turbo,
|
||
Messages: messages,
|
||
}
|
||
resp, err := c.client.CreateChatCompletion(c.ctx, req)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
answer = resp.Choices[0].Message.Content
|
||
for len(answer) > 0 {
|
||
if answer[0] == '\n' {
|
||
answer = answer[1:]
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
return answer, err
|
||
}
|
||
|
||
func NewGPT(ApiKey, UserId string) *ChatGPT {
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
go func() {
|
||
<-ctx.Done()
|
||
cancel()
|
||
}()
|
||
return &ChatGPT{
|
||
client: gogpt.NewClient(ApiKey),
|
||
ctx: ctx,
|
||
userId: UserId,
|
||
}
|
||
}
|
||
func (c *ChatGPT) Close() {
|
||
c.ctx.Done()
|
||
}
|