diff --git a/CHANGELOG.md b/CHANGELOG.md index eb552b8..49a552b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,19 @@ APIs change between minors without a deprecation period. ### Fixed +- **A project kept in the cluster could have its questions embedded outside + it.** The similar-question cache now asks the embedding model as the + question's project would, external fallbacks refused, and skips a model + that is itself outside. Its comparing no longer holds the cache's lock, + which every project's lookups waited on. +- **Restored values were left as placeholders in streamed tool calls**, in + completions and in reasoning, and a value with a quote in a whole answer's + tool-call arguments broke their JSON. +- **An outside guardrail service read what `redact` had taken out**, and a + streamed answer was sent to it twice. The docs now say a service stops a + streamed answer only under `stream: buffer_all`. +- **A replica found slow once could stay unused.** Its time to first token + is forgotten after five minutes without an answer, so it is tried again. - **Weights that run code got past the refusal in four ways.** A mirror version already on a node was used without being looked at again; a `?query` or `#fragment` on an `hf:` source hid the file's name; `.pt2` was diff --git a/docs/guardrails.md b/docs/guardrails.md index 77adfbd..42a1662 100644 --- a/docs/guardrails.md +++ b/docs/guardrails.md @@ -301,10 +301,15 @@ policies: string; empty sends `{"input": …}`). `blocked` and `reason` are dotted paths into the JSON answer, numbers for list items; true, a number above zero, a non-empty list or a string other than "false" is unsafe. On input the -service reads the latest user turn; on output the whole answer, once, at its -end — a streamed answer released a sentence at a time is then stopped or -flagged as the judge's is. The rule recorded is `service::`. -The service is told the text only, never who sent it. +service reads the latest user turn as the model will, with what `redact` +took out still out; on output the whole answer, once, at its end, as the +client reads it — with `restore`, the values put back. Being at the end, a +service can stop a streamed answer only where it is held back +(`stream: buffer_all`): released a sentence at a time, the client has read +it by then, and what the service finds ends the stream with +`content_filter`, flags it and is recorded. The rule recorded is +`service::`. The service is told the text only, never who +sent it. ## Where a policy applies diff --git a/internal/guard/gate.go b/internal/guard/gate.go index b6b03d5..e0333ef 100644 --- a/internal/guard/gate.go +++ b/internal/guard/gate.go @@ -418,7 +418,15 @@ func (g *Gate) Input(ctx context.Context, s Subject, body []byte) (Verdict, erro // Outside services next, then the judge: what a pattern already stopped // need not wait for either. if user != "" { - sf, serr := g.askServices(ctx, p, true, user) + // As the model will read it: what redact took out stays out of the + // outside service's sight too. + told := user + for _, rw := range rewrites { + if rw.seg.text != "" && strings.Contains(told, rw.seg.text) { + told = strings.ReplaceAll(told, rw.seg.text, rw.text) + } + } + sf, serr := g.askServices(ctx, p, true, told) if serr != nil && !v.Shadow { v.Took = g.now().Sub(start) return v, serr diff --git a/internal/guard/output.go b/internal/guard/output.go index 4103850..bf50c36 100644 --- a/internal/guard/output.go +++ b/internal/guard/output.go @@ -74,21 +74,23 @@ func (g *Gate) Output(ctx context.Context, s Subject) OutputFilter { // outFilter judges one answer. type outFilter struct { - ctx context.Context - g *Gate - p Policy - names string - shadow bool - schemas map[string]json.RawMessage - answer json.RawMessage - full strings.Builder // the whole answer, when it must fit a schema - user string // the request's last user turn, for the judge - sofar strings.Builder // the answer so far, for the judge - served strings.Builder // the answer so far, for outside services, asked once at its end - judged int // how much of it the judge has read - v Verdict - took time.Duration - stopped bool + // servicesAsked says the outside services have read the answer. + servicesAsked bool + ctx context.Context + g *Gate + p Policy + names string + shadow bool + schemas map[string]json.RawMessage + answer json.RawMessage + full strings.Builder // the whole answer, when it must fit a schema + user string // the request's last user turn, for the judge + sofar strings.Builder // the answer so far, for the judge + served strings.Builder // the answer so far, for outside services, asked once at its end + judged int // how much of it the judge has read + v Verdict + took time.Duration + stopped bool held []string // events not yet released sent bool // whether a chunk this filter made has gone out @@ -445,8 +447,13 @@ func (f *outFilter) flush(final bool) (bool, map[int]string, bool) { } if i == 0 { f.served.WriteString(head) - if final && f.askServices(f.served.String()).action == Block { - blocked = true + // Once: the finish chunk is final, and so is the end of the + // stream after it, which would ask again of the same answer. + if final && !f.servicesAsked { + f.servicesAsked = true + if f.askServices(f.served.String()).action == Block { + blocked = true + } } } if head != "" { diff --git a/internal/guard/service_test.go b/internal/guard/service_test.go index 8e23c69..af404e5 100644 --- a/internal/guard/service_test.go +++ b/internal/guard/service_test.go @@ -98,3 +98,44 @@ func TestAServiceIsCheckedWhenThePolicyIsRead(t *testing.T) { t.Errorf("at: %v %v", v, ok) } } + +// recorder is a service that finds nothing unsafe and keeps what it read. +func recorder(t *testing.T) (*httptest.Server, *[]string) { + var seen []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var in struct{ Input string } + _ = json.NewDecoder(r.Body).Decode(&in) + seen = append(seen, in.Input) + _, _ = w.Write([]byte(`{"flagged": false}`)) + })) + t.Cleanup(srv.Close) + return srv, &seen +} + +func TestAnOutsideServiceReadsWhatTheModelReadsNotWhatWasRedacted(t *testing.T) { + srv, seen := recorder(t) + g := gate(t, lib(map[string]Policy{"p": { + Input: Rules{PII: map[string]string{"email": "redact"}}, + Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged"}}, + }}, "p")) + if _, err := g.Input(context.Background(), Subject{}, chat("user", "mail ayse@example.com")); err != nil { + t.Fatal(err) + } + if len(*seen) != 1 || strings.Contains((*seen)[0], "ayse@") || !strings.Contains((*seen)[0], "[EMAIL_1]") { + t.Errorf("the service read %q", *seen) + } +} + +func TestAStreamedAnswerIsReadByAServiceOnce(t *testing.T) { + srv, seen := recorder(t) + g := gate(t, lib(map[string]Policy{"p": { + Services: []Service{{Name: "cls", URL: srv.URL, Blocked: "flagged", On: "output"}}, + }}, "p")) + got, _, _ := stream(g.Output(context.Background(), Subject{Stream: true}), "Hello ", "there.") + if got != "Hello there." { + t.Errorf("client got %q", got) + } + if len(*seen) != 1 || (*seen)[0] != "Hello there." { + t.Errorf("the service was asked %d times: %q", len(*seen), *seen) + } +} diff --git a/internal/router/cache.go b/internal/router/cache.go index 424b483..5f4c767 100644 --- a/internal/router/cache.go +++ b/internal/router/cache.go @@ -208,25 +208,37 @@ func (c *answerCache) remove(el *list.Element) { // nearest is the answer in group whose question is most like vec, when // that reaches least and the answer is younger than ttl, and how alike. +// The comparing is done outside the lock: a help desk's questions are one +// group of thousands, and every project's lookups wait on the lock. func (c *answerCache) nearest(group string, vec []float32, least float64, ttl time.Duration, now time.Time) (*cachedAnswer, float64) { c.mu.Lock() - defer c.mu.Unlock() - var best *list.Element - bestSim := least + cands := make([]*cachedAnswer, 0, len(c.groups[group])) for el := range c.groups[group] { - a := el.Value.(*cachedAnswer) + cands = append(cands, el.Value.(*cachedAnswer)) + } + c.mu.Unlock() + var best *cachedAnswer + bestSim := least + for _, a := range cands { if now.After(a.expires) || now.Sub(a.at) >= ttl || len(a.vec) != len(vec) { continue } if s := dot(a.vec, vec); s >= bestSim { - best, bestSim = el, s + best, bestSim = a, s } } if best == nil { return nil, 0 } - c.lru.MoveToFront(best) - return best.Value.(*cachedAnswer), bestSim + c.mu.Lock() + defer c.mu.Unlock() + // Pushed out meanwhile: not an answer to give. + el, ok := c.by[best.key] + if !ok || el.Value.(*cachedAnswer) != best { + return nil, 0 + } + c.lru.MoveToFront(el) + return best, bestSim } // answerCacheOf is a project's setting, read at most every requestLogTTL diff --git a/internal/router/endpoint/endpoint.go b/internal/router/endpoint/endpoint.go index 4057033..a1a9df6 100644 --- a/internal/router/endpoint/endpoint.go +++ b/internal/router/endpoint/endpoint.go @@ -170,6 +170,7 @@ type live struct { drainRate float64 maxRunning int ttft float64 // seconds, smoothed + ttftAt time.Time inflight int brk *breaker.Breaker @@ -252,6 +253,19 @@ func (r *Registry) Infos() []Info { return out } +// ttftForgotten is how long a replica's time to first token is believed +// without a new answer: one found slow and then given nothing would +// otherwise never be tried again, and its figure never mended. +const ttftForgotten = 5 * time.Minute + +// ttftNow is the smoothed time to first token, none once forgotten. +func (l *live) ttftNow(now time.Time) time.Duration { + if l.ttft == 0 || now.Sub(l.ttftAt) > ttftForgotten { + return 0 + } + return time.Duration(l.ttft * float64(time.Second)) +} + // ObserveTTFT records how long an answer streamed here took to its first token. func (r *Registry) ObserveTTFT(id string, d time.Duration) { if d <= 0 { @@ -264,7 +278,9 @@ func (r *Registry) ObserveTTFT(id string, d time.Duration) { return } const alpha = 0.2 // a few answers move it; one slow one does not - if l.ttft == 0 { + now := r.now() + defer func() { l.ttftAt = now }() + if l.ttft == 0 || now.Sub(l.ttftAt) > ttftForgotten { l.ttft = d.Seconds() } else { l.ttft = alpha*d.Seconds() + (1-alpha)*l.ttft @@ -320,7 +336,7 @@ func (r *Registry) Snapshots(model string) []Snapshot { Info: l.info, Stats: l.stats, DrainRate: l.drainRate, - TTFT: time.Duration(l.ttft * float64(time.Second)), + TTFT: l.ttftNow(now), Inflight: l.inflight, RunningCapacity: l.maxRunning, Stale: stale, diff --git a/internal/router/endpoint/endpoint_test.go b/internal/router/endpoint/endpoint_test.go index 999ecbe..b72b9af 100644 --- a/internal/router/endpoint/endpoint_test.go +++ b/internal/router/endpoint/endpoint_test.go @@ -319,3 +319,24 @@ func TestTTFTIsSmoothedPerEndpoint(t *testing.T) { t.Errorf("%v", s) } } + +func TestAReplicaFoundSlowIsTriedAgainLater(t *testing.T) { + // Slow once and then given nothing, its figure would never mend. + r := NewRegistry(Config{}) + now := time.Now() + r.now = func() time.Time { return now } + r.Replace([]Info{{ID: "a", Models: []string{"m"}}}) + r.ObserveTTFT("a", 10*time.Second) + if s := r.Snapshots("m"); s[0].TTFT != 10*time.Second { + t.Fatalf("%v", s[0].TTFT) + } + now = now.Add(6 * time.Minute) + if s := r.Snapshots("m"); s[0].TTFT != 0 { + t.Errorf("six minutes on, still %v", s[0].TTFT) + } + // The next answer starts the figure afresh rather than smoothing it in. + r.ObserveTTFT("a", time.Second) + if s := r.Snapshots("m"); s[0].TTFT != time.Second { + t.Errorf("after a new answer %v", s[0].TTFT) + } +} diff --git a/internal/router/restore.go b/internal/router/restore.go index d508755..1530dd9 100644 --- a/internal/router/restore.go +++ b/internal/router/restore.go @@ -22,27 +22,42 @@ import ( type restorer struct { values map[string]string longest int - // held is, per choice, the tail of the text so far that may be the start - // of a placeholder the next event completes: "[EMA" before "IL_1]". - held map[int]string + // held is, per choice and field, the tail of the text so far that may + // be the start of a placeholder the next event completes: "[EMA" before + // "IL_1]". Keyed by heldKey. + held map[heldKey]string // last is the most recent event, the shape a final release is sent in. last map[string]any } +// heldKey is where held text belongs: a choice, and in it the content, the +// reasoning, a completion's text, or a tool call's arguments (tool >= 0). +type heldKey struct { + choice int + field string + tool int +} + func newRestorer(values map[string]string) *restorer { - r := &restorer{values: values, held: map[int]string{}} + r := &restorer{values: values, held: map[heldKey]string{}} for ph := range values { r.longest = max(r.longest, len(ph)) } return r } -// text swaps every whole placeholder in s. -func (r *restorer) text(s string) string { +// text swaps every whole placeholder in s. In a tool call's arguments, +// which are JSON written into a string, a value goes in escaped for JSON: +// O"Brien would otherwise end the argument early. +func (r *restorer) text(s string, inJSON bool) string { if !strings.Contains(s, "[") { return s } for ph, v := range r.values { + if inJSON { + q, _ := json.Marshal(v) + v = string(q[1 : len(q)-1]) + } s = strings.ReplaceAll(s, ph, v) } return s @@ -50,8 +65,8 @@ func (r *restorer) text(s string) string { // split is s with its placeholders swapped, less a tail that could still // become one: what to send now, and what to hold for the next event. -func (r *restorer) split(s string) (send, hold string) { - s = r.text(s) +func (r *restorer) split(s string, inJSON bool) (send, hold string) { + s = r.text(s, inJSON) i := strings.LastIndexByte(s, '[') if i < 0 || strings.IndexByte(s[i:], ']') >= 0 || len(s)-i >= r.longest { return s, "" @@ -65,20 +80,43 @@ func (r *restorer) split(s string) (send, hold string) { return s, "" } -// body swaps the placeholders in a whole JSON answer. A placeholder is only -// ever inside a string there, so each value goes in escaped for one. +// body swaps the placeholders in a whole JSON answer: in every string, and +// in a tool call's arguments escaped for the JSON they are. func (r *restorer) body(b []byte) []byte { if !bytes.Contains(b, []byte("[")) { return b } - for ph, v := range r.values { - q, _ := json.Marshal(v) - b = bytes.ReplaceAll(b, []byte(ph), q[1:len(q)-1]) + d := json.NewDecoder(bytes.NewReader(b)) + d.UseNumber() + var v any + if err := d.Decode(&v); err != nil { + return b } - return b + out, err := json.Marshal(r.walk(v, "")) + if err != nil { + return b + } + return out } -// OnEvent swaps the placeholders in one streamed chat-completion event, +func (r *restorer) walk(v any, key string) any { + switch x := v.(type) { + case map[string]any: + for k, e := range x { + x[k] = r.walk(e, k) + } + case []any: + for i, e := range x { + x[i] = r.walk(e, key) + } + case string: + return r.text(x, key == "arguments") + } + return v +} + +// OnEvent swaps the placeholders in one streamed event — a chat delta's +// content, reasoning and tool calls' arguments, a completion's text — // holding back the start of one an event has split. func (r *restorer) OnEvent(data string) ([]string, bool) { if !strings.Contains(data, "[") && len(r.held) == 0 { @@ -93,9 +131,7 @@ func (r *restorer) OnEvent(data string) ([]string, bool) { choices, _ := ev["choices"].([]any) for i, c := range choices { ch, _ := c.(map[string]any) - delta, _ := ch["delta"].(map[string]any) - content, ok := delta["content"].(string) - if !ok { + if ch == nil { continue } idx := i @@ -104,16 +140,50 @@ func (r *restorer) OnEvent(data string) ([]string, bool) { idx = int(v) } } - send, hold := r.split(r.held[idx] + content) - // The last event of a choice releases everything. - if ch["finish_reason"] != nil { - send, hold = send+hold, "" + finished := ch["finish_reason"] != nil + delta, _ := ch["delta"].(map[string]any) + if delta == nil && finished && r.heldFor(idx) { + delta = map[string]any{} + ch["delta"] = delta } - delta["content"] = send - if hold == "" { - delete(r.held, idx) - } else { - r.held[idx] = hold + swap := func(container map[string]any, field string, k heldKey, inJSON bool) { + s, ok := container[field].(string) + if !ok && !(finished && r.held[k] != "") { + return + } + send, hold := r.split(r.held[k]+s, inJSON) + // The last event of a choice releases everything. + if finished { + send, hold = send+hold, "" + } + container[field] = send + if hold == "" { + delete(r.held, k) + } else { + r.held[k] = hold + } + } + if delta != nil { + swap(delta, "content", heldKey{idx, "content", -1}, false) + swap(delta, "reasoning_content", heldKey{idx, "reasoning_content", -1}, false) + calls, _ := delta["tool_calls"].([]any) + for j, tc := range calls { + call, _ := tc.(map[string]any) + fn, _ := call["function"].(map[string]any) + if fn == nil { + continue + } + t := j + if n, ok := call["index"].(json.Number); ok { + if v, err := n.Int64(); err == nil { + t = int(v) + } + } + swap(fn, "arguments", heldKey{idx, "arguments", t}, true) + } + } + if _, ok := ch["text"]; ok { + swap(ch, "text", heldKey{idx, "text", -1}, false) } } r.last = ev @@ -124,25 +194,44 @@ func (r *restorer) OnEvent(data string) ([]string, bool) { return []string{string(out)}, false } -// OnEnd releases what is still held, as one more event per choice. +func (r *restorer) heldFor(choice int) bool { + for k := range r.held { + if k.choice == choice { + return true + } + } + return false +} + +// OnEnd releases what is still held, as one more event for each. func (r *restorer) OnEnd() []string { if len(r.held) == 0 || r.last == nil { return nil } var out []string - for idx, text := range r.held { + for k, text := range r.held { ev := map[string]any{} - for k, v := range r.last { - if k != "choices" && k != "usage" { - ev[k] = v + for key, v := range r.last { + if key != "choices" && key != "usage" { + ev[key] = v } } - ev["choices"] = []any{map[string]any{"index": idx, "delta": map[string]any{"content": text}, "finish_reason": nil}} + ch := map[string]any{"index": k.choice, "finish_reason": nil} + switch k.field { + case "text": + ch["text"] = text + case "arguments": + ch["delta"] = map[string]any{"tool_calls": []any{map[string]any{"index": k.tool, + "function": map[string]any{"arguments": text}}}} + default: + ch["delta"] = map[string]any{k.field: text} + } + ev["choices"] = []any{ch} if b, err := json.Marshal(ev); err == nil { out = append(out, string(b)) } } - r.held = map[int]string{} + r.held = map[heldKey]string{} return out } diff --git a/internal/router/restore_test.go b/internal/router/restore_test.go index b48b5c8..80f5800 100644 --- a/internal/router/restore_test.go +++ b/internal/router/restore_test.go @@ -3,6 +3,7 @@ package router import ( "bytes" "context" + "encoding/json" "fmt" "net/http" "net/http/httptest" @@ -92,7 +93,7 @@ func TestARestorerHoldsOnlyWhatCouldBeAPlaceholder(t *testing.T) { {"end [", "end ", "["}, {"[X", "[X", ""}, } { - send, hold := r.split(c.in) + send, hold := r.split(c.in, false) if send != c.send || hold != c.hold { t.Errorf("split(%q) = %q, %q; want %q, %q", c.in, send, hold, c.send, c.hold) } @@ -148,3 +149,66 @@ func TestAnAnswerWithACallersValuesPutBackIsNeverReusedForAnother(t *testing.T) t.Errorf("x-llmgw-cache %s, want bypass for both", got) } } + +func TestAToolCallsArgumentsGetTheirValuesBackStreamedOrNot(t *testing.T) { + r := newRestorer(map[string]string{"[NAME_1]": `O"Brien`, "[EMAIL_1]": "a@example.com"}) + // Whole: the value goes into the arguments escaped for the JSON they are. + body := r.body([]byte(`{"choices":[{"message":{"content":"to [NAME_1]","tool_calls":[{"function":{"name":"send",` + + `"arguments":"{\"to\":\"[EMAIL_1]\",\"name\":\"[NAME_1]\"}"}}]}}]}`)) + var whole struct { + Choices []struct { + Message struct { + Content string + ToolCalls []struct{ Function struct{ Arguments string } } `json:"tool_calls"` + } + } + } + if err := json.Unmarshal(body, &whole); err != nil { + t.Fatalf("%v: %s", err, body) + } + var args map[string]string + if err := json.Unmarshal([]byte(whole.Choices[0].Message.ToolCalls[0].Function.Arguments), &args); err != nil || + args["to"] != "a@example.com" || args["name"] != `O"Brien` || whole.Choices[0].Message.Content != `to O"Brien` { + t.Fatalf("arguments %q (%v), content %q", whole.Choices[0].Message.ToolCalls[0].Function.Arguments, err, whole.Choices[0].Message.Content) + } + + // Streamed, the placeholder split across two events. + var argsText strings.Builder + for _, piece := range []string{`{"to":"[EMA`, `IL_1]"}`} { + ev, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"index": 0, "finish_reason": nil, + "delta": map[string]any{"tool_calls": []any{map[string]any{"index": 0, "function": map[string]any{"arguments": piece}}}}}}}) + out, _ := r.OnEvent(string(ev)) + for _, o := range out { + var e struct { + Choices []struct { + Delta struct { + ToolCalls []struct{ Function struct{ Arguments string } } `json:"tool_calls"` + } + } + } + _ = json.Unmarshal([]byte(o), &e) + for _, c := range e.Choices { + for _, tc := range c.Delta.ToolCalls { + argsText.WriteString(tc.Function.Arguments) + } + } + } + } + if argsText.String() != `{"to":"a@example.com"}` { + t.Errorf("streamed arguments read %q", argsText.String()) + } +} + +func TestHeldTextGoesOutWithTheFinishEvent(t *testing.T) { + // The answer ends on what could have begun a placeholder; the finish + // event, whose delta is empty, carries it rather than an event after. + r := newRestorer(map[string]string{"[EMAIL_1]": "a@example.com"}) + out1, _ := r.OnEvent(`{"choices":[{"index":0,"delta":{"content":"see [EMA"},"finish_reason":null}]}`) + out2, _ := r.OnEvent(`{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`) + if len(r.OnEnd()) != 0 { + t.Error("held text was released after the finish event") + } + if !strings.Contains(out1[0], `"see "`) || !strings.Contains(out2[0], `"content":"[EMA"`) { + t.Errorf("%s / %s", out1, out2) + } +} diff --git a/internal/router/similar.go b/internal/router/similar.go index f06676e..e919604 100644 --- a/internal/router/similar.go +++ b/internal/router/similar.go @@ -1,11 +1,14 @@ package router import ( + "bytes" "context" "encoding/json" "fmt" "math" + "net/http" "time" + "unicode/utf8" storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) @@ -24,7 +27,7 @@ import ( // similarityHeader says how alike a reused answer's question was. const similarityHeader = "x-llmgw-cache-similarity" -// maxSimilarQuestion bounds the questions compared: a longer one is a +// maxSimilarQuestion bounds the questions compared, in characters: a longer one is a // document pasted in, where a changed figure or name leaves the embedding // almost as it was. Those are reused only when they are the same. const maxSimilarQuestion = 2000 @@ -40,10 +43,15 @@ func (r *Router) similarLookup(ctx context.Context, req *parsedRequest, model st return nil, 0 } group, question, ok := similarGroup(req.ProjectID, req.Path, model, req.UpstreamBody) - if !ok || question == "" || len(question) > maxSimilarQuestion { + if !ok || question == "" || utf8.RuneCountInString(question) > maxSimilarQuestion { return nil, 0 } - vec, err := r.embedQuestion(ctx, question) + // A project whose requests stay in the cluster keeps its questions in + // it too: not embedded by a model outside it. + if _, m := r.resolve(r.cacheEmbedModel); req.inCluster && m.IsExternal() { + return nil, 0 + } + vec, err := r.embedQuestion(ctx, question, req.inCluster) if err != nil { r.log.Warn("a question could not be embedded for the cache; only the same request is reused", "model", r.cacheEmbedModel, "err", err) @@ -91,17 +99,34 @@ func similarGroup(project, path, model string, body []byte) (group, question str // embedQuestion has router.cache.embedding_model embed a question, on the // gateway's own account as a judge is asked, into a vector of unit length. -func (r *Router) embedQuestion(ctx context.Context, text string) ([]float32, error) { +func (r *Router) embedQuestion(ctx context.Context, text string, inCluster bool) ([]float32, error) { if _, m := r.resolve(r.cacheEmbedModel); m.QueryPrefix != "" { text = m.QueryPrefix + text } ctx, cancel := context.WithTimeout(ctx, similarTimeout) defer cancel() body, _ := json.Marshal(map[string]any{"model": r.cacheEmbedModel, "input": []string{text}}) - out, err := r.judgeOnce(ctx, "/v1/embeddings", body, nil, nil) + hreq, err := http.NewRequestWithContext(ctx, http.MethodPost, "/v1/embeddings", bytes.NewReader(body)) if err != nil { return nil, err } + req, err := r.parse(hreq, "/v1/embeddings") + if err != nil { + return nil, err + } + // As the gateway's own, as a judge is asked; and kept in the cluster + // when the question's project is, fallbacks included. + req.internal, req.inCluster = true, inCluster + req.RequestID = requestID("") + req.Header = http.Header{"X-Request-Id": []string{req.RequestID}} + w := &bufferWriter{header: http.Header{}} + if res := r.dispatch(ctx, w, req); res.err != nil { + return nil, res.err + } + if w.status != 0 && w.status != http.StatusOK { + return nil, fmt.Errorf("%s answered %d", r.cacheEmbedModel, w.status) + } + out := w.buf.Bytes() var res struct { Data []struct { Embedding []float32 `json:"embedding"`