rony-llm-agent/pkg/rag/embeddings/openai_compatible_test.go

package embeddings

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestNewOpenAICompatible_RequiresBaseURL(t *testing.T) {
	_, err := NewOpenAICompatible(OpenAICompatibleConfig{})
	if err == nil {
		t.Fatal("expected error when BaseURL is empty")
	}
}

func TestNewLlamaCpp_Defaults(t *testing.T) {
	e, err := NewLlamaCpp(LlamaCppConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if e.baseURL != "http://localhost:8080/v1" {
		t.Errorf("expected default base URL, got %q", e.baseURL)
	}
}

func TestNewLlamaCpp_Custom(t *testing.T) {
	e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: "http://custom:9000/v1", Model: "my-model"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if e.baseURL != "http://custom:9000/v1" {
		t.Errorf("expected custom base URL, got %q", e.baseURL)
	}
}

func TestOpenAICompatible_Embed_Success(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/embeddings" {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"data":[{"embedding":[0.1,0.2,0.3],"index":0}]}`))
	}))
	defer server.Close()

	e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: server.URL + "/v1"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	vector, err := e.Embed(context.Background(), "hola")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(vector) != 3 {
		t.Fatalf("expected 3 dimensions, got %d", len(vector))
	}
	if e.Dimensions() != 3 {
		t.Errorf("expected Dimensions() to learn 3 after a successful call, got %d", e.Dimensions())
	}
}

func TestOpenAICompatible_Embed_NotEnabled(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotImplemented)
		w.Write([]byte(`{"error":{"message":"This server does not support embeddings. Start it with ` + "`--embeddings`" + `"}}`))
	}))
	defer server.Close()

	e, err := NewLlamaCpp(LlamaCppConfig{BaseURL: server.URL + "/v1"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if _, err := e.Embed(context.Background(), "hola"); err == nil {
		t.Fatal("expected error when the server doesn't support embeddings")
	}
}

func TestOpenAICompatible_SendsBearerTokenWhenConfigured(t *testing.T) {
	var gotAuth string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"data":[{"embedding":[0.1],"index":0}]}`))
	}))
	defer server.Close()

	e, err := NewOpenAICompatible(OpenAICompatibleConfig{
		BaseURL: server.URL,
		APIKey:  "sk-test",
		Model:   "text-embedding-3-small",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if _, err := e.Embed(context.Background(), "hola"); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if gotAuth != "Bearer sk-test" {
		t.Fatalf("expected Authorization header to be sent, got %q", gotAuth)
	}
}

func TestOpenAICompatible_NoAuthHeaderWhenNoAPIKey(t *testing.T) {
	var gotAuth string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"data":[{"embedding":[0.1],"index":0}]}`))
	}))
	defer server.Close()

	e, err := NewOpenAICompatible(OpenAICompatibleConfig{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if _, err := e.Embed(context.Background(), "hola"); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if gotAuth != "" {
		t.Fatalf("expected no Authorization header without an API key, got %q", gotAuth)
	}
}

func TestOpenAICompatible_Dimensions_BeforeAnyCall(t *testing.T) {
	e, err := NewLlamaCpp(LlamaCppConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dims := e.Dimensions(); dims != 0 {
		t.Errorf("expected 0 dimensions before any successful call, got %d", dims)
	}
}