Add Ollama provider

ai-providers/ollama.ts

@@ -0,0 +1,71 @@
+import { ReadableStream, UnderlyingByteSource, ReadableByteStreamController } from 'stream/web'
+import { Ollama, ChatResponse } from 'ollama'
+import { AiProvider, StreamChunkCallback } from './provider'
+import { AiStreamEvent, encodeEvent } from './event'
+
+type OllamaStreamResponse = AsyncGenerator<ChatResponse>
+
+class OllamaByteSource implements UnderlyingByteSource {
+  type: 'bytes' = 'bytes'
+  response: OllamaStreamResponse
+  chunkCallback?: StreamChunkCallback
+
+  constructor (response: OllamaStreamResponse, chunkCallback?: StreamChunkCallback) {
+    this.response = response
+    this.chunkCallback = chunkCallback
+  }
+
+  async pull (controller: ReadableByteStreamController): Promise<void> {
+    const { done, value } = await this.response.next()
+    if (done !== undefined && done) {
+      controller.close()
+      return
+    }
+
+    let response = value.message.content
+    if (this.chunkCallback !== undefined) {
+      response = await this.chunkCallback(response)
+    }
+
+    const eventData: AiStreamEvent = {
+      event: 'content',
+      data: {
+        response
+      }
+    }
+    controller.enqueue(encodeEvent(eventData))
+  }
+}
+
+export class OllamaProvider implements AiProvider {
+  model: string
+  client: Ollama
+
+  constructor (host: string, model: string) {
+    this.model = model
+    this.client = new Ollama({ host })
+  }
+
+  async ask (prompt: string): Promise<string> {
+    const response = await this.client.chat({
+      model: this.model,
+      messages: [
+        { role: 'user', content: prompt }
+      ]
+    })
+
+    return response.message.content
+  }
+
+  async askStream (prompt: string, chunkCallback?: StreamChunkCallback | undefined): Promise<ReadableStream> {
+    const response = await this.client.chat({
+      model: this.model,
+      messages: [
+        { role: 'user', content: prompt }
+      ],
+      stream: true
+    })
+
+    return new ReadableStream(new OllamaByteSource(response, chunkCallback))
+  }
+}

config.d.ts

@@ -256,6 +256,12 @@ export interface AiWarpConfig {
             | "mistral-large-latest";
           apiKey: string;
         };
+      }
+    | {
+        ollama: {
+          host: string;
+          model: string;
+        };
       };
   promptDecorators?: {
     prefix?: string;

lib/generator.ts

@@ -78,6 +78,14 @@ class AiWarpGenerator extends ServiceGenerator {
           }
         }
         break
+      case 'ollama':
+        config.aiProvider = {
+          ollama: {
+            host: 'http://127.0.0.1:11434',
+            model: this.config.aiModel
+          }
+        }
+        break
       default:
         config.aiProvider = {
           openai: {

lib/schema.ts

@@ -68,6 +68,21 @@ const aiWarpSchema = {
           },
           required: ['mistral'],
           additionalProperties: false
+        },
+        {
+          properties: {
+            ollama: {
+              type: 'object',
+              properties: {
+                host: { type: 'string' },
+                model: { type: 'string' }
+              },
+              required: ['host', 'model'],
+              additionalProperties: false
+            }
+          },
+          required: ['ollama'],
+          additionalProperties: false
         }
       ]
     },

package.json

@@ -36,6 +36,7 @@
     "fast-json-stringify": "^5.13.0",
     "fastify-user": "^0.3.3",
     "json-schema-to-typescript": "^13.0.0",
+    "ollama": "^0.5.0",
     "openai": "^4.28.4",
     "snazzy": "^9.0.0",
     "ts-standard": "^12.0.2",

plugins/warp.ts

@@ -6,6 +6,7 @@ import { MistralProvider } from '../ai-providers/mistral.js'
 import { AiProvider, StreamChunkCallback } from '../ai-providers/provider'
 import { AiWarpConfig } from '../config'
 import createError from '@fastify/error'
+import { OllamaProvider } from '../ai-providers/ollama'
 
 const UnknownAiProviderError = createError('UNKNOWN_AI_PROVIDER', 'Unknown AI Provider')
 
@@ -16,6 +17,9 @@ function build (aiProvider: AiWarpConfig['aiProvider']): AiProvider {
   } else if ('mistral' in aiProvider) {
     const { model, apiKey } = aiProvider.mistral
     return new MistralProvider(model, apiKey)
+  } else if ('ollama' in aiProvider) {
+    const { host, model } = aiProvider.ollama
+    return new OllamaProvider(host, model)
   } else {
     throw new UnknownAiProviderError()
   }

tests/e2e/api.test.ts

@@ -3,7 +3,7 @@ import { before, after, describe, it } from 'node:test'
 import assert from 'node:assert'
 import { FastifyInstance } from 'fastify'
 import fastifyPlugin from 'fastify-plugin'
