dotfiles/emacs/.emacs.d/packages/llm/llm.el

;;; llm.el --- An Emacs interface to the LLM command-line tool  -*- lexical-binding: t; -*-

;; Copyright (C) 2024  Jeremy Isaac Dormitzer

;; Author: Jeremy Isaac Dormitzer <jeremydormitzer@hummingbird.co>
;; Version: 0.1
;; Package-Requires: ((emacs "24.3") (s "1.13") (markdown-mode "2.7"))
;; Keywords: tools

;; This program is free software; you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.

;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.

;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:

;; This package provides an Emacs interface to the LLM command-line tool.

;;; Code:
(require 's)
(require 'markdown-mode)

(defcustom llm-executable "llm"
  "Path to the llm executable."
  :type 'string
  :group 'llm)

(defcustom llm-model nil
  "The llm model to use."
  :type 'string
  :group 'llm)

(defcustom llm-max-tokens 5000
  "The maximum number of tokens to generate."
  :type 'integer
  :group 'llm)

(defun llm--ensure-executable ()
  "Ensure that the llm executable is available."
  (unless (executable-find llm-executable)
    (error
     "llm executable not found: see https://llm.datasette.io/en/stable/index.html for installation instructions")))

(defun llm--process-filter (proc string)
  (let* ((buffer (process-buffer proc))
         (window (get-buffer-window buffer))
         (string (replace-regexp-in-string "\r\n" "\n" string)))
    (when (buffer-live-p buffer)
      (with-current-buffer buffer
        (if (not (mark)) (push-mark))
        (exchange-point-and-mark) ;Use the mark to represent the cursor location
        (dolist (char (append string nil))
          (cond ((char-equal char ?\r)
                 (move-beginning-of-line 1))
                ((char-equal char ?\n)
                 (move-end-of-line 1) (newline))
                (t
                 (if (/= (point) (point-max)) ;Overwrite character
                     (delete-char 1))
                 (insert char))))
        (exchange-point-and-mark)))
    (if window
        (with-selected-window window
          (goto-char (point-max))))))

(define-derived-mode llm-mode markdown-mode "llm"
  "Major mode for LLM output.")

(define-key llm-mode-map
            (kbd "q") #'quit-window)

(when (fboundp #'evil-define-key)
  (evil-define-key 'normal llm-mode-map
    (kbd "q") #'quit-window))

(defun llm--run-async-process-sentinal (proc string)
  (with-current-buffer (process-buffer proc)
    (goto-char (point-max))
    (newline)
    (newline)
    (insert (format "[llm %s]" (s-trim string)))))

(defun llm--run-async (name buffer-name &rest llm-args)
  "Run llm with LLM-ARGS asynchronously.

The process is named NAME and runs in BUFFER-NAME."
  (llm--ensure-executable)
  (when-let ((existing-buffer (get-buffer buffer-name)))
    (kill-buffer existing-buffer))
  (let ((proc (make-process :name name
                            :buffer buffer-name
                            :command (cons llm-executable llm-args)
                            :filter #'llm--process-filter)))
    (with-current-buffer (process-buffer proc)
      (llm-mode))
    (set-process-sentinel proc #'llm--run-async-process-sentinal)))

(cl-defun llm--prompt-args (&key prompt system-prompt options extra-args)
  "Construct the arguments to prompt LLM with PROMPT."
  (let* ((opts (-mapcat (lambda (pair)
                          (list "-o" (car pair) (cdr pair)))
                        options))
         (sys (when system-prompt
                (list "-s" system-prompt)))
         (model (when llm-model
                  (list "-m" llm-model))))
    (append (list "prompt") model sys opts extra-args (list prompt))))

(cl-defun llm--prompt-async (&key prompt system-prompt options extra-args name buffer-name)
  "Prompt LLM asynchronously with PROMPT and other options."
  (let* ((name (or name "llm-prompt"))
         (buffer-name (or buffer-name (format "*%s*" name)))
         (args (llm--prompt-args :prompt prompt
                                 :system-prompt system-prompt
                                 :options options
                                 :extra-args extra-args)))
    (apply #'llm--run-async name buffer-name args)))

;;;###autoload
(cl-defun llm-call (callback &rest llm-args)
  "Call llm with LLM-ARGS and call CALLBACK with the result."
  (when-let ((buf (get-buffer " *llm-call*")))
    (kill-buffer buf))
  (let ((proc (apply #'llm--run-async "llm-call" " *llm-call*" llm-args)))
    (set-process-sentinel proc
                          (lambda (proc event)
                            (unless (string= event "finished\n")
                              (error "llm-call failed: %s" (s-trim event)))
                            (with-current-buffer (process-buffer proc)
                              (goto-char (point-min))
                              (funcall callback (s-trim
                                                 (buffer-substring-no-properties
                                                  (point)
                                                  (point-max)))))
                            (kill-buffer (process-buffer proc))))))

;;;###autoload
(defun llm-set-model (model)
  "Set the LLM model to MODEL."
  (interactive (list (let* ((model-strings
                             (split-string (shell-command-to-string
                                            (format "%s models" (executable-find llm-executable)))
                                           "\n" t " "))
                            (models (mapcar
                                     (lambda (s)
                                       (cons s
                                             (cadr
                                              (s-match ".*?: \\(.*?\\)\\(?:[[:blank:]]\\|$\\)" s))))
                                     model-strings))
                            (selected (completing-read "Model: " models)))
                       (alist-get selected models nil nil #'equal))))
  (setq llm-model model))

(defvar llm-model-options-alist
  `(("Meta-Llama-3-8B-Instruct" . (("max_tokens" . ,(number-to-string llm-max-tokens)))))
  "Alist mapping model names to options to pass to llm.")

(defun llm--model-options (&optional model)
  "Get the extra arguments for MODEL."
  (let ((model (or model llm-model)))
    (alist-get model llm-model-options-alist nil nil #'equal)))

;;;###autoload
(defun llm-prompt (query &optional system-prompt)
  "Prompt llm with the QUERY and optionally SYSTEM-PROMPT."
  (interactive (list (read-string "Query: " nil nil)
                     (when current-prefix-arg
                       (read-string "System prompt: " nil nil))))
  (llm--prompt-async :prompt query :options (llm--model-options))
  (switch-to-buffer "*llm-prompt*"))

;;;###autoload
(defun llm-prompt-buffer (system-prompt)
  "Prompt llm with the contents of the current buffer and the SYSTEM-PROMPT."
  (interactive "sSystem prompt: ")
  (llm--prompt-async :prompt (buffer-substring-no-properties (point-min) (point-max))
                     :system-prompt system-prompt
                     :options (llm--model-options)
                     :name "llm-prompt-buffer"
                     :buffer-name "*llm-prompt-buffer*")
  (switch-to-buffer "*llm-prompt-buffer*"))

(defun llm-prompt-region (system-prompt)
  "Prompt llm with the contents of the region and the SYSTEM-PROMPT."
  (interactive "sSystem prompt: ")
  (llm--prompt-async :prompt (buffer-substring-no-properties (region-beginning) (region-end))
                     :system-prompt system-prompt
                     :options (llm--model-options)
                     :name "llm-prompt-region"
                     :buffer-name "*llm-prompt-region*")
  (switch-to-buffer "*llm-prompt-region*"))

(defvar llm-chat-mode-map
  (make-sparse-keymap)
  "Keymap for `llm-chat-mode'.")

(defvar llm-chat-prompt-regexp "^> "
  "Regexp to match the prompt in `llm-chat-mode'.")

(define-derived-mode llm-chat-mode comint-mode "llm-chat"
  "Major mode for chatting with llm."
  (setq comint-prompt-regexp llm-chat-prompt-regexp)
  (setq comint-prompt-read-only t)
  (setq comint-process-echoes t))

(cl-defun llm--chat-args (&key system-prompt options)
  (let ((opts (-mapcat (lambda (pair)
                         (list "-o" (car pair) (cdr pair)))
                       options))
        (sys (when system-prompt
               (list "-s" system-prompt)))
        (model (when llm-model
                 (list "-m" llm-model))))
    (append (list "chat") model sys opts)))

;;;###autoload
(defun llm-chat (system-prompt &optional name)
  "Start a chat session with llm, prompting it with SYSTEM-PROMPT, naming the process and buffer NAME."
  (interactive (list (read-string "System prompt: " "You are a helpful AI assistant running inside the Emacs text editor.")
                     "llm-chat"))
  (let* ((name (or name "llm-chat"))
         (buffer-name (format "*%s*" name))
         (buffer (get-buffer-create buffer-name))
         (proc-alive (comint-check-proc buffer))
         (process (get-buffer-process buffer)))
    (unless proc-alive
      (with-current-buffer buffer
        (apply #'make-comint-in-buffer
               name
               buffer
               llm-executable
               nil
               (llm--chat-args :system-prompt system-prompt
                               :options (llm--model-options)))
        (llm-chat-mode)))
    (when buffer
      (pop-to-buffer buffer))))

;;;###autoload
(defun llm-doctor ()
  "Start a psychotherapy session with llm."
  (interactive)
  (llm-chat "You are an empathetic therapist." "llm-doctor"))

(provide 'llm)
;;; llm.el ends here