345 lines
17 KiB
Go
345 lines
17 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/auth"
|
|
"github.com/example/ollama-fair-gateway/internal/config"
|
|
"github.com/example/ollama-fair-gateway/internal/cost"
|
|
"github.com/example/ollama-fair-gateway/internal/metrics"
|
|
px "github.com/example/ollama-fair-gateway/internal/proxy"
|
|
"github.com/example/ollama-fair-gateway/internal/quota"
|
|
"github.com/example/ollama-fair-gateway/internal/scheduler"
|
|
"github.com/example/ollama-fair-gateway/internal/usage"
|
|
"github.com/example/ollama-fair-gateway/internal/worker"
|
|
)
|
|
|
|
func newPreflightServer(t *testing.T, capabilities []string, contextLength int64) *httptest.Server {
|
|
t.Helper()
|
|
backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/ps":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{map[string]any{"name": "codellama:7b", "model": "codellama:7b"}}})
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{map[string]any{"name": "codellama:7b", "model": "codellama:7b"}}})
|
|
case "/api/show":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"capabilities": capabilities, "model_info": map[string]any{"llama.context_length": contextLength}})
|
|
case "/api/chat":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"done": true, "message": map[string]any{"content": "ok"}, "prompt_eval_count": 2, "eval_count": 1})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
t.Cleanup(backend.Close)
|
|
cfg := &config.Config{
|
|
Server: config.ServerConfig{MaxBodyBytes: 1 << 20, MaxRequestDuration: config.Duration(time.Minute)},
|
|
Auth: config.AuthConfig{IPBypass: []config.IPBypassConfig{{CIDRs: []string{"127.0.0.1/32"}, Tenant: "t", Subject: "u"}}},
|
|
Scheduler: config.SchedulerConfig{GlobalConcurrency: 1, MaxQueue: 8, MaxQueuePerActor: 8, QueueTimeout: config.Duration(time.Second), DefaultTenantWeight: 1, DefaultActorWeight: 1, ComputePaths: []string{"/api/chat"}},
|
|
Cost: config.CostConfig{Default: config.ModelRate{InputCreditsPer1K: 1, OutputCreditsPer1K: 3, CachedInputFactor: 1}, DefaultMaxOutputTokens: 16},
|
|
Workers: []config.WorkerConfig{{Name: "w", URL: backend.URL, MaxConcurrent: 1, HealthInterval: config.Duration(time.Hour)}},
|
|
ModelCapabilities: config.ModelCapabilitiesConfig{Mode: "enforce", CacheTTL: config.Duration(time.Hour), ContextGuard: "reject"},
|
|
}
|
|
a, err := auth.New(context.Background(), cfg.Auth)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
wp := worker.New(cfg.Workers, "w")
|
|
wp.SetModelCapabilitiesConfig(cfg.ModelCapabilities)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
wp.Start(ctx)
|
|
met := metrics.New()
|
|
rec, _ := usage.New("", 100, time.Second, nil)
|
|
sv := New(cfg, Dependencies{Auth: a, Scheduler: scheduler.NewLocal(1, 8, 8), Quota: quota.Disabled{}, Estimator: cost.New(cfg.Cost), Workers: wp, Proxy: px.New(), Usage: rec, Metrics: met, Logger: slog.Default()})
|
|
front := httptest.NewServer(sv.Handler())
|
|
t.Cleanup(front.Close)
|
|
return front
|
|
}
|
|
|
|
func TestRejectsUnsupportedToolsBeforeOllama(t *testing.T) {
|
|
front := newPreflightServer(t, []string{"completion"}, 16384)
|
|
resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{"model":"codellama:7b","messages":[{"role":"user","content":"x"}],"tools":[{"type":"function","function":{"name":"f"}}]}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if !strings.Contains(string(b), "does not support required capability: tools") {
|
|
t.Fatalf("unexpected body: %s", b)
|
|
}
|
|
}
|
|
|
|
func TestAllowsSupportedTools(t *testing.T) {
|
|
front := newPreflightServer(t, []string{"completion", "tools"}, 16384)
|
|
resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{"model":"codellama:7b","messages":[{"role":"user","content":"x"}],"tools":[{"type":"function","function":{"name":"f"}}]}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != 200 {
|
|
b, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if got := resp.Header.Get("X-Gateway-Model-Capabilities"); !strings.Contains(got, "tools") {
|
|
t.Fatalf("capability header=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestRejectsExplicitContextOverModelMaximum(t *testing.T) {
|
|
front := newPreflightServer(t, []string{"completion"}, 4096)
|
|
resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{"model":"codellama:7b","messages":[{"role":"user","content":"x"}],"options":{"num_ctx":8192}}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "exceeds model context length") {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
}
|
|
|
|
func newContextGuardServer(t *testing.T, modelMax, loadedContext int64, parameters string, ctxCfg config.ContextPolicyConfig, workerLimits map[string]int64) (*httptest.Server, *atomic.Int64) {
|
|
t.Helper()
|
|
var inferenceCalls atomic.Int64
|
|
backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/ps":
|
|
models := []any{}
|
|
if loadedContext > 0 {
|
|
models = append(models, map[string]any{"name": "ctx:latest", "model": "ctx:latest", "context_length": loadedContext})
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": models})
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{map[string]any{"name": "ctx:latest", "model": "ctx:latest"}}})
|
|
case "/api/show":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"capabilities": []string{"completion", "vision"}, "model_info": map[string]any{"ctx.context_length": modelMax}, "parameters": parameters})
|
|
case "/api/chat", "/v1/chat/completions", "/v1/responses", "/api/generate":
|
|
inferenceCalls.Add(1)
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"done": true, "message": map[string]any{"content": "ok"}, "choices": []any{}, "usage": map[string]any{"prompt_tokens": 2, "completion_tokens": 1}, "prompt_eval_count": 2, "eval_count": 1})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
t.Cleanup(backend.Close)
|
|
cfg := &config.Config{
|
|
Server: config.ServerConfig{MaxBodyBytes: 1 << 20, MaxRequestDuration: config.Duration(time.Minute)},
|
|
Auth: config.AuthConfig{IPBypass: []config.IPBypassConfig{{CIDRs: []string{"127.0.0.1/32"}, Tenant: "t", Subject: "u"}}},
|
|
Scheduler: config.SchedulerConfig{GlobalConcurrency: 1, MaxQueue: 8, MaxQueuePerActor: 8, QueueTimeout: config.Duration(time.Second), DefaultTenantWeight: 1, DefaultActorWeight: 1, ComputePaths: []string{"/api/chat", "/api/generate", "/v1/chat/completions", "/v1/responses"}},
|
|
Cost: config.CostConfig{Default: config.ModelRate{InputCreditsPer1K: 1, OutputCreditsPer1K: 3, CachedInputFactor: 1}, DefaultMaxOutputTokens: 16},
|
|
Workers: []config.WorkerConfig{{Name: "w", URL: backend.URL, MaxConcurrent: 1, HealthInterval: config.Duration(time.Hour), ContextLimits: workerLimits}},
|
|
ModelCapabilities: config.ModelCapabilitiesConfig{Mode: "enforce", CacheTTL: config.Duration(time.Hour), ContextGuard: "reject", Context: ctxCfg},
|
|
}
|
|
a, err := auth.New(context.Background(), cfg.Auth)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
wp := worker.New(cfg.Workers, "w")
|
|
wp.SetModelCapabilitiesConfig(cfg.ModelCapabilities)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
wp.Start(ctx)
|
|
rec, _ := usage.New("", 100, time.Second, nil)
|
|
sv := New(cfg, Dependencies{Auth: a, Scheduler: scheduler.NewLocal(1, 8, 8), Quota: quota.Disabled{}, Estimator: cost.New(cfg.Cost), Workers: wp, Proxy: px.New(), Usage: rec, Metrics: metrics.New(), Logger: slog.Default()})
|
|
front := httptest.NewServer(sv.Handler())
|
|
t.Cleanup(front.Close)
|
|
return front, &inferenceCalls
|
|
}
|
|
|
|
func TestInvalidNumCtxRejectedBeforeBackend(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 0, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 15, VisionReserveTokensPerImage: 2048}, nil)
|
|
for _, raw := range []string{`-1`, `0`, `"8192"`, `1.5`, `null`} {
|
|
resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{"model":"ctx:latest","messages":[{"role":"user","content":"x"}],"options":{"num_ctx":`+raw+`}}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
b, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "options.num_ctx must be a positive integer") {
|
|
t.Fatalf("num_ctx=%s status=%d body=%s", raw, resp.StatusCode, b)
|
|
}
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestResponsesInstructionsAndMaxOutputRespectEffectiveContext(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 4096, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 15, VisionReserveTokensPerImage: 2048}, nil)
|
|
body := `{"model":"ctx:latest","instructions":"` + strings.Repeat("x", 8000) + `","input":"hello","max_output_tokens":3000}`
|
|
resp, err := http.Post(front.URL+"/v1/responses", "application/json", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "effective worker context") {
|
|
t.Fatalf("status=%d body=%s headers=%v", resp.StatusCode, b, resp.Header)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestOpenAIUsesModelfileNumCtxWhenModelNotLoaded(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 0, "num_ctx 16384\ntemperature 0.7", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 10, VisionReserveTokensPerImage: 2048}, nil)
|
|
body := `{"model":"ctx:latest","messages":[{"role":"user","content":"` + strings.Repeat("x", 12000) + `"}],"max_tokens":1024}`
|
|
resp, err := http.Post(front.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("status=%d body=%s headers=%v", resp.StatusCode, b, resp.Header)
|
|
}
|
|
if calls.Load() != 1 {
|
|
t.Fatalf("backend calls=%d", calls.Load())
|
|
}
|
|
if got := resp.Header.Get("X-Gateway-Model-Configured-Context"); got != "16384" {
|
|
t.Fatalf("configured context header=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestExplicitNumCtxHonorsGatewayAndWorkerCaps(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 4096, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 10, VisionReserveTokensPerImage: 2048}, map[string]int64{"*": 16384})
|
|
resp, err := http.Post(front.URL+"/api/chat", "application/json", strings.NewReader(`{"model":"ctx:latest","messages":[{"role":"user","content":"x"}],"options":{"num_ctx":20000}}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "cannot be served by any eligible worker") {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestVisionReserveParticipatesInContextGuard(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 4096, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 0, VisionReserveTokensPerImage: 3000}, nil)
|
|
body := `{"model":"ctx:latest","messages":[{"role":"user","content":[{"type":"text","text":"hello"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"max_tokens":1500}`
|
|
resp, err := http.Post(front.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || resp.Header.Get("X-Gateway-Context-Vision-Reserve") != "3000" {
|
|
t.Fatalf("status=%d reserve=%q body=%s", resp.StatusCode, resp.Header.Get("X-Gateway-Context-Vision-Reserve"), b)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestContextGuardRoutesToWorkerWithSufficientLoadedContext(t *testing.T) {
|
|
var smallCalls, largeCalls atomic.Int64
|
|
backend := func(ctxTokens int64, calls *atomic.Int64) *httptest.Server {
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/ps":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{map[string]any{"name": "ctx:latest", "model": "ctx:latest", "context_length": ctxTokens}}})
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{map[string]any{"name": "ctx:latest", "model": "ctx:latest"}}})
|
|
case "/api/show":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"capabilities": []string{"completion"}, "model_info": map[string]any{"ctx.context_length": 131072}})
|
|
case "/v1/chat/completions":
|
|
calls.Add(1)
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{}, "usage": map[string]any{"prompt_tokens": 3000, "completion_tokens": 100}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
}
|
|
small := backend(4096, &smallCalls)
|
|
defer small.Close()
|
|
large := backend(16384, &largeCalls)
|
|
defer large.Close()
|
|
cfg := &config.Config{
|
|
Server: config.ServerConfig{MaxBodyBytes: 1 << 20, MaxRequestDuration: config.Duration(time.Minute)},
|
|
Auth: config.AuthConfig{IPBypass: []config.IPBypassConfig{{CIDRs: []string{"127.0.0.1/32"}, Tenant: "t", Subject: "u"}}},
|
|
Scheduler: config.SchedulerConfig{GlobalConcurrency: 2, MaxQueue: 8, MaxQueuePerActor: 8, QueueTimeout: config.Duration(time.Second), DefaultTenantWeight: 1, DefaultActorWeight: 1, ComputePaths: []string{"/v1/chat/completions"}},
|
|
Cost: config.CostConfig{Default: config.ModelRate{InputCreditsPer1K: 1, OutputCreditsPer1K: 3, CachedInputFactor: 1}, DefaultMaxOutputTokens: 16},
|
|
Workers: []config.WorkerConfig{
|
|
{Name: "small", URL: small.URL, MaxConcurrent: 1, HealthInterval: config.Duration(time.Hour)},
|
|
{Name: "large", URL: large.URL, MaxConcurrent: 1, HealthInterval: config.Duration(time.Hour)},
|
|
},
|
|
ModelCapabilities: config.ModelCapabilitiesConfig{Mode: "enforce", CacheTTL: config.Duration(time.Hour), ContextGuard: "reject", Context: config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 15, VisionReserveTokensPerImage: 2048}},
|
|
}
|
|
a, err := auth.New(context.Background(), cfg.Auth)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
wp := worker.New(cfg.Workers, "small")
|
|
wp.SetModelCapabilitiesConfig(cfg.ModelCapabilities)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
wp.Start(ctx)
|
|
rec, _ := usage.New("", 100, time.Second, nil)
|
|
sv := New(cfg, Dependencies{Auth: a, Scheduler: scheduler.NewLocal(2, 8, 8), Quota: quota.Disabled{}, Estimator: cost.New(cfg.Cost), Workers: wp, Proxy: px.New(), Usage: rec, Metrics: metrics.New(), Logger: slog.Default()})
|
|
front := httptest.NewServer(sv.Handler())
|
|
defer front.Close()
|
|
body := `{"model":"ctx:latest","messages":[{"role":"user","content":"` + strings.Repeat("x", 12000) + `"}],"max_tokens":4096}`
|
|
resp, err := http.Post(front.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if got := resp.Header.Get("X-Gateway-Worker"); got != "large" {
|
|
t.Fatalf("worker=%q, want large; small=%d large=%d", got, smallCalls.Load(), largeCalls.Load())
|
|
}
|
|
if smallCalls.Load() != 0 || largeCalls.Load() != 1 {
|
|
t.Fatalf("unexpected backend calls small=%d large=%d", smallCalls.Load(), largeCalls.Load())
|
|
}
|
|
}
|
|
|
|
func TestOpenAICannotSpoofEffectiveContextWithOptionsNumCtx(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 4096, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 15, VisionReserveTokensPerImage: 2048}, nil)
|
|
resp, err := http.Post(front.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"ctx:latest","messages":[{"role":"user","content":"x"}],"options":{"num_ctx":16384}}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "only supported on native Ollama /api endpoints") {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestGenerateSuffixParticipatesInContextGuard(t *testing.T) {
|
|
front, calls := newContextGuardServer(t, 131072, 4096, "", config.ContextPolicyConfig{MaxRequestedTokens: 32768, DefaultWorkerTokens: 4096, EstimationMarginPercent: 10, VisionReserveTokensPerImage: 2048}, nil)
|
|
body := `{"model":"ctx:latest","prompt":"hello","suffix":"` + strings.Repeat("s", 14000) + `","options":{"num_predict":1024}}`
|
|
resp, err := http.Post(front.URL+"/api/generate", "application/json", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 400 || !strings.Contains(string(b), "effective worker context") {
|
|
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("backend called %d times", calls.Load())
|
|
}
|
|
}
|