You've already forked agentic-coding-workflow
130 lines
4.3 KiB
Go
130 lines
4.3 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestOpenAI_Embed_ReturnsVectors(t *testing.T) {
|
|
var gotInputs []string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var body struct {
|
|
Input []string `json:"input"`
|
|
Model string `json:"model"`
|
|
Dimensions int `json:"dimensions"`
|
|
}
|
|
_ = json.NewDecoder(r.Body).Decode(&body)
|
|
gotInputs = body.Input
|
|
require.Equal(t, PlatformEmbeddingDims, body.Dimensions)
|
|
// Echo two 3-dim vectors (dimension count here is irrelevant to the test).
|
|
resp := map[string]any{
|
|
"object": "list",
|
|
"data": []map[string]any{
|
|
{"object": "embedding", "index": 0, "embedding": []float64{0.1, 0.2, 0.3}},
|
|
{"object": "embedding", "index": 1, "embedding": []float64{0.4, 0.5, 0.6}},
|
|
},
|
|
"model": body.Model,
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(resp)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c, _ := NewOpenAI("k", srv.URL)
|
|
vecs, err := c.Embed(context.Background(), "text-embedding-3-small", []string{"alpha", "beta"})
|
|
require.NoError(t, err)
|
|
require.Len(t, vecs, 2)
|
|
require.Equal(t, []float32{0.1, 0.2, 0.3}, vecs[0])
|
|
require.Equal(t, []float32{0.4, 0.5, 0.6}, vecs[1])
|
|
require.Equal(t, []string{"alpha", "beta"}, gotInputs)
|
|
require.Equal(t, PlatformEmbeddingDims, c.EmbedDims("text-embedding-3-small"))
|
|
}
|
|
|
|
func TestOpenAI_Embed_RespectsResponseIndex(t *testing.T) {
|
|
// Server returns vectors out of order; Embed must place them by Index.
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
resp := map[string]any{
|
|
"object": "list",
|
|
"data": []map[string]any{
|
|
{"object": "embedding", "index": 1, "embedding": []float64{9, 9}},
|
|
{"object": "embedding", "index": 0, "embedding": []float64{1, 1}},
|
|
},
|
|
"model": "m",
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(resp)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c, _ := NewOpenAI("k", srv.URL)
|
|
vecs, err := c.Embed(context.Background(), "m", []string{"a", "b"})
|
|
require.NoError(t, err)
|
|
require.Equal(t, []float32{1, 1}, vecs[0])
|
|
require.Equal(t, []float32{9, 9}, vecs[1])
|
|
}
|
|
|
|
func TestOpenAI_Embed_EmptyInputs(t *testing.T) {
|
|
c, _ := NewOpenAI("k", "")
|
|
vecs, err := c.Embed(context.Background(), "m", nil)
|
|
require.NoError(t, err)
|
|
require.Nil(t, vecs)
|
|
}
|
|
|
|
func TestAnthropic_Embed_Unsupported(t *testing.T) {
|
|
c, _ := NewAnthropic("k", "")
|
|
_, err := c.Embed(context.Background(), "m", []string{"x"})
|
|
require.ErrorIs(t, err, ErrEmbeddingsUnsupported)
|
|
require.Equal(t, 0, c.EmbedDims("m"))
|
|
}
|
|
|
|
func TestFakeEmbedder_Deterministic(t *testing.T) {
|
|
e := NewFakeEmbedder()
|
|
v1, err := e.Embed(context.Background(), "m", []string{"hello"})
|
|
require.NoError(t, err)
|
|
v2, err := e.Embed(context.Background(), "m", []string{"hello"})
|
|
require.NoError(t, err)
|
|
require.Equal(t, v1[0], v2[0], "same input must yield identical vector")
|
|
require.Len(t, v1[0], PlatformEmbeddingDims)
|
|
|
|
other, err := e.Embed(context.Background(), "m", []string{"world"})
|
|
require.NoError(t, err)
|
|
require.NotEqual(t, v1[0], other[0], "different inputs must differ")
|
|
}
|
|
|
|
func TestFakeEmbedder_Error(t *testing.T) {
|
|
e := &FakeEmbedder{Dims: 4, Err: errors.New("boom")}
|
|
_, err := e.Embed(context.Background(), "m", []string{"x"})
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestRegistry_GetEmbedder_TypeAsserts(t *testing.T) {
|
|
// OpenAI client implements Embedder.
|
|
id := uuid.New()
|
|
reg := NewRegistry(&fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{
|
|
id: {ID: id, Provider: "openai", BaseURL: "http://localhost", APIKey: "k"},
|
|
}})
|
|
emb, err := reg.GetEmbedder(context.Background(), id)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, emb)
|
|
require.Equal(t, PlatformEmbeddingDims, emb.EmbedDims("m"))
|
|
}
|
|
|
|
func TestRegistry_GetEmbedder_AnthropicSupportedButUnsupportedEmbed(t *testing.T) {
|
|
// Anthropic client implements Embedder (Embed returns ErrEmbeddingsUnsupported).
|
|
id := uuid.New()
|
|
reg := NewRegistry(&fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{
|
|
id: {ID: id, Provider: "anthropic", BaseURL: "http://localhost", APIKey: "k"},
|
|
}})
|
|
emb, err := reg.GetEmbedder(context.Background(), id)
|
|
require.NoError(t, err)
|
|
_, eerr := emb.Embed(context.Background(), "m", []string{"x"})
|
|
require.ErrorIs(t, eerr, ErrEmbeddingsUnsupported)
|
|
}
|