1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798 |
- package channel
- import (
- "bytes"
- "context"
- "github.com/openai/openai-go"
- "github.com/openai/openai-go/option"
- "github.com/zeromicro/go-zero/core/logx"
- "io"
- "io/ioutil"
- "net/http"
- "strings"
- "time"
- )
- type ChatEngine struct {
- BaseURL string
- APIKey string
- IsFastGpt bool
- }
- func NewChatEngine(baseUrl string, apiKey string, isFastGpt bool) *ChatEngine {
- if baseUrl == "" {
- baseUrl = "https://api.openai.com/v1/"
- }
- if !strings.HasSuffix(baseUrl, "/") {
- baseUrl += "/"
- }
- return &ChatEngine{
- BaseURL: baseUrl,
- APIKey: apiKey,
- IsFastGpt: isFastGpt,
- }
- }
- func (o *ChatEngine) ChatCompletions(ctx context.Context, model openai.ChatModel, msg string, chatId string) (string, error) {
- client := openai.NewClient(
- option.WithBaseURL(o.BaseURL),
- option.WithAPIKey(o.APIKey),
- option.WithMiddleware(o.Logger),
- )
- options := make([]option.RequestOption, 0)
- if o.IsFastGpt {
- options = append(options, option.WithJSONSet("chatId", chatId))
- }
- param := openai.ChatCompletionNewParams{
- Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
- openai.UserMessage(msg),
- }),
- Model: openai.F(model),
- }
- chatCompletion, err := client.Chat.Completions.New(ctx, param, options...)
- if err != nil {
- return "", err
- }
- return chatCompletion.Choices[0].Message.Content, nil
- }
- func (o *ChatEngine) Logger(req *http.Request, next option.MiddlewareNext) (res *http.Response, err error) {
- // Before the request
- start := time.Now()
- //logx.Info("request:", req)
- // 读取并打印请求的 body
- var requestBody []byte
- if req.Body != nil {
- requestBody, err = ioutil.ReadAll(req.Body)
- if err != nil {
- logx.Error("读取请求 body 失败:", err)
- return nil, err
- }
- // 重新设置请求的 body
- req.Body = io.NopCloser(bytes.NewBuffer(requestBody))
- }
- logx.Info("request:", req)
- logx.Info("request body:", string(requestBody))
- // Forward the request to the next handler
- res, err = next(req)
- // Handle stuff after the request
- end := time.Now()
- //耗时
- logx.Info("resp: ", res, err, end.Sub(start))
- return res, err
- }
|