From 2ad8715d2e711edaf827cbda2a0b19c8b326d1b3 Mon Sep 17 00:00:00 2001
From: Aly Raffauf <alyraffauf@gmail.com>
Date: Tue, 23 Apr 2024 21:24:05 -0400
Subject: [PATCH] mauville: enable ollama

---
 hosts/mauville/default.nix               |  5 ++++
 nixosModules/services/default.nix        |  2 +-
 nixosModules/services/ollama/default.nix | 33 ++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 nixosModules/services/ollama/default.nix

diff --git a/hosts/mauville/default.nix b/hosts/mauville/default.nix
index f3b83de6..9efd42b7 100644
--- a/hosts/mauville/default.nix
+++ b/hosts/mauville/default.nix
@@ -48,6 +48,11 @@ in {
     };
     services = {
       binaryCache.enable = true;
+      ollama = {
+        enable = true;
+        gpu = "amd";
+        listenAddress = "0.0.0.0:11434";
+      };
     };
     system = {
       plymouth.enable = true;
diff --git a/nixosModules/services/default.nix b/nixosModules/services/default.nix
index d3502354..9d915266 100644
--- a/nixosModules/services/default.nix
+++ b/nixosModules/services/default.nix
@@ -4,5 +4,5 @@
   config,
   ...
 }: {
-  imports = [./binaryCache ./flatpak];
+  imports = [./binaryCache ./flatpak ./ollama];
 }
diff --git a/nixosModules/services/ollama/default.nix b/nixosModules/services/ollama/default.nix
new file mode 100644
index 00000000..63ac73ec
--- /dev/null
+++ b/nixosModules/services/ollama/default.nix
@@ -0,0 +1,33 @@
+{
+  pkgs,
+  lib,
+  config,
+  ...
+}: {
+  options = {
+    alyraffauf.services.ollama.enable = lib.mkEnableOption "Enable ollama interface for LLMs.";
+    alyraffauf.services.ollama.listenAddress = lib.mkOption {
+      description = "Listen Address for Ollama.";
+      default = "127.0.0.1:11434";
+      type = lib.types.str;
+    };
+    alyraffauf.services.ollama.gpu = lib.mkOption {
+      description = "Type of GPU for enabling GPU acceleration.";
+      default = null;
+      type = lib.types.str;
+    };
+  };
+
+  config = lib.mkIf config.alyraffauf.services.ollama.enable {
+    services.ollama = {
+      enable = true;
+      acceleration =
+        if config.alyraffauf.services.ollama.gpu == "amd"
+        then "rocm"
+        else if config.alyraffauf.services.ollama.gpu == "nvidia"
+        then "cuda"
+        else null;
+      listenAddress = config.alyraffauf.services.ollama.listenAddress;
+    };
+  };
+}