Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
4.6 kB · 131 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132package executor
import ( "context" "encoding/json" "io" "net/http" "strings" "testing" "time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson")
type kimiHomeModelCapabilityDispatcher struct{}
func (*kimiHomeModelCapabilityDispatcher) HeartbeatOK() bool { return true }
func (*kimiHomeModelCapabilityDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { return json.Marshal(map[string]any{ "model": "kimi-k3", "provider": "kimi", "model_info": map[string]any{ "id": "kimi-k3", "type": "kimi", "context_length": 1048576, "thinking": map[string]any{ "levels": []string{"low", "high"}, }, "user_defined": false, }, "auth": cliproxyauth.Auth{ ID: "home-kimi-k3-auth", Provider: "kimi", Status: cliproxyauth.StatusActive, Attributes: map[string]string{}, Metadata: map[string]any{ "access_token": "token", "expired": time.Now().Add(time.Hour).Format(time.RFC3339), }, }, })}
func (*kimiHomeModelCapabilityDispatcher) AbortAmbiguousDispatch() {}
func TestKimiHomeModelThinkingCapabilitiesOverrideLocalRegistry(t *testing.T) { tests := []struct { name string stream bool }{ {name: "non-streaming"}, {name: "streaming", stream: true}, }
for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { var upstreamBody []byte ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) { var errRead error upstreamBody, errRead = io.ReadAll(req.Body) if errRead != nil { return nil, errRead } if tt.stream { return &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader( "data: {\"id\":\"chatcmpl_test\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"k3\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n" + "data: [DONE]\n\n", )), }, nil } return &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader( `{"id":"chatcmpl_test","object":"chat.completion","created":1,"model":"k3","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, )), }, nil }))
cfg := &config.Config{RequestRetry: 1} cfg.Home.Enabled = true manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetConfig(cfg) manager.RegisterExecutor(NewKimiExecutor(cfg)) manager.PublishHomeDispatch(&kimiHomeModelCapabilityDispatcher{}, executionregistry.New(), 1)
payload := []byte(`{"model":"kimi-k3","reasoning_effort":"max","messages":[{"role":"user","content":"hello"}]}`) req := cliproxyexecutor.Request{Model: "kimi-k3", Payload: payload} opts := cliproxyexecutor.Options{ Stream: tt.stream, SourceFormat: sdktranslator.FormatOpenAI, ResponseFormat: sdktranslator.FormatOpenAI, OriginalRequest: payload, } if tt.stream { result, errExecute := manager.ExecuteStream(ctx, []string{"kimi"}, req, opts) if errExecute != nil { t.Fatalf("ExecuteStream() error = %v", errExecute) } for chunk := range result.Chunks { if chunk.Err != nil { t.Fatalf("stream chunk error = %v", chunk.Err) } } } else { if _, errExecute := manager.Execute(ctx, []string{"kimi"}, req, opts); errExecute != nil { t.Fatalf("Execute() error = %v", errExecute) } }
if got := gjson.GetBytes(upstreamBody, "thinking.type").String(); got != "enabled" { t.Fatalf("thinking.type = %q, want enabled; body=%s", got, upstreamBody) } if got := gjson.GetBytes(upstreamBody, "thinking.effort").String(); got != "high" { t.Fatalf("thinking.effort = %q, want high from Home model capabilities; body=%s", got, upstreamBody) } if effort := gjson.GetBytes(upstreamBody, "reasoning_effort"); effort.Exists() { t.Fatalf("reasoning_effort should be absent, got %s; body=%s", effort.Raw, upstreamBody) } }) }}