From 6bd00aaed4c5e988ba4a86d9ead8946fbf362197 Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Sat, 8 Jun 2024 07:00:01 -0400
Subject: [PATCH] [Go] proposal: vector DB Inits return actions

Vector DBs only register two actions: an indexer and a retriever.
Continue to allow users to look up these actions by name.
But also return them from Init.
---
 go/plugins/localvec/localvec.go | 14 ++++++++------
 go/samples/menu/main.go         | 10 +++++-----
 go/samples/rag/main.go          | 12 ++++++------
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go
index fcb06726d..27e00b411 100644
--- a/go/plugins/localvec/localvec.go
+++ b/go/plugins/localvec/localvec.go
@@ -46,9 +46,9 @@ type Config struct {
 }
 
 // Init initializes a new local vector database. This will register a new
-// indexer and retriever with genkit.
+// indexer and retriever with genkit, and return them.
 // This retriever may only be used by a single goroutine at a time.
-func Init(ctx context.Context, cfg Config) (err error) {
+func Init(ctx context.Context, cfg Config) (_ *ai.IndexerAction, _ *ai.RetrieverAction, err error) {
 	defer func() {
 		if err != nil {
 			err = fmt.Errorf("localvec.Init: %w", err)
@@ -56,19 +56,21 @@ func Init(ctx context.Context, cfg Config) (err error) {
 	}()
 	ds, err := newDocStore(cfg.Dir, cfg.Name, cfg.Embedder, cfg.EmbedderOptions)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
-	ai.DefineIndexer(provider, cfg.Name, ds.index)
-	ai.DefineRetriever(provider, cfg.Name, ds.retrieve)
-	return nil
+	ia := ai.DefineIndexer(provider, cfg.Name, ds.index)
+	ra := ai.DefineRetriever(provider, cfg.Name, ds.retrieve)
+	return ia, ra, nil
 }
 
 // Indexer returns the indexer with the given name.
+// The name must match the [Config.Name] value passed to [Init].
 func Indexer(name string) *ai.IndexerAction {
 	return ai.LookupIndexer(provider, name)
 }
 
 // Retriever returns the retriever with the given name.
+// The name must match the [Config.Name] value passed to [Init].
 func Retriever(name string) *ai.RetrieverAction {
 	return ai.LookupRetriever(provider, name)
 }
diff --git a/go/samples/menu/main.go b/go/samples/menu/main.go
index c4ce70e3c..9de9e166c 100644
--- a/go/samples/menu/main.go
+++ b/go/samples/menu/main.go
@@ -100,14 +100,14 @@ func main() {
 		log.Fatal(err)
 	}
 
-	const localvecName = "go-menu-items"
-	if err := localvec.Init(ctx, localvec.Config{
-		Name:     localvecName,
+	indexer, retriever, err := localvec.Init(ctx, localvec.Config{
+		Name:     "go-menu-items",
 		Embedder: vertexai.Embedder(embeddingGecko),
-	}); err != nil {
+	})
+	if err != nil {
 		log.Fatal(err)
 	}
-	if err := setup04(ctx, localvec.Indexer(localvecName), localvec.Retriever(localvecName), model); err != nil {
+	if err := setup04(ctx, indexer, retriever, model); err != nil {
 		log.Fatal(err)
 	}
 
diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go
index e120c34cf..46d9b55d5 100644
--- a/go/samples/rag/main.go
+++ b/go/samples/rag/main.go
@@ -91,11 +91,11 @@ func main() {
 		log.Fatal(err)
 	}
 
-	const localvecName = "simpleQa"
-	if err := localvec.Init(context.Background(), localvec.Config{
-		Name:     localvecName,
+	indexer, retriever, err := localvec.Init(context.Background(), localvec.Config{
+		Name:     "simpleQa",
 		Embedder: googleai.Embedder("embedding-001"),
-	}); err != nil {
+	})
+	if err != nil {
 		log.Fatal(err)
 	}
 
@@ -107,7 +107,7 @@ func main() {
 		indexerReq := &ai.IndexerRequest{
 			Documents: []*ai.Document{d1, d2, d3},
 		}
-		err := ai.Index(ctx, localvec.Indexer(localvecName), indexerReq)
+		err := ai.Index(ctx, indexer, indexerReq)
 		if err != nil {
 			return "", err
 		}
@@ -116,7 +116,7 @@ func main() {
 		retrieverReq := &ai.RetrieverRequest{
 			Document: dRequest,
 		}
-		response, err := ai.Retrieve(ctx, localvec.Retriever(localvecName), retrieverReq)
+		response, err := ai.Retrieve(ctx, retriever, retrieverReq)
 		if err != nil {
 			return "", err
 		}