[fix] Keep streamed tool calls out of response content
cyanseek
parent
f659b55761
commit
86aad50621
|
|
@ -129,9 +129,11 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
|
|||
|
||||
Thread(target=_generate).start()
|
||||
|
||||
full_text = ""
|
||||
emitted = 0
|
||||
pending = ""
|
||||
in_tool_call = False
|
||||
tool_count = 0
|
||||
thinking_ended = not bool(open_thinking)
|
||||
strip_newlines = False
|
||||
|
||||
while True:
|
||||
text = queue.get()
|
||||
|
|
@ -140,36 +142,58 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
|
|||
if isinstance(text, dict):
|
||||
yield json.dumps(text, ensure_ascii=False)
|
||||
continue
|
||||
full_text += text
|
||||
pending += text
|
||||
|
||||
if not thinking_ended:
|
||||
pos = full_text.find('</think>')
|
||||
if pos >= 0:
|
||||
thinking_ended = True
|
||||
new_r = full_text[emitted:pos]
|
||||
if new_r:
|
||||
yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
|
||||
emitted = pos + len('</think>')
|
||||
after = full_text[emitted:].lstrip('\n')
|
||||
emitted = len(full_text) - len(after)
|
||||
if after:
|
||||
yield json.dumps({"choices": [{"delta": {"content": after}}]}, ensure_ascii=False)
|
||||
emitted = len(full_text)
|
||||
while pending:
|
||||
if strip_newlines:
|
||||
pending = pending.lstrip('\n')
|
||||
if not pending:
|
||||
break
|
||||
strip_newlines = False
|
||||
if in_tool_call:
|
||||
end = pending.find('</tool_call>')
|
||||
if end < 0:
|
||||
break
|
||||
end += len('</tool_call>')
|
||||
block, pending = pending[:end], pending[end:]
|
||||
_, _, tool_calls = parse_response(block)
|
||||
if tool_calls:
|
||||
call = tool_calls[0]
|
||||
call['index'] = tool_count
|
||||
call['id'] = f"call_{int(time.time())}_{tool_count}"
|
||||
tool_count += 1
|
||||
yield json.dumps({"choices": [{"delta": {"tool_calls": [call]}}]}, ensure_ascii=False)
|
||||
else:
|
||||
yield json.dumps({"choices": [{"delta": {"content": block}}]}, ensure_ascii=False)
|
||||
in_tool_call = False
|
||||
continue
|
||||
|
||||
marker = '<tool_call>' if thinking_ended else '</think>'
|
||||
field = 'content' if thinking_ended else 'reasoning_content'
|
||||
pos = pending.find(marker)
|
||||
if pos < 0:
|
||||
# 保留可能跨 chunk 的标记前缀，其余文本立即发送。
|
||||
keep = next((n for n in range(len(marker) - 1, 0, -1)
|
||||
if pending.endswith(marker[:n])), 0)
|
||||
end = len(pending) - keep
|
||||
if end:
|
||||
yield json.dumps({"choices": [{"delta": {field: pending[:end]}}]}, ensure_ascii=False)
|
||||
pending = pending[end:]
|
||||
break
|
||||
if pos:
|
||||
yield json.dumps({"choices": [{"delta": {field: pending[:pos]}}]}, ensure_ascii=False)
|
||||
pending = pending[pos:]
|
||||
if thinking_ended:
|
||||
in_tool_call = True
|
||||
else:
|
||||
new_r = full_text[emitted:]
|
||||
if new_r:
|
||||
yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
|
||||
emitted = len(full_text)
|
||||
else:
|
||||
new_c = full_text[emitted:]
|
||||
if new_c:
|
||||
yield json.dumps({"choices": [{"delta": {"content": new_c}}]}, ensure_ascii=False)
|
||||
emitted = len(full_text)
|
||||
pending = pending[len(marker):]
|
||||
thinking_ended = True
|
||||
strip_newlines = True
|
||||
|
||||
_, _, tool_calls = parse_response(full_text)
|
||||
if tool_calls:
|
||||
yield json.dumps({"choices": [{"delta": {"tool_calls": tool_calls}}]}, ensure_ascii=False)
|
||||
yield json.dumps({"choices": [{"delta": {}, "finish_reason": "tool_calls" if tool_calls else "stop"}]}, ensure_ascii=False)
|
||||
if pending:
|
||||
field = 'content' if thinking_ended else 'reasoning_content'
|
||||
yield json.dumps({"choices": [{"delta": {field: pending}}]}, ensure_ascii=False)
|
||||
yield json.dumps({"choices": [{"delta": {}, "finish_reason": "tool_calls" if tool_count else "stop"}]}, ensure_ascii=False)
|
||||
|
||||
except Exception as e:
|
||||
yield json.dumps({"error": str(e)})
|
||||
|
|
|
|||
|
|
@ -0,0 +1,147 @@
|
|||
"""python -m unittest discover -s tests -p test_stream_tool_calls.py -v"""
|
||||
import json
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from fastapi.testclient import TestClient
|
||||
from openai import OpenAI
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from scripts import serve_openai_api as server
|
||||
|
||||
|
||||
TOOL = '<tool_call>{"name":"weather","arguments":{"city":"杭州"}}</tool_call>'
|
||||
TOOLS = [{'type': 'function', 'function': {'name': 'weather', 'parameters': {
|
||||
'type': 'object', 'properties': {'city': {'type': 'string'}}, 'required': ['city']}}}]
|
||||
|
||||
|
||||
class ScriptedModel:
|
||||
def __init__(self, chunks, tokenize=False):
|
||||
self.chunks, self.tokenize = chunks, tokenize
|
||||
|
||||
def generate(self, input_ids, streamer, **kwargs):
|
||||
if self.tokenize:
|
||||
streamer.put(input_ids.cpu())
|
||||
for token in streamer.tokenizer.encode(''.join(self.chunks), add_special_tokens=False):
|
||||
streamer.put(torch.tensor([token]))
|
||||
streamer.end()
|
||||
else:
|
||||
for chunk in self.chunks:
|
||||
streamer.on_finalized_text(chunk)
|
||||
streamer.on_finalized_text('', stream_end=True)
|
||||
|
||||
|
||||
class StreamToolTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.tokenizer = AutoTokenizer.from_pretrained(
|
||||
Path(__file__).resolve().parents[1] / 'model', local_files_only=True)
|
||||
|
||||
def setUp(self):
|
||||
globals_patch = patch.multiple(server, tokenizer=self.tokenizer, device='cpu',
|
||||
model=ScriptedModel([]), create=True)
|
||||
globals_patch.start()
|
||||
self.addCleanup(globals_patch.stop)
|
||||
self.http = TestClient(server.app)
|
||||
self.addCleanup(self.http.close)
|
||||
self.request = {'model': 'minimind', 'messages': [{'role': 'user', 'content': 'Hi'}]}
|
||||
|
||||
def frames(self, chunks, thinking=False):
|
||||
server.model = ScriptedModel(chunks)
|
||||
response = self.http.post('/v1/chat/completions', json={
|
||||
**self.request, 'stream': True, 'open_thinking': thinking})
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertIn('text/event-stream', response.headers['content-type'])
|
||||
return [json.loads(line[6:]) for line in response.text.splitlines()
|
||||
if line.startswith('data: ')]
|
||||
|
||||
def assert_output(self, frames, content, count, reasoning=''):
|
||||
self.assertFalse(any('error' in frame for frame in frames), frames)
|
||||
deltas = [frame['choices'][0]['delta'] for frame in frames]
|
||||
self.assertEqual(''.join(d.get('content', '') for d in deltas), content)
|
||||
self.assertEqual(''.join(d.get('reasoning_content', '') for d in deltas), reasoning)
|
||||
calls = [call for d in deltas for call in d.get('tool_calls', [])]
|
||||
self.assertEqual(len(calls), count)
|
||||
self.assertEqual([call['index'] for call in calls], list(range(count)))
|
||||
self.assertEqual(len({call['id'] for call in calls}), count)
|
||||
for call in calls:
|
||||
self.assertEqual(call['function']['name'], 'weather')
|
||||
self.assertEqual(json.loads(call['function']['arguments']), {'city': '杭州'})
|
||||
self.assertEqual(frames[-1]['choices'][0]['finish_reason'], 'tool_calls' if count else 'stop')
|
||||
|
||||
def test_split_markers_and_multiple_calls(self):
|
||||
text = 'Before ' + TOOL + ' between ' + TOOL + ' after'
|
||||
for chunks in ([text], list(text), ['Before <tool_', text[len('Before <tool_'):]]):
|
||||
with self.subTest(chunks=len(chunks)):
|
||||
self.assert_output(self.frames(chunks), 'Before between after', 2)
|
||||
|
||||
def test_thinking_boundary_and_tool_text_inside_reasoning(self):
|
||||
reasoning = 'Consider ' + TOOL + '\n'
|
||||
text = reasoning + '</think>\n\n' + TOOL + ' answer'
|
||||
self.assert_output(self.frames(list(text), thinking=True), ' answer', 1, reasoning)
|
||||
|
||||
def test_invalid_and_unfinished_calls_preserve_text(self):
|
||||
for text in ('plain < text', 'end <tool_', '<tool_call>{"name":',
|
||||
'<tool_call>bad json</tool_call>', '<tool_call>[]</tool_call>'):
|
||||
with self.subTest(text=text):
|
||||
self.assert_output(self.frames(list(text)), text, 0)
|
||||
invalid = '<tool_call>bad json</tool_call>'
|
||||
self.assert_output(self.frames([invalid, TOOL]), invalid, 1)
|
||||
|
||||
def test_empty_and_unfinished_thinking(self):
|
||||
self.assert_output(self.frames([]), '', 0)
|
||||
self.assert_output(self.frames(['reason </thi'], thinking=True), '', 0, 'reason </thi')
|
||||
|
||||
def test_plain_content_and_tool_call_arrive_before_generation_ends(self):
|
||||
def generate(input_ids, streamer, **kwargs):
|
||||
streamer.on_finalized_text('first ')
|
||||
self.assertTrue(continue_generation.wait(5))
|
||||
streamer.on_finalized_text(TOOL)
|
||||
self.assertTrue(finish_generation.wait(5))
|
||||
streamer.on_finalized_text(' last', stream_end=True)
|
||||
|
||||
from threading import Event
|
||||
continue_generation, finish_generation = Event(), Event()
|
||||
server.model.generate = generate
|
||||
stream = server.generate_stream_response(self.request['messages'], 0.7, 0.92, 128)
|
||||
self.addCleanup(continue_generation.set)
|
||||
self.addCleanup(finish_generation.set)
|
||||
first = json.loads(next(stream))
|
||||
if first['choices'][0]['delta'].get('role') == 'assistant':
|
||||
first = json.loads(next(stream))
|
||||
self.assertEqual(first['choices'][0]['delta']['content'], 'first ')
|
||||
continue_generation.set()
|
||||
call = json.loads(next(stream))
|
||||
self.assertEqual(call['choices'][0]['delta']['tool_calls'][0]['function']['name'], 'weather')
|
||||
finish_generation.set()
|
||||
self.assert_output([first, call] + [json.loads(frame) for frame in stream], 'first last', 1)
|
||||
|
||||
def test_real_text_streamer_and_openai_client(self):
|
||||
server.model = ScriptedModel(['Before ', TOOL, ' after'], tokenize=True)
|
||||
with OpenAI(api_key='local-test', base_url='http://testserver/v1',
|
||||
http_client=self.http, max_retries=0) as client:
|
||||
chunks = list(client.chat.completions.create(**self.request, tools=TOOLS, stream=True))
|
||||
self.assertEqual(''.join(c.choices[0].delta.content or '' for c in chunks), 'Before after')
|
||||
calls = [call for c in chunks for call in c.choices[0].delta.tool_calls or []]
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0].index, 0)
|
||||
self.assertEqual(json.loads(calls[0].function.arguments), {'city': '杭州'})
|
||||
|
||||
def test_tiny_minimind_generation(self):
|
||||
torch.manual_seed(42)
|
||||
server.model = MiniMindForCausalLM(MiniMindConfig(
|
||||
hidden_size=32, num_hidden_layers=1, num_attention_heads=4,
|
||||
num_key_value_heads=2, intermediate_size=64,
|
||||
vocab_size=len(self.tokenizer), max_position_embeddings=256)).eval()
|
||||
response = self.http.post('/v1/chat/completions', json={
|
||||
**self.request, 'stream': True, 'max_tokens': 4})
|
||||
frames = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith('data: ')]
|
||||
self.assertFalse(any('error' in frame for frame in frames), frames)
|
||||
self.assertEqual(frames[-1]['choices'][0]['finish_reason'], 'stop')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Loading…
Reference in New Issue