googleai: add embeddings to vertex (#546)

for #410

llms/googleai/palm/palm_llm.go

@@ -9,7 +9,7 @@ import (
 
 	"github.com/tmc/langchaingo/callbacks"
 	"github.com/tmc/langchaingo/llms"
-	"github.com/tmc/langchaingo/llms/googleai/palm/internal/palmclient"
+	"github.com/tmc/langchaingo/llms/googleai/internal/palmclient"
 )
 
 var (

llms/googleai/vertex/embeddings.go

@@ -0,0 +1,28 @@
+package vertex
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/tmc/langchaingo/llms/googleai/internal/palmclient"
+)
+
+// CreateEmbedding creates embeddings from texts.
+func (g *Vertex) CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error) {
+	embeddings, err := g.palmClient.CreateEmbedding(ctx, &palmclient.EmbeddingRequest{
+		Input: texts,
+	})
+	if err != nil {
+		return [][]float32{}, err
+	}
+
+	if len(embeddings) == 0 {
+		return nil, errors.New("empty response")
+	}
+	if len(texts) != len(embeddings) {
+		return embeddings, fmt.Errorf("returned %d embeddings for %d texts", len(embeddings), len(texts))
+	}
+
+	return embeddings, nil
+}

llms/googleai/vertex/new.go

@@ -5,21 +5,22 @@ package vertex
 
 import (
 	"context"
-	"log"
 
 	"cloud.google.com/go/vertexai/genai"
 	"github.com/tmc/langchaingo/callbacks"
 	"github.com/tmc/langchaingo/llms"
+	"github.com/tmc/langchaingo/llms/googleai/internal/palmclient"
 )
 
 // Vertex is a type that represents a Vertex AI API client.
 //
-// TODO: This isn't in common code; may need PaLM client for embeddings, etc.
-// Note the deltas: type of topk, candidate count.
+// Right now, the Vertex Gemini SDK doesn't support embeddings; therefore,
+// for embeddings we also hold a palmclient.
 type Vertex struct {
 	CallbacksHandler callbacks.Handler
 	client           *genai.Client
 	opts             options
+	palmClient       *palmclient.PaLMClient
 }
 
 var _ llms.Model = &Vertex{}
@@ -31,15 +32,20 @@ func NewVertex(ctx context.Context, opts ...Option) (*Vertex, error) {
 		opt(&clientOptions)
 	}
 
-	v := &Vertex{
-		opts: clientOptions,
+	client, err := genai.NewClient(ctx, clientOptions.cloudProject, clientOptions.cloudLocation)
+	if err != nil {
+		return nil, err
 	}
 
-	client, err := genai.NewClient(ctx, clientOptions.cloudProject, clientOptions.cloudLocation)
+	palmClient, err := palmclient.New(clientOptions.cloudProject) //nolint:contextcheck
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
 
-	v.client = client
+	v := &Vertex{
+		opts:       clientOptions,
+		client:     client,
+		palmClient: palmClient,
+	}
 	return v, nil
 }

llms/googleai/vertex/vertex.go

@@ -118,11 +118,6 @@ func convertCandidates(candidates []*genai.Candidate) (*llms.ContentResponse, er
 	return &contentResponse, nil
 }
 
-// CreateEmbedding creates embeddings from texts.
-func (g *Vertex) CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error) {
-	panic("not implemented")
-}
-
 // convertParts converts between a sequence of langchain parts and genai parts.
 func convertParts(parts []llms.ContentPart) ([]genai.Part, error) {
 	convertedParts := make([]genai.Part, 0, len(parts))

llms/googleai/vertex/vertex_test.go

@@ -89,3 +89,16 @@ func TestMultiContentTextStream(t *testing.T) {
 	assert.GreaterOrEqual(t, len(chunks), 2)
 	assert.Regexp(t, "(?i)dog|canid|canine", sb.String())
 }
+
+func TestEmbeddings(t *testing.T) {
+	t.Parallel()
+	llm := newClient(t)
+
+	texts := []string{"foo", "parrot"}
+	res, err := llm.CreateEmbedding(context.Background(), texts)
+	require.NoError(t, err)
+
+	assert.Equal(t, len(texts), len(res))
+	assert.NotEmpty(t, res[0])
+	assert.NotEmpty(t, res[1])
+}