homelab-frontend/internal/proxy/embeddings_and_rerank_test.go

package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"forgejo.riotpiao.com/rock/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, err := http.Post(
		server.URL+"/v1/embeddings",
		"application/json",
		bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"hello"}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/embeddings",
		"application/json",
		bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"test"}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/rerank",
		"application/json",
		bytes.NewReader([]byte(`{"model":"BAAI/bge-reranker-base","query":"test","texts":["a","b"]}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/rerank",
		"application/json",
		bytes.NewReader([]byte(`{"model":"BAAI/bge-reranker-base","query":"q","texts":["a"]}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/embeddings",
		"application/json",
		bytes.NewReader([]byte(`{"model":"unknown-embeddings","input":"test"}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/rerank",
		"application/json",
		bytes.NewReader([]byte(`{"model":"unknown-reranker","query":"q","texts":["a"]}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/embeddings",
		"application/json",
		bytes.NewReader([]byte(originalBody)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/rerank",
		"application/json",
		bytes.NewReader([]byte(originalBody)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	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, err := http.Post(
		server.URL+"/v1/embeddings",
		"application/json",
		bytes.NewReader([]byte(`{"model":"nomic-ai/nomic-embed-text-v2-moe","input":"test"}`)),
	)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusInternalServerError {
		t.Errorf("expected upstream error status 500, got %d", resp.StatusCode)
	}
}