-
-
Notifications
You must be signed in to change notification settings - Fork 58
/
org-ai-oobabooga.el
267 lines (225 loc) · 11.3 KB
/
org-ai-oobabooga.el
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
;;; org-ai-oobabooga.el --- Integration of oobabooga/text-generation-webui -*- lexical-binding: t; -*-
;; This file is NOT part of GNU Emacs.
;; org-ai.el 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.
;; org-ai.el 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 org-ai.el.
;; If not, see <https://www.gnu.org/licenses/>.
;;; Commentary:
;; This integrates org-ai with the API provided by
;; https://github.com/oobabooga/text-generation-webui to support locally run
;; text generation models.
;;; Code:
(require 'cl-lib)
(require 'gv)
(require 'json)
(require 'org)
(require 'org-element)
(require 'websocket)
(require 'org-ai-block)
(defcustom org-ai-oobabooga-websocket-url "ws://127.0.0.1:5005/api/v1/stream"
"URL to use for the websocket connection.
See https://github.com/oobabooga/text-generation-webui#api."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-system-prefix "PROMPT: "
"Prefix to use for system messages when assembling the full prompt."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-user-prefix "You: "
"Prefix to use for user messages when assembling the full prompt."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-assistant-prefix "Assistant: "
"Prefix to use for AI assistant messages when assembling the full prompt."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-default-system-prompt nil
"Default system prompt to use if none is specified.
If nil, no system prompt is used."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-stopping-strings (list org-ai-oobabooga-user-prefix)
"List of strings that will cause the stream to be stopped."
:type 'string
:group 'org-ai-oobabooga)
(defcustom org-ai-oobabooga-default-payload nil
"Default payload to add to all requests.
E.g. if you want to set the temperature and top_p values:
(setq org-ai-oobabooga-default-payload
'((:temperature . 0.7) (:top_p . 0.1)))"
:type 'alist
:group 'org-ai-oobabooga)
(defvar org-ai-oobabooga--current-request nil
"Holds the websocket while a request is in progress.")
(defvar org-ai-oobabooga--chat-got-first-response nil)
(make-variable-buffer-local 'org-ai-oobabooga--chat-got-first-response)
(defvar org-ai-oobabooga--chat-ai-role-inserted nil)
(make-variable-buffer-local 'org-ai-oobabooga--chat-ai-role-inserted)
(defun org-ai--merge-alist (a &rest bs)
"Merge the alists A and BS.
If a key exists in multiple alists, the value from the last alist
is used."
(let ((result (copy-alist a)))
(cl-loop for b in bs
do (cl-loop for (k . v) in b
do (if (assoc k result)
(setcdr (assoc k result) v)
(push (cons k v) result)))
finally return result)))
(defun org-ai-oobabooga-create-prompt-default (messages)
"Create a prompt from MESSAGES.
MESSAGES is a vector of (:role role :content content) plists.
:role can be 'system, 'user or 'assistant."
(let ((system-prefix org-ai-oobabooga-system-prefix)
(user-prefix org-ai-oobabooga-user-prefix)
(assistant-prefix org-ai-oobabooga-assistant-prefix))
(org-ai--stringify-chat-messages messages
:default-system-prompt org-ai-oobabooga-default-system-prompt
:user-prefix user-prefix
:assistant-prefix assistant-prefix
:system-prefix system-prefix)))
(defcustom org-ai-oobabooga-create-prompt-function 'org-ai-oobabooga-create-prompt-default
"Function to create the prompt that will be passed to the LLM.
MESSAGES is a vector of (:role role :content content) plists.
:role can be 'system, 'user or 'assistant.
Example:
(setq org-ai-oobabooga-create-prompt-function
(lambda (messages)
(cl-loop for msg across messages
for role = (plist-get msg :role)
for content = (plist-get msg :content)
collect (cond ((eql role 'system) (format \"SYSTEM %s\" content))
((eql role 'user) (format \"USER %s\" content))
((eql role 'assistant) (format \"ASSISTANT %s\" content)))
into result
finally return (string-join result \"\n\n\"))))"
:type 'function
:group 'org-ai-oobabooga)
(cl-defun org-ai-oobabooga-stream (&optional &key
messages context
max-tokens temperature top-p)
"Send a request to the server from within an org-ai block.
MESSAGES is a list of messages to send to the server.
CONTEXT is the context of the block that is being processed.
MAX-TOKENS is the max-tokens value to use.
TEMPERATURE is the temperature value to use.
TOP-P is the top-p value to use."
(let ((context (or context (org-ai-special-block)))
(buffer (current-buffer)))
(let* ((info (org-ai-get-block-info context))
(max-tokens (or max-tokens (alist-get :max-tokens info)))
(temperature (or temperature (alist-get :temperature info)))
(top-p (or top-p (alist-get :top-p info)))
(callback (lambda (result) (org-ai-oobabooga--insert-chat-completion-response
context buffer result)))
(prompt (funcall org-ai-oobabooga-create-prompt-function messages)))
(org-ai-oobabooga-stream-request :prompt prompt
:callback callback
:max-tokens max-tokens
:temperature temperature
:top-p top-p))))
(cl-defun org-ai-oobabooga-stream-request (&optional &key
prompt callback
max-tokens temperature top-p)
"Send a request to the oobabooga/text-generation-webui server.
PROMPT is the prompt to send.
CALLBACK is the function to call when a response is received.
MAX-TOKENS is the max-tokens value to use.
TEMPERATURE is the temperature value to use.
TOP-P is the top-p value to use."
(let ((msg (let ((json-object-type 'alist)
(payload `((:prompt . ,prompt)
(:stopping_strings . ,org-ai-oobabooga-stopping-strings))))
(when max-tokens (push `(:max_new_tokens . ,max-tokens) payload))
(when temperature (push `(:temperature . ,temperature) payload))
(when top-p (push `(:top_p . ,top-p) payload))
(when org-ai-oobabooga-default-payload
(setq payload (org-ai--merge-alist org-ai-oobabooga-default-payload payload)))
(json-encode payload)))
(ws (websocket-open org-ai-oobabooga-websocket-url
:on-message (lambda (_ws frame)
(funcall callback (websocket-frame-text frame)))
:on-close (lambda (_ws)
(setq org-ai-oobabooga--current-request nil)
(funcall callback nil)))))
(setq org-ai--current-insert-position-marker nil
org-ai-oobabooga--chat-got-first-response nil
org-ai-oobabooga--chat-ai-role-inserted nil
org-ai-oobabooga--current-request ws)
;; (message "[org-ai-oobabooga] SENDING %s" msg)
(websocket-send-text ws msg)))
(defun org-ai-oobabooga-stop ()
"Stop the current request."
(interactive)
(when org-ai-oobabooga--current-request
(websocket-close org-ai-oobabooga--current-request)))
(defun org-ai-oobabooga--read-response (response)
"Parse RESPONSE as json."
(let ((json-object-type 'alist)
(json-key-type 'keyword)
(json-array-type 'vector))
(condition-case _err
(json-read-from-string response)
(error
'((:error . (format "Failed to parse response: %s" response)))))))
(defun org-ai-oobabooga--insert-chat-completion-response (context buffer response)
"Handle stream response from oobabooga/text-generation-webui.
CONTEXT is the context of the block that was being processed.
BUFFER is the buffer that the block is in.
RESPONSE is the parsed json response from the server."
;; (message "[org-ai-oobabooga] RECEIVED %s" response)
(if response
;; process response
(with-current-buffer buffer
(let ((pos (or (and org-ai--current-insert-position-marker
(marker-position org-ai--current-insert-position-marker))
(org-element-property :contents-end context))))
(save-excursion
(goto-char pos)
;; make sure we have enough space at end of block, don't write on same line
(when (string-suffix-p "#+end_ai" (buffer-substring-no-properties (point) (line-end-position)))
(insert "\n")
(backward-char))
;; insert text
(let* ((response (org-ai-oobabooga--read-response response))
(err (alist-get :error response))
(type (alist-get :event response))
(text (alist-get :text response)))
(when err (error err))
(when (not org-ai-oobabooga--chat-ai-role-inserted)
(insert "\n[AI]: ")
(run-hook-with-args 'org-ai-after-chat-insertion-hook 'role "assistant")
(setq org-ai-oobabooga--chat-ai-role-inserted t))
(cond
((string= type "text_stream")
(when (or org-ai-oobabooga--chat-got-first-response
(not (string= (string-trim text) "")))
(when (and (not org-ai-oobabooga--chat-got-first-response)
(string-prefix-p "```" text))
;; start markdown codeblock responses on their own line
(insert "\n"))
(let ((text (if org-ai-oobabooga--chat-got-first-response
text
(string-trim-left text))))
(insert (decode-coding-string text 'utf-8))
(run-hook-with-args 'org-ai-after-chat-insertion-hook 'text text))
(setq org-ai-oobabooga--chat-got-first-response t)))
((string= type "stream_end")
(let ((text "\n\n[ME]: "))
(insert text)
(run-hook-with-args 'org-ai-after-chat-insertion-hook 'text text))))
(setq org-ai--current-insert-position-marker (point-marker))))))
;; insert new prompt and change position
(with-current-buffer buffer
(run-hook-with-args 'org-ai-after-chat-insertion-hook 'end "")
(org-element-cache-reset))))
;; -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
(provide 'org-ai-oobabooga)
;;; org-ai-oobabooga.el ends here