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

fix: handle ollama non-stream tool calls (#5865)

Browse Source 
* fix: handle ollama non-stream tool calls

* test: cover ollama non-stream tool call paths

---------

Co-authored-by: CaIon <i@caion.me>
This commit is contained in:
Mimi
2026-07-03 15:35:34 +08:00
committed by GitHub
co-authored by CaIon
parent 1dcb389d00
commit 0977965d93
2 changed files with 164 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
relay/channel/ollama/stream.go
+68 -25
View File
@@ -9,6 +9,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -24,15 +25,10 @@ type ollamaChatStreamChunk struct {
CreatedAt string `json:"created_at"`
// chat
Message *struct {
Role string `json:"role"`
Content string `json:"content"`
Thinking json.RawMessage `json:"thinking"`
ToolCalls []struct {
Function struct {
Name string `json:"name"`
Arguments interface{} `json:"arguments"`
} `json:"function"`
} `json:"tool_calls"`
Role string `json:"role"`
Content string `json:"content"`
Thinking json.RawMessage `json:"thinking"`
ToolCalls []OllamaToolCall `json:"tool_calls"`
} `json:"message"`
// generate
Response string `json:"response"`
@@ -46,6 +42,39 @@ type ollamaChatStreamChunk struct {
EvalDuration int64 `json:"eval_duration"`
}
func ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, includeIndex bool) ([]dto.ToolCallResponse, int) {
if len(toolCalls) == 0 {
return nil, startIndex
}
result := make([]dto.ToolCallResponse, 0, len(toolCalls))
for _, tc := range toolCalls {
var argBytes []byte
var err error
if tc.Function.Arguments == nil {
argBytes = []byte("{}")
} else {
argBytes, err = common.Marshal(tc.Function.Arguments)
if err != nil || len(argBytes) == 0 {
argBytes = []byte("{}")
}
}
tr := dto.ToolCallResponse{
ID: fmt.Sprintf("call_%d", startIndex),
Type: "function",
Function: dto.FunctionResponse{
Name: tc.Function.Name,
Arguments: string(argBytes),
},
}
if includeIndex {
tr.SetIndex(startIndex)
}
startIndex++
result = append(result, tr)
}
return result, startIndex
}
func toUnix(ts string) int64 {
if ts == "" {
return time.Now().Unix()
@@ -87,7 +116,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
continue
}
var chunk ollamaChatStreamChunk
if err := json.Unmarshal([]byte(line), &chunk); err != nil {
if err := common.Unmarshal([]byte(line), &chunk); err != nil {
logger.LogError(c, "ollama stream json decode error: "+err.Error()+" line="+line)
return usage, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
@@ -122,7 +151,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := json.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
@@ -132,16 +161,7 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
}
// tool calls
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
delta.Choices[0].Delta.ToolCalls = make([]dto.ToolCallResponse, 0, len(chunk.Message.ToolCalls))
for _, tc := range chunk.Message.ToolCalls {
// arguments -> string
argBytes, _ := json.Marshal(tc.Function.Arguments)
toolId := fmt.Sprintf("call_%d", toolCallIndex)
tr := dto.ToolCallResponse{ID: toolId, Type: "function", Function: dto.FunctionResponse{Name: tc.Function.Name, Arguments: string(argBytes)}}
tr.SetIndex(toolCallIndex)
toolCallIndex++
delta.Choices[0].Delta.ToolCalls = append(delta.Choices[0].Delta.ToolCalls, tr)
}
delta.Choices[0].Delta.ToolCalls, toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, toolCallIndex, true)
}
if data, err := common.Marshal(delta); err == nil {
_ = helper.StringData(c, string(data))
@@ -157,6 +177,9 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
if finishReason == "" {
finishReason = "stop"
}
if toolCallIndex > 0 {
finishReason = constant.FinishReasonToolCalls
}
// emit stop delta
if stop := helper.GenerateStopResponse(responseId, created, model, finishReason); stop != nil {
if data, err := common.Marshal(stop); err == nil {
@@ -197,6 +220,8 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
reasoningBuilder strings.Builder
lastChunk ollamaChatStreamChunk
parsedAny bool
toolCallIndex int
toolCalls []dto.ToolCallResponse
)
for _, ln := range lines {
ln = strings.TrimSpace(ln)
@@ -204,7 +229,7 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
continue
}
var ck ollamaChatStreamChunk
if err := json.Unmarshal([]byte(ln), &ck); err != nil {
if err := common.Unmarshal([]byte(ln), &ck); err != nil {
if len(lines) == 1 {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
@@ -217,7 +242,7 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := json.Unmarshal(ck.Message.Thinking, &thinkingContent); err == nil {
if err := common.Unmarshal(ck.Message.Thinking, &thinkingContent); err == nil {
reasoningBuilder.WriteString(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
@@ -230,11 +255,16 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
} else if ck.Response != "" {
aggContent.WriteString(ck.Response)
}
if ck.Message != nil && len(ck.Message.ToolCalls) > 0 {
var converted []dto.ToolCallResponse
converted, toolCallIndex = ollamaToolCallsToOpenAI(ck.Message.ToolCalls, toolCallIndex, false)
toolCalls = append(toolCalls, converted...)
}
}
if !parsedAny {
var single ollamaChatStreamChunk
if err := json.Unmarshal(body, &single); err != nil {
if err := common.Unmarshal(body, &single); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
lastChunk = single
@@ -244,7 +274,7 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := json.Unmarshal(single.Message.Thinking, &thinkingContent); err == nil {
if err := common.Unmarshal(single.Message.Thinking, &thinkingContent); err == nil {
reasoningBuilder.WriteString(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
@@ -253,6 +283,11 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
}
}
aggContent.WriteString(single.Message.Content)
if len(single.Message.ToolCalls) > 0 {
var converted []dto.ToolCallResponse
converted, toolCallIndex = ollamaToolCallsToOpenAI(single.Message.ToolCalls, toolCallIndex, false)
toolCalls = append(toolCalls, converted...)
}
} else {
aggContent.WriteString(single.Response)
}
@@ -269,8 +304,16 @@ func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
if finishReason == "" {
finishReason = "stop"
}
if len(toolCalls) > 0 {
finishReason = constant.FinishReasonToolCalls
}
msg := dto.Message{Role: "assistant", Content: contentPtr(content)}
if len(toolCalls) > 0 {
if rawToolCalls, err := common.Marshal(toolCalls); err == nil {
msg.ToolCalls = rawToolCalls
}
}
if rc := reasoningBuilder.String(); rc != "" {
msg.ReasoningContent = &rc
}
relay/channel/ollama/stream_test.go
+96
View File
@@ -0,0 +1,96 @@
package ollama
import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
raw string
}{
{
name: "compact json per-line parse path",
raw: `{"model":"llama3.1","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"Paris","days":0}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`,
},
{
name: "pretty json fallback parse path",
raw: `{
"model": "llama3.1",
"created_at": "2026-05-27T12:00:00Z",
"message": {
"role": "assistant",
"content": "",
"tool_calls": [
{
"function": {
"name": "get_weather",
"arguments": {
"city": "Paris",
"days": 0
}
}
}
]
},
"done": true,
"done_reason": "stop",
"prompt_eval_count": 5,
"eval_count": 7
}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(tt.raw)),
}
usage, apiErr := ollamaChatHandler(c, &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "fallback-model"},
}, resp)
require.Nil(t, apiErr)
require.NotNil(t, usage)
assert.Equal(t, 12, usage.TotalTokens)
var out dto.OpenAITextResponse
require.NoError(t, common.Unmarshal(w.Body.Bytes(), &out))
require.Len(t, out.Choices, 1)
assert.Equal(t, constant.FinishReasonToolCalls, out.Choices[0].FinishReason)
var toolCalls []dto.ToolCallResponse
require.NoError(t, common.Unmarshal(out.Choices[0].Message.ToolCalls, &toolCalls))
require.Len(t, toolCalls, 1)
assert.NotEmpty(t, toolCalls[0].ID)
assert.Equal(t, "function", toolCalls[0].Type)
assert.Equal(t, "get_weather", toolCalls[0].Function.Name)
assert.Nil(t, toolCalls[0].Index)
var args map[string]any
require.NoError(t, common.Unmarshal([]byte(toolCalls[0].Function.Arguments), &args))
assert.Equal(t, "Paris", args["city"])
assert.Equal(t, float64(0), args["days"])
})
}
}