diff --git a/scripts/serve_openai_api.py b/scripts/serve_openai_api.py index 09e7fc1..7073ecc 100644 --- a/scripts/serve_openai_api.py +++ b/scripts/serve_openai_api.py @@ -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('') - 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('') - 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('') + if end < 0: + break + end += len('') + 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 = '' if thinking_ended else '' + 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)}) diff --git a/tests/test_stream_tool_calls.py b/tests/test_stream_tool_calls.py new file mode 100644 index 0000000..e51332f --- /dev/null +++ b/tests/test_stream_tool_calls.py @@ -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 = '{"name":"weather","arguments":{"city":"杭州"}}' +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 \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 {"name":', + 'bad json', '[]'): + with self.subTest(text=text): + self.assert_output(self.frames(list(text)), text, 0) + invalid = 'bad json' + 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