From 3834c4e0db1929aa0e91fbfd9c12be01dbf976bb Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Fri, 14 Jun 2024 11:06:45 +0200
Subject: [PATCH] plugins/cmp: add cmp_ai source

---
 plugins/completion/cmp/sources.nix            |  1 +
 plugins/completion/cmp/sources/cmp-ai.nix     | 95 +++++++++++++++++++
 plugins/completion/cmp/sources/default.nix    |  1 +
 .../plugins/completion/cmp-ai.nix             | 37 ++++++++
 .../plugins/completion/cmp-all-sources.nix    |  5 +-
 5 files changed, 138 insertions(+), 1 deletion(-)
 create mode 100644 plugins/completion/cmp/sources/cmp-ai.nix
 create mode 100644 tests/test-sources/plugins/completion/cmp-ai.nix

diff --git a/plugins/completion/cmp/sources.nix b/plugins/completion/cmp/sources.nix
index f738c89d..baed7a5c 100644
--- a/plugins/completion/cmp/sources.nix
+++ b/plugins/completion/cmp/sources.nix
@@ -5,6 +5,7 @@
   calc = "cmp-calc";
   dap = "cmp-dap";
   cmdline = "cmp-cmdline";
+  cmp_ai = "cmp-ai";
   cmp-clippy = "cmp-clippy";
   cmp-cmdline-history = "cmp-cmdline-history";
   cmp_pandoc = "cmp-pandoc-nvim";
diff --git a/plugins/completion/cmp/sources/cmp-ai.nix b/plugins/completion/cmp/sources/cmp-ai.nix
new file mode 100644
index 00000000..3754a2f9
--- /dev/null
+++ b/plugins/completion/cmp/sources/cmp-ai.nix
@@ -0,0 +1,95 @@
+{
+  lib,
+  helpers,
+  config,
+  ...
+}:
+with lib;
+let
+  cfg = config.plugins.cmp-ai;
+in
+{
+  meta.maintainers = [ maintainers.GaetanLepage ];
+
+  options.plugins.cmp-ai = {
+    settings = helpers.mkSettingsOption {
+      description = "Options provided to the `require('cmp_ai.config'):setup` function.";
+
+      options = {
+        max_lines = helpers.defaultNullOpts.mkUnsignedInt 50 ''
+          How many lines of buffer context to use.
+        '';
+
+        run_on_every_keystroke = helpers.defaultNullOpts.mkBool true ''
+          Generate new completion items on every keystroke.
+        '';
+
+        provider = helpers.defaultNullOpts.mkStr "HF" ''
+          Which AI provider to use.
+
+          Check the [README](https://github.com/tzachar/cmp-ai/blob/main/README.md) to learn about
+          available options.
+        '';
+
+        provider_options = helpers.defaultNullOpts.mkAttrsOf types.anything { } ''
+          Options to forward to the provider.
+        '';
+
+        notify = helpers.defaultNullOpts.mkBool true ''
+          As some completion sources can be quit slow, setting this to `true` will trigger a
+          notification when a completion starts and ends using `vim.notify`.
+        '';
+
+        notify_callback = helpers.defaultNullOpts.mkLuaFn' {
+          description = ''
+            The default notify function uses `vim.notify`, but an override can be configured.
+          '';
+          pluginDefault = ''
+            function(msg)
+              vim.notify(msg)
+            end
+          '';
+          example = ''
+            function(msg)
+              require('notify').notify(msg, vim.log.levels.INFO, {
+                title = 'OpenAI',
+                render = 'compact',
+              })
+            end
+          '';
+        };
+
+        ignored_file_types = helpers.defaultNullOpts.mkAttrsOf' {
+          type = types.bool;
+          description = "Which filetypes to ignore.";
+          pluginDefault = { };
+          example = {
+            lua = true;
+            html = true;
+          };
+        };
+      };
+
+      example = {
+        max_lines = 1000;
+        provider = "HF";
+        notify = true;
+        notify_callback = ''
+          function(msg)
+            vim.notify(msg)
+          end
+        '';
+        run_on_every_keystroke = true;
+        ignored_file_types = {
+          lua = true;
+        };
+      };
+    };
+  };
+
+  config = mkIf cfg.enable {
+    extraConfigLua = ''
+      require('cmp_ai.config'):setup(${helpers.toLuaObject cfg.settings})
+    '';
+  };
+}
diff --git a/plugins/completion/cmp/sources/default.nix b/plugins/completion/cmp/sources/default.nix
index 3d4b0856..662a3738 100644
--- a/plugins/completion/cmp/sources/default.nix
+++ b/plugins/completion/cmp/sources/default.nix
@@ -23,6 +23,7 @@ in
   imports = [
     ./codeium-nvim.nix
     ./copilot-cmp.nix
+    ./cmp-ai.nix
     ./cmp-fish.nix
     ./cmp-git.nix
     ./cmp-tabby.nix
diff --git a/tests/test-sources/plugins/completion/cmp-ai.nix b/tests/test-sources/plugins/completion/cmp-ai.nix
new file mode 100644
index 00000000..4c9056c0
--- /dev/null
+++ b/tests/test-sources/plugins/completion/cmp-ai.nix
@@ -0,0 +1,37 @@
+{
+  empty = {
+    # We do not provide the required HF_API_KEY environment variable.
+    tests.dontRun = true;
+
+    plugins.cmp = {
+      enable = true;
+      settings.sources = [ { name = "cmp_ai"; } ];
+    };
+  };
+
+  example = {
+    # We do not provide the required HF_API_KEY environment variable.
+    tests.dontRun = true;
+
+    plugins = {
+      cmp = {
+        enable = true;
+        settings.sources = [ { name = "cmp_ai"; } ];
+      };
+      cmp-ai.settings = {
+        max_lines = 1000;
+        provider = "HF";
+        notify = true;
+        notify_callback = ''
+          function(msg)
+            vim.notify(msg)
+          end
+        '';
+        run_on_every_keystroke = true;
+        ignored_file_types = {
+          lua = true;
+        };
+      };
+    };
+  };
+}
diff --git a/tests/test-sources/plugins/completion/cmp-all-sources.nix b/tests/test-sources/plugins/completion/cmp-all-sources.nix
index c379970b..7677d4e7 100644
--- a/tests/test-sources/plugins/completion/cmp-all-sources.nix
+++ b/tests/test-sources/plugins/completion/cmp-all-sources.nix
@@ -14,7 +14,10 @@
         settings.sources =
           with pkgs.lib;
           let
-            disabledSources = optional (pkgs.stdenv.hostPlatform.system == "aarch64-linux") "cmp_tabnine";
+            disabledSources = [
+              # We do not provide the required HF_API_KEY environment variable.
+              "cmp_ai"
+            ] ++ optional (pkgs.stdenv.hostPlatform.system == "aarch64-linux") "cmp_tabnine";
 
             filterFunc = sourceName: !(elem sourceName disabledSources);