From 59c77108a55b86617c0e40a67ff2805e71e343c9 Mon Sep 17 00:00:00 2001
From: Asep Haryana <asepharyana@users.noreply.github.com>
Date: Sun, 26 Jul 2026 16:09:04 +0700
Subject: [PATCH] fix: skip leading EOS tokens in streaming + non-streaming

- Skip <|im_end|> generated as first token (prevents empty responses)
- Clean <think> tags as plain text in generated output
---
 src/infrastructure/llama/engine.rs | 10 +++++++++
 src/presentation/handler/chat.rs   | 33 ++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/src/infrastructure/llama/engine.rs b/src/infrastructure/llama/engine.rs
index 563453f..5fcd007 100644
--- a/src/infrastructure/llama/engine.rs
+++ b/src/infrastructure/llama/engine.rs
@@ -219,6 +219,16 @@ impl LlamaEngine {
 
         let mut current = inner.sample(sampler);
 
+        // Skip leading EOS tokens (like <|im_end|> as first token)
+        while output.is_empty() && self.model.is_eog_token(current) {
+            let pos = input_tokens.len() as i32 + output.len() as i32;
+            if let Err(e) = inner.decode(current, pos) {
+                tracing::info!("  Decode error: {e}");
+                break;
+            }
+            current = inner.sample(sampler);
+        }
+
         for _ in 0..max_tokens {
             if self.model.is_eog_token(current) {
                 break;
diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs
index e36f5e1..f746ef5 100644
--- a/src/presentation/handler/chat.rs
+++ b/src/presentation/handler/chat.rs
@@ -186,6 +186,39 @@ async fn handle_streaming(
                 break;
             }
 
+            // Skip leading EOS tokens (like <|im_end|> at start of generation)
+            if state.engine.is_eog(current) && text_buf.is_empty() {
+                // safety bound: don't skip more than 10
+                if count >= max_tokens || count > 10 {
+                    let chunk = serde_json::to_string(&SseChunk {
+                        id: chat_id.clone(),
+                        object: "chat.completion.chunk".into(),
+                        created,
+                        model: model_name.clone(),
+                        choices: vec![SseChoice {
+                            index: 0,
+                            delta: SseDelta {
+                                role: None,
+                                content: None,
+                                tool_calls: None,
+                            },
+                            finish_reason: Some("stop".into()),
+                        }],
+                    })
+                    .unwrap();
+                    let _ = tx.send(Ok(Event::default().data(chunk))).await;
+                    break;
+                }
+                count += 1;
+                let pos = input_tokens.len() as i32 + count as i32;
+                if let Err(e) = inner.decode(current, pos) {
+                    info!("  Decode error: {e}");
+                    break;
+                }
+                current = inner.sample(&mut sampler);
+                continue;
+            }
+
             if state.engine.is_eog(current) {
                 let reason = if has_tools && text_buf.contains("<tool_call>") {
                     "tool_calls"