fix(js/plugins/ollama): reconfigure ollama embedding

docs/plugins/ollama.md

@@ -122,3 +122,50 @@ const llmResponse = await generate({
   prompt: 'Tell me a joke.',
 });
 ```
+
+
+### Embeddings
+
+The Ollama plugin supports embeddings, allowing you to obtain vector representations of text content. To use embeddings, configure the plugin with an appropriate model that supports embeddings, and specify the `embedder` in your code.
+
+Example configuration:
+
+```js
+import { ollama } from 'genkitx-ollama';
+
+export default configureGenkit({
+  plugins: [
+    ollama({
+      embedders: [
+        {
+          name: 'nomic-embed-text', // specify the embedding model
+          dimensions: 768, // number of dimensions for the embedding
+        },
+      ],
+      serverAddress: 'http://127.0.0.1:11434', // default local address
+    }),
+  ],
+});
+```
+
+### Using Embeddings
+
+This example demonstrates generating a 768-dimensional embedding using the `nomic-embed-text` model.
+
+Make sure to run `ollama pull nomic-embed-text` to obtain this model, and `ollama serve` to serve the embedding model.
+
+Once configured, you can generate embeddings for your content using the `embed` function:
+
+```js
+import { embed } from '@genkit-ai/ai/embedder';
+
+// ....
+
+const result = await embed({
+  embedder: 'ollama/nomic-embed-text', // use the configured embedder
+  content: 'Hello, world!',
+  options: {
+    truncate: true, // optional: truncate content if necessary
+  },
+});
+```

js/plugins/ollama/package.json

@@ -17,7 +17,9 @@
     "compile": "tsup-node",
     "build:clean": "rimraf ./lib",
     "build": "npm-run-all build:clean check compile",
-    "build:watch": "tsup-node --watch"
+    "build:watch": "tsup-node --watch",
+    "test": "node --import tsx --test tests/*_test.ts",
+    "test:live": "LIVE_TEST=true node --import tsx --test tests/*_live_test.ts"
   },
   "repository": {
     "type": "git",

js/plugins/ollama/src/embeddings.ts

@@ -13,29 +13,85 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import { GenerateRequest } from '@genkit-ai/ai';
 import { defineEmbedder } from '@genkit-ai/ai/embedder';
+import { Document } from '@genkit-ai/ai/retriever';
 import { logger } from '@genkit-ai/core/logging';
 import z from 'zod';
 import { OllamaPluginParams } from './index.js';
 // Define the schema for Ollama embedding configuration
-export const OllamaEmbeddingConfigSchema = z.object({
-  modelName: z.string(),
-  serverAddress: z.string(),
-});
+export const OllamaEmbeddingConfigSchema = z
+  .object({
+    serverAddress: z.string().optional(),
+  })
+  .passthrough();
+
+export type EmbeddingModelDefinition = { name: string; dimensions: number };
+
 export type OllamaEmbeddingConfig = z.infer<typeof OllamaEmbeddingConfigSchema>;
+
 // Define the structure of the request and response for embedding
-interface OllamaEmbeddingInstance {
-  content: string;
-}
 interface OllamaEmbeddingPrediction {
   embedding: number[];
 }
+
 interface DefineOllamaEmbeddingParams {
   name: string;
   modelName: string;
   dimensions: number;
   options: OllamaPluginParams;
 }
+
+type RequestHeaders =
+  | Record<string, string>
+  | ((
+      params: { serverAddress: string; model: EmbeddingModelDefinition },
+      request: GenerateRequest
+    ) => Promise<Record<string, string> | void>);
+
+/**
+ * Helper function to create the Ollama embed request and headers.
+ */
+async function toOllamaEmbedRequest(
+  modelName: string,
+  document: Document,
+  serverAddress: string,
+  requestHeaders?: Record<string, string> | ((params: any, request: any) => any)
+): Promise<{
+  url: string;
+  requestPayload: any;
+  headers: Record<string, string>;
+}> {
+  const requestPayload = {
+    model: modelName,
+    prompt: document.text(),
+  };
+
+  // Determine headers
+  const extraHeaders = requestHeaders
+    ? typeof requestHeaders === 'function'
+      ? await requestHeaders(
+          {
+            serverAddress,
+            modelName,
+          },
+          document
+        )
+      : requestHeaders
+    : {};
+
+  const headers = {
+    'Content-Type': 'application/json',
+    ...extraHeaders, // Add any dynamic headers
+  };
+
+  return {
+    url: `${serverAddress}/api/embeddings`,
+    requestPayload,
+    headers,
+  };
+}
+
 export function defineOllamaEmbedder({
   name,
   modelName,
@@ -44,50 +100,53 @@ export function defineOllamaEmbedder({
 }: DefineOllamaEmbeddingParams) {
   return defineEmbedder(
     {
-      name,
+      name: `ollama/${name}`,
       configSchema: OllamaEmbeddingConfigSchema, // Use the Zod schema directly here
       info: {
-        // TODO: do we want users to be able to specify the label when they call this method directly?
-        label: 'Ollama Embedding - ' + modelName,
+        label: 'Ollama Embedding - ' + name,
         dimensions,
         supports: {
-          //  TODO: do any ollama models support other modalities?
           input: ['text'],
         },
       },
     },
-    async (input, _config) => {
-      const serverAddress = options.serverAddress;
+    async (input, config) => {
+      const serverAddress = config?.serverAddress || options.serverAddress;
+
       const responses = await Promise.all(
-        input.map(async (i) => {
-          const requestPayload = {
-            model: modelName,
-            prompt: i.text(),
-          };
+        input.map(async (doc) => {
+          // Generate the request and headers using the helper function
+          const { url, requestPayload, headers } = await toOllamaEmbedRequest(
+            modelName,
+            doc,
+            serverAddress,
+            options.requestHeaders
+          );
+
           let res: Response;
           try {
-            console.log('MODEL NAME: ', modelName);
-            res = await fetch(`${serverAddress}/api/embeddings`, {
+            res = await fetch(url, {
               method: 'POST',
-              headers: {
-                'Content-Type': 'application/json',
-              },
+              headers,
               body: JSON.stringify(requestPayload),
             });
           } catch (e) {
             logger.error('Failed to fetch Ollama embedding');
             throw new Error(`Error fetching embedding from Ollama: ${e}`);
           }
+
           if (!res.ok) {
             logger.error('Failed to fetch Ollama embedding');
             throw new Error(
               `Error fetching embedding from Ollama: ${res.statusText}`
             );
           }
+
           const responseData = (await res.json()) as OllamaEmbeddingPrediction;
           return responseData;
         })
       );
+
       return {
         embeddings: responses,
       };

js/plugins/ollama/src/index.ts

@@ -25,7 +25,7 @@ import {
 } from '@genkit-ai/ai/model';
 import { genkitPlugin, Plugin } from '@genkit-ai/core';
 import { logger } from '@genkit-ai/core/logging';
-import { defineOllamaEmbedder } from './embeddings';
+import { defineOllamaEmbedder, EmbeddingModelDefinition } from './embeddings';
 
 type ApiType = 'chat' | 'generate';
 
@@ -37,10 +37,16 @@ type RequestHeaders =
     ) => Promise<Record<string, string> | void>);
 
 type ModelDefinition = { name: string; type?: ApiType };
-type EmbeddingModelDefinition = { name: string; dimensions: number };
+
 export interface OllamaPluginParams {
-  models: ModelDefinition[];
-  embeddingModels?: EmbeddingModelDefinition[];
+  /*
+   * Models to be defined.
+   */
+  models?: ModelDefinition[];
+  /**
+   * Embedding models to be defined.
+   */
+  embedders?: EmbeddingModelDefinition[];
   /**
    *  ollama server address.
    */
@@ -54,12 +60,12 @@ export const ollama: Plugin<[OllamaPluginParams]> = genkitPlugin(
   async (params: OllamaPluginParams) => {
     const serverAddress = params?.serverAddress;
     return {
-      models: params.models.map((model) =>
+      models: params.models?.map((model) =>
         ollamaModel(model, serverAddress, params.requestHeaders)
       ),
-      embedders: params.embeddingModels?.map((model) =>
+      embedders: params.embedders?.map((model) =>
         defineOllamaEmbedder({
-          name: `${ollama}/model.name`,
+          name: model.name,
           modelName: model.name,
           dimensions: model.dimensions,
           options: params,

js/plugins/ollama/tests/embeddings_live_test.ts

@@ -14,10 +14,11 @@
  * limitations under the License.
  */
 import { embed } from '@genkit-ai/ai/embedder';
+import { configureGenkit } from '@genkit-ai/core';
 import assert from 'node:assert';
 import { describe, it } from 'node:test';
 import { defineOllamaEmbedder } from '../src/embeddings.js'; // Adjust the import path as necessary
-import { OllamaPluginParams } from '../src/index.js'; // Adjust the import path as necessary
+import { OllamaPluginParams, ollama } from '../src/index.js'; // Adjust the import path as necessary
 // Utility function to parse command-line arguments
 function parseArgs() {
   const args = process.argv.slice(2);
@@ -30,22 +31,64 @@ function parseArgs() {
   return { serverAddress, modelName };
 }
 const { serverAddress, modelName } = parseArgs();
-describe('defineOllamaEmbedder - Live Tests', () => {
-  const options: OllamaPluginParams = {
-    models: [{ name: modelName }],
-    serverAddress,
-  };
-  it('should successfully return embeddings', async () => {
-    const embedder = defineOllamaEmbedder({
-      name: 'live-test-embedder',
-      modelName: 'nomic-embed-text',
-      dimensions: 768,
-      options,
+if (process.env.LIVE_TEST) {
+  describe('Live Test: defineOllamaEmbedder', () => {
+    const options: OllamaPluginParams = {
+      models: [{ name: modelName }],
+      serverAddress,
+      requestHeaders: async () => {
+        // mock some async operation
+        await new Promise((resolve) => setTimeout(resolve, 200));
+
+        return {
+          'X-Custom-Header': 'custom-value',
+        };
+      },
+    };
+    it('live: should successfully return embeddings', async () => {
+      configureGenkit({
+        plugins: [
+          ollama({
+            serverAddress: 'http://127.0.0.1:11434', // default local address
+          }),
+        ],
+      });
+      const embedder = defineOllamaEmbedder({
+        name: 'live test embedder',
+        modelName: 'nomic-embed-text',
+        dimensions: 768,
+        options,
+      });
+      const result = await embed({
+        embedder,
+        content: 'Hello, world!',
+      });
+      assert.strictEqual(result.length, 768);
     });
-    const result = await embed({
-      embedder,
-      content: 'Hello, world!',
+  });
+
+  describe('E2E Test: Ollama Embedder', () => {
+    it('e2e: should successfully return embeddings using configureGenkit', async () => {
+      configureGenkit({
+        plugins: [
+          ollama({
+            embedders: [
+              { name: 'nomic-embed-text', dimensions: 768 }, // Use the existing embedder
+            ],
+            serverAddress: 'http://127.0.0.1:11434', // default local address
+          }),
+        ],
+      });
+
+      const result = await embed({
+        embedder: 'ollama/nomic-embed-text',
+        content: 'foo',
+        options: {
+          truncate: true,
+        },
+      });
+
+      assert.strictEqual(result.length, 768); // Assuming the embeddings should have 768 dimensions
     });
-    assert.strictEqual(result.length, 768);
   });
-});
+}