fix parsing logic

kanana_tool_calls/functionary_kanana_tool_parser.py

@@ -37,7 +37,6 @@ class BaseTemplate(ABC):
     def response_to_messages(self, generated_text):
         raise NotImplementedError
 
-
 class FunctionaryV3Llama31Template(BaseTemplate):
     def __init__(
             self,
@@ -144,7 +143,6 @@ def is_complete_json(input_str):
     except JSONDecodeError:
         return False
 
-
 @ToolParserManager.register_module(["functionary_v3_llama_31"])
 class FunctionaryV3Llama31ToolParser(ToolParser):
     def __init__(self, tokenizer: Union[PreTrainedTokenizerBase, AnyTokenizer]):
@@ -162,6 +160,8 @@ class FunctionaryV3Llama31ToolParser(ToolParser):
         self._python_tag_id = tokenizer.encode(self._python_tag,
                                              add_special_tokens=False)[0]
 
+        # added buffer for each tool call parser
+        self._buffer = ""
     def extract_tool_calls(
             self, model_output: str,
             request: ChatCompletionRequest) -> ExtractedToolCallInformation:
@@ -175,7 +175,7 @@ class FunctionaryV3Llama31ToolParser(ToolParser):
                                                 tool_calls=[],
                                                 content=result["content"])
 
-    # our template: <function=function_name>{"arg":"var"}<function>
+    # our template: <function=function_name>{"arg":"var"}</function>
     def extract_tool_calls_streaming(
         self,
         previous_text: str,
@@ -186,13 +186,36 @@ class FunctionaryV3Llama31ToolParser(ToolParser):
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> Union[DeltaMessage, None]:
+    
         # if current_text does not start with function tag (or python tag),
-        # stream right away as delta.content
         if not (current_text.startswith(self._python_tag)
                 or current_text.startswith(self._func_prefix)
                 or self._func_prefix.startswith(current_text)):
-            return DeltaMessage(content=delta_text)
-
+            # for cases like "The answer is <function="
+            # let current_text="<function="
+            if self._func_prefix in current_text:
+                idx = current_text.find(self._func_prefix)
+                current_text = current_text[idx:]
+                self._buffer = ""
+            # for cases like "The answer is <function"
+            # add delta_text to buffer to figure out whether to print or not later
+            elif delta_text.endswith("<")  or (current_text.endswith("<function") and delta_text.endswith("function")):
+                self._buffer += delta_text
+                return DeltaMessage(content=None)
+            # for cases that does not include "<function" at all,
+            # stream right away as delta.content                
+            else:
+                delta_text = self._buffer + delta_text
+                self._buffer = ""
+                return DeltaMessage(content=delta_text)
+
+        # for cases like "<" or "<function"
+        # add to delta_text to buffer 
+        if delta_text.endswith("<") or (current_text.endswith("<function") and delta_text.endswith("function")):
+            self._buffer += delta_text
+            
+            
+        # CHECK: this part not working (should use current_token_ids)
         # if current_text ends with stop token,
         # remove it from the text
         # CHECK: sometimes text is generated beyond <|eom_id|>
@@ -311,7 +334,7 @@ class FunctionaryV3Llama31ToolParser(ToolParser):
                 else:
                     delta = None
                 # re-set stuff pertaining to progress in the current tool
-                self.current_tool_id = len(tool_call_arr) - 1 # update current tool call # 1을 더하는게 아니라 이렇게 하네 ..
+                self.current_tool_id = len(tool_call_arr) - 1 # update current tool call 
                 self.current_tool_name_sent = False
                 self.streamed_args_for_tool.append("")
                 logger.debug("starting on new tool %d", self.current_tool_id)