Repository: promptfoo/promptfoo
Pull request #102 (2023-08-08 18:59:55)

Add ollama provider
Add support for https://github.com/jmorganca/ollama


Patch:
diff --git a/src/cache.ts b/src/cache.ts
index f347a9caf..eaed7882a 100644
--- a/src/cache.ts
+++ b/src/cache.ts
@@ -42,10 +42,11 @@ export function getCache() {
   return cacheInstance;
 }
 
-export async function fetchJsonWithCache(
+export async function fetchWithCache(
   url: RequestInfo,
   options: RequestInit = {},
   timeout: number,
+  format: 'json' | 'text' = 'json',
 ): Promise<{ data: any; cached: boolean }> {
   if (!enabled) {
     const resp = await fetchWithRetries(url, options, timeout);
@@ -75,7 +76,7 @@ export async function fetchJsonWithCache(
   // Fetch the actual data and store it in the cache
   const response = await fetchWithRetries(url, options, timeout);
   try {
-    const data = await response.json();
+    const data = format === 'json' ? await response.json() : await response.text();
     if (response.ok) {
       logger.debug(`Storing ${url} response in cache: ${JSON.stringify(data)}`);
       await cache.set(cacheKey, JSON.stringify(data));
diff --git a/src/providers.ts b/src/providers.ts
index e5c4680a2..b278078de 100644
--- a/src/providers.ts
+++ b/src/providers.ts
@@ -5,6 +5,7 @@ import { AnthropicCompletionProvider } from './providers/anthropic';
 import { ReplicateProvider } from './providers/replicate';
 import { LocalAiCompletionProvider, LocalAiChatProvider } from './providers/localai';
 import { LlamaProvider } from './providers/llama';
+import { OllamaProvider } from './providers/ollama';
 import { ScriptCompletionProvider } from './providers/scriptCompletion';
 import {
   AzureOpenAiChatCompletionProvider,
@@ -148,6 +149,9 @@ export async function loadApiProvider(
   if (providerPath === 'llama' || providerPath.startsWith('llama:')) {
     const modelName = providerPath.split(':')[1];
     return new LlamaProvider(modelName, context?.config);
+  } else if (providerPath.startsWith('ollama:')) {
+    const modelName = providerPath.split(':')[1];
+    return new OllamaProvider(modelName);
   } else if (providerPath?.startsWith('localai:')) {
     const options = providerPath.split(':');
     const modelType = options[1];
diff --git a/src/providers/azureopenai.ts b/src/providers/azureopenai.ts
index fcd527bfd..2507d47e0 100644
--- a/src/providers/azureopenai.ts
+++ b/src/providers/azureopenai.ts
@@ -1,5 +1,5 @@
 import logger from '../logger';
-import { fetchJsonWithCache } from '../cache';
+import { fetchWithCache } from '../cache';
 import { REQUEST_TIMEOUT_MS, parseChatPrompt } from './shared';
 
 import type { ApiProvider, ProviderEmbeddingResponse, ProviderResponse } from '../types.js';
@@ -61,7 +61,7 @@ export class AzureOpenAiEmbeddingProvider extends AzureOpenAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/openai/deployments/${this.deploymentName}/embeddings?api-version=2023-07-01-preview`,
         {
           method: 'POST',
@@ -171,7 +171,7 @@ export class AzureOpenAiCompletionProvider extends AzureOpenAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/openai/deployments/${this.deploymentName}/completions?api-version=2023-07-01-preview`,
         {
           method: 'POST',
@@ -258,7 +258,7 @@ export class AzureOpenAiChatCompletionProvider extends AzureOpenAiGenericProvide
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/openai/deployments/${this.deploymentName}/chat/completions?api-version=2023-07-01-preview`,
         {
           method: 'POST',
diff --git a/src/providers/llama.ts b/src/providers/llama.ts
index 494dc0910..6658030ff 100644
--- a/src/providers/llama.ts
+++ b/src/providers/llama.ts
@@ -1,4 +1,4 @@
-import { fetchJsonWithCache } from '../cache';
+import { fetchWithCache } from '../cache';
 import { REQUEST_TIMEOUT_MS } from './shared';
 
 import type { ApiProvider, ProviderResponse } from '../types.js';
@@ -65,7 +65,7 @@ export class LlamaProvider implements ApiProvider {
 
     let response;
     try {
-      response = await fetchJsonWithCache(
+      response = await fetchWithCache(
         `${process.env.LLAMA_BASE_URL || 'http://localhost:8080'}/completion`,
         {
           method: 'POST',
diff --git a/src/providers/localai.ts b/src/providers/localai.ts
index 2b19466d9..7313d7e74 100644
--- a/src/providers/localai.ts
+++ b/src/providers/localai.ts
@@ -1,5 +1,5 @@
 import logger from '../logger';
-import { fetchJsonWithCache } from '../cache';
+import { fetchWithCache } from '../cache';
 import { REQUEST_TIMEOUT_MS, parseChatPrompt } from './shared';
 
 import type { ApiProvider, ProviderResponse } from '../types.js';
@@ -40,7 +40,7 @@ export class LocalAiChatProvider extends LocalAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `${this.apiBaseUrl}/chat/completions`,
         {
           method: 'POST',
@@ -81,7 +81,7 @@ export class LocalAiCompletionProvider extends LocalAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `${this.apiBaseUrl}/completions`,
         {
           method: 'POST',
diff --git a/src/providers/ollama.ts b/src/providers/ollama.ts
new file mode 100644
index 000000000..4891ab76a
--- /dev/null
+++ b/src/providers/ollama.ts
@@ -0,0 +1,88 @@
+import logger from '../logger';
+import { fetchWithCache } from '../cache';
+
+import type { ApiProvider, ProviderResponse } from '../types.js';
+import { REQUEST_TIMEOUT_MS } from './shared';
+
+interface OllamaJsonL {
+  model: string;
+  created_at: string;
+  response?: string;
+  done: boolean;
+  context?: number[];
+  total_duration?: number;
+  load_duration?: number;
+  sample_count?: number;
+  sample_duration?: number;
+  prompt_eval_count?: number;
+  prompt_eval_duration?: number;
+  eval_count?: number;
+  eval_duration?: number;
+}
+
+export class OllamaProvider implements ApiProvider {
+  modelName: string;
+
+  constructor(modelName: string) {
+    this.modelName = modelName;
+  }
+
+  id(): string {
+    return `ollama:${this.modelName}`;
+  }
+
+  toString(): string {
+    return `[Ollama Provider ${this.modelName}]`;
+  }
+
+  async callApi(prompt: string): Promise<ProviderResponse> {
+    const params = {
+      model: this.modelName,
+      prompt,
+    };
+
+    logger.debug(`Calling Ollama API: ${JSON.stringify(params)}`);
+    let response;
+    try {
+      response = await fetchWithCache(
+        `${process.env.OLLAMA_BASE_URL || 'http://localhost:11434'}/api/generate`,
+        {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify(params),
+        },
+        REQUEST_TIMEOUT_MS,
+        'text',
+      );
+    } catch (err) {
+      return {
+        error: `API call error: ${String(err)}`,
+      };
+    }
+    logger.debug(`\tOllama API response: ${response.data}`);
+
+    try {
+      const output = response.data
+        .split('\n')
+        .map((line: string) => {
+          const parsed = JSON.parse(line) as OllamaJsonL;
+          if (parsed.response) {
+            return parsed.response;
+          }
+          return null;
+        })
+        .filter((s: string | null) => s !== null)
+        .join('');
+
+      return {
+        output,
+      };
+    } catch (err) {
+      return {
+        error: `API response error: ${String(err)}: ${JSON.stringify(response.data)}`,
+      };
+    }
+  }
+}
diff --git a/src/providers/openai.ts b/src/providers/openai.ts
index 782c37d08..f2540fe32 100644
--- a/src/providers/openai.ts
+++ b/src/providers/openai.ts
@@ -1,5 +1,5 @@
 import logger from '../logger';
-import { fetchJsonWithCache } from '../cache';
+import { fetchWithCache } from '../cache';
 import { REQUEST_TIMEOUT_MS, parseChatPrompt } from './shared';
 
 import type { ApiProvider, ProviderEmbeddingResponse, ProviderResponse } from '../types.js';
@@ -61,7 +61,7 @@ export class OpenAiEmbeddingProvider extends OpenAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/v1/embeddings`,
         {
           method: 'POST',
@@ -177,7 +177,7 @@ export class OpenAiCompletionProvider extends OpenAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/v1/completions`,
         {
           method: 'POST',
@@ -275,7 +275,7 @@ export class OpenAiChatCompletionProvider extends OpenAiGenericProvider {
     let data,
       cached = false;
     try {
-      ({ data, cached } = (await fetchJsonWithCache(
+      ({ data, cached } = (await fetchWithCache(
         `https://${this.apiHost}/v1/chat/completions`,
         {
           method: 'POST',
diff --git a/src/util.ts b/src/util.ts
index d9bd4dd39..c35878986 100644
--- a/src/util.ts
+++ b/src/util.ts
@@ -455,10 +455,12 @@ export function writeLatestResults(results: EvaluateSummary, config: Partial<Uni
         2,
       ),
     );
-    if (fs.existsSync(latestResultsPath) && fs.lstatSync(latestResultsPath).isSymbolicLink()) {
+
+    try {
       fs.unlinkSync(latestResultsPath);
-    }
+    } catch {}
     fs.symlinkSync(newResultsPath, latestResultsPath);
+
     cleanupOldResults();
   } catch (err) {
     logger.error(`Failed to write latest results to ${newResultsPath}:\n${err}`);


Test patch:
diff --git a/test/cache.test.ts b/test/cache.test.ts
index f6eb640ce..1b2c7ada1 100644
--- a/test/cache.test.ts
+++ b/test/cache.test.ts
@@ -1,10 +1,10 @@
-import { fetchJsonWithCache, disableCache, enableCache } from '../src/cache.js';
+import { fetchWithCache, disableCache, enableCache } from '../src/cache.js';
 import fetch, { Response } from 'node-fetch';
 
 jest.mock('node-fetch');
 const mockedFetch = fetch as jest.MockedFunction<typeof fetch>;
 
-describe('fetchJsonWithCache', () => {
+describe('fetchWithCache', () => {
   afterEach(() => {
     mockedFetch.mockReset();
   });
@@ -20,7 +20,7 @@ describe('fetchJsonWithCache', () => {
       json: () => Promise.resolve(response),
     } as Response);
 
-    const result = await fetchJsonWithCache(url, {}, 1000);
+    const result = await fetchWithCache(url, {}, 1000);
 
     expect(mockedFetch).toHaveBeenCalledTimes(1);
     expect(result).toEqual({ cached: false, data: response });
@@ -37,7 +37,7 @@ describe('fetchJsonWithCache', () => {
       json: () => Promise.resolve(response),
     } as Response);
 
-    const result = await fetchJsonWithCache(url, {}, 1000);
+    const result = await fetchWithCache(url, {}, 1000);
 
     expect(mockedFetch).toHaveBeenCalledTimes(1);
     expect(result).toEqual({ cached: false, data: response });
@@ -52,7 +52,7 @@ describe('fetchJsonWithCache', () => {
       json: () => Promise.resolve(response),
     } as Response);
 
-    const result = await fetchJsonWithCache(url, {}, 1000);
+    const result = await fetchWithCache(url, {}, 1000);
 
     expect(mockedFetch).toHaveBeenCalledTimes(0);
     expect(result).toEqual({ cached: true, data: response });
@@ -68,7 +68,7 @@ describe('fetchJsonWithCache', () => {
       json: () => Promise.resolve(response),
     } as Response);
 
-    const result = await fetchJsonWithCache(url, {}, 1000);
+    const result = await fetchWithCache(url, {}, 1000);
 
     expect(mockedFetch).toHaveBeenCalledTimes(1);
     expect(result).toEqual({ cached: false, data: response });
@@ -87,7 +87,7 @@ describe('fetchJsonWithCache', () => {
       json: () => Promise.resolve(response),
     } as Response);
 
-    const result = await fetchJsonWithCache(url, {}, 1000);
+    const result = await fetchWithCache(url, {}, 1000);
 
     expect(mockedFetch).toHaveBeenCalledTimes(1);
     expect(result).toEqual({ cached: false, data: response });
diff --git a/test/providers.test.ts b/test/providers.test.ts
index 84ba9c91e..eae0d924d 100644
--- a/test/providers.test.ts
+++ b/test/providers.test.ts
@@ -1,5 +1,5 @@
 import fetch from 'node-fetch';
-import { fetchJsonWithCache } from '../src/cache';
+import { fetchWithCache } from '../src/cache';
 
 import { OpenAiCompletionProvider, OpenAiChatCompletionProvider } from '../src/providers/openai';
 import { AnthropicCompletionProvider } from '../src/providers/anthropic';
@@ -12,6 +12,7 @@ import {
   AzureOpenAiChatCompletionProvider,
   AzureOpenAiCompletionProvider,
 } from '../src/providers/azureopenai';
+import { OllamaProvider } from '../src/providers/ollama';
 
 jest.mock('node-fetch', () => jest.fn());
 
@@ -167,6 +168,28 @@ describe('providers', () => {
     expect(result.output).toBe('Test output');
   });
 
+  test('OllamaProvider callApi', async () => {
+    const mockResponse = {
+      text: jest.fn()
+        .mockResolvedValue(`{"model":"llama2:13b","created_at":"2023-08-08T21:50:34.898068Z","response":"Gre","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:34.929199Z","response":"at","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:34.959989Z","response":" question","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:34.992117Z","response":"!","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:35.023658Z","response":" The","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:35.0551Z","response":" sky","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:35.086103Z","response":" appears","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:35.117166Z","response":" blue","done":false}
+{"model":"llama2:13b","created_at":"2023-08-08T21:50:41.695299Z","done":true,"context":[1,29871,1,13,9314],"total_duration":10411943458,"load_duration":458333,"sample_count":217,"sample_duration":154566000,"prompt_eval_count":11,"prompt_eval_duration":3334582000,"eval_count":216,"eval_duration":6905134000}`),
+    };
+    (fetch as unknown as jest.Mock).mockResolvedValue(mockResponse);
+
+    const provider = new OllamaProvider('llama');
+    const result = await provider.callApi('Test prompt');
+
+    expect(fetch).toHaveBeenCalledTimes(1);
+    expect(result.output).toBe('Great question! The sky appears blue');
+  });
+
   test('loadApiProvider with openai:chat', async () => {
     const provider = await loadApiProvider('openai:chat');
     expect(provider).toBeInstanceOf(OpenAiChatCompletionProvider);