rony-llm-agent/pkg/rag/backends/sqlitevec/sqlitevec_test.go

package sqlitevec_test

import (
	"context"
	"testing"

	"github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/sqlitevec"
)

func newTestBackend(t *testing.T) *sqlitevec.Backend {
	t.Helper()
	b, err := sqlitevec.New("")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	t.Cleanup(func() { b.Close() })
	return b
}

func TestBackend_UpsertAndSearchByVector(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", []float32{1, 0, 0}, "cómo desplegar a producción", map[string]string{"kind": "process"}); err != nil {
		t.Fatalf("upsert a: %v", err)
	}
	if err := b.Upsert(ctx, "b", []float32{0, 1, 0}, "receta de pan", nil); err != nil {
		t.Fatalf("upsert b: %v", err)
	}

	results, err := b.Search(ctx, "", []float32{1, 0, 0}, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("expected 2 results, got %d", len(results))
	}
	if results[0].ID != "a" {
		t.Fatalf("expected closest match to be 'a', got %q", results[0].ID)
	}
	if results[0].Content != "cómo desplegar a producción" {
		t.Fatalf("expected content to round-trip, got %q", results[0].Content)
	}
	if results[0].Metadata["kind"] != "process" {
		t.Fatalf("expected metadata to round-trip, got %v", results[0].Metadata)
	}
}

func TestBackend_UpsertReplacesExisting(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", []float32{1, 0}, "first version", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}
	if err := b.Upsert(ctx, "a", []float32{0, 1}, "second version", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}

	results, err := b.Search(ctx, "", []float32{0, 1}, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 1 {
		t.Fatalf("expected the upsert to replace, not duplicate; got %d results", len(results))
	}
	if results[0].Content != "second version" {
		t.Fatalf("expected replaced content, got %q", results[0].Content)
	}

	// The FTS5 side must also have been replaced, not duplicated.
	ftsResults, err := b.Search(ctx, "second version", nil, 5)
	if err != nil {
		t.Fatalf("fts search: %v", err)
	}
	if len(ftsResults) != 1 {
		t.Fatalf("expected 1 fts result after replace, got %d", len(ftsResults))
	}
}

func TestBackend_SearchTopK(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	for i, id := range []string{"a", "b", "c"} {
		vec := []float32{float32(i), 1, 1}
		if err := b.Upsert(ctx, id, vec, id, nil); err != nil {
			t.Fatalf("upsert %s: %v", id, err)
		}
	}

	results, err := b.Search(ctx, "", []float32{1, 1, 1}, 2)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("expected topK=2 results, got %d", len(results))
	}
}

func TestBackend_Forget(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", []float32{1, 0}, "content", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}
	if err := b.Forget(ctx, "a"); err != nil {
		t.Fatalf("forget: %v", err)
	}

	results, err := b.Search(ctx, "", []float32{1, 0}, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected no results after forget, got %d", len(results))
	}

	ftsResults, err := b.Search(ctx, "content", nil, 5)
	if err != nil {
		t.Fatalf("fts search: %v", err)
	}
	if len(ftsResults) != 0 {
		t.Fatalf("expected no fts results after forget, got %d", len(ftsResults))
	}
}

func TestBackend_ForgetAll(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	for _, id := range []string{"a", "b"} {
		if err := b.Upsert(ctx, id, []float32{1, 0}, id, nil); err != nil {
			t.Fatalf("upsert %s: %v", id, err)
		}
	}
	if err := b.ForgetAll(ctx); err != nil {
		t.Fatalf("forget all: %v", err)
	}

	results, err := b.Search(ctx, "", []float32{1, 0}, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected no results after forget all, got %d", len(results))
	}
}

func TestBackend_SearchEmpty(t *testing.T) {
	b := newTestBackend(t)

	results, err := b.Search(context.Background(), "", []float32{1, 0}, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected no results on empty store, got %d", len(results))
	}
}

func TestBackend_SearchByText_NoVectorFallsBackToFTS(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	// Upserted with an empty vector, as if the embedder had failed.
	if err := b.Upsert(ctx, "a", nil, "cómo desplegar a producción con Docker", map[string]string{"kind": "process"}); err != nil {
		t.Fatalf("upsert a: %v", err)
	}
	if err := b.Upsert(ctx, "b", nil, "receta de pan con masa madre", nil); err != nil {
		t.Fatalf("upsert b: %v", err)
	}

	results, err := b.Search(ctx, "desplegar producción", nil, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 1 {
		t.Fatalf("expected 1 fts match, got %d", len(results))
	}
	if results[0].ID != "a" {
		t.Fatalf("expected match to be 'a', got %q", results[0].ID)
	}
	if results[0].Metadata["kind"] != "process" {
		t.Fatalf("expected metadata to round-trip through fts path, got %v", results[0].Metadata)
	}
}

func TestBackend_SearchByText_NoMatches(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", nil, "receta de pan", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}

	results, err := b.Search(ctx, "algo completamente distinto", nil, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected no matches, got %d", len(results))
	}
}

func TestBackend_SearchByText_HandlesSpecialCharacters(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", nil, "usa docker-compose para levantar todo", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}

	// FTS5 operator characters in the query must not cause a syntax error.
	results, err := b.Search(ctx, `docker-compose "up" AND/OR *test*`, nil, 5)
	if err != nil {
		t.Fatalf("search should not error on special characters: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected at least one match despite special characters in the query")
	}
}

func TestBackend_SearchByText_EmptyQuery(t *testing.T) {
	b := newTestBackend(t)
	ctx := context.Background()

	if err := b.Upsert(ctx, "a", nil, "algo", nil); err != nil {
		t.Fatalf("upsert: %v", err)
	}

	results, err := b.Search(ctx, "   ", nil, 5)
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected no results for an empty query, got %d", len(results))
	}
}