From 8c1e1a47d300bdbfbd32417b5fd3e891c8856c7a Mon Sep 17 00:00:00 2001
From: uttarayan21 <email@uttarayan.me>
Date: Tue, 26 Nov 2024 16:08:13 +0530
Subject: [PATCH] feat: Added support for getting aichat api key from commands
 and files

---
 common/home.nix    | 16 ++------
 modules/aichat.nix | 98 ++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 86 insertions(+), 28 deletions(-)

diff --git a/common/home.nix b/common/home.nix
index adda95d5..749fc785 100644
--- a/common/home.nix
+++ b/common/home.nix
@@ -231,7 +231,6 @@ in {
       shellAliases =
         {
           g = "git";
-          aichat = "op plugin run -- aichat";
         }
         // lib.optionalAttrs pkgs.stdenv.isLinux {
           kmpv = "mpv --vo-kitty-use-shm=yes --vo=kitty --really-quiet";
@@ -244,15 +243,6 @@ in {
       # ${pkgs.spotify-player}/bin/spotify_player generate fish | source
       interactiveShellInit = ''
         ${pkgs.pfetch-rs}/bin/pfetch
-        function _aichat_fish
-            set -l _old (commandline)
-            if test -n $_old
-                echo -n "⌛"
-                commandline -f repaint
-                commandline (aichat -e $_old)
-            end
-        end
-        bind \co _aichat_fish
         ${lib.optionalString (device.isLinux && !device.isNix) "source /etc/profile.d/nix-daemon.fish"}
       '';
     };
@@ -352,20 +342,22 @@ in {
     home-manager = {enable = true;};
     aichat = {
       enable = true;
+      enableFishIntegration = true;
       settings = {
         clients = [
           {
             type = "openai-compatible";
             name = "llama";
             api_base = "https://llama.darksailor.dev/api/v1";
+            api_key_cmd = "op item get llama-api --fields label=credential --reveal";
             models = [
               {
-                name = "chat";
+                name = "qwen_2_5_1";
               }
             ];
           }
         ];
-        model = "llama:chat";
+        model = "llama:qwen_2_5_1";
       };
     };
   };
diff --git a/modules/aichat.nix b/modules/aichat.nix
index fc8b92d8..a2873b44 100644
--- a/modules/aichat.nix
+++ b/modules/aichat.nix
@@ -7,21 +7,71 @@
 with lib; let
   cfg = config.programs.aichat;
   yamlFormat = pkgs.formats.yaml {};
-  # configDir =
-  #   if pkgs.stdenv.isDarwin then
-  #     "${config.home.homeDirectory}Library/Application Support/aichat"
-  #   else
-  #     "${config.xdg.configHome}/aichat";
+  fishIntegration = ''
+    function _aichat_fish
+        set -l _old (commandline)
+        if test -n $_old
+            echo -n "⌛"
+            commandline -f repaint
+            commandline (aichat -e $_old)
+        end
+    end
+    bind \co _aichat_fish'';
+  nuIntegration = ''
+    def _aichat_nushell [] {
+        let _prev = (commandline)
+        if ($_prev != "") {
+            print '⌛'
+            commandline edit -r (aichat -e $_prev)
+        }
+    }
+
+    $env.config.keybindings = ($env.config.keybindings | append {
+            name: aichat_integration
+            modifier: control
+            keycode: char_o
+            mode: [emacs, vi_insert]
+            event:[
+                {
+                    send: executehostcommand,
+                    cmd: "_aichat_nushell"
+                }
+            ]
+        }
+    )
+  '';
+  bashIntegration = ''
+    _aichat_bash() {
+        if [[ -n "$READLINE_LINE" ]]; then
+            READLINE_LINE=$(aichat -e "$READLINE_LINE")
+            READLINE_POINT=''${#READLINE_LINE}
+        fi
+    }
+    bind -x '"\co": _aichat_bash'
+  '';
+  zshIntegration = ''
+    _aichat_zsh() {
+        if [[ -n "$BUFFER" ]]; then
+            local _old=$BUFFER
+            BUFFER+="⌛"
+            zle -I && zle redisplay
+            BUFFER=$(aichat -e "$_old")
+            zle end-of-line
+        fi
+    }
+    zle -N _aichat_zsh
+    bindkey '\co' _aichat_zsh
+  '';
 in {
   options = {
     programs.aichat = {
       enable = mkEnableOption "aichat";
-      package = mkOption {
-        type = types.package;
-        default = pkgs.aichat;
-        defaultText = literalExpression "pkgs.aichat";
-        description = "The aichat package to install.";
-      };
+      package = mkPackageOption pkgs "aichat" {};
+
+      enableFishIntegration = mkEnableOption "Fish integration" // {default = true;};
+      enableBashIntegration = mkEnableOption "Bash integration" // {default = true;};
+      enableZshIntegration = mkEnableOption "Zsh integration" // {default = true;};
+      enableNushellIntegration = mkEnableOption "Nushell integration" // {default = true;};
 
       settings = lib.mkOption {
         type = yamlFormat.type;
@@ -30,13 +80,29 @@ in {
     };
   };
 
-  config = {
-    home.packages = mkIf cfg.enable [cfg.package];
+  config = let
+    api_key_files = concatStringsSep " " (builtins.map (client: ''--run 'export ${lib.toUpper client.name}_API_KEY=`cat -v ${client.api_key_file}`' '') (builtins.filter (client: (builtins.hasAttr "api_key_file" client)) cfg.settings.clients));
+    api_key_cmds = concatStringsSep " " (builtins.map (client: ''--run 'export ${lib.toUpper client.name}_API_KEY=`${client.api_key_cmd}`' '') (builtins.filter (client: (builtins.hasAttr "api_key_cmd" client)) cfg.settings.clients));
+
+    aichat-wrapped = pkgs.symlinkJoin {
+      name = "aichat";
+      paths = [
+        cfg.package
+      ];
+      buildInputs = [pkgs.makeWrapper];
+      postBuild = ''
+        wrapProgram $out/bin/aichat ${api_key_files} ${api_key_cmds}
+      '';
+    };
+  in {
+    home.packages = mkIf cfg.enable [aichat-wrapped];
+
+    programs.fish.interactiveShellInit = mkIf cfg.enableFishIntegration fishIntegration;
+    programs.bash.initExtra = mkIf cfg.enableBashIntegration bashIntegration;
+    programs.zsh.initExtra = mkIf cfg.enableZshIntegration zshIntegration;
+    programs.nushell.extraConfig = mkIf cfg.enableNushellIntegration nuIntegration;
 
     xdg.configFile."aichat/config.yaml".source =
       yamlFormat.generate "config.yaml" cfg.settings;
-    # xdg.configFile = mkIf cfg.enable {
-    #   # "aichat/config.yaml".text = generators.toYAML {} cfg.settings;
-    # };
   };
 }