huangzhenpc/new-api-hunter
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 优化ollamaStreamHandler中的停止和最终使用响应逻辑

Browse Source
This commit is contained in:
somnifex
2025-09-15 23:23:53 +08:00
parent 503a2669d5
commit a9354df549
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
relay/channel/ollama/stream.go
View File

@@ -114,15 +114,23 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
continue
}
// done frame
// finalize once and break loop
usage.PromptTokens = chunk.PromptEvalCount
usage.CompletionTokens = chunk.EvalCount
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
finishReason := chunk.DoneReason
if finishReason == "" { finishReason = "stop" }
stop := helper.GenerateStopResponse(responseId, created, model, finishReason)
if data, err := common.Marshal(stop); err == nil { _ = helper.StringData(c, string(data)) }
final := helper.GenerateFinalUsageResponse(responseId, created, model, *usage)
if data, err := common.Marshal(final); err == nil { _ = helper.StringData(c, string(data)) }
// emit stop delta
if stop := helper.GenerateStopResponse(responseId, created, model, finishReason); stop != nil {
if data, err := common.Marshal(stop); err == nil { _ = helper.StringData(c, string(data)) }
}
// emit usage frame
if final := helper.GenerateFinalUsageResponse(responseId, created, model, *usage); final != nil {
if data, err := common.Marshal(final); err == nil { _ = helper.StringData(c, string(data)) }
}
// send [DONE]
helper.Done(c)
break
}
if err := scanner.Err(); err != nil && err != io.EOF { logger.LogError(c, "ollama stream scan error: "+err.Error()) }
return usage, nil