From d16597ca0b0626c362063be1ef3471019dcdcfe9 Mon Sep 17 00:00:00 2001 From: Admin Bot Date: Tue, 15 Sep 2026 09:52:55 +0900 Subject: [PATCH] feat: add upstreamModel config for model name mapping When the upstream LLM server expects a different model name than what clients send, the gateway now rewrites the 'model' field in the request body before forwarding. Config example: models: - name: "ornith:35b" # client-facing name address: "ornith-predictor:80" upstreamModel: "qwen2.5:72b" # what upstream expects Fixes 400 'model is required' errors when upstream model names differ. --- internal/config/config.go | 4 ++ internal/config/loader.go | 18 +++++---- internal/proxy/bodydispatch_test.go | 63 +++++++++++++++++++++++++++++ internal/proxy/router.go | 14 +++++++ 4 files changed, 91 insertions(+), 8 deletions(-) diff --git a/internal/config/config.go b/internal/config/config.go index 13bc39f..89e50ac 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -36,6 +36,10 @@ type ModelUpstream struct { Address string // Path is the upstream path for this model (e.g., "/v1/chat/completions"). Path string + // UpstreamModel is the model name to send to the upstream server. + // If empty, the client-provided model name (Name) is used as-is. + // Use this when the upstream expects a different model name than clients send. + UpstreamModel string // AuthRequired indicates whether this model requires JWT authentication. AuthRequired bool } diff --git a/internal/config/loader.go b/internal/config/loader.go index 344dfc6..e0d9ed3 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -37,10 +37,11 @@ type rawRoute struct { // rawModel represents a single model entry in the YAML configuration. type rawModel struct { - Name string `yaml:"name"` - Address string `yaml:"address"` - Path string `yaml:"path"` - AuthRequired *bool `yaml:"authRequired"` + Name string `yaml:"name"` + Address string `yaml:"address"` + Path string `yaml:"path"` + UpstreamModel string `yaml:"upstreamModel"` + AuthRequired *bool `yaml:"authRequired"` } // rawAdapter represents a service adapter in the YAML configuration. @@ -132,10 +133,11 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo authRequired = *rawModel.AuthRequired } models[rawModel.Name] = &ModelUpstream{ - Name: rawModel.Name, - Address: rawModel.Address, - Path: rawModel.Path, - AuthRequired: authRequired, + Name: rawModel.Name, + Address: rawModel.Address, + Path: rawModel.Path, + UpstreamModel: rawModel.UpstreamModel, + AuthRequired: authRequired, } } diff --git a/internal/proxy/bodydispatch_test.go b/internal/proxy/bodydispatch_test.go index 0e01828..eede0ac 100644 --- a/internal/proxy/bodydispatch_test.go +++ b/internal/proxy/bodydispatch_test.go @@ -381,3 +381,66 @@ func TestBodySizeCappedDispatch(t *testing.T) { t.Errorf("expected 200 for reasonable body, got %d", resp.StatusCode) } } + +// TestUpstreamModelRewrite verifies that the model field is rewritten when upstreamModel is set. +func TestUpstreamModelRewrite(t *testing.T) { + var receivedModel string + + upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var payload map[string]interface{} + json.Unmarshal(body, &payload) + receivedModel = payload["model"].(string) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{}`) + })) + defer upstreamServer.Close() + + upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://") + + cfg := &config.Config{ + Models: map[string]*config.ModelUpstream{ + // Client sends "ornith:35b", upstream expects "qwen2.5:72b-instruct" + "ornith:35b": { + Name: "ornith:35b", + Address: upstreamAddr, + UpstreamModel: "qwen2.5:72b-instruct", + }, + // No rewrite - upstream model same as client model + "reasoning": { + Name: "reasoning", + Address: upstreamAddr, + }, + }, + Routes: make(map[string]*config.Route), + } + + handler := New(cfg) + defer handler.Close() + + server := httptest.NewServer(handler) + defer server.Close() + + // Test 1: Model should be rewritten + requestBody := `{"model":"ornith:35b","messages":[{"role":"user","content":"hi"}]}` + resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody)) + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + + if receivedModel != "qwen2.5:72b-instruct" { + t.Errorf("expected upstream to receive model 'qwen2.5:72b-instruct', got '%s'", receivedModel) + } + + // Test 2: No rewrite when upstreamModel is empty + receivedModel = "" + requestBody = `{"model":"reasoning","messages":[{"role":"user","content":"hi"}]}` + resp, _ = http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody)) + resp.Body.Close() + + if receivedModel != "reasoning" { + t.Errorf("expected upstream to receive model 'reasoning', got '%s'", receivedModel) + } +} diff --git a/internal/proxy/router.go b/internal/proxy/router.go index d35ab58..da95efa 100644 --- a/internal/proxy/router.go +++ b/internal/proxy/router.go @@ -122,6 +122,20 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) { } } + // If upstream expects a different model name, rewrite the body + if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName { + payload["model"] = modelUpstream.UpstreamModel + newBody, err := json.Marshal(payload) + if err != nil { + return nil, &modelValidationError{ + Kind: "invalid_request", + Message: fmt.Sprintf("failed to rewrite model name: %v", err), + } + } + r.Body = io.NopCloser(bytes.NewReader(newBody)) + r.ContentLength = int64(len(newBody)) + } + // Determine the upstream path based on the request path upstreamPath := path if path == "/v1/rerank" {