Skip to content

Commit

Permalink
[Go] plugins/googlai: use batch embed RPC
Browse files Browse the repository at this point in the history
Use the BatchEmbedContents RPC so a single call can handle
multiple documents.
  • Loading branch information
jba committed Jul 7, 2024
1 parent 1546e28 commit 3b44d23
Showing 1 changed file with 11 additions and 7 deletions.
18 changes: 11 additions & 7 deletions go/plugins/googleai/googleai.go
Original file line number Diff line number Diff line change
Expand Up @@ -163,19 +163,23 @@ func DefineEmbedder(name string) *ai.Embedder {
// requires state.mu
func defineEmbedder(name string) *ai.Embedder {
return ai.DefineEmbedder(provider, name, func(ctx context.Context, input *ai.EmbedRequest) (*ai.EmbedResponse, error) {
// TODO: use the batch embedding API.
em := state.client.EmbeddingModel(name)
var res ai.EmbedResponse
// TODO: set em.TaskType from EmbedRequest.Options?
batch := em.NewBatch()
for _, doc := range input.Documents {
parts, err := convertParts(doc.Content)
if err != nil {
return nil, err
}
eres, err := em.EmbedContent(ctx, parts...)
if err != nil {
return nil, err
}
res.Embeddings = append(res.Embeddings, &ai.DocumentEmbedding{Embedding: eres.Embedding.Values})
batch.AddContent(parts...)
}
bres, err := em.BatchEmbedContents(ctx, batch)
if err != nil {
return nil, err
}
var res ai.EmbedResponse
for _, emb := range bres.Embeddings {
res.Embeddings = append(res.Embeddings, &ai.DocumentEmbedding{Embedding: emb.Values})
}
return &res, nil
})
Expand Down

0 comments on commit 3b44d23

Please sign in to comment.