huangzhenpc/newapi-yx-diy
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 7d6ba52d85
commit 176fd6eda1
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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