Comparing changes

v0.2.1 v0.3.0
2 commits 4 files changed
Changed files (4)
ext/elelem/llama/elelem.cpp
@@ -23,9 +23,29 @@ extern "C" {
     float temp;
     uint32_t seed;
     bool dead = false;
+    bool in_generate = false;
     std::vector<llama_token> cached;
   };
 
+  typedef void (*el_token_cb)(void *userdata, const char *piece);
+
+  // Returns the length of the prefix of `s` that ends on a complete UTF-8
+  // sequence, so callers don't flush a truncated multi-byte codepoint.
+  static size_t el_utf8_safe_len(const std::string &s) {
+    size_t n = s.size();
+    size_t back = 0;
+    while (back < n && back < 4 && ((unsigned char) s[n - 1 - back] & 0xC0) == 0x80) back++;
+    if (back == n) return 0;
+    unsigned char lead = (unsigned char) s[n - 1 - back];
+    size_t seq_len = 1;
+    if ((lead & 0xE0) == 0xC0) seq_len = 2;
+    else if ((lead & 0xF0) == 0xE0) seq_len = 3;
+    else if ((lead & 0xF8) == 0xF0) seq_len = 4;
+    else if (lead >= 0x80) seq_len = back + 1; // stray continuation/invalid lead; hold it
+    if (back + 1 < seq_len) return n - back - 1;
+    return n;
+  }
+
   static bool g_backend = false;
 
   static enum ggml_log_level el_log_threshold() {
@@ -154,7 +174,7 @@ extern "C" {
     return buf.c_str();
   }
 
-  static const char *el_generate_impl(void *handle, const char *messages_json, const char *tools_json, int max_tokens) {
+  static const char *el_generate_impl(void *handle, const char *messages_json, const char *tools_json, int max_tokens, el_token_cb on_token, void *userdata) {
     auto *h = (el_handle *) handle;
     const llama_vocab *vocab = llama_model_get_vocab(h->model);
     static thread_local std::string buf;
@@ -239,22 +259,40 @@ extern "C" {
     int budget = max_tokens > 0 ? std::min(max_tokens, headroom) : std::max(1, headroom / 2);
 
     std::string output;
+    std::string pending;
     char piece[512];
     int n_decoded = 0;
     bool eog = false;
     bool decode_failed = false;
+    auto last_flush = std::chrono::steady_clock::now();
     for (int t = 0; t < budget; t++) {
       llama_token id = llama_sampler_sample(smpl.get(), ctx, -1);
       if (llama_vocab_is_eog(vocab, id)) { eog = true; break; }
       int np = llama_token_to_piece(vocab, id, piece, (int32_t) sizeof(piece), 0, true);
-      if (np > 0) output.append(piece, np);
+      if (np > 0) {
+        output.append(piece, np);
+        if (on_token) pending.append(piece, np);
+      }
       n_decoded++;
+      if (on_token && !pending.empty()) {
+        auto now = std::chrono::steady_clock::now();
+        if (std::chrono::duration<double, std::milli>(now - last_flush).count() >= 50.0) {
+          size_t safe_len = el_utf8_safe_len(pending);
+          if (safe_len > 0) {
+            std::string chunk = pending.substr(0, safe_len);
+            on_token(userdata, chunk.c_str());
+            pending.erase(0, safe_len);
+          }
+          last_flush = now;
+        }
+      }
       if (t + 1 < budget && llama_decode(ctx, llama_batch_get_one(&id, 1)) != 0) {
         decode_failed = true;
         break;
       }
       if (t + 1 < budget) h->cached.push_back(id);
     }
+    if (on_token && !pending.empty()) on_token(userdata, pending.c_str());
     auto t3 = std::chrono::steady_clock::now();
 
     if (!decode_failed) guard.committed = true;
@@ -363,14 +401,22 @@ extern "C" {
     return buf.c_str();
   }
 
-  const char *el_generate(void *handle, const char *messages_json, const char *tools_json, int max_tokens) {
+  const char *el_generate(void *handle, const char *messages_json, const char *tools_json, int max_tokens, el_token_cb on_token, void *userdata) {
     static thread_local std::string errbuf;
     if (!handle) return el_error(errbuf, "null handle");
     if (!messages_json) return el_error(errbuf, "null messages");
-    if (((el_handle *) handle)->dead) return el_error(errbuf, "handle is unusable after a previous fatal error");
+    auto *h = (el_handle *) handle;
+    if (h->dead) return el_error(errbuf, "handle is unusable after a previous fatal error");
+    if (h->in_generate) return el_error(errbuf, "generate is not reentrant");
+
+    h->in_generate = true;
+    struct reentrancy_guard {
+      el_handle *h;
+      ~reentrancy_guard() { h->in_generate = false; }
+    } guard{h};
 
     try {
-      return el_generate_impl(handle, messages_json, tools_json, max_tokens);
+      return el_generate_impl(handle, messages_json, tools_json, max_tokens, on_token, userdata);
     } catch (const std::exception &e) {
       ((el_handle *) handle)->dead = true;
       return el_error(errbuf, e.what());
lib/elelem/llama/provider.rb
@@ -20,7 +20,7 @@ module Elelem
           lib = Fiddle.dlopen(SHIM)
           {
             open: Fiddle::Function.new(lib["el_open"], [V, I, I, I, F, I], V),
-            generate: Fiddle::Function.new(lib["el_generate"], [V, V, V, I], V),
+            generate: Fiddle::Function.new(lib["el_generate"], [V, V, V, I, V, V], V, need_gvl: true),
             close: Fiddle::Function.new(lib["el_close"], [V], Fiddle::TYPE_VOID)
           }
         end
@@ -35,7 +35,18 @@ module Elelem
       end
 
       def fetch(messages, tools = [], &block)
-        ptr = self.class.functions[:generate].call(@handle, JSON.generate(messages), JSON.generate(tools), @max_tokens)
+        streamed = false
+        on_token = Fiddle::Closure::BlockCaller.new(Fiddle::TYPE_VOID, [V, V]) do |_userdata, piece|
+          streamed = true
+          text = Fiddle::Pointer.new(piece).to_s.force_encoding(Encoding::UTF_8).scrub
+          block&.call(type: "thinking", text: text)
+        rescue Exception => e # rubocop:disable Lint/RescueException
+          Elelem.logger.warn("llama: streaming callback failed: #{e.message}")
+        end
+
+        ptr = self.class.functions[:generate].call(
+          @handle, JSON.generate(messages), JSON.generate(tools), @max_tokens, on_token, nil
+        )
         result = JSON.parse(Fiddle::Pointer.new(ptr).to_s)
 
         Elelem.logger.debug("llama: tool-call fallback used") if result["fallback"]
@@ -53,7 +64,7 @@ module Elelem
         end
         reasoning = result["reasoning"].to_s
         Elelem.logger.debug("llama: reasoning: #{reasoning}") unless reasoning.empty?
-        block&.call(type: "thinking", text: reasoning) unless reasoning.empty?
+        block&.call(type: "thinking", text: reasoning) unless reasoning.empty? || streamed
 
         content = result["content"].to_s
         if result["error"] && content.empty?
lib/elelem/llama/version.rb
@@ -2,6 +2,6 @@
 
 module Elelem
   module Llama
-    VERSION = "0.2.1"
+    VERSION = "0.3.0"
   end
 end
Gemfile.lock
@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    elelem-llama (0.2.1)
+    elelem-llama (0.3.0)
       elelem (~> 0.11)
       fiddle (~> 1.1)
       json (~> 3.0)
@@ -78,7 +78,7 @@ CHECKSUMS
   bigdecimal (4.1.2) sha256=53d217666027eab4280346fba98e7d5b66baaae1b9c3c1c0ffe89d48188a3fbd
   bundler (4.0.20) sha256=7978a8ac648767f5e635bc522445b79e80a52b907a39a36c2d8085ed6bc762ae
   elelem (0.11.0) sha256=552cafb092320e3896b8e07572d322fe2b94f0f10d965bf7b993b74517a6f744
-  elelem-llama (0.2.1)
+  elelem-llama (0.3.0)
   erb (6.0.7) sha256=c5ca6dc25b0ef974a44dc8f59fe847577122483b1968a38dec305c60bf91ee92
   fiddle (1.1.8) sha256=7fa8ee3627271497f3add5503acdbc3f40b32f610fc1cf49634f083ef3f32eee
   forwardable (1.4.0) sha256=f1cd40cc9812937980e1c76f1aa053660990a7c9b6a98fc37d945468afcce838