Fixes from review: the similar cache, restore, outside services, latency
CI / test (push) Successful in 9m48s
CI / test (push) Successful in 9m48s
- The similar-question cache embeds a question as its project would: a project kept in the cluster has external fallbacks refused, and an embedding model outside the cluster is not asked at all. The vectors are compared outside the cache's lock. The question's bound is in characters. - The restorer puts values back in streamed tool calls' arguments (escaped for the JSON they are), completions' text and reasoning, each held across events on its own, and releases a held tail with the finish event. A whole answer is walked as JSON, so a value with a quote no longer breaks tool-call arguments. - An outside guardrail service reads the user's turn as redacted, and a streamed answer once rather than at the finish chunk and again at the end. guardrails.md says it can stop a stream only under buffer_all. - A replica's smoothed time to first token is forgotten after five minutes without an answer: one found slow and then given nothing was never tried again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
fd94265e08
commit
54c6b66cb9
11 files changed
+372
-71
No files matched your search
@@ -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
|
||||
|
||||
+9
-4
@@ -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:<name>:<reason>`.
|
||||
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:<name>:<reason>`. The service is told the text only, never who
|
||||
sent it.
|
||||
|
||||
## Where a policy applies
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+24
-17
@@ -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 != "" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+123
-34
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"`
|
||||
|
||||
Reference in new issue
Block a user