feat(go/plugins/compat_oai): add Ollama Cloud support

go/plugins/compat_oai/ollamacloud/ollamacloud.go

@@ -0,0 +1,164 @@
+// Copyright 2025 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package ollamacloud
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/core/api"
+	"github.com/firebase/genkit/go/genkit"
+	"github.com/firebase/genkit/go/plugins/compat_oai"
+	"github.com/openai/openai-go/option"
+)
+
+const (
+	provider   = "ollamacloud"
+	apiBaseURL = "https://ollama.com"
+	apiVersion = "v1"
+)
+
+// supportedModels defines a curated set of Ollama Cloud models.
+// Model IDs are aligned with https://ollama.com/v1/models.
+var supportedModels = map[string]ai.ModelOptions{
+	// Large Language Models (text-only)
+	"gpt-oss:20b": {
+		Label:    "GPT-OSS 20B",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"gpt-oss:20b"},
+	},
+	"gpt-oss:120b": {
+		Label:    "GPT-OSS 120B",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"gpt-oss:120b"},
+	},
+	"qwen3-coder:480b": {
+		Label:    "Qwen3 Coder 480B",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"qwen3-coder:480b"},
+	},
+	"deepseek-v3.1:671b": {
+		Label:    "DeepSeek v3.1 671B",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"deepseek-v3.1:671b"},
+	},
+	"glm-4.6": {
+		Label:    "GLM-4.6",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"glm-4.6"},
+	},
+	"minimax-m2": {
+		Label:    "MiniMax M2",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"minimax-m2"},
+	},
+	"kimi-k2:1t": {
+		Label:    "Kimi K2 1T",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"kimi-k2:1t"},
+	},
+	"kimi-k2-thinking": {
+		Label:    "Kimi K2 Thinking",
+		Supports: &compat_oai.BasicText,
+		Versions: []string{"kimi-k2-thinking"},
+	},
+
+	// Multimodal Models (Vision + Text)
+	"qwen3-vl:235b-instruct": {
+		Label:    "Qwen3 VL 235B Instruct",
+		Supports: &compat_oai.Multimodal,
+		Versions: []string{"qwen3-vl:235b-instruct"},
+	},
+	"qwen3-vl:235b": {
+		Label:    "Qwen3 VL 235B",
+		Supports: &compat_oai.Multimodal,
+		Versions: []string{"qwen3-vl:235b"},
+	},
+}
+
+// OllamaCloud represents the Ollama Cloud plugin
+type OllamaCloud struct {
+	APIKey string
+	Opts   []option.RequestOption
+
+	openAICompatible *compat_oai.OpenAICompatible
+}
+
+// Name implements genkit.Plugin.
+func (o *OllamaCloud) Name() string {
+	return provider
+}
+
+// Init implements genkit.Plugin.
+func (o *OllamaCloud) Init(ctx context.Context) []api.Action {
+	apiKey := o.APIKey
+	if apiKey == "" {
+		apiKey = os.Getenv("OLLAMACLOUD_API_KEY")
+	}
+
+	if apiKey == "" {
+		panic("ollamacloud plugin initialization failed: API key is required")
+	}
+
+	if o.openAICompatible == nil {
+		o.openAICompatible = &compat_oai.OpenAICompatible{}
+	}
+
+	// Configure OpenAI-compatible client with Ollama Cloud settings
+	o.openAICompatible.Opts = []option.RequestOption{
+		option.WithAPIKey(apiKey),
+		option.WithBaseURL(fmt.Sprintf("%s/%s", apiBaseURL, apiVersion)),
+	}
+	if len(o.Opts) > 0 {
+		o.openAICompatible.Opts = append(o.openAICompatible.Opts, o.Opts...)
+	}
+
+	o.openAICompatible.Provider = provider
+	compatActions := o.openAICompatible.Init(ctx)
+
+	var actions []api.Action
+	actions = append(actions, compatActions...)
+
+	// Define available models
+	for model, opts := range supportedModels {
+		actions = append(actions, o.DefineModel(model, opts).(api.Action))
+	}
+
+	return actions
+}
+
+// Model returns the ai.Model with the given name.
+func (o *OllamaCloud) Model(g *genkit.Genkit, name string) ai.Model {
+	return o.openAICompatible.Model(g, api.NewName(provider, name))
+}
+
+// DefineModel defines a model with the given ID and options.
+func (o *OllamaCloud) DefineModel(id string, opts ai.ModelOptions) ai.Model {
+	return o.openAICompatible.DefineModel(provider, id, opts)
+}
+
+// ListActions implements genkit.Plugin.
+func (o *OllamaCloud) ListActions(ctx context.Context) []api.ActionDesc {
+	return o.openAICompatible.ListActions(ctx)
+}
+
+// ResolveAction implements genkit.Plugin.
+func (o *OllamaCloud) ResolveAction(atype api.ActionType, name string) api.Action {
+	return o.openAICompatible.ResolveAction(atype, name)
+}

go/plugins/compat_oai/ollamacloud/ollamacloud_live_test.go

@@ -0,0 +1,211 @@
+// Copyright 2025 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package ollamacloud
+
+import (
+	"context"
+	"math"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/genkit"
+	"github.com/openai/openai-go/option"
+)
+
+func TestPlugin(t *testing.T) {
+	apiKey := os.Getenv("OLLAMACLOUD_API_KEY")
+	if apiKey == "" {
+		t.Skip("Skipping test: OLLAMACLOUD_API_KEY environment variable not set")
+	}
+
+	ctx := context.Background()
+	ollamaCloud := &OllamaCloud{
+		APIKey: apiKey,
+		Opts: []option.RequestOption{
+			option.WithAPIKey(apiKey),
+		},
+	}
+
+	g := genkit.Init(ctx,
+		genkit.WithDefaultModel("ollamacloud/gpt-oss:20b"),
+		genkit.WithPlugins(ollamaCloud))
+
+	gablorkenTool := genkit.DefineTool(g, "gablorken", "use when need to calculate a gablorken",
+		func(ctx *ai.ToolContext, input struct {
+			Value float64
+			Over  float64
+		}) (float64, error) {
+			return math.Pow(input.Value, input.Over), nil
+		})
+
+	t.Log("ollamacloud plugin initialized")
+
+	t.Run("basic completion", func(t *testing.T) {
+		t.Log("generating basic completion response")
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithPrompt("What is the capital of France?"))
+		if err != nil {
+			t.Fatal("error generating basic completion response: ", err)
+		}
+		t.Logf("basic completion response: %+v", resp)
+		out := resp.Message.Content[0].Text
+		if !strings.Contains(strings.ToLower(out), "paris") {
+			t.Errorf("got %q, expecting it to contain 'Paris'", out)
+		}
+		// Verify usage statistics are present
+		if resp.Usage == nil || resp.Usage.TotalTokens == 0 {
+			t.Error("Expected non-zero usage statistics")
+		}
+	})
+
+	t.Run("streaming", func(t *testing.T) {
+		var streamedOutput string
+		chunks := 0
+		final, err := genkit.Generate(ctx, g,
+			ai.WithPrompt("Write a short paragraph about artificial intelligence."),
+			ai.WithStreaming(func(ctx context.Context, chunk *ai.ModelResponseChunk) error {
+				chunks++
+				for _, content := range chunk.Content {
+					streamedOutput += content.Text
+				}
+				return nil
+			}))
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Verify streaming worked
+		if chunks <= 1 {
+			t.Error("Expected multiple chunks for streaming")
+		}
+		// Verify the final output matches streamed content
+		finalOutput := ""
+		for _, content := range final.Message.Content {
+			finalOutput += content.Text
+		}
+		if streamedOutput != finalOutput {
+			t.Errorf("Streaming output doesn't match final output\nStreamed: %s\nFinal: %s",
+				streamedOutput, finalOutput)
+		}
+		t.Logf("streaming response: %+v", finalOutput)
+	})
+
+	t.Run("media part", func(t *testing.T) {
+		image := "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAHIAAABUAQMAAABk5vEVAAAABlBMVEX///8AAABVwtN+" +
+			"AAAAI0lEQVR4nGNgGHaA/z8UHIDwOWASDqP8Uf7w56On/1FAQwAAVM0exw1hqwkAAAAASUVORK5CYII="
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithModelName("ollamacloud/qwen3-vl:235b-instruct"),
+			ai.WithMessages(
+				ai.NewUserMessage(
+					ai.NewMediaPart("image/png", image),
+					ai.NewTextPart("Is there a rectangle in the picture? Yes or not."),
+				),
+			),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+		text := resp.Message.Content[0].Text
+		if !strings.Contains(strings.ToLower(text), "yes") {
+			t.Errorf("got %q, expecting it to contain 'yes'", text)
+		}
+	})
+
+	t.Run("system message", func(t *testing.T) {
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithPrompt("What are you?"),
+			ai.WithSystem("You are a helpful math tutor who loves numbers."))
+		if err != nil {
+			t.Fatal(err)
+		}
+		out := resp.Message.Content[0].Text
+		if !strings.Contains(strings.ToLower(out), "math") {
+			t.Errorf("got %q, expecting response to mention being a math tutor", out)
+		}
+		t.Logf("system message response: %+v", out)
+	})
+
+	t.Run("tool usage with basic completion", func(t *testing.T) {
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithModelName("ollamacloud/qwen3-coder:480b"),
+			ai.WithPrompt("Use the gablorken tool to calculate the gablorken of 2 over 3. Set Value=2 and Over=3 as numbers (not strings) and answer with the numeric result."),
+			ai.WithTools(gablorkenTool))
+		if err != nil {
+			t.Fatal(err)
+		}
+		out := resp.Message.Content[0].Text
+		const want = "8"
+		if !strings.Contains(out, want) {
+			t.Errorf("got %q, expecting it to contain %q", out, want)
+		}
+		t.Logf("tool usage with basic completion response: %+v", out)
+	})
+
+	t.Run("tool usage with streaming", func(t *testing.T) {
+		var streamedOutput string
+		chunks := 0
+		final, err := genkit.Generate(ctx, g,
+			ai.WithModelName("ollamacloud/qwen3-coder:480b"),
+			ai.WithPrompt("Use the gablorken tool to calculate the gablorken of 2 over 3. Set Value=2 and Over=3 as numbers (not strings) and answer with the numeric result."),
+			ai.WithTools(gablorkenTool),
+			ai.WithStreaming(func(ctx context.Context, chunk *ai.ModelResponseChunk) error {
+				chunks++
+				for _, content := range chunk.Content {
+					streamedOutput += content.Text
+				}
+				return nil
+			}))
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Verify streaming worked
+		if chunks <= 1 {
+			t.Error("Expected multiple chunks for streaming")
+		}
+		// Verify the final output matches streamed content
+		finalOutput := ""
+		for _, content := range final.Message.Content {
+			finalOutput += content.Text
+		}
+		if streamedOutput != finalOutput {
+			t.Errorf("Streaming output doesn't match final output\nStreamed: %s\nFinal: %s",
+				streamedOutput, finalOutput)
+		}
+		const want = "8"
+		if !strings.Contains(finalOutput, want) {
+			t.Errorf("got %q, expecting it to contain %q", finalOutput, want)
+		}
+		t.Logf("tool usage with streaming response: %+v", finalOutput)
+	})
+
+	t.Run("invalid config type", func(t *testing.T) {
+		// Try to use a string as config instead of *openai.ChatCompletionNewParams
+		config := "not a config"
+		_, err := genkit.Generate(ctx, g,
+			ai.WithPrompt("Write a short sentence about artificial intelligence."),
+			ai.WithConfig(config),
+		)
+		if err == nil {
+			t.Fatal("expected error for invalid config type")
+		}
+		if !strings.Contains(err.Error(), "unexpected config type: string") {
+			t.Errorf("got error %q, want error containing 'unexpected config type: string'", err.Error())
+		}
+		t.Logf("invalid config type error: %v", err)
+	})
+}