mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-02 20:40:52 +08:00
fix(aistudio): strip usage metadata from non-final stream chunks
This commit is contained in:
@@ -150,13 +150,15 @@ func (e *AistudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
|
|||||||
case wsrelay.MessageTypeStreamChunk:
|
case wsrelay.MessageTypeStreamChunk:
|
||||||
if len(event.Payload) > 0 {
|
if len(event.Payload) > 0 {
|
||||||
appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload))
|
appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload))
|
||||||
if detail, ok := parseGeminiStreamUsage(event.Payload); ok {
|
filtered := filterAistudioUsageMetadata(event.Payload)
|
||||||
|
if detail, ok := parseGeminiStreamUsage(filtered); ok {
|
||||||
reporter.publish(ctx, detail)
|
reporter.publish(ctx, detail)
|
||||||
}
|
}
|
||||||
}
|
lines := sdktranslator.TranslateStream(ctx, body.toFormat, opts.SourceFormat, req.Model, bytes.Clone(opts.OriginalRequest), translatedReq, bytes.Clone(filtered), ¶m)
|
||||||
lines := sdktranslator.TranslateStream(ctx, body.toFormat, opts.SourceFormat, req.Model, bytes.Clone(opts.OriginalRequest), translatedReq, bytes.Clone(event.Payload), ¶m)
|
for i := range lines {
|
||||||
for i := range lines {
|
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
|
||||||
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
|
}
|
||||||
|
break
|
||||||
}
|
}
|
||||||
case wsrelay.MessageTypeStreamEnd:
|
case wsrelay.MessageTypeStreamEnd:
|
||||||
return
|
return
|
||||||
@@ -281,3 +283,62 @@ func (e *AistudioExecutor) buildEndpoint(model, action, alt string) string {
|
|||||||
}
|
}
|
||||||
return base
|
return base
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// filterAistudioUsageMetadata removes usageMetadata from intermediate SSE events so that
|
||||||
|
// only the terminal chunk retains token statistics.
|
||||||
|
func filterAistudioUsageMetadata(payload []byte) []byte {
|
||||||
|
if len(payload) == 0 {
|
||||||
|
return payload
|
||||||
|
}
|
||||||
|
|
||||||
|
lines := bytes.Split(payload, []byte("\n"))
|
||||||
|
modified := false
|
||||||
|
for idx, line := range lines {
|
||||||
|
trimmed := bytes.TrimSpace(line)
|
||||||
|
if len(trimmed) == 0 || !bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
dataIdx := bytes.Index(line, []byte("data:"))
|
||||||
|
if dataIdx < 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
rawJSON := bytes.TrimSpace(line[dataIdx+5:])
|
||||||
|
cleaned, changed := stripUsageMetadataFromJSON(rawJSON)
|
||||||
|
if !changed {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
var rebuilt []byte
|
||||||
|
rebuilt = append(rebuilt, line[:dataIdx]...)
|
||||||
|
rebuilt = append(rebuilt, []byte("data:")...)
|
||||||
|
if len(cleaned) > 0 {
|
||||||
|
rebuilt = append(rebuilt, ' ')
|
||||||
|
rebuilt = append(rebuilt, cleaned...)
|
||||||
|
}
|
||||||
|
lines[idx] = rebuilt
|
||||||
|
modified = true
|
||||||
|
}
|
||||||
|
if !modified {
|
||||||
|
return payload
|
||||||
|
}
|
||||||
|
return bytes.Join(lines, []byte("\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
// stripUsageMetadataFromJSON drops usageMetadata when no finishReason is present.
|
||||||
|
func stripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) {
|
||||||
|
jsonBytes := bytes.TrimSpace(rawJSON)
|
||||||
|
if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) {
|
||||||
|
return rawJSON, false
|
||||||
|
}
|
||||||
|
finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason")
|
||||||
|
if finishReason.Exists() && finishReason.String() != "" {
|
||||||
|
return rawJSON, false
|
||||||
|
}
|
||||||
|
if !gjson.GetBytes(jsonBytes, "usageMetadata").Exists() {
|
||||||
|
return rawJSON, false
|
||||||
|
}
|
||||||
|
cleaned, err := sjson.DeleteBytes(jsonBytes, "usageMetadata")
|
||||||
|
if err != nil {
|
||||||
|
return rawJSON, false
|
||||||
|
}
|
||||||
|
return cleaned, true
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user