[Go] change embedder plugins to use Init and New

go/plugins/localvec/localvec.go

@@ -34,6 +34,16 @@ import (
 	"github.com/firebase/genkit/go/core/logger"
 )
 
+// Init registers all the actions in this package with [ai]'s Register calls.
+func Init(ctx context.Context, dir, name string, embedder ai.Embedder, embedderOptions any) error {
+	r, err := New(ctx, dir, name, embedder, embedderOptions)
+	if err != nil {
+		return err
+	}
+	ai.RegisterRetriever("devLocalVectorStore/"+name, r)
+	return nil
+}
+
 // New returns a new local vector database. This will register a new
 // retriever with genkit, and also return it.
 // This retriever may only be used by a single goroutine at a time.
@@ -43,7 +53,6 @@ func New(ctx context.Context, dir, name string, embedder ai.Embedder, embedderOp
 	if err != nil {
 		return nil, err
 	}
-	ai.RegisterRetriever("devLocalVectorStore/"+name, r)
 	return r, nil
 }
 

go/plugins/pinecone/genkit.go

@@ -38,8 +38,20 @@ const defaultTextKey = "_content"
 
 // Init registers all the actions in this package with [ai]'s Register calls.
 //
+// The arguments are as for [New].
+func Init(ctx context.Context, apiKey, host string, embedder ai.Embedder, embedderOptions any, textKey string) error {
+	r, err := New(ctx, apiKey, host, embedder, embedderOptions, textKey)
+	if err != nil {
+		return err
+	}
+	ai.RegisterRetriever("pinecone", r)
+	return nil
+}
+
+// New returns an [ai.Retriever] that uses Pinecone.
+//
 // apiKey is the API key to use to access Pinecone.
-// If it is the empty string, it is read from the PINECONE_API_INDEX
+// If it is the empty string, it is read from the PINECONE_API_KEY
 // environment variable.
 //
 // host is the controller host name. This implies which index to use.
@@ -50,17 +62,7 @@ const defaultTextKey = "_content"
 //
 // The textKey parameter is the metadata key to use to store document text
 // in Pinecone; the default is "_content".
