From 82cafa527fa90ee5195a3be2cb0bf2e9e47998e9 Mon Sep 17 00:00:00 2001 From: Eike Haller <58111764+eksrha@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:16:26 +0200 Subject: [PATCH] feat(router): accept the advertised embedding model id, report backend limits - strip the "model" field from /v1/embeddings when it names an id the router advertises; TEI rejects anything but its internal path and warned on every request before falling back - /v1/models now includes limits (max_batch_tokens, max_client_batch_size, max_input_length) read from the backend's /info, so clients can size batches --- README.md | 8 +++ internal/proxy/proxy.go | 105 +++++++++++++++++++++++++++++++++-- internal/proxy/proxy_test.go | 89 +++++++++++++++++++++++++++++ 3 files changed, 196 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index f855456..94ee8a6 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,14 @@ curl http://localhost:8080/v1/audio/transcriptions \ curl http://localhost:8080/v1/models ``` +`/v1/models` also reports the request limits of the embedding and reranker +backends (`limits.max_batch_tokens`, `limits.max_client_batch_size`, +`limits.max_input_length`), read from the running backend, so clients can size +batches without trial and error. The model id it advertises +(`BAAI/bge-m3`) can be sent back as `model` on `/v1/embeddings`; the router +strips it before forwarding, because TEI only accepts its internal model path +or no model at all. + ## Architecture ### Init Container Pattern (Embedding + Reranker) diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index cba8a6b..a60aa1d 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -1,6 +1,7 @@ package proxy import ( + "bytes" "context" "encoding/json" "fmt" @@ -120,6 +121,9 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { r.handleModels(w, req) case strings.HasPrefix(path, "/v1/embeddings") || path == "/embed" || path == "/embed_sparse": + if strings.HasPrefix(path, "/v1/embeddings") { + r.stripAdvertisedModel(req, r.cfg.Embedding) + } r.dispatch(w, req, r.cfg.Embedding, "embedding") case strings.HasPrefix(path, "/rerank") || strings.HasPrefix(path, "/v1/rerank"): @@ -219,10 +223,95 @@ func (r *Router) dispatch(w http.ResponseWriter, req *http.Request, b config.Bac // modelObj mirrors the OpenAI /v1/models entry shape. type modelObj struct { - ID string `json:"id"` - Object string `json:"object"` - Created int64 `json:"created"` - OwnedBy string `json:"owned_by"` + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + OwnedBy string `json:"owned_by"` + Limits *modelLimits `json:"limits,omitempty"` +} + +// modelLimits are the request-size limits of the running backend, read from +// its /info endpoint so clients can size batches without trial and error. +type modelLimits struct { + MaxBatchTokens int `json:"max_batch_tokens,omitempty"` // tokens per forward batch + MaxClientBatchSize int `json:"max_client_batch_size,omitempty"` // inputs per request + MaxInputLength int `json:"max_input_length,omitempty"` // tokens per single input +} + +// backendLimits asks a TEI backend for its limits. It returns nil when the +// backend has no /info endpoint or reports nothing usable. +func backendLimits(client *http.Client, b config.Backend) *modelLimits { + if b.Name != "embedding" && b.Name != "reranker" { + return nil + } + resp, err := client.Get(b.BaseURL + "/info") + if err != nil { + return nil + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return nil + } + var info struct { + MaxBatchTokens int `json:"max_batch_tokens"` + MaxClientBatchSize int `json:"max_client_batch_size"` + MaxInputLength int `json:"max_input_length"` + } + if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<16)).Decode(&info); err != nil { + return nil + } + l := &modelLimits{ + MaxBatchTokens: info.MaxBatchTokens, + MaxClientBatchSize: info.MaxClientBatchSize, + MaxInputLength: info.MaxInputLength, + } + if *l == (modelLimits{}) { + return nil + } + return l +} + +// stripAdvertisedModel removes the "model" field from a JSON request body +// when it names one of the ids this backend advertises on /v1/models. +// TEI serves exactly one model and rejects ids other than its internal path +// (logging a warning on every call before falling back); an absent field is +// the form it accepts, so clients can use the advertised id as-is. +func (r *Router) stripAdvertisedModel(req *http.Request, b config.Backend) { + if req.Body == nil || req.Method != http.MethodPost || + !strings.HasPrefix(req.Header.Get("Content-Type"), "application/json") { + return + } + body, err := io.ReadAll(req.Body) + _ = req.Body.Close() + setBody := func(data []byte) { + req.Body = io.NopCloser(bytes.NewReader(data)) + req.ContentLength = int64(len(data)) + req.Header.Set("Content-Length", strconv.Itoa(len(data))) + } + if err != nil { + setBody(body) + return + } + var fields map[string]json.RawMessage + if json.Unmarshal(body, &fields) != nil { + setBody(body) + return + } + var model string + if raw, ok := fields["model"]; !ok || json.Unmarshal(raw, &model) != nil { + setBody(body) + return + } + for _, m := range b.Models { + if model == m { + delete(fields, "model") + if out, err := json.Marshal(fields); err == nil { + setBody(out) + return + } + } + } + setBody(body) } type modelsResp struct { @@ -259,8 +348,12 @@ func (r *Router) handleModels(w http.ResponseWriter, req *http.Request) { valid := r.validateModels(b, upstreamResp.Data) if valid == nil { - models = append(models, r.staticModels(b)...) - continue + valid = r.staticModels(b) + } + if limits := backendLimits(client, b); limits != nil { + for i := range valid { + valid[i].Limits = limits + } } models = append(models, valid...) } diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 9539451..35dbb3a 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "reflect" "strings" "testing" "time" @@ -580,3 +581,91 @@ func TestUnknownPath_404(t *testing.T) { t.Errorf("error = %q, want not_found", errResp["error"]) } } + +func TestEmbeddings_AdvertisedModelIDIsStripped(t *testing.T) { + cases := []struct { + name string + body string + wantBody string // JSON, compared after normalisation + }{ + {"advertised id", `{"model":"BAAI/bge-m3","input":["a","b"]}`, `{"input":["a","b"]}`}, + {"no model", `{"input":"a"}`, `{"input":"a"}`}, + {"backend path", `{"model":"/model","input":"a"}`, `{"model":"/model","input":"a"}`}, + {"other id kept", `{"model":"other","input":"a"}`, `{"model":"other","input":"a"}`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var got string + backend := mockBackend(t, map[string]http.HandlerFunc{ + "/v1/embeddings": func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + got = string(b) + if int64(len(b)) != r.ContentLength { + t.Errorf("Content-Length = %d, body = %d bytes", r.ContentLength, len(b)) + } + jsonOK(`{"object":"list","data":[]}`)(w, r) + }, + }) + cfg := buildConfig( + config.Backend{Name: "embedding", BaseURL: backend.URL, Enabled: true, Timeout: 5 * time.Second, Models: []string{"BAAI/bge-m3"}}, + config.Backend{Name: "reranker"}, + config.Backend{Name: "whisper"}, + ) + router := buildRouter(t, cfg) + req := httptest.NewRequest(http.MethodPost, "/v1/embeddings", strings.NewReader(tc.body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + var gotM, wantM map[string]any + if err := json.Unmarshal([]byte(got), &gotM); err != nil { + t.Fatalf("backend got invalid JSON %q: %v", got, err) + } + _ = json.Unmarshal([]byte(tc.wantBody), &wantM) + if !reflect.DeepEqual(gotM, wantM) { + t.Errorf("backend body = %s, want %s", got, tc.wantBody) + } + }) + } +} + +func TestModels_EmbeddingLimits(t *testing.T) { + embedBackend := mockBackend(t, map[string]http.HandlerFunc{ + "/v1/models": jsonOK(modelsJSON("/model")), // TEI reports its internal path + "/info": jsonOK(`{"max_batch_tokens":4096,"max_client_batch_size":32,"max_input_length":8192}`), + }) + cfg := buildConfig( + config.Backend{Name: "embedding", BaseURL: embedBackend.URL, Enabled: true, Timeout: 5 * time.Second, Models: []string{"BAAI/bge-m3"}}, + config.Backend{Name: "reranker"}, + config.Backend{Name: "whisper"}, + ) + router := buildRouter(t, cfg) + + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/models", nil)) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + var resp struct { + Data []struct { + ID string `json:"id"` + Limits struct { + MaxBatchTokens int `json:"max_batch_tokens"` + MaxClientBatchSize int `json:"max_client_batch_size"` + MaxInputLength int `json:"max_input_length"` + } `json:"limits"` + } `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("invalid JSON: %v", err) + } + if len(resp.Data) != 1 || resp.Data[0].ID != "BAAI/bge-m3" { + t.Fatalf("models = %+v, want one entry BAAI/bge-m3", resp.Data) + } + l := resp.Data[0].Limits + if l.MaxBatchTokens != 4096 || l.MaxClientBatchSize != 32 || l.MaxInputLength != 8192 { + t.Errorf("limits = %+v, want 4096/32/8192", l) + } +}