From 76fb7879e87fbecea984fba08557adb3bd331d4c Mon Sep 17 00:00:00 2001 From: yetone Date: Sun, 4 Oct 2026 16:24:08 +0800 Subject: [PATCH] =?UTF-8?q?gateway:=20/v1/embeddings=20and=20/v1/rerank=20?= =?UTF-8?q?go=20to=20the=20model's=20provider=20(#765,=20dbydd:=20a=20loca?= =?UTF-8?q?l=20RAG=20stack=20had=20to=20send=20its=20embeddings=20around?= =?UTF-8?q?=20magpie,=20with=20a=20second=20base=20URL=20and=20key,=20outs?= =?UTF-8?q?ide=20the=20usage=20ledger)=20=E2=80=94=20magpie=20answered=20b?= =?UTF-8?q?oth=20with=20its=20catch-all=20404.=20retrieve=20(embeddings.go?= =?UTF-8?q?)=20takes=20the=20body=20as=20the=20agent=20sent=20it,=20resolv?= =?UTF-8?q?es=20its=20model=20as=20a=20turn's=20is,=20puts=20the=20provide?= =?UTF-8?q?r's=20own=20model=20id=20in,=20masks=20secrets=20as=20the=20set?= =?UTF-8?q?tings=20say=20(a=20document=20a=20reranker=20gives=20back=20has?= =?UTF-8?q?=20them=20again),=20and=20posts=20it=20to=20the=20provider's=20?= =?UTF-8?q?OpenAI-style=20base=20+=20/embeddings=20or=20/rerank,=20signed?= =?UTF-8?q?=20as=20the=20provider=20signs;=20the=20vendor's=20answer=20com?= =?UTF-8?q?es=20back=20as=20it=20said=20it,=20its=20refusal=20with=20its?= =?UTF-8?q?=20code=20and=20words.=20The=20call=20is=20on=20the=20Requests?= =?UTF-8?q?=20page=20and=20in=20the=20usage=20ledger=20(operation=20embedd?= =?UTF-8?q?ings=20/=20rerank)=20with=20the=20tokens=20the=20answer=20count?= =?UTF-8?q?ed=20(prompt=5Ftokens,=20total=5Ftokens,=20or=20Cohere's=20bill?= =?UTF-8?q?ed=20input).=20/embeddings=20and=20/rerank=20answer=20too,=20an?= =?UTF-8?q?d=20the=20404's=20list=20and=20GET=20/=20name=20them.=20With=20?= =?UTF-8?q?the=20owner's=20real=20keys=20under=20a=20sandbox=20HOME:=20aih?= =?UTF-8?q?ubmix/text-embedding-3-small=20with=20dimensions=204=20returned?= =?UTF-8?q?=20two=204-d=20vectors=20and=203=20tokens,=20aihubmix/jina-rera?= =?UTF-8?q?nker-m0=20ranked=20the=20magpie=20document=20first=20(37=20toke?= =?UTF-8?q?ns,=20recorded=20as=20rerank),=20and=20OpenRouter's=20402=20for?= =?UTF-8?q?=20no=20credits=20came=20back=20as=20its=20own=20words.=20TestE?= =?UTF-8?q?mbeddingsGoToTheModelsProvider,=20TestRerankGoesToTheModelsProv?= =?UTF-8?q?ider,=20TestRetrievalFailures=20and=20TestRetrievalRedacts=20fa?= =?UTF-8?q?il=20on=20the=20old=20code=20(404).=20go=20vet,=20builds=20(dar?= =?UTF-8?q?win,=20GOOS=3Dlinux,=20GOOS=3Dwindows)=20and=20go=20test=20-tag?= =?UTF-8?q?s=20nogui=20-count=3D1=20./...=20under=20a=20temp=20HOME=20pass?= =?UTF-8?q?=20(TestClientGoneRecordedAsCanceled=20flaked=20once=20under=20?= =?UTF-8?q?load=20and=20passes=20alone).?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/gateway/embeddings.go | 164 ++++++++++++++++++++++++++++ internal/gateway/embeddings_test.go | 163 +++++++++++++++++++++++++++ internal/gateway/gateway.go | 8 +- 3 files changed, 333 insertions(+), 2 deletions(-) create mode 100644 internal/gateway/embeddings.go create mode 100644 internal/gateway/embeddings_test.go diff --git a/internal/gateway/embeddings.go b/internal/gateway/embeddings.go new file mode 100644 index 00000000..267f7412 --- /dev/null +++ b/internal/gateway/embeddings.go @@ -0,0 +1,164 @@ +package gateway + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/yetone/magpie/internal/provider" + "github.com/yetone/magpie/internal/usage" +) + +// retrieveTimeout bounds an embeddings or rerank call: a vendor answers +// one in seconds, a long batch in a minute or two. +const retrieveTimeout = 3 * time.Minute + +// retrieve serves the retrieval APIs (#765): POST /v1/embeddings in +// OpenAI's shape and POST /v1/rerank in the shape Cohere, Jina and Voyage +// share. Neither is a chat, so neither is translated: the body goes to the +// model's provider at its OpenAI-style base + path as the agent sent it, +// the model id the provider's own, and the vendor's answer comes back as +// it said it. The call is recorded and its tokens counted like a turn's. +func (s *Server) retrieve(path, operation string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + body, admitted := s.requestBody(w, r, provider.Chat) + if !admitted { + return + } + start := time.Now() + var req map[string]json.RawMessage + if json.Unmarshal(body, &req) != nil || req == nil { + writeError(w, provider.Chat, 400, "the body isn't a JSON object") + return + } + var asked string + json.Unmarshal(req["model"], &asked) + if asked == "" { + writeError(w, provider.Chat, 400, "name the model: \"model\" is a magpie model id, as /v1/models lists them") + return + } + who := callerOf(r) + call := Call{Time: start, From: provider.Chat, Agent: who.agent, Via: who.via, Model: asked} + usage.Saw(agentOf(r)) + p, model, ok := provider.Resolve(asked) + if !ok { + msg := fmt.Sprintf("magpie knows no model %q", asked) + if off, isOff := provider.SwitchedOff(asked); isOff { + msg = switchedOff(off, asked) + } + call.Status, call.Error = 404, msg + writeError(w, provider.Chat, 404, msg) + s.record(call) + return + } + call.Provider, call.To = p.ID, provider.Chat + fail := func(code int, msg string) { + call.Status, call.Error, call.Millis = code, msg, time.Since(start).Milliseconds() + writeError(w, provider.Chat, code, msg) + s.record(call) + } + base := strings.TrimRight(p.Base(provider.Chat), "/") + if base == "" { + fail(400, fmt.Sprintf("%s has no OpenAI-style API for %s", p.Name, path)) + return + } + req["model"], _ = json.Marshal(model) + out, _ := json.Marshal(req) + var done func() + w, out, done = redacted(w, out) + defer done() + ctx, cancel := context.WithTimeout(r.Context(), retrieveTimeout) + defer cancel() + b, code, err := s.retrieveFrom(ctx, p, base+path, out) + call.Millis = time.Since(start).Milliseconds() + call.Status = code + in := retrievedTokens(b) + call.Usage.Input = in + providerKeyID, providerKeyName := "", "" + if p.Account == nil && p.Key != "" { + providerKeyID, providerKeyName = provider.KeyID(p.Key), p.KeyName + } + appendUsage(r, usage.Record{Operation: operation, Time: start, Agent: call.Agent, Via: call.Via, Provider: p.ID, Host: p.Where(), Model: model, Requested: asked, ProviderKeyID: providerKeyID, ProviderKeyName: providerKeyName, ProviderAccount: accountOf(p), + Input: in, Millis: call.Millis, Status: code, Session: sessionOf(r.Header)}) + if err != nil { + call.Error = err.Error() + s.record(call) + writeError(w, provider.Chat, code, err.Error()) + return + } + s.record(call) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + w.Write(b) + } +} + +// retrieveFrom posts body to url as the provider signs its requests, and +// reads the answer; a failure's code is the vendor's, with its message. +func (s *Server) retrieveFrom(ctx context.Context, p provider.Provider, url string, body []byte) ([]byte, int, error) { + ctx = p.Via(ctx) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, 500, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + if p.IsRemoteMagpie() { + passOnCaller(ctx, req) + } + if err := p.Sign(ctx, req, provider.Chat, body); err != nil { + return nil, 502, err + } + res, err := p.Do(s.client, req) + if err != nil { + if ctx.Err() != nil { + return nil, 504, fmt.Errorf("%s didn't answer in %s", p.Name, retrieveTimeout) + } + return nil, 502, err + } + defer res.Body.Close() + b, err := io.ReadAll(io.LimitReader(res.Body, 256<<20)) + if err != nil { + return nil, 502, err + } + if res.StatusCode >= 300 { + hint := "" + if res.StatusCode == http.StatusNotFound || res.StatusCode == http.StatusMethodNotAllowed { + hint = " — " + p.Name + " may not serve this API" + } + return b, res.StatusCode, fmt.Errorf("%s: %d %s%s%s", provider.HostOf(url), res.StatusCode, http.StatusText(res.StatusCode), vendorSaid(vendorMessage(b)), hint) + } + return b, res.StatusCode, nil +} + +// retrievedTokens is what an answer says it counted: OpenAI's +// prompt_tokens, Jina's and Voyage's total_tokens, or Cohere's billed +// input tokens. +func retrievedTokens(b []byte) int { + var a struct { + Usage struct { + Prompt int `json:"prompt_tokens"` + Input int `json:"input_tokens"` + Total int `json:"total_tokens"` + } `json:"usage"` + Meta struct { + Billed struct { + Input int `json:"input_tokens"` + } `json:"billed_units"` + } `json:"meta"` + } + if json.Unmarshal(b, &a) != nil { + return 0 + } + for _, n := range []int{a.Usage.Prompt, a.Usage.Input, a.Usage.Total, a.Meta.Billed.Input} { + if n > 0 { + return n + } + } + return 0 +} diff --git a/internal/gateway/embeddings_test.go b/internal/gateway/embeddings_test.go new file mode 100644 index 00000000..bab50dd9 --- /dev/null +++ b/internal/gateway/embeddings_test.go @@ -0,0 +1,163 @@ +package gateway + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" + "github.com/yetone/magpie/internal/settings" + "github.com/yetone/magpie/internal/usage" +) + +// library is a vendor with an embeddings model and a reranker; what each +// path was sent is kept. +type library struct { + mu sync.Mutex + sent map[string][]string +} + +func (l *library) ServeHTTP(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + l.mu.Lock() + if l.sent == nil { + l.sent = map[string][]string{} + } + l.sent[r.URL.Path] = append(l.sent[r.URL.Path], string(body)) + l.mu.Unlock() + if r.Header.Get("Authorization") != "Bearer key" { + w.WriteHeader(401) + return + } + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/v1/embeddings": + if strings.Contains(string(body), "chat-model") { + w.WriteHeader(400) + io.WriteString(w, `{"error":{"message":"chat-model is not an embedding model"}}`) + return + } + io.WriteString(w, `{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,-0.2]},{"object":"embedding","index":1,"embedding":[0.3,0.4]}],"model":"embed-1","usage":{"prompt_tokens":5,"total_tokens":5}}`) + case "/v1/rerank": + // Jina's shape, the documents given back + io.WriteString(w, `{"model":"rerank-1","results":[{"index":1,"relevance_score":0.9,"document":{"text":`+strings.SplitN(strings.SplitN(string(body), `"documents":[`, 2)[1], ",", 2)[0]+`}},{"index":0,"relevance_score":0.1}],"usage":{"total_tokens":12}}`) + default: + w.WriteHeader(404) + } +} + +func (l *library) got(path string) []string { + l.mu.Lock() + defer l.mu.Unlock() + return append([]string(nil), l.sent[path]...) +} + +func shelved(t *testing.T) (*Server, *library) { + t.Helper() + fresh(t) + l := &library{} + up := httptest.NewServer(l) + t.Cleanup(up.Close) + if err := provider.Save(provider.Provider{ID: "lib", Name: "Lib", Chat: up.URL + "/v1", Key: "key", Models: []string{"embed-1", "rerank-1", "chat-model"}}); err != nil { + t.Fatal(err) + } + return New(), l +} + +func postRetrieval(t *testing.T, s *Server, path, body string) (int, string) { + t.Helper() + rec := httptest.NewRecorder() + req := httptest.NewRequest("POST", path, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + s.Handler().ServeHTTP(rec, req) + return rec.Code, rec.Body.String() +} + +// #765: an embeddings request goes to the model's provider as it was sent, +// the model its own id, and the vendor's answer comes back; the call is +// in the usage ledger with its tokens. +func TestEmbeddingsGoToTheModelsProvider(t *testing.T) { + s, l := shelved(t) + code, raw := postRetrieval(t, s, "/v1/embeddings", `{"model":"lib/embed-1","input":["hello","world"],"encoding_format":"float","dimensions":2}`) + if code != 200 || !strings.Contains(raw, `"embedding":[0.1,-0.2]`) || !strings.Contains(raw, `"prompt_tokens":5`) { + t.Fatalf("%d %s", code, raw) + } + sent := l.got("/v1/embeddings") + if len(sent) != 1 || !strings.Contains(sent[0], `"model":"embed-1"`) || !strings.Contains(sent[0], `"input":["hello","world"]`) || !strings.Contains(sent[0], `"dimensions":2`) || !strings.Contains(sent[0], `"encoding_format":"float"`) { + t.Fatalf("vendor was sent %v", sent) + } + // the bare /embeddings an SDK given magpie's root as its base posts to + if code, raw := postRetrieval(t, s, "/embeddings", `{"model":"embed-1","input":"hi"}`); code != 200 { + t.Fatalf("bare path: %d %s", code, raw) + } + var recs []usage.Record + for i := 0; i < 50; i++ { + if recs = usage.Load(time.Time{}); len(recs) >= 2 { + break + } + time.Sleep(20 * time.Millisecond) + } + if len(recs) < 1 || recs[0].Operation != "embeddings" || recs[0].Provider != "lib" || recs[0].Model != "embed-1" || recs[0].Requested != "lib/embed-1" || recs[0].Input != 5 || recs[0].Status != 200 { + t.Fatalf("usage %+v", recs) + } +} + +func TestRerankGoesToTheModelsProvider(t *testing.T) { + s, l := shelved(t) + code, raw := postRetrieval(t, s, "/v1/rerank", `{"model":"lib/rerank-1","query":"magpie","documents":["a crow","a magpie"],"top_n":2,"return_documents":true}`) + if code != 200 || !strings.Contains(raw, `"relevance_score":0.9`) { + t.Fatalf("%d %s", code, raw) + } + sent := l.got("/v1/rerank") + if len(sent) != 1 || !strings.Contains(sent[0], `"model":"rerank-1"`) || !strings.Contains(sent[0], `"query":"magpie"`) || !strings.Contains(sent[0], `"top_n":2`) { + t.Fatalf("vendor was sent %v", sent) + } +} + +// What the vendor refuses comes back with its code and its words; a model +// magpie doesn't know is a 404 before any vendor is asked. +func TestRetrievalFailures(t *testing.T) { + s, l := shelved(t) + if code, raw := postRetrieval(t, s, "/v1/embeddings", `{"model":"lib/chat-model","input":"hi"}`); code != 400 || !strings.Contains(raw, "chat-model is not an embedding model") { + t.Fatalf("vendor refusal: %d %s", code, raw) + } + if code, raw := postRetrieval(t, s, "/v1/embeddings", `{"model":"nobody/embed","input":"hi"}`); code != 404 || !strings.Contains(raw, `nobody/embed`) { + t.Fatalf("unknown model: %d %s", code, raw) + } + if code, raw := postRetrieval(t, s, "/v1/embeddings", `{"input":"hi"}`); code != 400 || !strings.Contains(raw, "name the model") { + t.Fatalf("no model: %d %s", code, raw) + } + if code, _ := postRetrieval(t, s, "/v1/embeddings", `not json`); code != 400 { + t.Fatalf("not json: %d", code) + } + if n := len(l.got("/v1/embeddings")); n != 1 { + t.Fatalf("vendor asked %d times, want 1", n) + } +} + +// With secrets kept, what is embedded or ranked has them masked on the +// way out, and a document the reranker gives back has them again. +func TestRetrievalRedacts(t *testing.T) { + s, l := shelved(t) + if err := settings.Save(settings.Settings{Redact: true}); err != nil { + t.Fatal(err) + } + secret := "sk-ant-api03-" + strings.Repeat("Ab3x", 20) + if code, raw := postRetrieval(t, s, "/v1/embeddings", `{"model":"lib/embed-1","input":"key `+secret+`"}`); code != 200 { + t.Fatalf("%d %s", code, raw) + } + if sent := l.got("/v1/embeddings"); len(sent) != 1 || strings.Contains(sent[0], secret) { + t.Fatalf("vendor was sent %v", sent) + } + code, raw := postRetrieval(t, s, "/v1/rerank", `{"model":"lib/rerank-1","query":"q","documents":["`+secret+`","b"],"return_documents":true}`) + if code != 200 || !strings.Contains(raw, secret) { + t.Fatalf("%d %s", code, raw) + } + if sent := l.got("/v1/rerank"); len(sent) != 1 || strings.Contains(sent[0], secret) { + t.Fatalf("vendor was sent %v", sent) + } +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 5e32ace2..0e02c3d0 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -469,6 +469,10 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /images/generations", s.images(false)) mux.HandleFunc("POST /v1/images/edits", s.images(true)) mux.HandleFunc("POST /images/edits", s.images(true)) + mux.HandleFunc("POST /v1/embeddings", s.retrieve("/embeddings", "embeddings")) + mux.HandleFunc("POST /embeddings", s.retrieve("/embeddings", "embeddings")) + mux.HandleFunc("POST /v1/rerank", s.retrieve("/rerank", "rerank")) + mux.HandleFunc("POST /rerank", s.retrieve("/rerank", "rerank")) mux.HandleFunc("POST /v1/videos", s.videosCreate) mux.HandleFunc("POST /videos", s.videosCreate) mux.HandleFunc("GET /v1/videos/{id}", s.videosGet) @@ -481,14 +485,14 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /v1beta/models", s.geminiModels) mux.HandleFunc("POST /v1beta/models/{call...}", s.gemini) mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { - writeError(w, provider.Chat, http.StatusNotFound, "magpie serves /v1/chat/completions, /v1/responses, /v1/messages, /v1/systemone, /v1/images/generations, /v1/images/edits, /v1/videos and /v1beta/models/*") + writeError(w, provider.Chat, http.StatusNotFound, "magpie serves /v1/chat/completions, /v1/responses, /v1/messages, /v1/systemone, /v1/images/generations, /v1/images/edits, /v1/videos, /v1/embeddings, /v1/rerank and /v1beta/models/*") }) return s.counted(callerGuard(withCaller(keyLimited(mux)))) } func (s *Server) info(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"name": "magpie", "version": Version, "models": len(provider.Catalog()), "window": Window, - "apis": []string{"/v1/chat/completions", "/v1/responses", "/v1/messages", "/v1/systemone", "/v1beta/models/{model}:generateContent", "/v1/images/generations", "/v1/images/edits", "/v1/videos", "/v1/magpie/quotas", "/v1/magpie/quotas/history", "/v1/magpie/route"}}) + "apis": []string{"/v1/chat/completions", "/v1/responses", "/v1/messages", "/v1/systemone", "/v1beta/models/{model}:generateContent", "/v1/images/generations", "/v1/images/edits", "/v1/videos", "/v1/embeddings", "/v1/rerank", "/v1/magpie/quotas", "/v1/magpie/quotas/history", "/v1/magpie/route"}}) } // quotas is what is left of every subscription, plan and key magpie has,