allow one to specify width of vector in echo embedding

this will be useful for testing as not all embedding models emit a vector of the
same size

Author: jaffee

Diff for: pkg/aicli/echoai.go

@@ -2,6 +2,8 @@ package aicli
 
 import (
 	"io"
+	"strconv"
+	"strings"
 
 	"github.com/pkg/errors"
 )
@@ -31,10 +33,30 @@ func (c *Echo) GetEmbedding(req *EmbeddingRequest) ([]Embedding, error) {
 	if len(req.Inputs) == 0 {
 		return nil, errors.New("emtpy list of inputs")
 	}
+	// parse req.Model to see what size vector to return. Format is name_num.
+	split := strings.Split(req.Model, "_")
+	var vectorSize int
+	switch len(split) {
+	case 0:
+		return nil, errors.New("empty model name")
+	case 1:
+		vectorSize = 1
+	default:
+		num, err := strconv.Atoi(split[1])
+		if err != nil {
+			return nil, errors.Wrap(err, "parsing model vector size")
+		}
+		vectorSize = num
+	}
+	if vectorSize <= 0 {
+		return nil, errors.Errorf("vector size must be greater than 0, got %d", vectorSize)
+	}
 	ret := make([]Embedding, len(req.Inputs))
 	for i := range ret {
-		ret[i].Embedding = []float32{0.42}
+		ret[i].Embedding = make([]float32, vectorSize)
+		for j := range ret[i].Embedding {
+			ret[i].Embedding[j] = 0.42
+		}
 	}
 	return ret, nil
-
 }

Diff for: pkg/aicli/echoai_test.go

@@ -44,3 +44,41 @@ func TestEcho(t *testing.T) {
 		})
 	}
 }
+
+func TestEchoEmbeddings(t *testing.T) {
+	c := &aicli.Echo{}
+	cases := []struct {
+		model  string
+		exp    []float32
+		expErr string
+	}{
+		{
+			model: "basic",
+			exp:   []float32{0.42},
+		},
+		{
+			model: "numbered_2",
+			exp:   []float32{0.42, 0.42},
+		},
+		{
+			model:  "numbered_-11",
+			expErr: "size must be greater than 0",
+		},
+	}
+
+	for i, tst := range cases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			er := &aicli.EmbeddingRequest{
+				Model:  tst.model,
+				Inputs: []string{"hello"},
+			}
+			emb, err := c.GetEmbedding(er)
+			if tst.expErr != "" {
+				require.Contains(t, err.Error(), tst.expErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tst.exp, emb[0].Embedding)
+			}
+		})
+	}
+}