Files
CLIProxyAPI/internal/api/modules/amp/fallback_handlers_test.go
Luis Pater ed5ec5b55c feat(amp): enhance model mapping and Gemini thinking configuration
This commit introduces several improvements to the AMP (Advanced Model Proxy) module:

- **Model Mapping Logic:** The `FallbackHandler` now uses a more robust approach for model mapping. It includes the extraction and preservation of dynamic "thinking suffixes" (e.g., `(xhigh)`) during mapping, ensuring that these configurations are correctly applied to the mapped model. A new `resolveMappedModel` function centralizes this logic for cleaner code.
- **ModelMapper Verification:** The `ModelMapper` in `model_mapping.go` now verifies that the target model of a mapping has available providers *after* normalizing it. This prevents mappings to non-existent or unresolvable models.
- **Gemini Thinking Configuration Cleanup:** In `gemini_thinking.go`, unnecessary `generationConfig.thinkingConfig.include_thoughts` and `generationConfig.thinkingConfig.thinkingBudget` fields are now deleted from the request body when applying Gemini thinking levels. This prevents potential conflicts or redundant configurations.
- **Testing:** A new test case `TestModelMapper_MapModel_TargetWithThinkingSuffix` has been added to `model_mapping_test.go` to specifically cover the preservation of thinking suffixes during model mapping.
2025-12-20 22:19:35 +08:00

74 lines
2.0 KiB
Go

package amp
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"net/http/httputil"
"testing"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
)
func TestFallbackHandler_ModelMapping_PreservesThinkingSuffixAndRewritesResponse(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"},
})
defer reg.UnregisterClient("test-client-amp-fallback")
mapper := NewModelMapper([]config.AmpModelMapping{
{From: "gpt-5.2", To: "test/gpt-5.2"},
})
fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy { return nil }, mapper, nil)
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,
})
}
r := gin.New()
r.POST("/chat/completions", fallback.WrapHandler(handler))
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)
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)
}
}