100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package ollama
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
|
|
"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
|
|
wantID 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":[{"id":"call_upstream","function":{"name":"get_weather","arguments":{"city":"Paris","days":0}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`,
|
|
wantID: "call_upstream",
|
|
},
|
|
{
|
|
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
|
|
}`,
|
|
wantID: "call_0",
|
|
},
|
|
}
|
|
|
|
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.Equal(t, tt.wantID, 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"])
|
|
})
|
|
}
|
|
}
|