-import { MOCK_CONTENT_RESPONSE, buildExpectedStreamBodyString } from '../utils/mocks'
+import { MOCK_CONTENT_RESPONSE, OLLAMA_MOCK_HOST, buildExpectedStreamBodyString } from '../utils/mocks'
 import { AiWarpConfig } from '../../config'
 import { buildAiWarpApp } from '../utils/stackable'
 
@@ -24,6 +24,15 @@ const providers: Provider[] = [
       }
     }
   },
+  {
+    name: 'Ollama',
+    config: {
+      ollama: {
+        host: OLLAMA_MOCK_HOST,
+        model: 'some-model'
+      }
+    }
+  },
   {
     name: 'Mistral',
     config: {

tests/e2e/index.ts

@@ -1,7 +1,8 @@
 import './api.test'
 import './rate-limiting.test'
 import './auth.test'
-import { mockMistralApi, mockOpenAiApi } from '../utils/mocks'
+import { mockMistralApi, mockOllama, mockOpenAiApi } from '../utils/mocks'
 
 mockOpenAiApi()
 mockMistralApi()
+mockOllama()

tests/types/schema.test-d.ts

@@ -40,6 +40,15 @@ expectAssignable<AiWarpConfig>({
   }
 })
 
+expectAssignable<AiWarpConfig>({
+  aiProvider: {
+    ollama: {
+      host: '',
+      model: 'some-model'
+    }
+  }
+})
+
 expectAssignable<AiWarpConfig>({
   $schema: './stackable.schema.json',
   service: {

tests/unit/ai-providers.test.ts

@@ -4,13 +4,15 @@ import assert from 'node:assert'
 import { MistralProvider } from '../../ai-providers/mistral'
 import { OpenAiProvider } from '../../ai-providers/open-ai'
 import { AiProvider } from '../../ai-providers/provider'
-import { MOCK_CONTENT_RESPONSE, buildExpectedStreamBodyString } from '../utils/mocks'
+import { MOCK_CONTENT_RESPONSE, OLLAMA_MOCK_HOST, buildExpectedStreamBodyString } from '../utils/mocks'
+import { OllamaProvider } from '../../ai-providers/ollama'
 
 const expectedStreamBody = buildExpectedStreamBodyString()
 
 const providers: AiProvider[] = [
   new OpenAiProvider('gpt-3.5-turbo', ''),
-  new MistralProvider('open-mistral-7b', '')
+  new MistralProvider('open-mistral-7b', ''),
+  new OllamaProvider(OLLAMA_MOCK_HOST, 'some-model')
 ]
 
 for (const provider of providers) {

tests/unit/index.ts

@@ -1,6 +1,7 @@
 import './generator.test'
 import './ai-providers.test'
-import { mockMistralApi, mockOpenAiApi } from '../utils/mocks'
+import { mockMistralApi, mockOllama, mockOpenAiApi } from '../utils/mocks'
 
 mockOpenAiApi()
 mockMistralApi()
+mockOllama()

tests/utils/mocks.ts

@@ -188,3 +188,65 @@ export function mockMistralApi (): void {
     }
   }).persist()
 }
+
+export const OLLAMA_MOCK_HOST = 'http://127.0.0.1:41434'
+let isOllamaMocked = false
+
+/**
+ * @see https://github.com/ollama/ollama/blob/9446b795b58e32c8b248a76707780f4f96b6434f/docs/api.md
+ */
+export function mockOllama (): void {
+  if (isOllamaMocked) {
+    return
+  }
+
+  isOllamaMocked = true
+
+  establishMockAgent()
+
+  const pool = mockAgent.get(OLLAMA_MOCK_HOST)
+  pool.intercept({
+    path: '/api/chat',
+    method: 'POST'
+  }).reply(200, (opts) => {
+    if (typeof opts.body !== 'string') {
+      throw new Error(`body is not a string (${typeof opts.body})`)
+    }
+
+    const body = JSON.parse(opts.body)
+
+    let response = ''
+    if (body.stream === true) {
+      for (let i = 0; i < MOCK_STREAMING_CONTENT_CHUNKS.length; i++) {
+        response += JSON.stringify({
+          model: 'llama2',
+          created_at: '2023-08-04T08:52:19.385406455-07:00',
+          message: {
+            role: 'assistant',
+            content: MOCK_STREAMING_CONTENT_CHUNKS[i],
+            images: null
+          },
+          done: i === MOCK_STREAMING_CONTENT_CHUNKS.length - 1
+        })
+        response += '\n'
+      }
+    } else {
+      response += JSON.stringify({
+        model: 'llama2',
+        created_at: '2023-08-04T19:22:45.499127Z',
+        message: {
+          role: 'assistant',
+          content: MOCK_CONTENT_RESPONSE,
+          images: null
+        },
+        done: true
+      })
+    }
+
+    return response
+  }, {
+    headers: {
+      'content-type': 'application/json'
+    }
+  }).persist()
+}