Skip to content

fix: handle ollama non-stream tool calls #5865

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Calcium-Ion merged 2 commits into from
Jul 3, 2026
Merged

fix: handle ollama non-stream tool calls #5865

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5f183a4
fix: handle ollama non-stream tool calls
stevenjoezhang May 27, 2026
8b8a08c
test: cover ollama non-stream tool call paths
Calcium-Ion Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 68 additions & 25 deletions relay/channel/ollama/stream.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"`
Expand All @@ -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()
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand All @@ -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))
Expand All @@ -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 {
Expand Down Expand Up @@ -197,14 +220,16 @@ 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)
if ln == "" {
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)
}
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
}
Expand All @@ -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
}
Expand Down
72 changes: 72 additions & 0 deletions relay/channel/ollama/stream_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
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"
)

func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)

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}`
resp := &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(raw)),
}

usage, apiErr := ollamaChatHandler(c, &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "fallback-model"},
}, resp)
if apiErr != nil {
t.Fatalf("ollamaChatHandler returned error: %v", apiErr)
}
if usage.TotalTokens != 12 {
t.Fatalf("unexpected usage total: got %d", usage.TotalTokens)
}

var out dto.OpenAITextResponse
if err := common.Unmarshal(w.Body.Bytes(), &out); err != nil {
t.Fatalf("failed to decode response: %v", err)
}
if len(out.Choices) != 1 {
t.Fatalf("expected one choice, got %d", len(out.Choices))
}
if out.Choices[0].FinishReason != constant.FinishReasonToolCalls {
t.Fatalf("expected finish_reason %q, got %q", constant.FinishReasonToolCalls, out.Choices[0].FinishReason)
}

var toolCalls []dto.ToolCallResponse
if err := common.Unmarshal(out.Choices[0].Message.ToolCalls, &toolCalls); err != nil {
t.Fatalf("failed to decode tool calls: %v", err)
}
if len(toolCalls) != 1 {
t.Fatalf("expected one tool call, got %d", len(toolCalls))
}
if toolCalls[0].ID == "" || toolCalls[0].Type != "function" || toolCalls[0].Function.Name != "get_weather" {
t.Fatalf("unexpected tool call: %+v", toolCalls[0])
}
if toolCalls[0].Index != nil {
t.Fatalf("non-stream tool call should not include index: %+v", toolCalls[0].Index)
}

var args map[string]any
if err := common.Unmarshal([]byte(toolCalls[0].Function.Arguments), &args); err != nil {
t.Fatalf("failed to decode tool call arguments: %v", err)
}
if args["city"] != "Paris" || args["days"] != float64(0) {
t.Fatalf("unexpected tool call arguments: %+v", args)
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved