diff --git a/embeddings/cohere.go b/embeddings/cohere.go index 5b38d8e..5ee8116 100644 --- a/embeddings/cohere.go +++ b/embeddings/cohere.go @@ -9,6 +9,9 @@ import ( "time" ) +// cohereEndpoint is overridable in tests. +var cohereEndpoint = "https://api.cohere.com/v2/embed" + // cohere implements the Cohere embedding API. // Supported models: embed-english-v3.0, embed-multilingual-v3.0. type cohere struct { @@ -39,7 +42,7 @@ func (p *cohere) Embed(ctx context.Context, text string) ([]float32, error) { if err != nil { return nil, fmt.Errorf("cohere: marshal: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.cohere.com/v2/embed", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", cohereEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("cohere: create request: %w", err) } @@ -98,7 +101,7 @@ func (p *cohere) EmbedBatch(ctx context.Context, texts []string) ([][]float32, e if err != nil { return nil, fmt.Errorf("cohere: marshal batch: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.cohere.com/v2/embed", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", cohereEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("cohere: create request: %w", err) } @@ -159,7 +162,7 @@ func (p *cohere) EmbedWithMode(ctx context.Context, text string, mode EmbedMode) if err != nil { return nil, fmt.Errorf("cohere: marshal: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.cohere.com/v2/embed", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", cohereEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("cohere: create request: %w", err) } diff --git a/embeddings/memo_test.go b/embeddings/memo_test.go index 1c3f100..4efc759 100644 --- a/embeddings/memo_test.go +++ b/embeddings/memo_test.go @@ -119,6 +119,60 @@ func (p *namedProvider) Name() string { return p.name } // TestMemoizedProvider_ModelSwapInvalidates pins the core #1 fix: a memo built // for one model must not serve its vectors to a different model. Since the memo // is namespaced by Name(), each model has its own key space. +func TestEmbeddingMemoLen(t *testing.T) { + t.Parallel() + m := NewEmbeddingMemo(10) + if m.Len() != 0 { + t.Fatalf("expected empty memo, got %d entries", m.Len()) + } + m.Put("a", []float32{1}) + if m.Len() != 1 { + t.Fatalf("expected 1 entry, got %d", m.Len()) + } + m.Put("b", []float32{2}) + if m.Len() != 2 { + t.Fatalf("expected 2 entries, got %d", m.Len()) + } +} + +func TestEmbeddingMemoPutModeUpdate(t *testing.T) { + t.Parallel() + m := NewEmbeddingMemo(10) + m.Put("a", []float32{1, 2, 3}) + if m.Len() != 1 { + t.Fatalf("expected 1 entry after put, got %d", m.Len()) + } + // Update existing entry + m.Put("a", []float32{4, 5, 6}) + if m.Len() != 1 { + t.Fatalf("expected still 1 entry after update, got %d", m.Len()) + } + vec, ok := m.Get("a") + if !ok || vec[0] != 4 { + t.Errorf("expected updated vector [4 5 6], got %v", vec) + } +} + +func TestEmbeddingMemoNewWithZeroEntries(t *testing.T) { + t.Parallel() + m := NewEmbeddingMemoNS("test", 0) + if m.max != 1024 { + t.Errorf("expected default 1024 max for zero input, got %d", m.max) + } +} + +func TestMemoizedProviderNameAndDims(t *testing.T) { + t.Parallel() + inner := NewLocal() + p := NewMemoizedProvider(inner, 10) + if p.Name() != inner.Name() { + t.Errorf("expected Name=%q, got %q", inner.Name(), p.Name()) + } + if p.Dims() != inner.Dims() { + t.Errorf("expected Dims=%d, got %d", inner.Dims(), p.Dims()) + } +} + func TestMemoizedProvider_ModelSwapInvalidates(t *testing.T) { t.Parallel() ctx := context.Background() diff --git a/embeddings/provider.go b/embeddings/provider.go index a814305..ccece80 100644 --- a/embeddings/provider.go +++ b/embeddings/provider.go @@ -18,6 +18,10 @@ import ( // between OpenAI and Voyage is ever needed. var httpClient = &http.Client{Timeout: 30 * time.Second} +// Endpoint URLs are package-level vars so tests can override them with httptest servers. +var openAIEndpoint = "https://api.openai.com/v1/embeddings" +var voyageEndpoint = "https://api.voyageai.com/v1/embeddings" + // maxRetries is the maximum number of rate-limit retries before giving up. const maxRetries = 5 @@ -71,7 +75,7 @@ func (p *openAI) Embed(ctx context.Context, text string) ([]float32, error) { if err != nil { return nil, fmt.Errorf("openai: marshal: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.openai.com/v1/embeddings", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", openAIEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("openai: create request: %w", err) } @@ -147,7 +151,7 @@ func (p *openAI) EmbedBatch(ctx context.Context, texts []string) ([][]float32, e if err != nil { return nil, fmt.Errorf("openai: marshal batch: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.openai.com/v1/embeddings", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", openAIEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("openai: create request: %w", err) } @@ -252,7 +256,7 @@ func (p *voyage) Embed(ctx context.Context, text string) ([]float32, error) { if err != nil { return nil, fmt.Errorf("voyage: marshal: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.voyageai.com/v1/embeddings", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", voyageEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("voyage: create request: %w", err) } @@ -312,7 +316,7 @@ func (p *voyage) EmbedBatch(ctx context.Context, texts []string) ([][]float32, e if err != nil { return nil, fmt.Errorf("voyage: marshal batch: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.voyageai.com/v1/embeddings", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", voyageEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("voyage: create request: %w", err) } @@ -376,7 +380,7 @@ func (p *voyage) EmbedWithMode(ctx context.Context, text string, mode EmbedMode) if err != nil { return nil, fmt.Errorf("voyage: marshal: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", "https://api.voyageai.com/v1/embeddings", bytes.NewReader(body)) + req, err := http.NewRequestWithContext(ctx, "POST", voyageEndpoint, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("voyage: create request: %w", err) } diff --git a/embeddings/provider_test.go b/embeddings/provider_test.go new file mode 100644 index 0000000..f1d19d9 --- /dev/null +++ b/embeddings/provider_test.go @@ -0,0 +1,717 @@ +package embeddings + +import ( + "context" + "encoding/json" + "math" + "net/http" + "net/http/httptest" + "testing" +) + +func TestNewLocal(t *testing.T) { + p := NewLocal() + if p == nil { + t.Fatal("NewLocal returned nil") + } +} + +func TestLocalStubName(t *testing.T) { + p := NewLocal() + if p.Name() != "local:stub" { + t.Errorf("expected 'local:stub', got %q", p.Name()) + } +} + +func TestLocalStubDims(t *testing.T) { + p := NewLocal() + if p.Dims() != 128 { + t.Errorf("expected 128 dims, got %d", p.Dims()) + } +} + +func TestLocalStubEmbed(t *testing.T) { + p := NewLocal() + ctx := context.Background() + + vec, err := p.Embed(ctx, "hello world") + if err != nil { + t.Fatalf("Embed returned error: %v", err) + } + if len(vec) != 128 { + t.Errorf("expected 128-dim vector, got %d", len(vec)) + } + + // Deterministic + vec2, _ := p.Embed(ctx, "hello world") + for i := range vec { + if vec[i] != vec2[i] { + t.Errorf("expected deterministic embedding, diff at index %d: %f vs %f", i, vec[i], vec2[i]) + break + } + } + + // Different text produces different vector + vec3, _ := p.Embed(ctx, "different text") + same := true + for i := range vec { + if vec[i] != vec3[i] { + same = false + break + } + } + if same { + t.Error("different inputs should produce different vectors") + } +} + +func TestLocalStubEmbedEmpty(t *testing.T) { + p := NewLocal() + ctx := context.Background() + + vec, err := p.Embed(ctx, "") + if err != nil { + t.Fatalf("Embed empty string returned error: %v", err) + } + if len(vec) != 128 { + t.Errorf("expected 128-dim vector for empty string, got %d", len(vec)) + } +} + +func TestLocalStubEmbedBatch(t *testing.T) { + p := NewLocal() + ctx := context.Background() + + vecs, err := p.EmbedBatch(ctx, []string{"a", "b", "c"}) + if err != nil { + t.Fatalf("EmbedBatch returned error: %v", err) + } + if len(vecs) != 3 { + t.Errorf("expected 3 vectors, got %d", len(vecs)) + } + for i, v := range vecs { + if len(v) != 128 { + t.Errorf("vector %d: expected 128 dims, got %d", i, len(v)) + } + } +} + +func TestLocalStubEmbedBatchEmpty(t *testing.T) { + p := NewLocal() + ctx := context.Background() + + vecs, err := p.EmbedBatch(ctx, []string{}) + if err != nil { + t.Fatalf("EmbedBatch empty returned error: %v", err) + } + if len(vecs) != 0 { + t.Errorf("expected 0 vectors, got %d", len(vecs)) + } +} + +func TestLocalStubEmbedWithMode(t *testing.T) { + p := NewLocal() + ctx := context.Background() + + vec, err := p.EmbedWithMode(ctx, "test", ModeQuery) + if err != nil { + t.Fatalf("EmbedWithMode returned error: %v", err) + } + if len(vec) != 128 { + t.Errorf("expected 128-dim vector, got %d", len(vec)) + } +} + +func TestNewOpenAI(t *testing.T) { + p := NewOpenAI("sk-test", "text-embedding-3-small") + if p == nil { + t.Fatal("NewOpenAI returned nil") + } + if p.Dims() != 1536 { + t.Errorf("expected 1536 dims, got %d", p.Dims()) + } + if p.Name() != "openai:text-embedding-3-small" { + t.Errorf("expected 'openai:text-embedding-3-small', got %q", p.Name()) + } +} + +func TestNewOpenAILarge(t *testing.T) { + p := NewOpenAI("sk-test", "text-embedding-3-large") + if p.Dims() != 3072 { + t.Errorf("expected 3072 dims for large model, got %d", p.Dims()) + } +} + +func TestNewOpenAIEmptyKey(t *testing.T) { + p := NewOpenAI("", "text-embedding-3-small") + if p == nil { + t.Fatal("NewOpenAI with empty key returned nil") + } +} + +func TestNewCohere(t *testing.T) { + p := NewCohere("co-test", "embed-english-v3.0") + if p == nil { + t.Fatal("NewCohere returned nil") + } + if p.Dims() != 1024 { + t.Errorf("expected 1024 dims, got %d", p.Dims()) + } + if p.Name() != "cohere:embed-english-v3.0" { + t.Errorf("expected 'cohere:embed-english-v3.0', got %q", p.Name()) + } +} + +func TestNewCohereDefaultModel(t *testing.T) { + p := NewCohere("co-test", "") + if p.Name() != "cohere:embed-english-v3.0" { + t.Errorf("expected default model 'cohere:embed-english-v3.0', got %q", p.Name()) + } +} + +func TestNewVoyage(t *testing.T) { + p := NewVoyage("voy-test", "voyage-code-3") + if p == nil { + t.Fatal("NewVoyage returned nil") + } + if p.Dims() != 1024 { + t.Errorf("expected 1024 dims, got %d", p.Dims()) + } + if p.Name() != "voyage:voyage-code-3" { + t.Errorf("expected 'voyage:voyage-code-3', got %q", p.Name()) + } +} + +func TestNewVoyageDefaultModel(t *testing.T) { + p := NewVoyage("voy-test", "") + if p.Name() != "voyage:voyage-code-3" { + t.Errorf("expected default model 'voyage:voyage-code-3', got %q", p.Name()) + } +} + +func TestCosine(t *testing.T) { + a := []float32{1, 0, 0} + b := []float32{1, 0, 0} + c := Cosine(a, b) + if math.Abs(float64(c-1.0)) > 0.0001 { + t.Errorf("expected cosine ~1.0 for identical vectors, got %f", c) + } +} + +func TestCosineOrthogonal(t *testing.T) { + a := []float32{1, 0} + b := []float32{0, 1} + c := Cosine(a, b) + if math.Abs(float64(c)) > 0.0001 { + t.Errorf("expected cosine ~0 for orthogonal vectors, got %f", c) + } +} + +func TestCosineOpposite(t *testing.T) { + a := []float32{1, 0} + b := []float32{-1, 0} + c := Cosine(a, b) + if math.Abs(float64(c-(-1.0))) > 0.0001 { + t.Errorf("expected cosine ~-1 for opposite vectors, got %f", c) + } +} + +func TestCosineDifferentLengths(t *testing.T) { + a := []float32{1, 0} + b := []float32{1, 0, 0} + c := Cosine(a, b) + if c != 0 { + t.Errorf("expected 0 for different-length vectors, got %f", c) + } +} + +func TestCosineZeroVector(t *testing.T) { + a := []float32{0, 0} + b := []float32{1, 0} + c := Cosine(a, b) + if c != 0 { + t.Errorf("expected 0 for zero vector, got %f", c) + } +} + +func TestNormalize(t *testing.T) { + v := []float32{3, 4} + n := normalize(v) + expectedNorm := math.Sqrt(9 + 16) + if math.Abs(float64(n[0])-3.0/expectedNorm) > 0.0001 { + t.Errorf("unexpected normalized value[0]: %f", n[0]) + } + if math.Abs(float64(n[1])-4.0/expectedNorm) > 0.0001 { + t.Errorf("unexpected normalized value[1]: %f", n[1]) + } +} + +func TestNormalizeZeroVector(t *testing.T) { + v := []float32{0, 0} + n := normalize(v) + if len(n) != 2 || n[0] != 0 || n[1] != 0 { + t.Errorf("expected zero vector unchanged, got %v", n) + } +} + +func TestNormalizeAlreadyNormalized(t *testing.T) { + v := []float32{1, 0} + n := normalize(v) + if math.Abs(float64(n[0]-1.0)) > 0.0001 || math.Abs(float64(n[1])) > 0.0001 { + t.Errorf("expected unchanged unit vector, got %v", n) + } +} + +func TestEmbedModeValues(t *testing.T) { + if ModeDocument != 0 { + t.Errorf("expected ModeDocument=0, got %d", ModeDocument) + } + if ModeQuery != 1 { + t.Errorf("expected ModeQuery=1, got %d", ModeQuery) + } +} + +func TestMaxRetries(t *testing.T) { + if maxRetries != 5 { + t.Errorf("expected maxRetries=5, got %d", maxRetries) + } +} + +func TestHttpClientTimeout(t *testing.T) { + if httpClient.Timeout == 0 { + t.Error("expected httpClient to have a timeout") + } +} + +// --- HTTP mock tests for OpenAI, Voyage, and Cohere --- + +func TestOpenAIEmbed_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer sk-test" { + t.Error("missing or bad Authorization header") + } + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{0.1, 0.2, 0.3}, "index": 0}, + }, + }) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + vec, err := p.Embed(context.Background(), "hello") + if err != nil { + t.Fatalf("Embed failed: %v", err) + } + if len(vec) != 3 || vec[0] != 0.1 || vec[1] != 0.2 || vec[2] != 0.3 { + t.Errorf("unexpected vector: %v", vec) + } +} + +func TestOpenAIEmbed_EmptyResponse(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"data": []map[string]any{}}) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for empty response") + } +} + +func TestOpenAIEmbed_NonRetryableError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]string{"message": "invalid request"}, + }) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for 400") + } +} + +func TestOpenAIEmbedBatch_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{1, 0}, "index": 1}, + {"embedding": []float32{0, 1}, "index": 0}, + }, + }) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + vecs, err := p.EmbedBatch(context.Background(), []string{"a", "b"}) + if err != nil { + t.Fatalf("EmbedBatch failed: %v", err) + } + if len(vecs) != 2 { + t.Fatalf("expected 2 vectors, got %d", len(vecs)) + } + // Should be reordered by index: [0]={0,1}, [1]={1,0} + if vecs[0][0] != 0 || vecs[0][1] != 1 { + t.Errorf("expected vecs[0]={0,1}, got %v", vecs[0]) + } + if vecs[1][0] != 1 || vecs[1][1] != 0 { + t.Errorf("expected vecs[1]={1,0}, got %v", vecs[1]) + } +} + +func TestOpenAIEmbedBatch_Empty(t *testing.T) { + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + vecs, err := p.EmbedBatch(context.Background(), nil) + if err != nil { + t.Fatalf("EmbedBatch(nil) failed: %v", err) + } + if vecs != nil { + t.Errorf("expected nil for empty batch, got %v", vecs) + } +} + +func TestOpenAIEmbedWithMode(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{0.5, 0.5}, "index": 0}, + }, + }) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + vec, err := p.EmbedWithMode(context.Background(), "test", ModeQuery) + if err != nil { + t.Fatalf("EmbedWithMode failed: %v", err) + } + if len(vec) != 2 { + t.Errorf("expected 2 dims, got %d", len(vec)) + } +} + +func TestVoyageEmbed_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{0.1, 0.2}}, + }, + }) + })) + defer ts.Close() + orig := voyageEndpoint + voyageEndpoint = ts.URL + defer func() { voyageEndpoint = orig }() + + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + vec, err := p.Embed(context.Background(), "hello") + if err != nil { + t.Fatalf("Embed failed: %v", err) + } + if len(vec) != 2 || vec[0] != 0.1 || vec[1] != 0.2 { + t.Errorf("unexpected vector: %v", vec) + } +} + +func TestVoyageEmbed_Error(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]any{ + "detail": "internal error", + }) + })) + defer ts.Close() + orig := voyageEndpoint + voyageEndpoint = ts.URL + defer func() { voyageEndpoint = orig }() + + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for 500") + } +} + +func TestVoyageEmbed_EmptyResponse(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"data": []map[string]any{}}) + })) + defer ts.Close() + orig := voyageEndpoint + voyageEndpoint = ts.URL + defer func() { voyageEndpoint = orig }() + + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for empty response") + } +} + +func TestVoyageEmbedBatch_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{1, 0}}, + {"embedding": []float32{0, 1}}, + }, + }) + })) + defer ts.Close() + orig := voyageEndpoint + voyageEndpoint = ts.URL + defer func() { voyageEndpoint = orig }() + + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + vecs, err := p.EmbedBatch(context.Background(), []string{"a", "b"}) + if err != nil { + t.Fatalf("EmbedBatch failed: %v", err) + } + if len(vecs) != 2 { + t.Fatalf("expected 2 vectors, got %d", len(vecs)) + } +} + +func TestVoyageEmbedBatch_Empty(t *testing.T) { + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + vecs, err := p.EmbedBatch(context.Background(), nil) + if err != nil { + t.Fatalf("EmbedBatch(nil) failed: %v", err) + } + if vecs != nil { + t.Errorf("expected nil for empty batch, got %v", vecs) + } +} + +func TestVoyageEmbedWithMode(t *testing.T) { + var inputType string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + InputType string `json:"input_type"` + } + json.NewDecoder(r.Body).Decode(&body) + inputType = body.InputType + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{0.5}}, + }, + }) + })) + defer ts.Close() + orig := voyageEndpoint + voyageEndpoint = ts.URL + defer func() { voyageEndpoint = orig }() + + p := NewVoyage("voy-test", "voyage-code-3").(*voyage) + + // Document mode + _, err := p.EmbedWithMode(context.Background(), "doc", ModeDocument) + if err != nil { + t.Fatalf("EmbedWithMode document failed: %v", err) + } + if inputType != "search_document" { + t.Errorf("expected search_document, got %q", inputType) + } + + // Query mode + _, err = p.EmbedWithMode(context.Background(), "query", ModeQuery) + if err != nil { + t.Fatalf("EmbedWithMode query failed: %v", err) + } + if inputType != "search_query" { + t.Errorf("expected search_query, got %q", inputType) + } +} + +func TestCohereEmbed_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "embeddings": [][]float32{{0.1, 0.2, 0.3}}, + }) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + vec, err := p.Embed(context.Background(), "hello") + if err != nil { + t.Fatalf("Embed failed: %v", err) + } + if len(vec) != 3 || vec[0] != 0.1 || vec[1] != 0.2 || vec[2] != 0.3 { + t.Errorf("unexpected vector: %v", vec) + } +} + +func TestCohereEmbed_Error(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(401) + json.NewEncoder(w).Encode(map[string]any{ + "message": "unauthorized", + }) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for 401") + } +} + +func TestCohereEmbed_EmptyResponse(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"embeddings": [][]float32{}}) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + _, err := p.Embed(context.Background(), "hello") + if err == nil { + t.Fatal("expected error for empty embeddings") + } +} + +func TestCohereEmbedBatch_Success(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "embeddings": [][]float32{{1, 0}, {0, 1}}, + }) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + vecs, err := p.EmbedBatch(context.Background(), []string{"a", "b"}) + if err != nil { + t.Fatalf("EmbedBatch failed: %v", err) + } + if len(vecs) != 2 { + t.Fatalf("expected 2 vectors, got %d", len(vecs)) + } +} + +func TestCohereEmbedBatch_Empty(t *testing.T) { + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + vecs, err := p.EmbedBatch(context.Background(), nil) + if err != nil { + t.Fatalf("EmbedBatch(nil) failed: %v", err) + } + if vecs != nil { + t.Errorf("expected nil for empty batch, got %v", vecs) + } +} + +func TestCohereEmbedBatch_EmptyResponse(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"embeddings": [][]float32{}}) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + _, err := p.EmbedBatch(context.Background(), []string{"a"}) + if err == nil { + t.Fatal("expected error for empty batch response") + } +} + +func TestCohereEmbedWithMode(t *testing.T) { + var inputType string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + InputType string `json:"input_type"` + } + json.NewDecoder(r.Body).Decode(&body) + inputType = body.InputType + json.NewEncoder(w).Encode(map[string]any{ + "embeddings": [][]float32{{0.5}}, + }) + })) + defer ts.Close() + orig := cohereEndpoint + cohereEndpoint = ts.URL + defer func() { cohereEndpoint = orig }() + + p := NewCohere("co-test", "embed-english-v3.0").(*cohere) + + _, err := p.EmbedWithMode(context.Background(), "doc", ModeDocument) + if err != nil { + t.Fatalf("EmbedWithMode document failed: %v", err) + } + if inputType != "search_document" { + t.Errorf("expected search_document, got %q", inputType) + } + + _, err = p.EmbedWithMode(context.Background(), "query", ModeQuery) + if err != nil { + t.Fatalf("EmbedWithMode query failed: %v", err) + } + if inputType != "search_query" { + t.Errorf("expected search_query, got %q", inputType) + } +} + +func TestOpenAIEmbed_RetryThenSuccess(t *testing.T) { + attempts := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts++ + if attempts == 1 { + w.WriteHeader(429) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]string{"message": "Rate limit exceeded. Please try again in 0.001s"}, + }) + return + } + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"embedding": []float32{0.5}, "index": 0}, + }, + }) + })) + defer ts.Close() + orig := openAIEndpoint + openAIEndpoint = ts.URL + defer func() { openAIEndpoint = orig }() + + p := NewOpenAI("sk-test", "text-embedding-3-small").(*openAI) + vec, err := p.Embed(context.Background(), "hello") + if err != nil { + t.Fatalf("Embed should have retried and succeeded: %v", err) + } + if len(vec) != 1 || vec[0] != 0.5 { + t.Errorf("unexpected vector: %v", vec) + } +} diff --git a/portablegraph/portable_graph.go b/portablegraph/portable_graph.go index efb9677..4bc7a92 100644 --- a/portablegraph/portable_graph.go +++ b/portablegraph/portable_graph.go @@ -131,7 +131,6 @@ func (g *PortableGraph) ToGraphSpec() *graphcontracts.GraphSpec { nodes = append(nodes, graphcontracts.NodeSpec{ ID: id, - Type: graphcontracts.NodeTypeSystem, Name: node.Name, Config: config, })