Add aimenu for AI-powered imenu outlines

This commit introduces aimenu, a new package that generates imenu-like outlines using a language model. The new package includes functions for generating and handling outlines, managing cache, and interacting with the user to select headers.
This commit is contained in:
jdormit 2024-08-11 22:55:35 -04:00 committed by Jeremy Dormitzer
parent c622de4d80
commit 71861b026a
2 changed files with 225 additions and 0 deletions

View File

@ -216,6 +216,15 @@ Assistant: Write a short story about a dragon who discovers a hidden talent that
(embark-general-map (embark-general-map
"?" #'gptel-quick)) "?" #'gptel-quick))
(use-package aimenu
:straight `(:local-repo ,(expand-file-name "packages/aimenu" user-emacs-directory))
:defer t
:config
(setq aimenu-backend gptel-backend-ollama
aimenu-model "gemma2:2b")
:general
(ai-map "i" #'aimenu))
(defun gptel-commit-message () (defun gptel-commit-message ()
"Generate a commit message via gptel." "Generate a commit message via gptel."
(interactive) (interactive)

View File

@ -0,0 +1,216 @@
;;; aimenu.el --- imenu-like outline generated by an LLM -*- lexical-binding: t; -*-
;; Copyright (C) 2024 Jeremy Dormitzer
;; Author: Jeremy Dormitzer <jeremy.dormitzer@gmail.com>
;; Keywords: tools
;; Package-Requires: ((emacs "25.1") (gptel "0.9.0"))
;; 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 imenu-like outline interface for any buffer, powered by an LLM.
;;; Code:
(require 'json)
(defvar aimenu-gptel-backend (gptel-make-openai "aimenu-openai"
:key #'gptel-api-key-from-auth-source
:models '("gpt-4o-mini"))
"The gptel backend to use for aimenu requests.")
(defvar aimenu-gptel-model "gpt-4o-mini"
"The gptel model to use for aimenu requests.")
(defvar aimenu-outline-cache (make-hash-table :test 'equal)
"Cache for storing outline responses based on prompt hash.")
;;;###autoload
(defun aimenu-bust-cache ()
"Bust the aimenu cache."
(interactive)
(clrhash aimenu-outline-cache))
(defun aimenu-hash-string (str)
"Return a hash for the given string STR using SHA-256."
(secure-hash 'sha256 str))
(defun aimenu-get-buffer-with-line-numbers (buffer)
"Return the contents of BUFFER with line numbers added to each line."
(with-temp-buffer
(insert-buffer-substring buffer)
(goto-char (point-min))
(let ((line-number 1))
(while (not (eobp))
(insert (format "%d: " line-number))
(forward-line 1)
(setq line-number (1+ line-number))))
(buffer-string)))
(defun aimenu-parse-and-strip-response (response)
"Parse and strip markdown code fences from LLM RESPONSE."
(with-temp-buffer
(insert response)
(goto-char (point-min))
(while (re-search-forward "```\\(json\\)*" nil t)
(replace-match ""))
(buffer-string)))
(defun aimenu-header-completions (outline headers buffer)
(lambda (string pred action)
(if (eq action 'metadata)
(list 'metadata
(cons 'display-sort-function
(lambda (headers)
(sort headers
(lambda (a b) (< (alist-get a outline) (alist-get b outline))))))
(cons 'annotation-function
(lambda (header)
(let* ((line-number (alist-get header outline nil nil #'equal))
(line (with-current-buffer buffer
(save-excursion
(goto-char (point-min))
(forward-line line-number)
(thing-at-point 'line)))))
(format "%s%s: %s"
(propertize " " 'display '(space :align-to center))
line-number
(string-trim line))))))
(complete-with-action action headers string pred))))
(defun aimenu-handle-outline-response (outline)
"Handle the outline response by prompting the user to select a header.
OUTLINE is the parsed JSON response."
(let* ((current-buffer (current-buffer))
(headers (mapcar #'car outline))
(chosen-header (completing-read "Choose an outline header: " (aimenu-header-completions outline headers current-buffer)))
(line-number (alist-get chosen-header outline nil nil #'equal)))
(funcall-interactively #'goto-line line-number)))
;;;###autoload
(defun aimenu ()
"Generate an outline of the current buffer using an LLM."
(interactive)
(let* ((buffer-contents (aimenu-get-buffer-with-line-numbers (current-buffer)))
(prompt-hash (aimenu-hash-string buffer-contents))
(cached-response (gethash prompt-hash aimenu-outline-cache))
(gptel-backend aimenu-gptel-backend)
(gptel-model aimenu-gptel-model))
(if cached-response
(aimenu-handle-outline-response cached-response)
(gptel-request
buffer-contents
:system "Create an outline of the buffer. Return a JSON object where the keys are the outline headers and the values are the line numbers that correspond to each outline header. Reply only in valid JSON without any code fences or additional text.
Here are some examples of your task:
USER:
1: * Outline header 1
2: Some content here.
3: ** Subheader 1.1
4: More content.
5: ** Subheader 1.2
6: Even more content.
7: * Outline header 2
8: And some more content.
9: ** Subheader 2.1
10: Still more content.
11: *** Outline header 3
12: Final bit of content.
ASSISTANT:
{
\"Outline header 1\": 1,
\"Subheader 1.1\": 3,
\"Subheader 1.2\": 5,
\"Outline header 2\": 7,
\"Subheader 2.1\": 9,
\"Outline header 3\": 11
}
USER:
1: # Main Header
2: Some introductory text.
3: ## Subheader 1
4: More details here.
5: ## Subheader 2
6: Additional information.
7: ### Sub-subheader 2.1
8: Finer details.
ASSISTANT:
{
\"Main Header\": 1,
\"Subheader 1\": 3,
\"Subheader 2\": 5,
\"Sub-subheader 2.1\": 7
}
USER:
1: def main_function():
2: # Main function logic
3: pass
4:
5: def helper_function():
6: # Helper function logic
7: pass
ASSISTANT:
{
\"main_function\": 1,
\"helper_function\": 5
}
USER:
1: main_section:
2: key1: value1
3: key2: value2
4: sub_section1:
5: sub_key1: sub_value1
6: sub_section2:
7: sub_key2: sub_value2
8: sub_sub_section:
9: sub_sub_key: sub_sub_value
ASSISTANT:
{
\"main_section\": 1,
\"sub_section1\": 4,
\"sub_section2\": 6,
\"sub_sub_section\": 8
}"
:callback (lambda (response info)
(if response
(let* ((stripped-response (aimenu-parse-and-strip-response response))
(response-json (ignore-errors (let ((json-key-type 'string))
(json-read-from-string stripped-response))))
(outline (if (json-alist-p response-json)
response-json
(error "Invalid response format: %s" stripped-response)))
(buffer (plist-get info :buffer)))
(puthash (aimenu-hash-string buffer-contents)
outline
aimenu-outline-cache)
(with-current-buffer buffer
(aimenu-handle-outline-response outline))
(pop-to-buffer buffer))
(message "gptel-request failed with message: %s" (plist-get info :status))))))))
(provide 'aimenu)
;;; aimenu.el ends here