Ilker ManapandClaude Opus 5.5
dbcecd1e85
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
81 lines
2.3 KiB
Go
81 lines
2.3 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
|
)
|
|
|
|
// Judge sends a chat completion to a judge model and returns its answer: for
|
|
// the guardrails, which ask a model on the cluster about the traffic they
|
|
// judge. It goes past admission — a judge call queued behind the tenants it
|
|
// is judging would wait on them — and past accounting and the guardrails
|
|
// themselves, but through the scheduler and the replicas like any request.
|
|
func (r *Router) Judge(ctx context.Context, model, path string, body []byte) ([]byte, error) {
|
|
if path != "/v1/chat/completions" && path != "/v1/completions" {
|
|
return nil, fmt.Errorf("a judge is asked at /v1/chat/completions or /v1/completions, not %s", path)
|
|
}
|
|
hreq, err := http.NewRequestWithContext(ctx, http.MethodPost, path, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req, err := r.parse(hreq, path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.internal = true
|
|
req.Stream = false
|
|
req.RequestID = requestID("")
|
|
req.Header = http.Header{"X-Request-Id": []string{req.RequestID}}
|
|
w := &bufferWriter{header: http.Header{}}
|
|
res := r.dispatch(ctx, w, req)
|
|
if res.err != nil {
|
|
return nil, res.err
|
|
}
|
|
if w.status != 0 && w.status != http.StatusOK {
|
|
return nil, fmt.Errorf("the judge %s answered %d", model, w.status)
|
|
}
|
|
return w.buf.Bytes(), nil
|
|
}
|
|
|
|
// GuardFormat is how a judge model's answers are read, from its catalog
|
|
// entry: for guard.Gate.SetJudge.
|
|
func (r *Router) GuardFormat(model string) (string, []string, bool) {
|
|
_, entry := r.resolve(model)
|
|
if !entry.IsGuard() || entry.Guard == nil {
|
|
return "", nil, false
|
|
}
|
|
return entry.Guard.Format, entry.Guard.Categories, true
|
|
}
|
|
|
|
var _ guard.Caller = (*Router)(nil).Judge
|
|
|
|
// bufferWriter is a ResponseWriter that keeps what is written, for a request
|
|
// the gateway makes of itself.
|
|
type bufferWriter struct {
|
|
header http.Header
|
|
status int
|
|
buf bytes.Buffer
|
|
}
|
|
|
|
func (b *bufferWriter) Header() http.Header { return b.header }
|
|
func (b *bufferWriter) WriteHeader(code int) {
|
|
if b.status == 0 {
|
|
b.status = code
|
|
}
|
|
}
|
|
func (b *bufferWriter) Write(p []byte) (int, error) {
|
|
if b.status == 0 {
|
|
b.status = http.StatusOK
|
|
}
|
|
if b.buf.Len()+len(p) > 1<<20 {
|
|
return 0, errors.New("a judge's answer over 1 MiB")
|
|
}
|
|
return b.buf.Write(p)
|
|
}
|
|
func (b *bufferWriter) Flush() {}
|