From dfb8acf757019bec099bf482c70bb3c8b4357b31 Mon Sep 17 00:00:00 2001
From: ken koch <kkoch986@gmail.com>
Date: Wed, 20 Mar 2024 23:46:40 -0400
Subject: [PATCH] feat: add ollama support

---
 Dockerfile         |  2 +-
 Makefile           |  5 ++-
 README.md          |  3 +-
 docker-compose.yml | 13 ++++++-
 index.ts           | 92 ++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 111 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 6fbe9b4..959c800 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,7 +3,7 @@ FROM node:18-alpine AS build
 WORKDIR /srv
 COPY package*.json /srv/
 RUN npm ci
-# COPY tsconfig.json /srv/
+
 COPY index.ts /srv/index.ts
 RUN npx tsc index.ts
 RUN npm ci --production
diff --git a/Makefile b/Makefile
index b2006ae..f17513b 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,9 @@
 GIT_SHA=$(shell git rev-parse --short=8 HEAD)
 
-all: index.js
+pull-models:
+	curl http://localhost:11434/api/pull -d '{"name": "neural-chat"}'
+
+all: index.js pull-models
 	node index.js
 
 index.js: index.ts
diff --git a/README.md b/README.md
index b8f6abb..a95ddde 100644
--- a/README.md
+++ b/README.md
@@ -2,9 +2,10 @@
 
 This module holds the service for generating responses in the interactive
 pipeline.
-It currently offers 2 backends:
+It currently offers 3 backends:
 
 * [Character.AI](https://beta.character.ai/) (using [the Node Unoffical API](https://github.com/realcoloride/node_characterai))
+* [OLlama](https://ollama.com/)
 * An "Echo" backend which just replies with exactly what you send it
 
 ## Interface
diff --git a/docker-compose.yml b/docker-compose.yml
index cced0c2..1216e9a 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -5,7 +5,18 @@ services:
     ports:
       - "3000:3000"
     environment:
-      BACKEND_TYPE: "characterai"
+      BACKEND_TYPE: "ollama"
       SERVER_PORT: 3000
       CHARACTER_AI_TOKEN: "${CHARACTER_AI_TOKEN}"
       CHARACTER_AI_CHARACTER_ID: "1AYQkwEQ83I3JKxMeNvctG4m7kQL1zTPJuGugAVsT_k"
+  ollama:
+    image: ollama/ollama:latest
+    ports:
+      - "11434:11434"
+    # deploy:
+    #   resources:
+    #     reservations:
+    #       devices:
+    #         - driver: nvidia
+    #           count: 1
+    #           capabilities: [gpu]
diff --git a/index.ts b/index.ts
index a4eed1c..57da511 100644
--- a/index.ts
+++ b/index.ts
@@ -13,11 +13,88 @@ type Response = {
   text: string,
 };
 
+type Message = {
+  role: string,
+  content: string,
+};
+
 interface Backend {
   start(): Promise<Response>,
   receive(message:string): Promise<Response>,
 };
 
+// OLlama backend is backed by ollama
+// the assumption is that relevant models have already been pulled
+class OLlamaBackend implements Backend {
+  private baseURL:string;
+  private model:string;
+  private systemPrompts:string[];
+  private initialPrompt: string;
+  private messages: Message[];
+
+  async initialize(baseURL: string, model: string, systemPrompts: string[], initialPrompt: string) {
+    this.model = model;
+    this.baseURL = baseURL;
+    this.systemPrompts = systemPrompts;
+    this.initialPrompt = initialPrompt;
+    // LATER: we could pull the model at this point, but that seems like not necessary at this point
+  }
+
+  async sendPrompt(message:string): Promise<Response> {
+    if (message != "") {
+      this.messages.push({
+        role: 'user',
+        content: message,
+      });
+    }
+    const p = await fetch(this.baseURL + "api/chat", {
+      method: "POST",
+      body: JSON.stringify({
+        "model": this.model,
+        // LATER: support streaming calls?
+        "stream": false,
+        "options": {
+          "temperature": 1
+        },
+        "messages": this.messages,
+      }),
+    })
+    .then((response) => response.json())
+    .then((resp: any) => {
+      // TODO: error handling
+      this.messages.push(resp.message);
+      const i = resp.created_at;
+      return {
+        id: i,
+        characterId: this.model,
+        chatId: i,
+        imageRelativePath: "",
+        srcName: "ollama-" + this.model,
+        text: resp.message.content,
+      };
+    });
+    return p;
+  }
+
+  async start(): Promise<Response> {
+    // send the basic system prompts and return the response
+    const messages = this.systemPrompts.map((m) => ({
+      "role": "system",
+      "content": m,
+    }));
+    messages.push({
+      "role": "user",
+      "content": this.initialPrompt,
+    });
+    this.messages = messages;
+    return this.sendPrompt("");
+  }
+
+  async receive(message: string): Promise<Response> {
+    return this.sendPrompt(message);
+  }
+}
+
 // Echo backend is just a testing backend that replies with exactly
 // the message it was sent
 class EchoBackend implements Backend {
@@ -141,6 +218,21 @@ class CharacterAIBackend implements Backend {
     case "echo":
       backend = new EchoBackend();
       break ;
+    case "ollama":
+      const ol = new OLlamaBackend();
+    // TODO: better parameterize all of this...
+      await ol.initialize(
+        "http://ollama:11434/",
+        "neural-chat",
+        [
+          "you are a zoltar inspired fortune teller with a tiki / bayou flair",
+          "your name is Makani",
+          "you should always respond to prompts with a whimsical style",
+        ],
+        "please introduce yourself",
+      );
+      backend = ol;
+      break ;
     case "characterai":
       const cai = new CharacterAIBackend();
       await cai.initialize(token, characterId)