diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go index 3dbcc59b7..93ba44977 100644 --- a/go/plugins/googleai/googleai.go +++ b/go/plugins/googleai/googleai.go @@ -26,6 +26,7 @@ import ( "github.com/firebase/genkit/go/ai" "github.com/firebase/genkit/go/internal" + "github.com/firebase/genkit/go/plugins/internal/gemini" "github.com/firebase/genkit/go/plugins/internal/uri" "github.com/google/generative-ai-go/genai" "google.golang.org/api/iterator" @@ -41,24 +42,10 @@ var state struct { } var ( - basicText = ai.ModelCapabilities{ - Multiturn: true, - Tools: true, - SystemRole: true, - Media: false, - } - - multimodal = ai.ModelCapabilities{ - Multiturn: true, - Tools: true, - SystemRole: true, - Media: true, - } - knownCaps = map[string]ai.ModelCapabilities{ - "gemini-1.0-pro": basicText, - "gemini-1.5-pro": multimodal, - "gemini-1.5-flash": multimodal, + "gemini-1.0-pro": gemini.BasicText, + "gemini-1.5-pro": gemini.Multimodal, + "gemini-1.5-flash": gemini.Multimodal, } knownEmbedders = []string{"text-embedding-004", "embedding-001"} diff --git a/go/plugins/internal/gemini/gemini.go b/go/plugins/internal/gemini/gemini.go new file mode 100644 index 000000000..9ede1b8f0 --- /dev/null +++ b/go/plugins/internal/gemini/gemini.go @@ -0,0 +1,37 @@ +// 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 gemini contains code that is common to both the googleai and vertexai plugins. +// Most most cannot be shared in this way because the import paths are different. +package gemini + +import "github.com/firebase/genkit/go/ai" + +var ( + // BasicText describes model capabilities for text-only Gemini models. + BasicText = ai.ModelCapabilities{ + Multiturn: true, + Tools: true, + SystemRole: true, + Media: false, + } + + // Multimodal describes model capabilities for multimodal Gemini models. + Multimodal = ai.ModelCapabilities{ + Multiturn: true, + Tools: true, + SystemRole: true, + Media: true, + } +) diff --git a/go/plugins/vertexai/vertexai.go b/go/plugins/vertexai/vertexai.go index 93a2988d5..00072aae9 100644 --- a/go/plugins/vertexai/vertexai.go +++ b/go/plugins/vertexai/vertexai.go @@ -25,6 +25,7 @@ import ( "cloud.google.com/go/vertexai/genai" "github.com/firebase/genkit/go/ai" "github.com/firebase/genkit/go/internal" + "github.com/firebase/genkit/go/plugins/internal/gemini" "github.com/firebase/genkit/go/plugins/internal/uri" "google.golang.org/api/iterator" "google.golang.org/api/option" @@ -33,24 +34,10 @@ import ( const provider = "vertexai" var ( - basicText = ai.ModelCapabilities{ - Multiturn: true, - Tools: true, - SystemRole: true, - Media: false, - } - - multimodal = ai.ModelCapabilities{ - Multiturn: true, - Tools: true, - SystemRole: true, - Media: true, - } - knownCaps = map[string]ai.ModelCapabilities{ - "gemini-1.0-pro": basicText, - "gemini-1.5-pro": multimodal, - "gemini-1.5-flash": multimodal, + "gemini-1.0-pro": gemini.BasicText, + "gemini-1.5-pro": gemini.Multimodal, + "gemini-1.5-flash": gemini.Multimodal, } knownEmbedders = []string{