From 1f1cf3a6be3842af9f07cc6937dd6683511e7ddd Mon Sep 17 00:00:00 2001
From: Daylin Morgan <me@dayl.in>
Date: Thu, 8 Aug 2024 13:53:51 -0500
Subject: [PATCH] add llm package + plugins

---
 hosts/othalan/default.nix                 |  1 +
 hosts/othalan/system.nix                  |  6 +---
 modules/llm/default.nix                   | 24 ++++++++++++++++
 modules/oizys.nix                         |  2 ++
 pkgs/llm-plugins/llm-claude-3/default.nix | 34 +++++++++++++++++++++++
 pkgs/llm-plugins/llm-ollama/default.nix   | 33 ++++++++++++++++++++++
 6 files changed, 95 insertions(+), 5 deletions(-)
 create mode 100644 modules/llm/default.nix
 create mode 100644 pkgs/llm-plugins/llm-claude-3/default.nix
 create mode 100644 pkgs/llm-plugins/llm-ollama/default.nix

diff --git a/hosts/othalan/default.nix b/hosts/othalan/default.nix
index b8d15f0..ade5b61 100644
--- a/hosts/othalan/default.nix
+++ b/hosts/othalan/default.nix
@@ -22,6 +22,7 @@
       "tex"
       "zig"
     ];
+    llm = enabled;
   };
 
   services.restic.backups.gdrive = {
diff --git a/hosts/othalan/system.nix b/hosts/othalan/system.nix
index 9bf3841..97997cd 100644
--- a/hosts/othalan/system.nix
+++ b/hosts/othalan/system.nix
@@ -1,8 +1,4 @@
-{
-  pkgs,
-  enabled,
-  ...
-}:
+{ pkgs, enabled, ... }:
 {
   networking.networkmanager = enabled;
   services.fwupd = enabled;
diff --git a/modules/llm/default.nix b/modules/llm/default.nix
new file mode 100644
index 0000000..7a72f0f
--- /dev/null
+++ b/modules/llm/default.nix
@@ -0,0 +1,24 @@
+{
+  pkgs,
+  config,
+  enabled,
+  mkOizysModule,
+  ...
+}:
+
+let
+  inherit (pkgs) python3Packages;
+  llm-ollama = python3Packages.callPackage ../../pkgs/llm-plugins/llm-ollama { };
+  llm-claude3 = python3Packages.callPackage ../../pkgs/llm-plugins/llm-claude-3 { };
+  llm = (
+    pkgs.llm.withPlugins [
+      llm-ollama
+      llm-claude3
+    ]
+  );
+in
+
+mkOizysModule config "llm" {
+  services.ollama = enabled;
+  environment.systemPackages = [ llm ];
+}
diff --git a/modules/oizys.nix b/modules/oizys.nix
index 516d5a2..3ed2589 100644
--- a/modules/oizys.nix
+++ b/modules/oizys.nix
@@ -38,6 +38,8 @@ in
 
     nix-ld
     restic
+
+    llm
   ];
 
   options.oizys = {
diff --git a/pkgs/llm-plugins/llm-claude-3/default.nix b/pkgs/llm-plugins/llm-claude-3/default.nix
new file mode 100644
index 0000000..28e0804
--- /dev/null
+++ b/pkgs/llm-plugins/llm-claude-3/default.nix
@@ -0,0 +1,34 @@
+{
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+  wheel,
+
+  # deps
+  anthropic,
+  ...
+}:
+
+buildPythonPackage rec {
+  pname = "llm-claude-3";
+  version = "0.4";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "simonw";
+    repo = "llm-claude-3";
+    rev = version;
+    hash = "sha256-5qF5BK319PNzB4XsLdYvtyq/SxBDdHJ9IoKWEnvNRp4=";
+  };
+
+  nativeBuildInputs = [
+    setuptools
+    wheel
+  ];
+
+  dependencies = [ anthropic ];
+
+  dontCheckRuntimeDeps = true;
+}
diff --git a/pkgs/llm-plugins/llm-ollama/default.nix b/pkgs/llm-plugins/llm-ollama/default.nix
new file mode 100644
index 0000000..09b1e5e
--- /dev/null
+++ b/pkgs/llm-plugins/llm-ollama/default.nix
@@ -0,0 +1,33 @@
+{
+  buildPythonPackage,
+  fetchPypi,
+
+  # build-system
+  setuptools,
+
+  # deps
+  ollama,
+  pydantic,
+
+  ...
+}:
+buildPythonPackage rec {
+  pname = "llm-ollama";
+  version = "0.5.0";
+  pyproject = true;
+  src = fetchPypi {
+    inherit version;
+    pname = "llm_ollama";
+    hash = "sha256-M3FF9fAZ2rr+toKoz/rLRPZxB7LIHqmZQXJBdKR4fVk=";
+  };
+
+  dependencies = [
+    ollama
+    pydantic
+  ];
+
+  build-system = [ setuptools ];
+
+  # will only be used in environment with llm installed
+  dontCheckRuntimeDeps = true;
+}