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,13 @@
  * 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';
 
-async function toOllamaEmbedRequest(
+export async function toOllamaEmbedRequest(
   modelName: string,
   dimensions: number,
   documents: Document[],
@@ -59,11 +61,13 @@ 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<any> {
+  return embedder(
     {
       name: `ollama/${name}`,
       info: {
@@ -75,13 +79,14 @@ export function defineOllamaEmbedder(
         },
       },
     },
-    async (input, config) => {
-      const serverAddress = config?.serverAddress || options.serverAddress;
+    async (request, config) => {
+      const serverAddress =
+        options.serverAddress || DEFAULT_OLLAMA_SERVER_ADDRESS;
 
       const { url, requestPayload, headers } = await toOllamaEmbedRequest(
         modelName,
         dimensions,
-        input,
+        request.input,
         serverAddress,
         options.requestHeaders
       );

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,
@@ -56,7 +63,7 @@ import type {
 export type { OllamaPluginParams };
 
 export type OllamaPlugin = {
-  (params?: OllamaPluginParams): GenkitPlugin;
+  (params?: OllamaPluginParams): GenkitPluginV2;
 
   model(
     name: string,
@@ -65,59 +72,60 @@ 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 ollamaPlugin(params?: OllamaPluginParams): GenkitPluginV2 {
+  if (!params) {
+    params = {};
+  }
+  if (!params.serverAddress) {
+    params.serverAddress = DEFAULT_OLLAMA_SERVER_ADDRESS;
+  }
+  const serverAddress = params.serverAddress;
 
-const GENERIC_MODEL_INFO = {
-  supports: {
-    multiturn: true,
-    media: true,
-    tools: true,
-    toolChoice: true,
-    systemRole: true,
-    constrained: 'all',
-  },
-} as ModelInfo;
+  return genkitPluginV2({
+    name: 'ollama',
+    init() {
+      const actions: ResolvableAction[] = [];
 
-const DEFAULT_OLLAMA_SERVER_ADDRESS = 'http://localhost:11434';
+      if (params?.models) {
+        for (const model of params.models) {
+          actions.push(
+            createOllamaModel(model, serverAddress, params.requestHeaders)
+          );
+        }
+      }
 
-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?.embedders && params.serverAddress) {
+        for (const embedder of params.embedders) {
+          actions.push(
+            defineOllamaEmbedder({
+              name: embedder.name,
+              modelName: embedder.name,
+              dimensions: embedder.dimensions,
+              options: { ...params, serverAddress },
+            })
+          );
+        }
+      }
 
-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
-    );
-  }
+      return actions;
+    },
+    async resolve(actionType, actionName) {
+      // dynamically resolve models, for embedders user must provide dimensions.
+      if (actionType === 'model') {
+        return await createOllamaModel(
+          {
+            name: actionName,
+          },
+          serverAddress,
+          params?.requestHeaders
+        );
+      }
+      return undefined;
+    },
+    async list() {
+      return await listActions(serverAddress, params?.requestHeaders);
+    },
+  });
 }
 
 async function listActions(
@@ -131,39 +139,13 @@ async function listActions(
       ?.filter((m) => m.model && !m.model.includes('embed'))
       .map((m) =>
         modelActionMetadata({
-          name: `ollama/${m.model}`,
+          name: m.model,
           info: GENERIC_MODEL_INFO,
         })
       ) || []
   );
 }
 
-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);
-    },
-    async (ai, actionType, actionName) => {
-      resolveAction(
-        ai,
-        actionType,
-        actionName,
-        serverAddress,
-        params?.requestHeaders
-      );
-    },
-    async () => await listActions(serverAddress, params?.requestHeaders)
-  );
-}
-
 async function listLocalModels(
   serverAddress: string,
   requestHeaders?: RequestHeaders
@@ -217,26 +199,25 @@ export const OllamaConfigSchema = GenerationCommonConfigSchema.extend({
     .optional(),
 });
 
-function defineOllamaModel(
-  ai: Genkit,
-  model: ModelDefinition,
+function createOllamaModel(
+  modelDef: ModelDefinition,
   serverAddress: string,
   requestHeaders?: RequestHeaders
 ) {
-  return ai.defineModel(
+  return model(
     {
-      name: `ollama/${model.name}`,
-      label: `Ollama - ${model.name}`,
+      name: 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;
+        request.config as any;
       const options: Record<string, any> = { ...rest };
       if (topP !== undefined) {
         options.top_p = topP;
@@ -250,29 +231,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,16 +278,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({
-            index: 0,
+          opts.sendChunk({
             content: message.content,
           });
           textResponse += message.content[0].text;
@@ -329,7 +309,7 @@ function defineOllamaModel(
 
       return {
         message,
-        usage: getBasicUsageStats(input.messages, message),
+        usage: getBasicUsageStats(request.messages, message),
         finishReason: 'stop',
       } as GenerateResponseData;
     }
@@ -500,7 +480,7 @@ function toGenkitToolRequest(tool_calls: OllamaToolCall[]): ToolRequestPart[] {
   }));
 }
 
-function readChunks(reader) {
+function readChunks(reader: ReadableStreamDefaultReader<Uint8Array>) {
   return {
     async *[Symbol.asyncIterator]() {
       let readResult = await reader.read();