Add Streaming Tool Calling support

nemotron_toolcall_parser_no_streaming.py

@@ -44,6 +44,13 @@ class NemotronJSONToolParser(ToolParser):
 
         self.tool_call_regex = re.compile(r"<TOOLCALL>(.*?)</TOOLCALL>", re.DOTALL)
 
+        self.reasoning_end_token: str = "</think>"
+        self.special_token_buffer: str = ""
+        self.parsing_special_token: bool = False
+        self.parsing_tool_call: bool = False
+        self.parsing_reasoning: bool = True
+
+
     def extract_tool_calls(
         self,
         model_output: str,
@@ -95,6 +102,41 @@ class NemotronJSONToolParser(ToolParser):
                     tool_calls=[],
                     content=model_output,
                 )
+            
+
+    def _parse_tool_xml(self,xml:str):
+        tool_calls = []
+        str_tool_calls = self.tool_call_regex.findall(xml)[0].strip()
+        if not str_tool_calls.startswith("["):
+            str_tool_calls = "[" + str_tool_calls
+        if not str_tool_calls.endswith("]"):
+            str_tool_calls = "]" + str_tool_calls
+        json_tool_calls = json.loads(str_tool_calls)
+        for i,tool_call in enumerate(json_tool_calls):
+            try:
+                tool_calls.append(DeltaToolCall(
+                    type="function",
+                    index=i,
+                    id=f"{tool_call["name"]}-{i}",
+                    function=DeltaFunctionCall(name=tool_call["name"],arguments=json.dumps(tool_call["arguments"], ensure_ascii=False) if isinstance(tool_call["arguments"], dict) else tool_call["arguments"])
+                ))
+            except:
+                continue
+        return tool_calls
+
+    def _split(self,k:str,v:str):
+        idx = v.find(k)
+        return v[:idx],v[idx:]
+    
+    def _splitr(self,k:str,v:str):
+        idx = v.find(k) + len(k)
+        return v[:idx],v[idx:]
+    
+    def _partial_match(self,s1:str,s2:str):
+        for i in range(min(len(s1),len(s2))):
+            if s1[i] != s2[i]:
+                return False
+        return True
 
     def extract_tool_calls_streaming(
         self,
@@ -106,5 +148,65 @@ class NemotronJSONToolParser(ToolParser):
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> Union[DeltaMessage, None]:
+        
+        """
+        As a stopgap, this script also implements reasoning parsing until vllm accepts reasoning plugins.
+
+        Handles streaming tool call detection and parsing for Nemotron format:
+        <TOOLCALL>[{"name": "toolName", "arguments": {...}}]</TOOLCALL>
+
+        This code runs once per token, its not very efficient. 
+        """
 
-        raise NotImplementedError("Tool calling is not supported in streaming mode!")
+
+        def send_message(content:str):
+            if self.parsing_reasoning:
+                return DeltaMessage(reasoning_content=content)
+            return DeltaMessage(content=content)
+
+
+
+        if not self.parsing_special_token:
+            if "<" in delta_text:
+                self.parsing_special_token = True
+                before,after = self._split("<",delta_text)
+                self.special_token_buffer += after
+                return send_message(before) #DeltaMessage(content=before)
+            else:
+                return send_message(delta_text) #DeltaMessage(content=delta_text)
+        else:
+            self.special_token_buffer += delta_text
+
+            partial_match_tool_token = self._partial_match(self.tool_call_start_token,self.special_token_buffer)
+            partial_match_reasoning_token = self._partial_match(self.reasoning_end_token,self.special_token_buffer)
+
+            if partial_match_tool_token and partial_match_reasoning_token:
+                #ambiguous token so we continue for now
+                return None
+            elif partial_match_tool_token:
+                #the above code is responsible for kicking us out if we arent on track so we can just wait till we find a end token
+                if self.tool_call_end_token in self.special_token_buffer:
+                    before,after = self._splitr(self.tool_call_end_token,self.special_token_buffer)
+                    tool_calls = self._parse_tool_xml(before)
+                    self.special_token_buffer = ""
+                    self.parsing_special_token = False
+                    return DeltaMessage(tool_calls=tool_calls, content=after)
+                
+            elif partial_match_reasoning_token:
+                if self.reasoning_end_token in self.special_token_buffer:
+                    before,after = self._splitr(self.reasoning_end_token,self.special_token_buffer)
+                    self.special_token_buffer = ""
+                    self.parsing_reasoning = False
+                    self.parsing_special_token = False
+                    return DeltaMessage(reasoning_content=before, content=after)
+                    
+            else:
+                #neither token matched so its a dud
+                content = self.special_token_buffer
+                self.special_token_buffer = ""
+                self.parsing_tool_call = False
+                return DeltaMessage(content=content)
+            
+            
+            
+            return None