From 16a965d56cd00d888625dcaf73664e1814bf3d5c Mon Sep 17 00:00:00 2001
From: JMARyA <jmarya@hydrar.de>
Date: Sun, 29 Dec 2024 10:01:54 +0100
Subject: [PATCH] update + ollama

---
 src/config.rs         |  8 ++++++++
 src/install/drives.rs |  2 ++
 src/install/mod.rs    |  6 ++++++
 src/install/ollama.rs | 24 ++++++++++++++++++++++++
 4 files changed, 40 insertions(+)
 create mode 100644 src/install/ollama.rs

diff --git a/src/config.rs b/src/config.rs
index cdc2190..4cb1c80 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -13,6 +13,14 @@ pub struct InstallConfig {
     pub user: Vec<UserConfig>,
     /// SSH Configuration
     pub ssh: Option<SSHConfig>,
+    /// Ollama AI Config
+    pub ai: Option<OllamaConfig>
+}
+
+#[derive(Debug, Deserialize)]
+pub struct OllamaConfig {
+    pub models: Option<Vec<String>>,
+    pub gpu: bool
 }
 
 #[derive(Debug, Deserialize)]
diff --git a/src/install/drives.rs b/src/install/drives.rs
index fb3d021..c9413be 100644
--- a/src/install/drives.rs
+++ b/src/install/drives.rs
@@ -2,6 +2,8 @@ use crate::{config::DriveConfig, run_command};
 
 use super::str_vec;
 
+// TODO : Add support for using entire block device
+
 pub fn format_drives(conf: &DriveConfig, encrypted: bool) {
     // EFI (BOOT)
     run_command(
diff --git a/src/install/mod.rs b/src/install/mod.rs
index 8d257f3..5f3dd4e 100644
--- a/src/install/mod.rs
+++ b/src/install/mod.rs
@@ -10,6 +10,7 @@ use boot::setup_bootloader;
 use drives::{format_drives, mount_drives};
 use first_boot::{first_boot_values, genfstab};
 use kernel::setup_mkinitcpio;
+use ollama::setup_ollama;
 use security::{setup_secure_boot, setup_tpm_unlock};
 use ssh::setup_ssh;
 use user::setup_users;
@@ -24,6 +25,7 @@ pub mod security;
 pub mod ssh;
 pub mod user;
 pub mod zram;
+pub mod ollama;
 
 use crate::{
     config::InstallConfig,
@@ -118,6 +120,10 @@ pub fn install(conf: InstallConfig) {
         // TODO : Enable docker
     }
 
+    if let Some(ai) = conf.ai {
+        setup_ollama(&ai);
+    }
+
     setup_zram();
     setup_mkinitcpio(&conf.drive);
     setup_secure_boot();
diff --git a/src/install/ollama.rs b/src/install/ollama.rs
new file mode 100644
index 0000000..389ed49
--- /dev/null
+++ b/src/install/ollama.rs
@@ -0,0 +1,24 @@
+use crate::{config::OllamaConfig, pkg::install_pkgs, run_command};
+
+use super::str_vec;
+
+
+pub fn setup_ollama(conf: &OllamaConfig) {
+    if conf.gpu {
+        install_pkgs(&["ollama-cuda"]);
+    } else {
+        install_pkgs(&["ollama"]);
+    }
+
+    run_command(&str_vec(vec![
+        "arch-chroot",
+        "/mnt",
+        "systemctl",
+        "enable",
+        "ollama.service"
+    ]), None, false);
+
+    for model in conf.models.clone().unwrap_or_default() {
+        // TODO : Pull models
+    }
+}
\ No newline at end of file