From f2d06d08b2a5f5c34438c732b7e763c256ed7d73 Mon Sep 17 00:00:00 2001 From: Ian Lance Taylor Date: Wed, 29 May 2024 10:34:07 -0700 Subject: [PATCH 1/2] [Go] change embedder plugins to use Init and New Change the pinecone and localvec plugins to use a New function that returns an ai.Retriever, and an Init function that registers one with genkit. This makes the pinecone and localvec plugins consistent with each other, and with other plugins such as googleai. --- go/plugins/localvec/localvec.go | 11 ++++++++++- go/plugins/pinecone/genkit.go | 26 ++++++++++++++------------ go/plugins/pinecone/genkit_test.go | 2 +- go/plugins/pinecone/pinecone.go | 2 +- 4 files changed, 26 insertions(+), 15 deletions(-) diff --git a/go/plugins/localvec/localvec.go b/go/plugins/localvec/localvec.go index cc2374bb05..ecc256f799 100644 --- a/go/plugins/localvec/localvec.go +++ b/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 } diff --git a/go/plugins/pinecone/genkit.go b/go/plugins/pinecone/genkit.go index c709d83251..c84b3b5d65 100644 --- a/go/plugins/pinecone/genkit.go +++ b/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 diff --git a/go/plugins/pinecone/genkit_test.go b/go/plugins/pinecone/genkit_test.go index 4c72fa437c..73f01eb854 100644 --- a/go/plugins/pinecone/genkit_test.go +++ b/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) } diff --git a/go/plugins/pinecone/pinecone.go b/go/plugins/pinecone/pinecone.go index 999990febf..99671aa41d 100644 --- a/go/plugins/pinecone/pinecone.go +++ b/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) From 7757ff5c1e37f110cc017ed11525521bdacc2104 Mon Sep 17 00:00:00 2001 From: Ian Lance Taylor Date: Wed, 29 May 2024 10:39:03 -0700 Subject: [PATCH 2/2] Fix build. --- go/samples/menu/main.go | 93 ++++++++++++++++++++++++++ go/samples/menu/s01.go | 86 ++++++++++++++++++++++++ go/samples/menu/s02.go | 103 +++++++++++++++++++++++++++++ go/samples/menu/testdata/menu.json | 97 +++++++++++++++++++++++++++ go/samples/rag/main.go | 2 +- 5 files changed, 380 insertions(+), 1 deletion(-) create mode 100644 go/samples/menu/main.go create mode 100644 go/samples/menu/s01.go create mode 100644 go/samples/menu/s02.go create mode 100644 go/samples/menu/testdata/menu.json diff --git a/go/samples/menu/main.go b/go/samples/menu/main.go new file mode 100644 index 0000000000..49c18a6c4f --- /dev/null +++ b/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) + } +} diff --git a/go/samples/menu/s01.go b/go/samples/menu/s01.go new file mode 100644 index 0000000000..4a2634bc14 --- /dev/null +++ b/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) + } +} diff --git a/go/samples/menu/s02.go b/go/samples/menu/s02.go new file mode 100644 index 0000000000..c677bac884 --- /dev/null +++ b/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 + }, +) diff --git a/go/samples/menu/testdata/menu.json b/go/samples/menu/testdata/menu.json new file mode 100644 index 0000000000..d46739205d --- /dev/null +++ b/go/samples/menu/testdata/menu.json @@ -0,0 +1,97 @@ +[ + { + "title": "Mozzarella Sticks", + "price": 8, + "description": "Crispy fried mozzarella sticks served with marinara sauce." + }, + { + "title": "Chicken Wings", + "price": 10, + "description": "Crispy fried chicken wings tossed in your choice of sauce." + }, + { + "title": "Nachos", + "price": 12, + "description": "Crispy tortilla chips topped with melted cheese, chili, sour cream, and salsa." + }, + { + "title": "Onion Rings", + "price": 7, + "description": "Crispy fried onion rings served with ranch dressing." + }, + { + "title": "French Fries", + "price": 5, + "description": "Crispy fried french fries." + }, + { + "title": "Mashed Potatoes", + "price": 6, + "description": "Creamy mashed potatoes." + }, + { + "title": "Coleslaw", + "price": 4, + "description": "Homemade coleslaw." + }, + { + "title": "Classic Cheeseburger", + "price": 12, + "description": "A juicy beef patty topped with melted American cheese, lettuce, tomato, and onion on a toasted bun." + }, + { + "title": "Bacon Cheeseburger", + "price": 14, + "description": "A classic cheeseburger with the addition of crispy bacon." + }, + { + "title": "Mushroom Swiss Burger", + "price": 15, + "description": "A beef patty topped with sautéed mushrooms, melted Swiss cheese, and a creamy horseradish sauce." + }, + { + "title": "Chicken Sandwich", + "price": 13, + "description": "A crispy chicken breast on a toasted bun with lettuce, tomato, and your choice of sauce." + }, + { + "title": "Pulled Pork Sandwich", + "price": 14, + "description": "Slow-cooked pulled pork on a toasted bun with coleslaw and barbecue sauce." + }, + { + "title": "Reuben Sandwich", + "price": 15, + "description": "Thinly sliced corned beef, Swiss cheese, sauerkraut, and Thousand Island dressing on rye bread." + }, + { + "title": "House Salad", + "price": 8, + "description": "Mixed greens with your choice of dressing." + }, + { + "title": "Caesar Salad", + "price": 9, + "description": "Romaine lettuce with croutons, Parmesan cheese, and Caesar dressing." + }, + { + "title": "Greek Salad", + "price": 10, + "description": "Mixed greens with feta cheese, olives, tomatoes, cucumbers, and red onions." + }, + { + "title": "Chocolate Lava Cake", + "price": 8, + "description": "A warm, gooey chocolate cake with a molten chocolate center." + }, + { + "title": "Apple Pie", + "price": 7, + "description": "A classic apple pie with a flaky crust and warm apple filling." + }, + { + "title": "Cheesecake", + "price": 8, + "description": "A creamy cheesecake with a graham cracker crust." + } +] diff --git a/go/samples/rag/main.go b/go/samples/rag/main.go index d049b48d9d..ef9331b6a8 100644 --- a/go/samples/rag/main.go +++ b/go/samples/rag/main.go @@ -125,7 +125,7 @@ func main() { var sb strings.Builder for _, d := range response.Documents { - sb.WriteString(d.Content[0].Text()) + sb.WriteString(d.Content[0].Text) sb.WriteByte('\n') }