[Go] factor out common code for Gemini models (#614)

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"}

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,
+ }
+)

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{