go如何實現openai伺服器推送

卓能文發表於2024-06-30

最近在玩openai,但opeanai的響應是一個漫長的過程,不能讓瀏覽器端無限制等待,必須將openai的返回資訊實時推送到客戶端,讓使用者獲得一個滿意的互動過程。直接上原始碼:

package chat

import (
	"context"
	"errors"
	"io"
	"os"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/joho/godotenv"

	v1 "SmartPaper/api/chat/v1"

	openai "github.com/sashabaranov/go-openai"
)

func (c *ChatV1) Chat(ctx context.Context, req *v1.ChatReq) (res *v1.ChatRes, err error) {
	err = godotenv.Load(".env")
	if err != nil {
		return res, err
	}

	OPENAI_API_KEY := os.Getenv("OPENAI_API_KEY")
	OPENAI_BASE_URL := os.Getenv("OPENAI_BASE_URL")
	config := openai.DefaultConfig(OPENAI_API_KEY)
	config.BaseURL = OPENAI_BASE_URL
	clientai := openai.NewClientWithConfig(config) // openai.NewClient(OPENAI_API_KEY)

	reqai := openai.ChatCompletionRequest{
		Model: openai.GPT4o,
		// MaxTokens: 128*1024 - 1,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleUser,
				Content: req.Prompt,
			},
		},
		Stream: true,
	}
	stream, err := clientai.CreateChatCompletionStream(ctx, reqai)
	if err != nil {
		return res, err
	}
	defer stream.Close()

	response := g.RequestFromCtx(ctx).Response
	for {
		responseai, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return res, nil
		}

		if err != nil {
			return res, err
		}

		if len(responseai.Choices) > 0 {
			response.Write(responseai.Choices[0].Delta.Content)
			response.Flush()
		}
	}
}

相關文章