Fix llm sse requests and add debugging

Author: ovistoica

src/simulflow/processors/google.clj

@@ -67,9 +67,7 @@
    :params (schema/->describe-parameters GoogleLLMConfigSchema)
    :workload :io})
 
-(defn init!
-  [params]
-  (uai/init-llm-processor! GoogleLLMConfigSchema params :gemini))
+(def init! (partial uai/init-llm-processor! GoogleLLMConfigSchema :gemini))
 
 (defn transition
   [state transition]

src/simulflow/processors/openai.clj

@@ -141,9 +141,7 @@
    :params (schema/->describe-parameters OpenAILLMConfigSchema)
    :workload :io})
 
-(defn init!
-  [params]
-  (uai/init-llm-processor! OpenAILLMConfigSchema params :openai))
+(def init! (partial uai/init-llm-processor! OpenAILLMConfigSchema :openai))
 
 (defn transition
   [state transition]

src/simulflow/utils/openai.clj

@@ -71,19 +71,27 @@
 ;; Common processor functions
 (defn vthread-pipe-response-with-interrupt
   "Common function to handle streaming response from LLM"
-  [{:keys [in-ch out-ch interrupt-chan on-end]}]
+  [{:keys [in-ch out-ch interrupt-ch on-end] :as data}]
+  (t/log! {:msg "Piping out result" :data data :level :trace :id :llm-processor})
   (vthread-loop []
-    (if-let [[chunk c] (a/alts!! [in-ch interrupt-chan])]
-      (when (= c in-ch)
-        (a/>!! out-ch chunk)
-        (recur))
-      ;; No more data or interruption, call on-end and exit
-      (when (fn? on-end)
-        (on-end)))))
+    (let [[val port] (a/alts!! [interrupt-ch in-ch] :priority true)]
+      (t/log! {:level :trace
+               :msg "Piping current val"
+               :data {:val val
+                      :chan (if (= port in-ch) :in-ch :interrupt-ch)}
+               :id :llm-processor})
+
+      (if (and (= port in-ch) val)
+        (do
+          (a/>!! out-ch val)
+          (recur))
+        ;; No more data or interruption, call on-end and exit
+        (when (fn? on-end)
+          (on-end))))))
 
 (defn init-llm-processor!
   "Common initialization function for OpenAI-compatible LLM processors"
-  [schema params log-id]
+  [schema log-id params]
   (let [parsed-config (schema/parse-with-defaults schema params)
         llm-write (a/chan 100)
         llm-read (a/chan 1024)
@@ -92,22 +100,26 @@
     (vthread-loop []
       (when-let [command (a/<!! llm-write)]
         (try
-          (t/log! {:level :debug :id log-id :data command} "Processing request command")
-          (assert (#{:command/sse-request :command/interrupt-request} command)
+          (t/log! {:level :debug :id log-id :data command} "Processing command")
+          (assert (#{:command/sse-request :command/interrupt-request} (:command/kind command))
                   "LLM processor only supports SSE request or interrupt request commands")
           ;; Execute the command and handle the streaming response
-          (cond
-            (= command :command/sse-request)
+          (case (:command/kind command)
+
+            :command/sse-request
             (do
+              (t/log! {:level :debug :id log-id :msg "Making SSE request"})
               (reset! request-in-progress? true)
               (vthread-pipe-response-with-interrupt {:in-ch (command/handle-command command)
                                                      :out-ch llm-read
                                                      :interrupt-ch interrupt-ch
                                                      :on-end #(reset! request-in-progress? false)}))
 
-            (= command :command/interrupt-request)
+            :command/interrupt-request
             (when @request-in-progress?
-              (a/>!! interrupt-ch command)))
+              (a/>!! interrupt-ch command))
+
+            nil)
           (catch Exception e
             (t/log! {:level :error :id log-id :error e} "Error processing command")))
         (recur)))

src/simulflow/utils/request.clj

@@ -7,7 +7,8 @@
    [hato.client :as http]
    [hato.middleware :as hm]
    [simulflow.async :refer [vthread vthread-loop]]
-   [simulflow.utils.core :as u])
+   [simulflow.utils.core :as u]
+   [taoensso.telemere :as t])
   (:import
    (java.io InputStream)))
 
@@ -23,12 +24,19 @@
           (recur))))))
 
 (defn- parse-openai-event [raw-event]
-  (let [data-idx (string/index-of raw-event "{")
+  (let [_ (t/log! {:level :trace :msg "SSE parsing raw event" :data {:raw-event-preview (take 100 raw-event)}})
+        data-idx (string/index-of raw-event "{")
         done-idx (string/index-of raw-event "[DONE]")]
+    (t/log! {:level :trace :msg "SSE parsing indices" :data {:data-idx data-idx :done-idx done-idx}})
     (if done-idx
-      :done
-      (-> (subs raw-event data-idx)
-          (u/parse-if-json :throw-on-error? true)))))
+      (do
+        (t/log! {:level :trace :msg "SSE parsing DONE event"})
+        :done)
+      (let [json-str (subs raw-event data-idx)
+            _ (t/log! {:level :trace :msg "SSE parsing JSON string" :data {:json-preview (take 100 json-str)}})
+            parsed (u/parse-if-json json-str :throw-on-error? true)]
+        (t/log! {:level :trace :msg "SSE parsing result" :data {:parsed-preview (take 100 (str parsed))}})
+        parsed))))
 
 ; Per this discussion: https://community.openai.com/t/clarification-for-max-tokens/19576
 ; if the max_tokens is not provided, the response will try to use all the available
