feat(router): accept the advertised embedding model id, report backend limits

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)

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...)
 	}

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)
+	}
+}