Merge 86aad50621 into f659b55761

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 = ""
+        pending = ""
-        emitted = 0
+        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:
+            while pending:
-                pos = full_text.find('</think>')
+                if strip_newlines:
-                if pos >= 0:
+                    pending = pending.lstrip('\n')
-                    thinking_ended = True
+                    if not pending:
-                    new_r = full_text[emitted:pos]
+                        break
-                    if new_r:
+                    strip_newlines = False
-                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                if in_tool_call:
-                    emitted = pos + len('</think>')
+                    end = pending.find('</tool_call>')
-                    after = full_text[emitted:].lstrip('\n')
+                    if end < 0:
-                    emitted = len(full_text) - len(after)
+                        break
-                    if after:
+                    end += len('</tool_call>')
-                        yield json.dumps({"choices": [{"delta": {"content": after}}]}, ensure_ascii=False)
+                    block, pending = pending[:end], pending[end:]
-                        emitted = len(full_text)
+                    _, _, 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:]
+                    pending = pending[len(marker):]
-                    if new_r:
+                    thinking_ended = True
-                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                    strip_newlines = True
-                        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)
 
-        _, _, tool_calls = parse_response(full_text)
+        if pending:
-        if tool_calls:
+            field = 'content' if thinking_ended else 'reasoning_content'
-            yield json.dumps({"choices": [{"delta": {"tool_calls": tool_calls}}]}, ensure_ascii=False)
+            yield json.dumps({"choices": [{"delta": {field: pending}}]}, ensure_ascii=False)
-        yield json.dumps({"choices": [{"delta": {}, "finish_reason": "tool_calls" if tool_calls else "stop"}]}, 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)})

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 = '<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()