From 3b20a6b6d734027f270a18f96b2f9bf81e81596b Mon Sep 17 00:00:00 2001
From: Kirill Radzikhovskyy <kirillrdy@gmail.com>
Date: Sat, 19 Oct 2024 10:10:06 +1100
Subject: [PATCH 1/3] python312Packages.llama-cpp-python: init at 0.3.1

---
 .../llama-cpp-python/default.nix              | 96 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 98 insertions(+)
 create mode 100644 pkgs/development/python-modules/llama-cpp-python/default.nix

diff --git a/pkgs/development/python-modules/llama-cpp-python/default.nix b/pkgs/development/python-modules/llama-cpp-python/default.nix
new file mode 100644
index 0000000000000..9a2f9458a8757
--- /dev/null
+++ b/pkgs/development/python-modules/llama-cpp-python/default.nix
@@ -0,0 +1,96 @@
+{
+  lib,
+  buildPythonPackage,
+  cmake,
+  fetchFromGitHub,
+  gitUpdater,
+  ninja,
+  pathspec,
+  pyproject-metadata,
+  pytestCheckHook,
+  pythonOlder,
+  scikit-build-core,
+
+  config,
+  cudaSupport ? config.cudaSupport,
+  cudaPackages ? { },
+
+  diskcache,
+  jinja2,
+  numpy,
+  typing-extensions,
+  scipy,
+  huggingface-hub,
+}:
+
+buildPythonPackage rec {
+  pname = "llama-cpp-python";
+  version = "0.3.1";
+  pyproject = true;
+
+  disabled = pythonOlder "3.7";
+
+  src = fetchFromGitHub {
+    owner = "abetlen";
+    repo = "llama-cpp-python";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-eO1zvNJZBE5BCnbgbh00tFIRWBCWor1lIsrLXs/HFds=";
+    fetchSubmodules = true;
+  };
+
+  dontUseCmakeConfigure = true;
+  SKBUILD_CMAKE_ARGS = lib.strings.concatStringsSep ";" (
+    lib.optionals cudaSupport [
+      "-DGGML_CUDA=on"
+      "-DCUDAToolkit_ROOT=${lib.getDev cudaPackages.cuda_nvcc}"
+      "-DCMAKE_CUDA_COMPILER=${lib.getExe cudaPackages.cuda_nvcc}"
+    ]
+  );
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    pathspec
+    pyproject-metadata
+    scikit-build-core
+  ];
+
+  buildInputs = lib.optionals cudaSupport (
+    with cudaPackages;
+    [
+      cuda_cudart # cuda_runtime.h
+      cuda_cccl # <thrust/*>
+      libcublas # cublas_v2.h
+    ]
+  );
+
+  propagatedBuildInputs = [
+    diskcache
+    jinja2
+    numpy
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+    scipy
+    huggingface-hub
+  ];
+
+  disabledTests = [
+    # tries to download model from huggingface-hub
+    "test_real_model"
+    "test_real_llama"
+  ];
+
+  pythonImportsCheck = [ "llama_cpp" ];
+
+  passthru.updateScript = gitUpdater { rev-prefix = "v"; };
+
+  meta = {
+    description = "Python bindings for llama.cpp";
+    homepage = "https://github.com/abetlen/llama-cpp-python";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ kirillrdy ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index fc2904ec43ac3..8608ab2b1c3e3 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -7324,6 +7324,8 @@ self: super: with self; {
 
   lizard = callPackage ../development/python-modules/lizard { };
 
+  llama-cpp-python = callPackage ../development/python-modules/llama-cpp-python { };
+
   llama-cloud = callPackage ../development/python-modules/llama-cloud { };
 
   llama-index = callPackage ../development/python-modules/llama-index { };

From 417356674681d03e835a5876c1043d10d85bbbc7 Mon Sep 17 00:00:00 2001
From: Kirill Radzikhovskyy <kirillrdy@gmail.com>
Date: Sat, 26 Oct 2024 06:50:47 +1100
Subject: [PATCH 2/3] python312Packages.llama-cpp-python: add passthru.test to
 build with CUDA support

---
 pkgs/development/python-modules/llama-cpp-python/default.nix | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pkgs/development/python-modules/llama-cpp-python/default.nix b/pkgs/development/python-modules/llama-cpp-python/default.nix
index 9a2f9458a8757..30eaf7ee6b773 100644
--- a/pkgs/development/python-modules/llama-cpp-python/default.nix
+++ b/pkgs/development/python-modules/llama-cpp-python/default.nix
@@ -10,6 +10,7 @@
   pytestCheckHook,
   pythonOlder,
   scikit-build-core,
+  llama-cpp-python,
 
   config,
   cudaSupport ? config.cudaSupport,
@@ -86,6 +87,7 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "llama_cpp" ];
 
   passthru.updateScript = gitUpdater { rev-prefix = "v"; };
+  passthru.tests.llama-cpp-python = llama-cpp-python.override { cudaSupport = true; };
 
   meta = {
     description = "Python bindings for llama.cpp";

From c9ca8b6f54fc982e6a5ac68576a84989e169745e Mon Sep 17 00:00:00 2001
From: Kirill Radzikhovskyy <kirillrdy@gmail.com>
Date: Sun, 27 Oct 2024 04:45:16 +1100
Subject: [PATCH 3/3] python312Packages.llama-cpp-python: use stdenv from
 cudaPackages

---
 .../python-modules/llama-cpp-python/default.nix       | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/llama-cpp-python/default.nix b/pkgs/development/python-modules/llama-cpp-python/default.nix
index 30eaf7ee6b773..4c0391858e4aa 100644
--- a/pkgs/development/python-modules/llama-cpp-python/default.nix
+++ b/pkgs/development/python-modules/llama-cpp-python/default.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   cmake,
   fetchFromGitHub,
@@ -23,14 +24,18 @@
   scipy,
   huggingface-hub,
 }:
-
-buildPythonPackage rec {
-  pname = "llama-cpp-python";
+let
   version = "0.3.1";
+in
+buildPythonPackage {
+  pname = "llama-cpp-python";
+  inherit version;
   pyproject = true;
 
   disabled = pythonOlder "3.7";
 
+  stdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv;
+
   src = fetchFromGitHub {
     owner = "abetlen";
     repo = "llama-cpp-python";