Golang 使用 Gin 请求和转发 Openai API 使用 event-stream 格式

Qiang 发布在技术 1

04/17/2023 更新:

两个要点:

  1. gin 不得使用 Gzip 压缩数据;
    如果使用 Gzip 中间件,那么所有数据流会在事件结束时一起发送给请求方,而不是刷新时发送。参见此问答评论。
  2. 回复的数据前缀需要与客户端(前端)保持一致,如 datamessage
    如果前后端采用的数据结构不一致,会导致无法正常接收数据,比如默认使用 data: <message> 格式发送,那么客户端应接收 data 数据,如果使用 message: <message> 格式发送,应接收 message 数据。

以下为原回答:

使用了 go-gpt3sse 两个包,主要流程大致流程如下。

func main() {
	router.Any("/openai/completions", proxy, completionsOpenAI)
}

func completionsOpenAI(c *gin.Context) {
	// 定义请求 Json 数据结构
	var data gogpt.CompletionRequest
	// 获取请求的 Json 数据,并转化为对象
	if err := c.ShouldBindJSON(&data); err != nil {
		c.String(http.StatusBadRequest, err.Error())
		return
	}

	// 创建一个 go-gpt3 客户端,指定 api key
	cc := gogpt.NewClient(config.AiKey)

	// 创建一个 stream 会话
	stream, err := cc.CreateCompletionStream(ctx, data)
	if err != nil {
		// 创建失败,则打印错误信息并返回
		fmt.Printf("CompletionStream error: %v\n", err)
		return
	}
	// 在函数结束后,关闭 stream 流
	defer stream.Close()

	// 循环获取 stream 流数据
	for {
		// 接收一次响应数据
		response, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			// 流式传输结束,接收到 `[DONE]` 消息
			fmt.Println("Stream finished")
			return
		}

		if err != nil {
			// 接收数据失败,返回
			fmt.Printf("Stream error: %v\n", err)
			return
		}

		fmt.Printf("Stream response: %v\n", response)
		// 向请求响应体里写 `event-stream` 格式的数据
		sse.Encode(c.Writer, sse.Event{
			Data: data,
		})
		// 刷新数据,以通知请求端
		c.Writer.Flush()
	}
}

路过了部分异常等状态设置和响应体的 Header 设置等,请自行根据业务需要配置。

TOP
前往 GitHub Discussion 评论