gateway: /v1/embeddings and /v1/rerank go to the model's provider (#765, dbydd: a local RAG stack had to send its embeddings around magpie, with a second base URL and key, outside the usage ledger) — magpie answered both with its catch-all 404. retrieve (embeddings.go) takes the body as the agent sent it, resolves its model as a turn's is, puts the provider's own model id in, masks secrets as the settings say (a document a reranker gives back has them again), and posts it to the provider's OpenAI-style base + /embeddings or /rerank, signed as the provider signs; the vendor's answer comes back as it said it, its refusal with its code and words. The call is on the Requests page and in the usage ledger (operation embeddings / rerank) with the tokens the answer counted (prompt_tokens, total_tokens, or Cohere's billed input). /embeddings and /rerank answer too, and the 404's list and GET / name them. With the owner's real keys under a sandbox HOME: aihubmix/text-embedding-3-small with dimensions 4 returned two 4-d vectors and 3 tokens, aihubmix/jina-reranker-m0 ranked the magpie document first (37 tokens, recorded as rerank), and OpenRouter's 402 for no credits came back as its own words. TestEmbeddingsGoToTheModelsProvider, TestRerankGoesToTheModelsProvider, TestRetrievalFailures and TestRetrievalRedacts fail on the old code (404). go vet, builds (darwin, GOOS=linux, GOOS=windows) and go test -tags nogui -count=1 ./... under a temp HOME pass (TestClientGoneRecordedAsCanceled flaked once under load and passes alone).
This commit is contained in:
yetone committed
1 parent
a7cd091ac6
commit
76fb7879e8
3 files changed
+333
-2
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user