feat: Implement think tag handling

TilmanGriesel · Jan 28, 2025 · dff2934 · dff2934
1 parent f7c9e81
commit dff2934
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/services/web/src/static/js/chatService.js b/services/web/src/static/js/chatService.js
@@ -86,6 +86,8 @@ export class ChatService {
     const decoder = new TextDecoder("utf-8");
     let buffer = "";
     let isFirstChunk = true;
+    let isInThinkBlock = false;
+    let thinkBuffer = "";
 
     try {
       while (true) {
@@ -112,14 +114,40 @@ export class ChatService {
               return;
             }
             if (data.chunk) {
-              // trim the first chunk to remove potential leading whitespace
+              // process chunk
               const chunk = isFirstChunk ? data.chunk.trimLeft() : data.chunk;
               if (chunk.length > 0) {
-                onChunk(chunk);
+                // handle think tags
+                for (let i = 0; i < chunk.length; i++) {
+                  if (chunk.slice(i).startsWith("<think>")) {
+                    isInThinkBlock = true;
+                    i += "<think>".length - 1;
+                    continue;
+                  }
+                  if (chunk.slice(i).startsWith("</think>")) {
+                    isInThinkBlock = false;
+                    // process think content
+                    console.log("Think block content:", thinkBuffer);
+                    thinkBuffer = ""; // clear buffer
+                    i += "</think>".length - 1;
+                    continue;
+                  }
+
+                  if (isInThinkBlock) {
+                    thinkBuffer += chunk[i];
+                  } else {
+                    onChunk(chunk[i]);
+                  }
+                }
                 isFirstChunk = false;
               }
             }
             if (data.done) {
+              if (thinkBuffer) {
+                console.log("Final think block content:", thinkBuffer);
+                thinkBuffer = "";
+              }
+
               const ttsText = data.full_response.llm.replies[0];
               console.log(ttsText);
               window.postMessage(

diff --git a/tools/embed/src/core/document_embedder.py b/tools/embed/src/core/document_embedder.py
@@ -48,7 +48,6 @@ def __init__(
         self.embedding_pipeline = None
         self.embedding_dimension = None
 
-        # Validate HuggingFace configuration if selected
         if self.provider == ModelProvider.HUGGINGFACE and not self.hf_api_key:
             raise ValueError(
                 "HuggingFace API key is required when using HuggingFace provider"