cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
67 lines
3.5 KiB
Go
67 lines
3.5 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
hostdto "github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/model"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestRelayAttemptSwitchSafety(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name, reason, previous, conversation string
|
|
count int
|
|
pin, cancel, written, disabled, want bool
|
|
}{
|
|
{name: "explicit mismatch", reason: "model_mismatch", count: 1, want: true},
|
|
{name: "token disabled", reason: "model_mismatch", count: 1, disabled: true},
|
|
{name: "third channel allowed", reason: "model_mismatch", count: 2, want: true},
|
|
{name: "third channel final", reason: "model_mismatch", count: 3},
|
|
{name: "missing model", reason: "missing_model", count: 1},
|
|
{name: "confirmation timeout", reason: "confirmation_timeout", count: 1},
|
|
{name: "pinned", reason: "model_mismatch", count: 1, pin: true},
|
|
{name: "cancelled", reason: "model_mismatch", count: 1, cancel: true},
|
|
{name: "already delivered", reason: "model_mismatch", count: 1, written: true},
|
|
{name: "upstream response id", reason: "model_mismatch", count: 1, previous: "resp_private"},
|
|
{name: "upstream conversation", reason: "model_mismatch", count: 1, conversation: `"conv_private"`},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(ctx)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "a"}}})
|
|
if tc.pin {
|
|
service.GetChannelConstraints(c).AddPin(hostdto.ChannelPin{ChannelId: 1, Source: hostdto.PinSourceToken, Rank: hostdto.PinRankToken, RetryMode: hostdto.PinRetrySingleAttempt})
|
|
}
|
|
if tc.cancel {
|
|
cancel()
|
|
}
|
|
if tc.written {
|
|
_, _ = c.Writer.WriteString("already sent")
|
|
}
|
|
info := &relaycommon.RelayInfo{DurableAttemptBilling: true, UpstreamRequests: tc.count, Request: &dto.OpenAIResponsesRequest{PreviousResponseID: tc.previous, Conversation: []byte(tc.conversation)}, ResponseGuard: &relaycommon.ResponseModelGuard{Attempts: []*relaycommon.ResponseGuardAttempt{{Reason: tc.reason}}}}
|
|
err := types.NewError(errors.New("guard rejected"), relaycommon.ResponseGuardErrorCode)
|
|
info.AstraAutoSwitch = !tc.disabled
|
|
require.Equal(t, tc.want, canSwitchRejectedAstra(c, info, err))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRelayAttemptGuardResult(t *testing.T) {
|
|
require.Equal(t, "mismatch_rejected", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "model_mismatch", Number: 1}))
|
|
require.Equal(t, "switched_success", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 2}))
|
|
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 1}))
|
|
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra", ReturnedModel: "gpt-6-astra"}))
|
|
require.Equal(t, "unchecked", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra"}))
|
|
require.Equal(t, "transport_error", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "transport_error", Number: 1}))
|
|
}
|