refactor: standardize dataTag processing across response translators
docker-image / docker (push) Has been cancelled
goreleaser / goreleaser (push) Has been cancelled

- Unified `dataTag` initialization by removing spaces after `data:`.
- Replaced manual slicing with `bytes.TrimSpace` for consistent and robust handling of JSON payloads.
This commit is contained in:
Luis Pater
2025-09-21 11:16:03 +08:00
parent 83a1fa618d
commit e5a6fd2d4f
11 changed files with 38 additions and 51 deletions
+3 -3
View File
@@ -554,7 +554,7 @@ func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName st
rawJSON, _ = sjson.SetBytes(rawJSON, "project", c.GetProjectID())
rawJSON, _ = sjson.SetBytes(rawJSON, "model", modelName)
dataTag := []byte("data: ")
dataTag := []byte("data:")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
@@ -619,7 +619,7 @@ func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName st
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, line[6:], &param)
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, bytes.TrimSpace(line[5:]), &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
@@ -630,7 +630,7 @@ func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName st
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
dataChan <- line[6:]
dataChan <- bytes.TrimSpace(line[5:])
}
c.AddAPIResponseData(ctx, line)
}
+3 -3
View File
@@ -298,7 +298,7 @@ func (c *GeminiClient) SendRawMessageStream(ctx context.Context, modelName strin
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
dataTag := []byte("data: ")
dataTag := []byte("data:")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
@@ -342,7 +342,7 @@ func (c *GeminiClient) SendRawMessageStream(ctx context.Context, modelName strin
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, line[6:], &param)
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, bytes.TrimSpace(line[5:]), &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
@@ -353,7 +353,7 @@ func (c *GeminiClient) SendRawMessageStream(ctx context.Context, modelName strin
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
dataChan <- line[6:]
dataChan <- bytes.TrimSpace(line[5:])
}
c.AddAPIResponseData(ctx, line)
}
+6 -19
View File
@@ -291,9 +291,8 @@ func (c *OpenAICompatibilityClient) SendRawMessageStream(ctx context.Context, mo
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
dataTag := []byte("data: ")
dataUglyTag := []byte("data:") // Some APIs providers don't add space after "data:", fuck for them all
doneTag := []byte("data: [DONE]")
dataTag := []byte("data:")
doneTag := []byte("[DONE]")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
@@ -332,19 +331,10 @@ func (c *OpenAICompatibilityClient) SendRawMessageStream(ctx context.Context, mo
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
if bytes.Equal(line, doneTag) {
if bytes.Equal(bytes.TrimSpace(line[5:]), doneTag) {
break
}
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, line[6:], &param)
for i := 0; i < len(lines); i++ {
c.AddAPIResponseData(ctx, line)
dataChan <- []byte(lines[i])
}
} else if bytes.HasPrefix(line, dataUglyTag) {
if bytes.Equal(line, doneTag) {
break
}
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, line[5:], &param)
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, originalRequestRawJSON, rawJSON, bytes.TrimSpace(line[5:]), &param)
for i := 0; i < len(lines); i++ {
c.AddAPIResponseData(ctx, line)
dataChan <- []byte(lines[i])
@@ -356,13 +346,10 @@ func (c *OpenAICompatibilityClient) SendRawMessageStream(ctx context.Context, mo
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
if bytes.Equal(line, doneTag) {
if bytes.Equal(bytes.TrimSpace(line[5:]), doneTag) {
break
}
c.AddAPIResponseData(newCtx, line[6:])
dataChan <- line[6:]
} else if bytes.HasPrefix(line, dataUglyTag) {
c.AddAPIResponseData(newCtx, line[5:])
c.AddAPIResponseData(newCtx, bytes.TrimSpace(line[5:]))
dataChan <- line[5:]
}
}
+6 -6
View File
@@ -215,8 +215,8 @@ func (c *QwenClient) SendRawMessageStream(ctx context.Context, modelName string,
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
dataTag := []byte("data: ")
doneTag := []byte("data: [DONE]")
dataTag := []byte("data:")
doneTag := []byte("[DONE]")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
@@ -264,7 +264,7 @@ func (c *QwenClient) SendRawMessageStream(ctx context.Context, modelName string,
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
lines := translator.Response(handlerType, c.Type(), ctx, modelName, originalRequestRawJSON, rawJSON, line[6:], &param)
lines := translator.Response(handlerType, c.Type(), ctx, modelName, originalRequestRawJSON, rawJSON, bytes.TrimSpace(line[5:]), &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
@@ -274,9 +274,9 @@ func (c *QwenClient) SendRawMessageStream(ctx context.Context, modelName string,
} else {
for scanner.Scan() {
line := scanner.Bytes()
if !bytes.HasPrefix(line, doneTag) {
if bytes.HasPrefix(line, dataTag) {
dataChan <- line[6:]
if bytes.HasPrefix(line, dataTag) {
if !bytes.Equal(bytes.TrimSpace(line[5:]), doneTag) {
dataChan <- bytes.TrimSpace(line[5:])
}
}
c.AddAPIResponseData(ctx, line)