diff --git a/relay/channel/ollama/stream.go b/relay/channel/ollama/stream.go index ee3a8034..b57761fa 100644 --- a/relay/channel/ollama/stream.go +++ b/relay/channel/ollama/stream.go @@ -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 } diff --git a/relay/channel/ollama/stream_test.go b/relay/channel/ollama/stream_test.go new file mode 100644 index 00000000..8d28396c --- /dev/null +++ b/relay/channel/ollama/stream_test.go @@ -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"]) + }) + } +}