embeddings_test.go

package blmstudio_test

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	blmstudio "github.com/bearaujus/blmstudio"
)

func TestEmbed_Happy(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.EmbeddingResponse{
			Data: []blmstudio.EmbeddingObject{{
				Embedding: []float64{0.1, 0.2, 0.3},
				Index:     0,
			}},
			Model: "embed-model",
		})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	resp, err := c.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: "embed-model",
		Input: blmstudio.NewEmbeddingStringsInput([]string{"hello world"}),
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(resp.Data) == 0 {
		t.Fatal("expected at least one embedding")
	}
	if len(resp.Data[0].Embedding) != 3 {
		t.Errorf("got %d dims, want 3", len(resp.Data[0].Embedding))
	}
}

func TestEmbed_Single(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.EmbeddingResponse{
			Data: []blmstudio.EmbeddingObject{{Embedding: []float64{1.0}, Index: 0}},
		})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	resp, err := c.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: "embed-model",
		Input: blmstudio.NewEmbeddingStringInput("single"),
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(resp.Data) == 0 {
		t.Fatal("expected embedding data")
	}
}

func TestEmbed_Empty(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.EmbeddingResponse{})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	// Zero-value EmbeddingInput has nil raw → triggers ErrEmptyText
	_, err := c.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: "embed-model",
	})
	if !errors.Is(err, blmstudio.ErrEmptyText) {
		t.Errorf("got %v, want ErrEmptyText", err)
	}
}

func TestEmbed_4xx(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 422, map[string]any{"error": "invalid input", "message": "bad text"})
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: "embed-model",
		Input: blmstudio.NewEmbeddingStringInput("text"),
	})
	if _, ok := blmstudio.IsAPIError(err); !ok {
		t.Errorf("expected APIError, got %v", err)
	}
}

func TestEmbed_MalformedJSON(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		_, _ = w.Write([]byte("not-json"))
	}))
	defer srv.Close()

	c := newTestClient(t, srv)
	_, err := c.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: "embed-model",
		Input: blmstudio.NewEmbeddingStringInput("text"),
	})
	if err == nil {
		t.Fatal("expected error for malformed JSON, got nil")
	}
}

func TestEmbed_Timeout(t *testing.T) {
	t.Parallel()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(500 * time.Millisecond)
		writeJSON(w, 200, blmstudio.EmbeddingResponse{})
	}))
	defer srv.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	c := newTestClient(t, srv)
	_, err := c.Embed(ctx, &blmstudio.EmbeddingRequest{
		Model: "embed-model",
		Input: blmstudio.NewEmbeddingStringInput("text"),
	})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("got %v, want DeadlineExceeded", err)
	}
}

func BenchmarkEmbed(b *testing.B) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, 200, blmstudio.EmbeddingResponse{
			Data: []blmstudio.EmbeddingObject{{Embedding: make([]float64, 384), Index: 0}},
		})
	}))
	defer srv.Close()

	c := blmstudio.New(blmstudio.WithBaseURL(srv.URL))
	input := &blmstudio.EmbeddingRequest{
		Model: "bench-model",
		Input: blmstudio.NewEmbeddingStringInput("benchmark text for embedding"),
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = c.Embed(context.Background(), input)
	}
}