mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-18 04:10:51 +08:00
feat(routing): implement unified model routing with OAuth and API key providers
- Added a new routing package to manage provider registration and model resolution. - Introduced Router, Executor, and Provider interfaces to handle different provider types. - Implemented OAuthProvider and APIKeyProvider to support OAuth and API key authentication. - Enhanced DefaultModelMapper to include OAuth model alias handling and fallback mechanisms. - Updated context management in API handlers to preserve fallback models. - Added tests for routing logic and provider selection. - Enhanced Claude request conversion to handle reasoning content based on thinking mode.
This commit is contained in:
@@ -2,7 +2,7 @@ package amp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/http/httputil"
|
||||
@@ -10,64 +10,152 @@ import (
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestFallbackHandler_ModelMapping_PreservesThinkingSuffixAndRewritesResponse(t *testing.T) {
|
||||
// Characterization tests for fallback_handlers.go
|
||||
// These tests capture existing behavior before refactoring to routing layer
|
||||
|
||||
func TestFallbackHandler_WrapHandler_LocalProvider_NoMapping(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client-amp-fallback", "codex", []*registry.ModelInfo{
|
||||
{ID: "test/gpt-5.2", OwnedBy: "openai", Type: "codex"},
|
||||
// Setup: model that has local providers
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
body := `{"model": "gemini-2.5-pro", "messages": [{"role": "user", "content": "hello"}]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/provider/anthropic/v1/messages", bytes.NewReader([]byte(body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
|
||||
// Handler that should be called (not proxy)
|
||||
handlerCalled := false
|
||||
handler := func(c *gin.Context) {
|
||||
handlerCalled = true
|
||||
c.JSON(200, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
// Create fallback handler
|
||||
fh := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
return nil // no proxy
|
||||
})
|
||||
defer reg.UnregisterClient("test-client-amp-fallback")
|
||||
|
||||
// Execute
|
||||
wrapped := fh.WrapHandler(handler)
|
||||
wrapped(c)
|
||||
|
||||
// Assert: handler should be called directly (no mapping needed)
|
||||
assert.True(t, handlerCalled, "handler should be called for local provider")
|
||||
assert.Equal(t, 200, w.Code)
|
||||
}
|
||||
|
||||
func TestFallbackHandler_WrapHandler_MappingApplied(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// Setup: model that needs mapping
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
body := `{"model": "claude-opus-4-5-20251101", "messages": [{"role": "user", "content": "hello"}]}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/provider/anthropic/v1/messages", bytes.NewReader([]byte(body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
|
||||
// Handler to capture rewritten body
|
||||
var capturedBody []byte
|
||||
handler := func(c *gin.Context) {
|
||||
capturedBody, _ = io.ReadAll(c.Request.Body)
|
||||
c.JSON(200, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
// Create fallback handler with mapper
|
||||
mapper := NewModelMapper([]config.AmpModelMapping{
|
||||
{From: "claude-opus-4-5-20251101", To: "claude-opus-4-5-thinking"},
|
||||
})
|
||||
// TODO: Setup oauth aliases for testing
|
||||
|
||||
fh := NewFallbackHandlerWithMapper(
|
||||
func() *httputil.ReverseProxy { return nil },
|
||||
mapper,
|
||||
func() bool { return false },
|
||||
)
|
||||
|
||||
// Execute
|
||||
wrapped := fh.WrapHandler(handler)
|
||||
wrapped(c)
|
||||
|
||||
// Assert: body should be rewritten
|
||||
assert.Contains(t, string(capturedBody), "claude-opus-4-5-thinking")
|
||||
|
||||
// Assert: context should have mapped model
|
||||
mappedModel, exists := c.Get(MappedModelContextKey)
|
||||
assert.True(t, exists, "MappedModelContextKey should be set")
|
||||
assert.NotEmpty(t, mappedModel)
|
||||
}
|
||||
|
||||
func TestFallbackHandler_WrapHandler_ThinkingSuffixPreserved(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
// Model with thinking suffix
|
||||
body := `{"model": "claude-opus-4-5-20251101(xhigh)", "messages": []}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/provider/anthropic/v1/messages", bytes.NewReader([]byte(body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
|
||||
var capturedBody []byte
|
||||
handler := func(c *gin.Context) {
|
||||
capturedBody, _ = io.ReadAll(c.Request.Body)
|
||||
c.JSON(200, gin.H{"status": "ok"})
|
||||
}
|
||||
|
||||
mapper := NewModelMapper([]config.AmpModelMapping{
|
||||
{From: "gpt-5.2", To: "test/gpt-5.2"},
|
||||
{From: "claude-opus-4-5-20251101", To: "claude-opus-4-5-thinking"},
|
||||
})
|
||||
|
||||
fh := NewFallbackHandlerWithMapper(
|
||||
func() *httputil.ReverseProxy { return nil },
|
||||
mapper,
|
||||
func() bool { return false },
|
||||
)
|
||||
|
||||
fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy { return nil }, mapper, nil)
|
||||
wrapped := fh.WrapHandler(handler)
|
||||
wrapped(c)
|
||||
|
||||
// Assert: thinking suffix should be preserved
|
||||
assert.Contains(t, string(capturedBody), "(xhigh)")
|
||||
}
|
||||
|
||||
func TestFallbackHandler_WrapHandler_NoProvider_NoMapping_ProxyEnabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
body := `{"model": "unknown-model", "messages": []}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/provider/anthropic/v1/messages", bytes.NewReader([]byte(body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
|
||||
// Note: Proxy test needs proper setup with reverse proxy
|
||||
|
||||
handler := func(c *gin.Context) {
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"model": req.Model,
|
||||
"seen_model": req.Model,
|
||||
})
|
||||
t.Error("handler should not be called when proxy is available")
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
r.POST("/chat/completions", fallback.WrapHandler(handler))
|
||||
// TODO: Setup proxy properly
|
||||
fh := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
// Return mock proxy
|
||||
return nil
|
||||
})
|
||||
|
||||
reqBody := []byte(`{"model":"gpt-5.2(xhigh)"}`)
|
||||
req := httptest.NewRequest(http.MethodPost, "/chat/completions", bytes.NewReader(reqBody))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
wrapped := fh.WrapHandler(handler)
|
||||
wrapped(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("Expected status 200, got %d", w.Code)
|
||||
}
|
||||
|
||||
var resp struct {
|
||||
Model string `json:"model"`
|
||||
SeenModel string `json:"seen_model"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("Failed to parse response JSON: %v", err)
|
||||
}
|
||||
|
||||
if resp.Model != "gpt-5.2(xhigh)" {
|
||||
t.Errorf("Expected response model gpt-5.2(xhigh), got %s", resp.Model)
|
||||
}
|
||||
if resp.SeenModel != "test/gpt-5.2(xhigh)" {
|
||||
t.Errorf("Expected handler to see test/gpt-5.2(xhigh), got %s", resp.SeenModel)
|
||||
}
|
||||
// Assert: proxy should be called when no local provider
|
||||
// Note: This test needs proxy setup to work properly
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user