394 lines
11 KiB
Go
394 lines
11 KiB
Go
|
package proxy
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Riotpiaole/homelab-frontend/internal/config"
|
||
|
|
)
|
||
|
|
|
||
|
|
// TestEmbeddingsPassthroughNoRewrite verifies /v1/embeddings is not rewritten
|
||
|
|
func TestEmbeddingsPassthroughNoRewrite(t *testing.T) {
|
||
|
|
embeddingsCalled := false
|
||
|
|
upstreamPath := ""
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
embeddingsCalled = true
|
||
|
|
upstreamPath = r.URL.Path
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, `{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}]}`)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe": {
|
||
|
|
Name: "nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/embeddings",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"hello"}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if !embeddingsCalled {
|
||
|
|
t.Errorf("upstream embeddings service was not called")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify path is NOT rewritten (should stay /v1/embeddings)
|
||
|
|
if upstreamPath != "/v1/embeddings" {
|
||
|
|
t.Errorf("expected upstream path /v1/embeddings, got %s", upstreamPath)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestEmbeddingsResponsePassthrough verifies response body is unmodified
|
||
|
|
func TestEmbeddingsResponsePassthrough(t *testing.T) {
|
||
|
|
expectedResponse := `{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}]}`
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, expectedResponse)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe": {
|
||
|
|
Name: "nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/embeddings",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"test"}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
body, _ := io.ReadAll(resp.Body)
|
||
|
|
|
||
|
|
if string(body) != expectedResponse {
|
||
|
|
t.Errorf("response was modified. Expected:\n%s\n\nGot:\n%s", expectedResponse, string(body))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestRerankPathRewrite verifies /v1/rerank is rewritten to /rerank
|
||
|
|
func TestRerankPathRewrite(t *testing.T) {
|
||
|
|
upstreamPath := ""
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
upstreamPath = r.URL.Path
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, `{"results":[{"index":0,"score":0.9}]}`)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"BAAI/bge-reranker-base": {
|
||
|
|
Name: "BAAI/bge-reranker-base",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/rerank",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"BAAI/bge-reranker-base","query":"test","texts":["a","b"]}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
// Verify path IS rewritten to /rerank
|
||
|
|
if upstreamPath != "/rerank" {
|
||
|
|
t.Errorf("expected upstream path /rerank, got %s", upstreamPath)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestRerankResponsePassthrough verifies response is unmodified
|
||
|
|
func TestRerankResponsePassthrough(t *testing.T) {
|
||
|
|
expectedResponse := `{"results":[{"index":0,"score":0.95},{"index":1,"score":0.85}]}`
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, expectedResponse)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"BAAI/bge-reranker-base": {
|
||
|
|
Name: "BAAI/bge-reranker-base",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/rerank",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"BAAI/bge-reranker-base","query":"q","texts":["a"]}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
body, _ := io.ReadAll(resp.Body)
|
||
|
|
|
||
|
|
if string(body) != expectedResponse {
|
||
|
|
t.Errorf("response was modified")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestEmbeddingsUnknownModel returns error for unknown model
|
||
|
|
func TestEmbeddingsUnknownModel(t *testing.T) {
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe": {
|
||
|
|
Name: "nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
Address: "localhost:9000",
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/embeddings",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"unknown-embeddings","input":"test"}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected 400 for unknown embeddings model, got %d", resp.StatusCode)
|
||
|
|
}
|
||
|
|
|
||
|
|
ct := resp.Header.Get("Content-Type")
|
||
|
|
if !strings.Contains(ct, "application/problem+json") {
|
||
|
|
t.Errorf("expected problem+json for unknown model")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestRerankerUnknownModel returns error for unknown model
|
||
|
|
func TestRerankerUnknownModel(t *testing.T) {
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"BAAI/bge-reranker-base": {
|
||
|
|
Name: "BAAI/bge-reranker-base",
|
||
|
|
Address: "localhost:9000",
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/rerank",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"unknown-reranker","query":"q","texts":["a"]}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected 400 for unknown reranker model, got %d", resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestEmbeddingsBodyForwarded verifies body is byte-identical to upstream
|
||
|
|
func TestEmbeddingsBodyForwarded(t *testing.T) {
|
||
|
|
receivedBody := ""
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
body, _ := io.ReadAll(r.Body)
|
||
|
|
receivedBody = string(body)
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, `{"object":"list","data":[]}`)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe": {
|
||
|
|
Name: "nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
originalBody := `{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"test data with special chars: \u0001"}`
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/embeddings",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(originalBody)),
|
||
|
|
)
|
||
|
|
resp.Body.Close()
|
||
|
|
|
||
|
|
// The received body should match the original (though may have different formatting)
|
||
|
|
var orig, received map[string]interface{}
|
||
|
|
json.Unmarshal([]byte(originalBody), &orig)
|
||
|
|
json.Unmarshal([]byte(receivedBody), &received)
|
||
|
|
|
||
|
|
if orig["model"] != received["model"] || orig["input"] != received["input"] {
|
||
|
|
t.Errorf("body was not forwarded correctly")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestRerankerBodyForwarded verifies body is byte-identical to upstream
|
||
|
|
func TestRerankerBodyForwarded(t *testing.T) {
|
||
|
|
receivedBody := ""
|
||
|
|
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
body, _ := io.ReadAll(r.Body)
|
||
|
|
receivedBody = string(body)
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
fmt.Fprint(w, `{"results":[]}`)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"BAAI/bge-reranker-base": {
|
||
|
|
Name: "BAAI/bge-reranker-base",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
originalBody := `{"model":"BAAI/bge-reranker-base","query":"test","texts":["a","b","c"]}`
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/rerank",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(originalBody)),
|
||
|
|
)
|
||
|
|
resp.Body.Close()
|
||
|
|
|
||
|
|
var orig, received map[string]interface{}
|
||
|
|
json.Unmarshal([]byte(originalBody), &orig)
|
||
|
|
json.Unmarshal([]byte(receivedBody), &received)
|
||
|
|
|
||
|
|
if orig["model"] != received["model"] || orig["query"] != received["query"] {
|
||
|
|
t.Errorf("body was not forwarded correctly")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestUpstreamErrorStatusSurfaced verifies upstream errors are returned as-is
|
||
|
|
func TestUpstreamErrorStatusSurfaced(t *testing.T) {
|
||
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusInternalServerError)
|
||
|
|
fmt.Fprint(w, `{"error":"upstream failure"}`)
|
||
|
|
}))
|
||
|
|
defer upstreamServer.Close()
|
||
|
|
|
||
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
||
|
|
|
||
|
|
cfg := &config.Config{
|
||
|
|
Models: map[string]*config.ModelUpstream{
|
||
|
|
"nomic-ai/nomic-embed-text-v2-moe": {
|
||
|
|
Name: "nomic-ai/nomic-embed-text-v2-moe",
|
||
|
|
Address: upstreamAddr,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Routes: make(map[string]*config.Route),
|
||
|
|
}
|
||
|
|
|
||
|
|
handler := New(cfg)
|
||
|
|
defer handler.Close()
|
||
|
|
|
||
|
|
server := httptest.NewServer(handler)
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
resp, _ := http.Post(
|
||
|
|
server.URL+"/v1/embeddings",
|
||
|
|
"application/json",
|
||
|
|
bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"test"}`)),
|
||
|
|
)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusInternalServerError {
|
||
|
|
t.Errorf("expected upstream error status 500, got %d", resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|