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.
This commit is contained in:
@@ -36,6 +36,10 @@ type ModelUpstream struct {
|
|||||||
Address string
|
Address string
|
||||||
// Path is the upstream path for this model (e.g., "/v1/chat/completions").
|
// Path is the upstream path for this model (e.g., "/v1/chat/completions").
|
||||||
Path string
|
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 indicates whether this model requires JWT authentication.
|
||||||
AuthRequired bool
|
AuthRequired bool
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -40,6 +40,7 @@ type rawModel struct {
|
|||||||
Name string `yaml:"name"`
|
Name string `yaml:"name"`
|
||||||
Address string `yaml:"address"`
|
Address string `yaml:"address"`
|
||||||
Path string `yaml:"path"`
|
Path string `yaml:"path"`
|
||||||
|
UpstreamModel string `yaml:"upstreamModel"`
|
||||||
AuthRequired *bool `yaml:"authRequired"`
|
AuthRequired *bool `yaml:"authRequired"`
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -135,6 +136,7 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo
|
|||||||
Name: rawModel.Name,
|
Name: rawModel.Name,
|
||||||
Address: rawModel.Address,
|
Address: rawModel.Address,
|
||||||
Path: rawModel.Path,
|
Path: rawModel.Path,
|
||||||
|
UpstreamModel: rawModel.UpstreamModel,
|
||||||
AuthRequired: authRequired,
|
AuthRequired: authRequired,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -381,3 +381,66 @@ func TestBodySizeCappedDispatch(t *testing.T) {
|
|||||||
t.Errorf("expected 200 for reasonable body, got %d", resp.StatusCode)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
|
// Determine the upstream path based on the request path
|
||||||
upstreamPath := path
|
upstreamPath := path
|
||||||
if path == "/v1/rerank" {
|
if path == "/v1/rerank" {
|
||||||
|
|||||||
Reference in New Issue
Block a user