@@ -48,33 +56,47 @@
   [{:keys [request params]}]
   (let [close? (:stream/close? params)
         parse-event (or (:parse-event params) parse-openai-event)
-        event-stream ^InputStream (:body (http/request (merge request
-                                                              params
-                                                              {:as :stream})))
+        result (http/request (merge request
+                             params
+                             {:as :stream}))
+        _ (t/log! {:level :trace :msg "SSE Request result" :data result})
+        event-stream ^InputStream (:body result)
         buffer-size (calc-buffer-size params)
         events (a/chan (a/buffer buffer-size) (map parse-event))]
     (vthread
       (try
         (loop [byte-coll []]
           (let [byte-arr (byte-array (max 1 (.available event-stream)))
                 bytes-read (.read event-stream byte-arr)]
+            (t/log! {:level :trace :msg "SSE read bytes" :data {:bytes-read bytes-read :available (.available event-stream)}})
 
             (if (neg? bytes-read)
 
               ;; Input stream closed, exiting read-loop
-              nil
+              (do
+                (t/log! {:level :trace :msg "SSE stream closed"})
+                nil)
 
               (let [next-byte-coll (concat byte-coll (seq byte-arr))
                     data (slurp (byte-array next-byte-coll))]
+                (t/log! {:level :trace :msg "SSE data chunk" :data {:data-length (count data) :data-preview (take 200 data)}})
                 (if-let [es (not-empty (re-seq event-mask data))]
-                  (if (every? true? (map #(a/>!! events %) es))
-                    (recur (drop (apply + (map #(count (.getBytes ^String %)) es))
-                                 next-byte-coll))
-
-                    ;; Output stream closed, exiting read-loop
-                    nil)
-
-                  (recur next-byte-coll))))))
+                  (do
+                    (t/log! {:level :trace :msg "SSE parsed events" :data {:event-count (count es) :events (map #(take 100 %) es)}})
+                    (if (every? true? (map #(do
+                                            (t/log! {:level :trace :msg "SSE putting event in channel" :data {:event (take 100 %)}})
+                                            (a/>!! events %)) es))
+                      (recur (drop (apply + (map #(count (.getBytes ^String %)) es))
+                                   next-byte-coll))
+
+                      ;; Output stream closed, exiting read-loop
+                      (do
+                        (t/log! {:level :trace :msg "SSE output stream closed during event sending"})
+                        nil)))
+
+                  (do
+                    (t/log! {:level :trace :msg "SSE no events found in data chunk"})
+                    (recur next-byte-coll)))))))
         (finally
           (when close?
             (a/close! events))