Просмотр исходного кода

feat: openai stream completion

Tienson Qin 2 лет назад
Родитель
Сommit
7047c42309
4 измененных файлов с 50 добавлено и 4 удалено
  1. 1 0
      package.json
  2. 42 2
      src/main/frontend/modules/ai/openai.cljs
  3. 2 2
      src/main/frontend/modules/ai/protocol.cljs
  4. 5 0
      yarn.lock

+ 1 - 0
package.json

@@ -133,6 +133,7 @@
         "remove-accents": "0.4.2",
         "sanitize-filename": "1.6.3",
         "send-intent": "3.0.11",
+        "sse.js": "^0.6.1",
         "threads": "1.6.5",
         "url": "^0.11.0",
         "yargs-parser": "20.2.4"

+ 42 - 2
src/main/frontend/modules/ai/openai.cljs

@@ -1,7 +1,9 @@
 (ns frontend.modules.ai.openai
   (:require [frontend.modules.ai.protocol :as protocol]
             [frontend.util :as util]
-            [cljs-bean.core :as bean]))
+            [cljs-bean.core :as bean]
+            [clojure.string :as string]
+            ["sse.js" :refer [SSE]]))
 
 (defrecord OpenAI [token]
   protocol/AI
@@ -20,6 +22,41 @@
                        (map #(get-in % [:message :content]))))
                 (fn [failed-resp]
                   failed-resp)))
+  (ask-stream [_this q {:keys [model on-message on-finished]
+                        :or {model "gpt-3.5-turbo"}}]
+    (let [*buffer (atom "")
+          sse ^js (SSE. "https://api.openai.com/v1/chat/completions"
+                        (bean/->js
+                         {:headers {:Content-Type "application/json"
+                                    :authorization (str "Bearer " token)}
+                          :method "POST"
+                          :payload (js/JSON.stringify
+                                    (bean/->js {:model model
+                                                :stream true
+                                                :stop ["\n\n"]
+                                                :messages [{:role "user"
+                                                            :content q}]}))}))]
+      (.addEventListener sse "message"
+                         (fn [e]
+                           (let [data (.-data e)]
+                             (if (and (string? data)
+                                      (= data "[DONE]"))
+                               (do
+                                 (.close sse)
+                                 (when on-finished (on-finished @*buffer)))
+                               (try
+                                 (let [result (-> (bean/->clj (js/JSON.parse data))
+                                                  :choices
+                                                  first)
+                                       content (get-in result [:delta :content])]
+                                   (when content
+                                     (swap! *buffer str content)
+                                     (when on-message (on-message @*buffer))))
+                                 (catch :default e
+                                   (prn "OpenAI request failed: " e)
+                                   (.close sse)))))))
+
+      (.stream sse)))
   (summarize [this content opts])
   (translate [this content opts])
   (generate-text [this description opts])
@@ -32,5 +69,8 @@
 
   (protocol/ask open-ai "What's logseq?" {})
 
-  ;; {:result {:status 200, :success true, :body {:id "chatcmpl-6s5E1oW3DFvzuds9zWDIIMW49hxSo", :object "chat.completion", :created 1678347817, :model "gpt-3.5-turbo-0301", :usage {:prompt_tokens 12, :completion_tokens 119, :total_tokens 131}, :choices [{:message {:role "assistant", :content "\n\nLogseq is a personal knowledge management and note-taking tool that is used for creating and organizing ideas and notes in a hierarchical and linked form. This tool is designed to help people to capture and organize their thoughts, to-do lists, project plans, and other ideas in a way that is easy to manage and access. It offers a range of features such as graph visualization, markdown formatting, task management, smart blocks, and integrations with other tools like Roam Research, Obsidian, and Notion. The software can be run locally on your computer or accessed through a web browser."}, :finish_reason "stop", :index 0}]}, :headers {"openai-organization" "user-wszz12mmidc9hm2pniid1miq", "content-type" "application/json", "access-control-allow-origin" "*", "content-length" "882", "openai-version" "2020-10-01", "strict-transport-security" "max-age=15724800; includeSubDomains", "openai-processing-ms" "4902", "date" "Thu, 09 Mar 2023 07:43:42 GMT", "x-request-id" "fe837c73411e1f277b3672a8c255607f", "openai-model" "gpt-3.5-turbo-0301", "cache-control" "no-cache, must-revalidate"}, :trace-redirects ["https://api.openai.com/v1/chat/completions" "https://api.openai.com/v1/chat/completions"], :error-code :no-error, :error-text ""}}
+  (protocol/ask-stream open-ai "What's logseq?" {:on-message (fn [message]
+                                                               (prn "received: " message))
+                                                 :on-finished (fn [message]
+                                                                (prn "finished: " message))})
   )

+ 2 - 2
src/main/frontend/modules/ai/protocol.cljs

@@ -1,12 +1,12 @@
 (ns frontend.modules.ai.protocol)
 
 (defprotocol AI
-  ;; TODO: thread questions
   (ask [this q opts])
+  (ask-stream [this q opts])
   ;; (index-graph [this repo])
+  ;; encode && embedding
   (summarize [this content opts])
   (translate [this content opts])
-  ;; draft email/etc.
   (generate-text [this description opts])
   (generate-image [this description opts])
   (speech-to-text [this audio opts])

+ 5 - 0
yarn.lock

@@ -6735,6 +6735,11 @@ sprintf-js@^1.1.2:
   resolved "https://registry.yarnpkg.com/sprintf-js/-/sprintf-js-1.1.2.tgz#da1765262bf8c0f571749f2ad6c26300207ae673"
   integrity sha512-VE0SOVEHCk7Qc8ulkWw3ntAzXuqf7S2lvwQaDLRnUeIEaKNQJzV6BwmLKhOqT61aGhfUMrXeaBk+oDGCzvhcug==
 
+sse.js@^0.6.1:
+  version "0.6.1"
+  resolved "https://registry.yarnpkg.com/sse.js/-/sse.js-0.6.1.tgz#1de304f51a5b870e78c0ad72c3b1d7b682a733a1"
+  integrity sha512-peXG6GnWqF5hnubhMw0WfB6rqQy7z7LaMBT067vqgQwC3gKz8JGFzexBSV80FqZ9JoUDwo3Xt5nxkrGrgbPrtA==
+
 stable@^0.1.8:
   version "0.1.8"
   resolved "https://registry.yarnpkg.com/stable/-/stable-0.1.8.tgz#836eb3c8382fe2936feaf544631017ce7d47a3cf"