streaming chatgpt in clojure

gpt.clj

(ns app.gpt
  (:require [clj-http.client :as http]
            [clojure.repl]
            [cheshire.core :as json]
            [clojure.pprint :refer [pprint]]
            [taoensso.timbre :as log]
            [clojure.core.async :as a :refer [<! >! go]]
            [clojure.string :as str]
            [clojure.java.io :as io])
  (:import [java.io InputStream]))

(log/refer-timbre)

(def event-mask (re-pattern (str "data: (.+?)\n\n")))

(defn- parse-event [raw-event]
  (let [json-str (second raw-event)] 
    (if (= json-str "DONE")
      {:status :done}
      (json/decode json-str true))))

(def DEFAULT_MODEL "gpt-4")

(defn auth-headers []
  {"Authorization" (str "Bearer " (System/getenv "OPENAPI_KEY"))})

(defn chat-streaming [params]
  (let [response (http/post "https://api.openai.com/v1/chat/completions"
                            {:headers (auth-headers)
                             :form-params (merge {:model DEFAULT_MODEL
                                                  :stream true} params)
                             :content-type :json
                             :accept :json
                             :as :stream})
        event-stream ^InputStream (:body response)
        events (a/chan (a/sliding-buffer 10))]
    (a/thread
      (loop [accum ""]
        (let [byte-array (byte-array 4096)
              bytes-read (.read event-stream byte-array)
              new-data (if (neg? bytes-read)
                         accum
                         (str accum (String. byte-array 0 bytes-read "UTF-8")))
              matches (re-seq event-mask new-data)
              remaining-data (str/replace new-data event-mask "")]
          (doseq [match matches]
            (let [event (parse-event match)]
              (if (and (map? event) (= (:status event) :done))
                (do (info "Received DONE token, closing stream")
                    (.close event-stream))
                (a/>!! events event))))
          (if (neg? bytes-read)
            (do (info "Input stream closed, exiting read-loop")
                (.close event-stream))
            (recur remaining-data)))))
    events))

(defn log-chat-stream [params]
  (let [events (chat-streaming params)]
    (go
      (loop []
        (when-let [event (<! events)]
          (if (= (:status event) :done)
            (info "Stream ended with DONE token.")
            (do
              (info "Received event:" event)
              (recur))))))))

;; Example usage
(comment
  (log-chat-stream {:messages [{:role "system" :content "You are a helpful assistant."}
                               {:role "user" :content "Tell me a joke."}]}))