From f1cd9102957d1afa48c41cc5103b4298aec4533b Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 12 Jun 2024 10:21:48 -0400 Subject: [PATCH 1/4] chore: renamed google ai and vertex ai provider ids --- go/plugins/googleai/googleai.go | 2 +- go/plugins/vertexai/vertexai.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go index dd4aa42ec3..9b8128b823 100644 --- a/go/plugins/googleai/googleai.go +++ b/go/plugins/googleai/googleai.go @@ -28,7 +28,7 @@ import ( "google.golang.org/api/option" ) -const provider = "google-genai" +const provider = "googleai" // Config provides configuration options for the Init function. type Config struct { diff --git a/go/plugins/vertexai/vertexai.go b/go/plugins/vertexai/vertexai.go index 6069a386ec..08e85ad3ac 100644 --- a/go/plugins/vertexai/vertexai.go +++ b/go/plugins/vertexai/vertexai.go @@ -27,7 +27,7 @@ import ( "google.golang.org/api/option" ) -const provider = "google-genai" +const provider = "vertexai" // Config provides configuration options for the Init function. type Config struct { From 0bbb8481ffaa6a5653921954c388192d5e1762ba Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 12 Jun 2024 21:22:37 -0400 Subject: [PATCH 2/4] feat: [GO] maintain an internal lookup of known model capabilities --- go/plugins/googleai/googleai.go | 28 +++++++++++++++++++++++----- js/flow/tests/durable_test.ts | 2 +- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go index 9b8128b823..3f2f0b5ff8 100644 --- a/go/plugins/googleai/googleai.go +++ b/go/plugins/googleai/googleai.go @@ -28,7 +28,23 @@ import ( "google.golang.org/api/option" ) -const provider = "googleai" +const ( + provider = "googleai" +) + +var ( + defaultModelCapabilities = ai.ModelCapabilities{ + Multiturn: true, + } + basicText = ai.ModelCapabilities{ + Multiturn: true, + Tools: true, + SystemRole: true, + } + knownModelsCapabilities = map[string]ai.ModelCapabilities{ + "gemini-1.0-pro": basicText, + } +) // Config provides configuration options for the Init function. type Config struct { @@ -90,11 +106,13 @@ func Init(ctx context.Context, cfg Config) (err error) { } func defineModel(name string, client *genai.Client) { + capabilities, ok := knownModelsCapabilities[name] + if !ok { + capabilities = defaultModelCapabilities + } meta := &ai.ModelMetadata{ - Label: "Google AI - " + name, - Supports: ai.ModelCapabilities{ - Multiturn: true, - }, + Label: "Google AI - " + name, + Supports: capabilities, } g := generator{model: name, client: client} ai.DefineModel(provider, name, meta, g.generate) diff --git a/js/flow/tests/durable_test.ts b/js/flow/tests/durable_test.ts index 5d0ace5c30..79e8843f2a 100644 --- a/js/flow/tests/durable_test.ts +++ b/js/flow/tests/durable_test.ts @@ -69,7 +69,7 @@ describe('durable', () => { name: 'testFlow', inputSchema: z.string(), outputSchema: z.string(), - experimentalDurable: true, + experimentalDurable: true, }, async (input) => { const response = await interrupt( From d492b916639692ce7fec070fd751d67166db10ef Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 18 Jun 2024 08:40:35 -0400 Subject: [PATCH 3/4] gemini-1.5-flash --- go/plugins/googleai/googleai.go | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go index 6482d8e957..6b231c9230 100644 --- a/go/plugins/googleai/googleai.go +++ b/go/plugins/googleai/googleai.go @@ -36,9 +36,17 @@ var ( Multiturn: true, Tools: true, SystemRole: true, + Media: false, + } + multimodal = ai.ModelCapabilities{ + Multiturn: true, + Tools: true, + SystemRole: true, + Media: true, } knownModelsCapabilities = map[string]ai.ModelCapabilities{ - "gemini-1.0-pro": basicText, + "gemini-1.0-pro": basicText, + "gemini-1.5-flash": multimodal, } ) From b6c53db566acef37a41bbcce760130830f4751cf Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 18 Jun 2024 08:53:38 -0400 Subject: [PATCH 4/4] config --- go/plugins/googleai/googleai.go | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go index 6b231c9230..229a87f22c 100644 --- a/go/plugins/googleai/googleai.go +++ b/go/plugins/googleai/googleai.go @@ -58,7 +58,7 @@ type Config struct { APIKey string // Generative models to provide. // If empty, a complete list will be obtained from the service. - Models []string + Models map[string]*ai.ModelCapabilities // Embedding models to provide. // If empty, a complete list will be obtained from the service. Embedders []string @@ -80,7 +80,10 @@ func Init(ctx context.Context, cfg Config) (err error) { return err } - needModels := len(cfg.Models) == 0 + needModels := cfg.Models == nil || len(cfg.Models) == 0 + if needModels { + cfg.Models = map[string]*ai.ModelCapabilities{} + } needEmbedders := len(cfg.Embedders) == 0 if needModels || needEmbedders { iter := client.ListModels(ctx) @@ -95,15 +98,15 @@ func Init(ctx context.Context, cfg Config) (err error) { // Model names are of the form "models/name". name := path.Base(mi.Name) if needModels && slices.Contains(mi.SupportedGenerationMethods, "generateContent") { - cfg.Models = append(cfg.Models, name) + cfg.Models[name] = nil } if needEmbedders && slices.Contains(mi.SupportedGenerationMethods, "embedContent") { cfg.Embedders = append(cfg.Embedders, name) } } } - for _, name := range cfg.Models { - defineModel(name, client) + for name, c := range cfg.Models { + defineModel(name, client, c) } for _, name := range cfg.Embedders { defineEmbedder(name, client) @@ -111,14 +114,19 @@ func Init(ctx context.Context, cfg Config) (err error) { return nil } -func defineModel(name string, client *genai.Client) { - capabilities, ok := knownModelsCapabilities[name] - if !ok { - capabilities = defaultModelCapabilities +func defineModel(name string, client *genai.Client, capabilities *ai.ModelCapabilities) { + c := defaultModelCapabilities + if capabilities == nil { + foundCapability, ok := knownModelsCapabilities[name] + if ok { + c = foundCapability + } + } else { + c = *capabilities } meta := &ai.ModelMetadata{ Label: "Google AI - " + name, - Supports: capabilities, + Supports: c, } g := generator{model: name, client: client} ai.DefineModel(provider, name, meta, g.generate)