-func Init(ctx context.Context, apiKey, host string, embedder ai.Embedder, embedderOptions any, textKey string) error {
-	r, err := newRetriever(ctx, apiKey, host, embedder, embedderOptions, textKey)
-	if err != nil {
-		return err
-	}
-	ai.RegisterRetriever("pinecone", r)
-	return nil
-}
-
-// newRetriever returns a new ai.Retriever to register.
-func newRetriever(ctx context.Context, apiKey, host string, embedder ai.Embedder, embedderOptions any, textKey string) (ai.Retriever, error) {
+func New(ctx context.Context, apiKey, host string, embedder ai.Embedder, embedderOptions any, textKey string) (ai.Retriever, error) {
 	client, err := NewClient(ctx, apiKey)
 	if err != nil {
 		return nil, err

go/plugins/pinecone/genkit_test.go

@@ -72,7 +72,7 @@ func TestGenkit(t *testing.T) {
 	embedder.Register(d2, v2)
 	embedder.Register(d3, v3)
 
-	r, err := newRetriever(ctx, *testAPIKey, indexData.Host, embedder, nil, "")
+	r, err := New(ctx, *testAPIKey, indexData.Host, embedder, nil, "")
 	if err != nil {
 		t.Fatal(err)
 	}

go/plugins/pinecone/pinecone.go

@@ -64,7 +64,7 @@ type Client struct {
 // NewClient builds a Client.
 //
 // apiKey is the API key to use to access Pinecone.
-// If it is the empty string, it is read from the PINECONE_API_INDEX
+// If it is the empty string, it is read from the PINECONE_API_KEY
 // environment variable.
 func NewClient(ctx context.Context, apiKey string) (*Client, error) {
 	key, err := resolveAPIKey(apiKey)

go/samples/menu/main.go

@@ -0,0 +1,93 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/firebase/genkit/go/genkit"
+	"github.com/firebase/genkit/go/plugins/vertexai"
+	"github.com/invopop/jsonschema"
+)
+
+// menuItem is the data model for an item on the menu.
+type menuItem struct {
+	Title       string  `json:"title" jsonschema_description:"The name of the menu item"`
+	Description string  `json:"description" jsonschema_description:"Details including ingredients and preparation"`
+	Price       float64 `json:"price" jsonschema_description:"Price in dollars"`
+}
+
+// menuItemSchema is the JSON schema for a menuItem.
+var menuItemSchema = jsonschema.Reflect(menuItem{})
+
+// menuQuestionInput is a question about the menu.
+type menuQuestionInput struct {
+	Question string `json:"question"`
+}
+
+// menuQuestionInputSchema is the JSON schema for a menuQuestionInput.
+var menuQuestionInputSchema = jsonschema.Reflect(menuQuestionInput{})
+
+// answerOutput is an answer to a question.
+type answerOutput struct {
+	Answer string `json:"answer"`
+}
+
+// answerOutputSchema is the JSON schema for an answerOutput.
+var answerOutputSchema = jsonschema.Reflect(answerOutput{})
+
+// dataMenuQuestionInput is a question about the menu,
+// where the menu is provided in the JSON data.
+type dataMenuQuestionInput struct {
+	MenuData []menuItem `json:"menuData"`
+	Question string     `json:"question"`
+}
+
+// dataMenuQuestionInputSchema is the JSON schema for a dataMenuQuestionInput.
+var dataMenuQuestionInputSchema = jsonschema.Reflect(dataMenuQuestionInput{})
+
+// textMenuQuestionInput is for a question about the menu,
+// where the menu is provided as unstructured text.
+type textMenuQuestionInput struct {
+	MenuText string `json:"menuText"`
+	Question string `json:"question"`
+}
+
+// textMenuQuestionInputSchema is the JSON schema for a textMenuQuestionInput.
+var textMenuQuestionInputSchema = jsonschema.Reflect(textMenuQuestionInput{})
+
+func main() {
+	projectID := os.Getenv("GCLOUD_PROJECT")
+	if projectID == "" {
+		fmt.Fprintln(os.Stderr, "menu example requires setting GCLOUD_PROJECT in the environment.")
+		os.Exit(1)
+	}
+
+	location := "us-central1"
+	if envLocation := os.Getenv("GCLOUD_LOCATION"); envLocation != "" {
+		location = envLocation
+	}
+
+	if err := vertexai.Init(context.Background(), "gemini-1.0-pro", projectID, location); err != nil {
+		log.Fatal(err)
+	}
+
+	if err := genkit.StartFlowServer(""); err != nil {
+		log.Fatal(err)
+	}
+}

go/samples/menu/s01.go

@@ -0,0 +1,86 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"log"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/plugins/dotprompt"
+)
+
+var s01VanillaPrompt *dotprompt.Prompt
+
+func init() {
+	var err error
+	s01VanillaPrompt, err = dotprompt.Define("s01_vanillaPrompt",
+		`You are acting as a helpful AI assistant named "Walt" that can answer 
+		 questions about the food available on the menu at Walt's Burgers.
+		 Customer says: ${input.question}`,
+		&dotprompt.Config{
+			Model: "google-vertexai/gemini-1.0-pro",
+			InputSchema: menuQuestionInputSchema,
+		},
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+var s01StaticMenuDotPrompt *dotprompt.Prompt
+
+func init() {
+	var err error
+	s01StaticMenuDotPrompt, err = dotprompt.Define("s01_staticMenuDotPrompt",
+		`You are acting as a helpful AI assistant named "Walt" that can answer 
+		 questions about the food available on the menu at Walt's Burgers. 
+		 Here is today's menu:
+
+		 - The Regular Burger $12
+		   The classic charbroiled to perfection with your choice of cheese
+
+		 - The Fancy Burger $13
+		   Classic burger topped with bacon & Blue Cheese
+
+		 - The Bacon Burger $13
+		   Bacon cheeseburger with your choice of cheese.
+
+		 - Everything Burger $14
+		   Heinz 57 sauce, American cheese, bacon, fried egg & crispy onion bits
+
+		 - Chicken Breast Sandwich $12
+		   Tender juicy chicken breast on a brioche roll.
+		   Grilled, blackened, or fried
+
+		 Our fresh 1/2 lb. beef patties are made using choice cut
+		 brisket, short rib & sirloin. Served on a toasted
+		 brioche roll with chips. Served with lettuce, tomato & pickles. 
+		 Onions upon request. Substitute veggie patty $2
+
+		 Answer this customer's question, in a concise and helpful manner,
+		 as long as it is about food.
+
+		 Question:
+		 {{question}} ?`,
+		&dotprompt.Config{
+			Model: "google-vertexai/gemini-1.0-pro",
+			InputSchema:  menuQuestionInputSchema,
+			OutputFormat: ai.OutputFormatText,
+		},
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+}

go/samples/menu/s02.go

@@ -0,0 +1,103 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/genkit"
+	"github.com/firebase/genkit/go/plugins/dotprompt"
+)
+
+var s02DataMenuPrompt *dotprompt.Prompt
+
+func init() {
+	var err error
+	s02DataMenuPrompt, err = dotprompt.Define("s02_dataMenu",
+		`You are acting as a helpful AI assistant named Walt that can answer 
+		 questions about the food available on the menu at Walt's Burgers. 
+
+		 Answer this customer's question, in a concise and helpful manner,
+		 as long as it is about food on the menu or something harmless like sports.
+		 Use the tools available to answer menu questions.
+		 DO NOT INVENT ITEMS NOT ON THE MENU.
+
+		 Question:
+		 {{question}} ?`,
+		&dotprompt.Config{
+			Model: "google-vertexai/gemini-1.0-pro",
+			InputSchema: menuQuestionInputSchema,
+			OutputFormat: ai.OutputFormatText,
+			Tools: []*ai.ToolDefinition{
+				menuToolDef,
+			},
+		},
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+var menuToolDef = &ai.ToolDefinition{
+	Name:        "todaysMenu",
+	OutputSchema: map[string]any{
+		"menuData": []menuItem{},
+	},
+}
+
+func menu(ctx context.Context, input map[string]any) (map[string]any, error) {
+	fmt.Println("********** called menu **********")
+
+	f, err := os.Open("testdata/menu.json")
+	if err != nil {
+		return nil, err
+	}
+	decoder := json.NewDecoder(f)
+	var m map[string]any
+	if err := decoder.Decode(&m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+func init() {
+	ai.RegisterTool("menu", menuToolDef, nil, menu)
+}
+
+var s02MenuQuestionFlow = genkit.DefineFlow("s02_menuQuestion",
+	func(ctx context.Context, input *menuQuestionInput, _ genkit.NoStream) (*answerOutput, error) {
+		vars, err := s02DataMenuPrompt.BuildVariables(input)
+		if err != nil {
+			return nil, err
+		}
+		ai := &dotprompt.ActionInput{Variables: vars}
+
+		resp, err := s02DataMenuPrompt.Execute(ctx, ai)
+		if err != nil {
+			return nil, err
+		}
+
+		text, err := resp.Text()
+		if err != nil {
+			return nil, fmt.Errorf("s02MenuQuestionFlow: %v", err)
+		}
+		return &answerOutput{Answer: text}, nil
+	},
+)