65 lines
1.6 KiB
Go
65 lines
1.6 KiB
Go
package embedding
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
|
)
|
|
|
|
func TestManager_SelectsAPIWhenConfigured(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
cfg := config.EmbeddingConfig{
|
|
APIURL: "http://localhost:9999/v1/embeddings",
|
|
APIKey: "secret",
|
|
Model: "text-embedding-3-small",
|
|
}
|
|
m := New(cfg)
|
|
if m.Embedder() == nil {
|
|
t.Fatal("Embedder() returned nil")
|
|
}
|
|
if _, ok := m.Embedder().(*APIEmbedder); !ok {
|
|
t.Fatalf("Embedder() = %T, want *APIEmbedder", m.Embedder())
|
|
}
|
|
}
|
|
|
|
func TestManager_SelectsNGramByDefault(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
m := New(config.EmbeddingConfig{})
|
|
emb, ok := m.Embedder().(*NGramEmbedder)
|
|
if !ok {
|
|
t.Fatalf("Embedder() = %T, want *NGramEmbedder", m.Embedder())
|
|
}
|
|
if emb.Dim() != defaultNGramDim {
|
|
t.Errorf("default ngram dim = %d, want %d", emb.Dim(), defaultNGramDim)
|
|
}
|
|
}
|
|
|
|
func TestManager_BlankAPIURLSelectsNGram(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
cfg := config.EmbeddingConfig{APIURL: " "} // whitespace counts as unset
|
|
m := New(cfg)
|
|
if _, ok := m.Embedder().(*NGramEmbedder); !ok {
|
|
t.Fatalf("Embedder() = %T, want *NGramEmbedder for blank URL", m.Embedder())
|
|
}
|
|
}
|
|
|
|
func TestManager_EmbedDelegatesToSelectedEmbedder(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
m := New(config.EmbeddingConfig{})
|
|
vectors, err := m.Embed(context.Background(), []string{"hello"})
|
|
if err != nil {
|
|
t.Fatalf("Embed() error: %v", err)
|
|
}
|
|
if len(vectors) != 1 {
|
|
t.Fatalf("got %d vectors, want 1", len(vectors))
|
|
}
|
|
if len(vectors[0]) != defaultNGramDim {
|
|
t.Errorf("vector length = %d, want %d", len(vectors[0]), defaultNGramDim)
|
|
}
|
|
}
|