feat: reduce the number of stream memory copies

service/aiproxy/common/render/event.go

@@ -0,0 +1,51 @@
+package render
+
+import (
+	"net/http"
+
+	"github.com/labring/sealos/service/aiproxy/common/conv"
+)
+
+var (
+	contentType = []string{"text/event-stream"}
+	noCache     = []string{"no-cache"}
+)
+
+type OpenAISSE struct {
+	Data string
+}
+
+const (
+	nn   = "\n\n"
+	data = "data: "
+)
+
+var (
+	nnBytes   = conv.StringToBytes(nn)
+	dataBytes = conv.StringToBytes(data)
+)
+
+func (r OpenAISSE) Render(w http.ResponseWriter) error {
+	r.WriteContentType(w)
+
+	for _, bytes := range [][]byte{
+		dataBytes,
+		conv.StringToBytes(r.Data),
+		nnBytes,
+	} {
+		// nosemgrep: go.lang.security.audit.xss.no-direct-write-to-responsewriter.no-direct-write-to-responsewriter
+		if _, err := w.Write(bytes); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (r OpenAISSE) WriteContentType(w http.ResponseWriter) {
+	header := w.Header()
+	header["Content-Type"] = contentType
+
+	if _, exist := header["Cache-Control"]; !exist {
+		header["Cache-Control"] = noCache
+	}
+}

service/aiproxy/common/render/render.go

@@ -3,11 +3,9 @@ package render
 import (
 	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/gin-gonic/gin"
 	json "github.com/json-iterator/go"
-	"github.com/labring/sealos/service/aiproxy/common"
 	"github.com/labring/sealos/service/aiproxy/common/conv"
 )
 
@@ -20,9 +18,7 @@ func StringData(c *gin.Context, str string) {
 	if c.IsAborted() {
 		return
 	}
-	str = strings.TrimPrefix(str, "data:")
-	// str = strings.TrimSuffix(str, "\r")
-	c.Render(-1, common.CustomEvent{Data: "data: " + strings.TrimSpace(str)})
+	c.Render(-1, OpenAISSE{Data: str})
 	c.Writer.Flush()
 }
 
@@ -37,7 +33,8 @@ func ObjectData(c *gin.Context, object any) error {
 	if err != nil {
 		return fmt.Errorf("error marshalling object: %w", err)
 	}
-	StringData(c, conv.BytesToString(jsonData))
+	c.Render(-1, OpenAISSE{Data: conv.BytesToString(jsonData)})
+	c.Writer.Flush()
 	return nil
 }
 

service/aiproxy/relay/adaptor/aws/claude/main.go

@@ -202,7 +202,7 @@ func StreamHandler(meta *meta.Meta, c *gin.Context) (*relaymodel.ErrorWithStatus
 	c.Stream(func(_ io.Writer) bool {
 		event, ok := <-stream.Events()
 		if !ok {
-			render.StringData(c, "[DONE]")
+			render.Done(c)
 			return false
 		}
 

service/aiproxy/relay/adaptor/aws/llama3/main.go

@@ -14,7 +14,6 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
 	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
 	"github.com/gin-gonic/gin"
-	"github.com/labring/sealos/service/aiproxy/common"
 	"github.com/labring/sealos/service/aiproxy/common/random"
 	"github.com/labring/sealos/service/aiproxy/common/render"
 	"github.com/labring/sealos/service/aiproxy/middleware"
@@ -209,7 +208,7 @@ func StreamHandler(meta *meta.Meta, c *gin.Context) (*relaymodel.ErrorWithStatus
 	c.Stream(func(_ io.Writer) bool {
 		event, ok := <-stream.Events()
 		if !ok {
-			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+			render.Done(c)
 			return false
 		}
 

service/aiproxy/relay/adaptor/coze/adaptor.go

@@ -82,9 +82,9 @@ func (a *Adaptor) DoRequest(_ *meta.Meta, _ *gin.Context, req *http.Request) (*h
 func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *relaymodel.Usage, err *relaymodel.ErrorWithStatusCode) {
 	var responseText *string
 	if utils.IsStreamResponse(resp) {
-		err, responseText = StreamHandler(c, resp)
+		err, responseText = StreamHandler(meta, c, resp)
 	} else {
-		err, responseText = Handler(c, resp, meta.InputTokens, meta.ActualModel)
+		err, responseText = Handler(meta, c, resp)
 	}
 	if responseText != nil {
 		usage = openai.ResponseText2Usage(*responseText, meta.ActualModel, meta.InputTokens)

service/aiproxy/relay/adaptor/coze/main.go

@@ -14,6 +14,7 @@ import (
 	"github.com/labring/sealos/service/aiproxy/relay/adaptor/coze/constant/messagetype"
 	"github.com/labring/sealos/service/aiproxy/relay/adaptor/openai"
 	"github.com/labring/sealos/service/aiproxy/relay/constant"
+	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	"github.com/labring/sealos/service/aiproxy/relay/model"
 )
 
@@ -85,7 +86,7 @@ func ResponseCoze2OpenAI(cozeResponse *Response) *openai.TextResponse {
 	return &fullTextResponse
 }
 
-func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *string) {
+func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *string) {
 	defer resp.Body.Close()
 
 	log := middleware.GetLogger(c)
@@ -96,7 +97,6 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 	scanner.Split(bufio.ScanLines)
 
 	common.SetEventStreamHeaders(c)
-	var modelName string
 
 	for scanner.Scan() {
 		data := scanner.Bytes()
@@ -124,7 +124,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 		for _, choice := range response.Choices {
 			responseText += conv.AsString(choice.Delta.Content)
 		}
-		response.Model = modelName
+		response.Model = meta.OriginModel
 		response.Created = createdTime
 
 		_ = render.ObjectData(c, response)
@@ -139,7 +139,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 	return nil, &responseText
 }
 
-func Handler(c *gin.Context, resp *http.Response, _ int, modelName string) (*model.ErrorWithStatusCode, *string) {
+func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *string) {
 	defer resp.Body.Close()
 
 	log := middleware.GetLogger(c)
@@ -159,7 +159,7 @@ func Handler(c *gin.Context, resp *http.Response, _ int, modelName string) (*mod
 		}, nil
 	}
 	fullTextResponse := ResponseCoze2OpenAI(&cozeResponse)
-	fullTextResponse.Model = modelName
+	fullTextResponse.Model = meta.OriginModel
 	jsonResponse, err := json.Marshal(fullTextResponse)
 	if err != nil {
 		return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil