feat(js/plugins/ollama): migrate ollama plugin to v2 plugin API

js/plugins/ollama/src/constants.ts

@@ -0,0 +1,34 @@
+/**
+ * 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.
+ */
+
+import { ModelInfo } from 'genkit/model';
+
+export const ANY_JSON_SCHEMA: Record<string, any> = {
+  $schema: 'http://json-schema.org/draft-07/schema#',
+};
+
+export const GENERIC_MODEL_INFO = {
+  supports: {
+    multiturn: true,
+    media: true,
+    tools: true,
+    toolChoice: true,
+    systemRole: true,
+    constrained: 'all',
+  },
+} as ModelInfo;
+
+export const DEFAULT_OLLAMA_SERVER_ADDRESS = 'http://localhost:11434';

js/plugins/ollama/src/embeddings.ts

@@ -13,11 +13,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import type { Document, EmbedderAction, Genkit } from 'genkit';
+import type { Document, EmbedderAction } from 'genkit';
+import { embedder } from 'genkit/plugin';
 import type { EmbedRequest, EmbedResponse } from 'ollama';
+import { DEFAULT_OLLAMA_SERVER_ADDRESS } from './constants.js';
 import type { DefineOllamaEmbeddingParams, RequestHeaders } from './types.js';
+import { OllamaEmbedderConfigSchema } from './types.js';
 
-async function toOllamaEmbedRequest(
+export async function toOllamaEmbedRequest(
   modelName: string,
   dimensions: number,
   documents: Document[],
@@ -59,13 +62,18 @@ async function toOllamaEmbedRequest(
   };
 }
 
-export function defineOllamaEmbedder(
-  ai: Genkit,
-  { name, modelName, dimensions, options }: DefineOllamaEmbeddingParams
-): EmbedderAction<any> {
-  return ai.defineEmbedder(
+export function defineOllamaEmbedder({
+  name,
+  modelName,
+  dimensions,
+  options,
+}: DefineOllamaEmbeddingParams): EmbedderAction<
+  typeof OllamaEmbedderConfigSchema
+> {
+  return embedder(
     {
       name: `ollama/${name}`,
+      configSchema: OllamaEmbedderConfigSchema,
       info: {
         label: 'Ollama Embedding - ' + name,
         dimensions,
@@ -75,9 +83,11 @@ export function defineOllamaEmbedder(
         },
       },
     },
-    async (input, config) => {
-      const serverAddress = config?.serverAddress || options.serverAddress;
-
+    async ({ input, options: requestOptions }, config) => {
+      const serverAddress =
+        requestOptions?.serverAddress ||
+        options.serverAddress ||
+        DEFAULT_OLLAMA_SERVER_ADDRESS;
       const { url, requestPayload, headers } = await toOllamaEmbedRequest(
         modelName,
         dimensions,

js/plugins/ollama/src/index.ts

@@ -20,7 +20,6 @@ import {
   z,
   type ActionMetadata,
   type EmbedderReference,
-  type Genkit,
   type ModelReference,
   type ToolRequest,
   type ToolRequestPart,
@@ -35,11 +34,19 @@ import {
   type GenerateRequest,
   type GenerateResponseData,
   type MessageData,
-  type ModelInfo,
   type ToolDefinition,
 } from 'genkit/model';
-import { genkitPlugin, type GenkitPlugin } from 'genkit/plugin';
-import type { ActionType } from 'genkit/registry';
+import {
+  genkitPluginV2,
+  model,
+  type GenkitPluginV2,
+  type ResolvableAction,
+} from 'genkit/plugin';
+import {
+  ANY_JSON_SCHEMA,
+  DEFAULT_OLLAMA_SERVER_ADDRESS,
+  GENERIC_MODEL_INFO,
+} from './constants.js';
 import { defineOllamaEmbedder } from './embeddings.js';
 import type {
   ApiType,
@@ -51,12 +58,13 @@ import type {
   OllamaTool,
   OllamaToolCall,
   RequestHeaders,
+  ResolveActionOptions,
 } from './types.js';
 
 export type { OllamaPluginParams };
 
 export type OllamaPlugin = {
-  (params?: OllamaPluginParams): GenkitPlugin;
+  (params?: OllamaPluginParams): GenkitPluginV2;
 
   model(
     name: string,
@@ -65,59 +73,50 @@ export type OllamaPlugin = {
   embedder(name: string, config?: Record<string, any>): EmbedderReference;
 };
 
-const ANY_JSON_SCHEMA: Record<string, any> = {
-  $schema: 'http://json-schema.org/draft-07/schema#',
-};
+function initializer(serverAddress: string, params: OllamaPluginParams = {}) {
+  const actions: ResolvableAction[] = [];
 
-const GENERIC_MODEL_INFO = {
-  supports: {
-    multiturn: true,
-    media: true,
-    tools: true,
-    toolChoice: true,
-    systemRole: true,
-    constrained: 'all',
-  },
-} as ModelInfo;
-
-const DEFAULT_OLLAMA_SERVER_ADDRESS = 'http://localhost:11434';
-
-async function initializer(
-  ai: Genkit,
-  serverAddress: string,
-  params?: OllamaPluginParams
-) {
-  params?.models?.map((model) =>
-    defineOllamaModel(ai, model, serverAddress, params?.requestHeaders)
-  );
-  params?.embedders?.map((model) =>
-    defineOllamaEmbedder(ai, {
-      name: model.name,
-      modelName: model.name,
-      dimensions: model.dimensions,
-      options: params!,
-    })
-  );
+  if (params?.models) {
+    for (const model of params.models) {
+      actions.push(
+        defineOllamaModel(model, serverAddress, params.requestHeaders)
+      );
+    }
+  }
+
+  if (params?.embedders && params.serverAddress) {
+    for (const embedder of params.embedders) {
+      actions.push(
+        defineOllamaEmbedder({
+          name: embedder.name,
+          modelName: embedder.name,
+          dimensions: embedder.dimensions,
+          options: params,
+        })
+      );
+    }
+  }
+
+  return actions;
 }
 
-function resolveAction(
-  ai: Genkit,
-  actionType: ActionType,
-  actionName: string,
-  serverAddress: string,
-  requestHeaders?: RequestHeaders
-) {
-  // We can only dynamically resolve models, for embedders user must provide dimensions.
-  if (actionType === 'model') {
-    defineOllamaModel(
-      ai,
-      {
-        name: actionName,
-      },
-      serverAddress,
-      requestHeaders
-    );
+function resolveAction({
+  params,
+  actionType,
+  actionName,
+  serverAddress,
+}: ResolveActionOptions) {
+  switch (actionType) {
+    case 'model':
+      return defineOllamaModel(
+        {
+          name: actionName,
+        },
+        serverAddress,
+        params?.requestHeaders
+      );
   }
+  return undefined;
 }
 
 async function listActions(
@@ -138,30 +137,21 @@ async function listActions(
   );
 }
 
-function ollamaPlugin(params?: OllamaPluginParams): GenkitPlugin {
-  if (!params) {
-    params = {};
-  }
-  if (!params.serverAddress) {
-    params.serverAddress = DEFAULT_OLLAMA_SERVER_ADDRESS;
-  }
-  const serverAddress = params.serverAddress;
-  return genkitPlugin(
-    'ollama',
-    async (ai: Genkit) => {
-      await initializer(ai, serverAddress, params);
+function ollamaPlugin(params: OllamaPluginParams = {}): GenkitPluginV2 {
+  const serverAddress = params.serverAddress || DEFAULT_OLLAMA_SERVER_ADDRESS;
+
+  return genkitPluginV2({
+    name: 'ollama',
+    init() {
+      return initializer(serverAddress, params);
     },
-    async (ai, actionType, actionName) => {
-      resolveAction(
-        ai,
-        actionType,
-        actionName,
-        serverAddress,
-        params?.requestHeaders
-      );
+    resolve(actionType, actionName) {
+      return resolveAction({ params, actionType, actionName, serverAddress });
     },
-    async () => await listActions(serverAddress, params?.requestHeaders)
-  );
+    async list() {
+      return await listActions(serverAddress, params?.requestHeaders);
+    },
+  });
 }
 
 async function listLocalModels(
@@ -218,26 +208,25 @@ export const OllamaConfigSchema = GenerationCommonConfigSchema.extend({
 });
 
 function defineOllamaModel(
-  ai: Genkit,
-  model: ModelDefinition,
+  modelDef: ModelDefinition,
   serverAddress: string,
   requestHeaders?: RequestHeaders
 ) {
-  return ai.defineModel(
+  return model<typeof OllamaConfigSchema>(
     {
-      name: `ollama/${model.name}`,
-      label: `Ollama - ${model.name}`,
+      name: `ollama/${modelDef.name}`,
+      label: `Ollama - ${modelDef.name}`,
       configSchema: OllamaConfigSchema,
       supports: {
-        multiturn: !model.type || model.type === 'chat',
+        multiturn: !modelDef.type || modelDef.type === 'chat',
         systemRole: true,
-        tools: model.supports?.tools,
+        tools: modelDef.supports?.tools,
       },
     },
-    async (input, streamingCallback) => {
+    async (request, opts) => {
       const { topP, topK, stopSequences, maxOutputTokens, ...rest } =
-        input.config as any;
-      const options: Record<string, any> = { ...rest };
+        request.config || {};
+      const options = { ...rest };
       if (topP !== undefined) {
         options.top_p = topP;
       }
@@ -250,29 +239,29 @@ function defineOllamaModel(
       if (maxOutputTokens !== undefined) {
         options.num_predict = maxOutputTokens;
       }
-      const type = model.type ?? 'chat';
-      const request = toOllamaRequest(
-        model.name,
-        input,
+      const type = modelDef.type ?? 'chat';
+      const ollamaRequest = toOllamaRequest(
+        modelDef.name,
+        request,
         options,
         type,
-        !!streamingCallback
+        opts?.streamingRequested
       );
-      logger.debug(request, `ollama request (${type})`);
+      logger.debug(ollamaRequest, `ollama request (${type})`);
 
       const extraHeaders = await getHeaders(
         serverAddress,
         requestHeaders,
-        model,
-        input
+        modelDef,
+        request
       );
       let res;
       try {
         res = await fetch(
           serverAddress + (type === 'chat' ? '/api/chat' : '/api/generate'),
           {
             method: 'POST',
-            body: JSON.stringify(request),
+            body: JSON.stringify(ollamaRequest),
             headers: {
               'Content-Type': 'application/json',
               ...extraHeaders,
@@ -297,15 +286,15 @@ function defineOllamaModel(
 
       let message: MessageData;
 
-      if (streamingCallback) {
+      if (opts.streamingRequested) {
         const reader = res.body.getReader();
         const textDecoder = new TextDecoder();
         let textResponse = '';
         for await (const chunk of readChunks(reader)) {
           const chunkText = textDecoder.decode(chunk);
           const json = JSON.parse(chunkText);
           const message = parseMessage(json, type);
-          streamingCallback({
+          opts.sendChunk({
             index: 0,
             content: message.content,
           });
@@ -329,7 +318,7 @@ function defineOllamaModel(
 
       return {
         message,
-        usage: getBasicUsageStats(input.messages, message),
+        usage: getBasicUsageStats(request.messages, message),
         finishReason: 'stop',
       } as GenerateResponseData;
     }
@@ -500,7 +489,7 @@ function toGenkitToolRequest(tool_calls: OllamaToolCall[]): ToolRequestPart[] {
   }));
 }
 
-function readChunks(reader) {
+function readChunks(reader: ReadableStreamDefaultReader<Uint8Array>) {
   return {
     async *[Symbol.asyncIterator]() {
       let readResult = await reader.read();
@@ -536,6 +525,7 @@ function isValidOllamaTool(tool: ToolDefinition): boolean {
 }
 
 export const ollama = ollamaPlugin as OllamaPlugin;
+
 ollama.model = (
   name: string,
   config?: z.infer<typeof OllamaConfigSchema>