org-ai-md.el

;;; org-ai-md.el --- Description -*- lexical-binding: t; -*-
;;
;; Copyright (C) 2023 Manuel Odendahl
;;
;; Author: Manuel Odendahl <wesen@ruinwesen.com>
;; Maintainer: Manuel Odendahl <wesen@ruinwesen.com>
;; Created: July 06, 2023
;; Modified: July 06, 2023
;; Version: 0.0.1
;; Keywords: abbrev bib c calendar comm convenience data docs emulations extensions faces files frames games hardware help hypermedia i18n internal languages lisp local maint mail matching mouse multimedia news outlines processes terminals tex tools unix vc wp
;; Homepage: https://github.com/wesen/md
;; Package-Requires: ((emacs "27.1"))
;;
;; This file is not part of GNU Emacs.
;;
;;; Commentary:
;;
;;  Description
;;
;;; Code:


(require 's)
(require 'seq)
(require 'org-element)
(require 'org-ai)

(defvar org-ai-md-markdown-code-regexp "^```\\(.*\\)\n\\(\\(.\\|\n\\)*?\\)```")

(defun org-ai-md-get-match (s location)
  "S is a pair (START . END). GET-MATCH returns the string between START and END."
  (let ((start (car location))
        (end (cdr location)))
    (substring s start end)))

(defun org-ai-md-get-regexp-matches (s subexp-level)
  "ORG-AI-MD-GET-REGEXP-MATCHES call S-MATCHED-POSITIONS-ALL with the given SUBEXP-LEVEL
and returns the substrings of the matches."
  (let ((matches (s-matched-positions-all org-ai-md-markdown-code-regexp s subexp-level)))
    (print matches)
    (mapcar (lambda (match)
              (org-ai-md-get-match s match))
            matches)))

(defun org-ai-md-get-code-blocks (s)
  "MD-GET-CODE-BLOCKS returns a list of (LANGUAGE . CODE) pairs."
  ;; get all matches at subexp level 1 (language) and 2 (code)
  (let ((languages (org-ai-md-get-regexp-matches s 1))
        (codes (org-ai-md-get-regexp-matches s 2)))
    ;; zip the 2 lists together
        (seq-mapn (lambda (language code)
                    (cons (s-trim language) code))
                  languages
                  codes)))

(defun org-ai-md-find-special-block (element block-type)
  "Finds the first special block of type BLOCK-TYPE in ELEMENT."
  (let ((type (org-element-type element))
        (block-type- (org-element-property :type element)))
    (if (eq type 'special-block)
        (if (string= (s-upcase block-type-) (s-upcase block-type))
            element
          nil)
      (let ((parent (org-element-property :parent element)))
        (when parent (org-ai-md-find-special-block parent block-type))))))

(defun org-ai-md-get-org-block-content ()
  "Gets the content of the surrounding #+BEGIN_AI / #+END_AI block."
  (save-excursion
    (let* ((ai-block (org-ai-md-find-special-block (org-element-at-point) "AI")))
      (if ai-block
          (buffer-substring (org-element-property :contents-begin ai-block)
                            (org-element-property :contents-end ai-block))
        ""))))

(defun org-ai-md-src-org-preamble (language)
  "If LANGUAGE is of the format LANG (NAME) then output a named source block,
otherwise just a normal source block."
  (let ((matches (s-match "^\\(.*\\) (\\(.*\\))$" language)))
    (if matches
        (format "#+NAME: %s\n#+BEGIN_SRC %s" (nth 2 matches) (nth 1 matches))
      (format "#+BEGIN_SRC %s" language))))

(defun org-ai-md-extract-source-blocks ()
  (interactive)
  (save-excursion
    (let* ((org-block-content (org-ai-md-get-org-block-content))
           (code-blocks (org-ai-md-get-code-blocks org-block-content)))
      ;; Go to the end of buffer
      (goto-char (point-max))
      (newline)
      ;; Insert each block as an org source block
      (seq-each (lambda (code-block)
                  (insert (format "%s\n%s\n#+END_SRC\n\n"
                                  (org-ai-md-src-org-preamble (car code-block))
                                  (s-trim-right (cdr code-block)))))
                code-blocks))))

(defvar org-ai-md--role-to-md
  '((user . "**ME**:")
    (system . "**SYS**:")
    (assistant . "**AI**:")))

(defvar org-ai-md--role-to-token
  '((user . "[ME]:")
    (system . "[SYS]:")
    (assistant . "[AI]:")))

(defun org-ai-md--get-role-markdown (role &optional role-to-md)
  (let ((role-to-md (unless role-to-md org-ai-md--role-to-md)))
    (alist-get role role-to-md)))

(defun org-ai-md--block-content-to-md (content &optional role-to-md)
  "Converts CONTENT to markdown. CONTENT is a list of chat messages of the format
(:role ROLE :content \"foobar\"). The output is a blockquoted block with
[ROLE] CONTENT."
  (let ((content
         (seq-map (lambda (message)
                    (let ((m  (format "> %s %s"
                                      ;; lookup role in role-to-md
                                      (org-ai-md--get-role-markdown (plist-get message :role) role-to-md)
                                      (plist-get message :content))))
                      (s-join "\n> " (s-split "\n" m))))
                  content)))
    (s-join "\n\n" content)))

(defun org-ai-md--remove-until (pred list &optional include)
  "Removes elements from LIST until PRED is true. If INCLUDE is true, then the
first element for which PRED is true is also removed."
  (let ((l (seq-copy list)))
    (while (not (funcall pred (seq-first l)))
      (setq l (seq-subseq l 1)))
    (if include (seq-subseq l 1) l)))

(defun org-ai-md--messages-to-org-ai (messages)
  "Converts MESSAGES to org-ai format."
  (let ((messages (seq-copy messages)))
    (s-join "\n\n" (seq-map (lambda (message)
                            (format "%s %s" (alist-get (plist-get message :role) org-ai-md--role-to-token)
                                    (plist-get message :content)))
                          messages))))

(defun org-ai-md--strip-last-assistant (content)
  "Strips the last assistant message from CONTENT, returning a list of messages."
  (let* ((messages (nreverse (org-ai--collect-chat-messages content))))
    (nreverse (org-ai-md--remove-until (lambda (message)
                                         (not (eq (plist-get message :role) 'assistant)))
                                       messages
                                       nil))))

(defun org-ai-md--insert-block-after (content &optional org-mode-header)
  "Inserts CONTENT as a new org-ai block after the current block. It also
adds the necessary header entries from ORG-MODE-HEADER."
  (let* ((block (org-ai-md-find-special-block (org-element-at-point) "AI"))
         (block-end (org-element-property :end block))
         (header-string (if org-mode-header
                            (s-join " " (seq-map (lambda (header)
                                                    (format "%s %s" (car header) (cdr header)))
                                                  org-mode-header))
                          "")))
    (progn (goto-char block-end)
           (insert (format "\n\n#+begin_ai %s\n%s\n\n#+end_ai\n\n" header-string content)))))

(defun org-ai-md-regenerate ()
  "Takes the current org-ai block, strips the last assistant entry, and reinserts
it as a new org-ai block after the current block."
  (interactive)
  (let* ((content (org-ai-get-block-content))
         ;; BLOCK-INFO is an alist of key-value pairs of the org-mode headers
         (block-info (org-ai-get-block-info))
         (messages (org-ai-md--strip-last-assistant content))
         (ai (org-ai-md--messages-to-org-ai messages)))
    (progn (org-ai-md--insert-block-after ai block-info)
           (org-previous-block 1)
           (goto-char (org-element-property :end (org-element-at-point)))
           (previous-line 3)
           (message "Regenerated block."))))

(defun org-ai-md-current-block-to-md ()
  "Convert the current block to markdown and stores it in the kill ring."
  (interactive)
  (let* ((content (org-ai--collect-chat-messages (org-ai-get-block-content)))
         (md (org-ai-md--block-content-to-md content)))
    (progn (kill-new md)
           (message "Copied to kill ring: %s" md))))

(provide 'org-ai-md)
;;; org-ai-md.el ends here