Merge 15b2dd4df4 into f659b55761

scripts/serve_openai_api.py

@@ -7,6 +7,7 @@ import sys
 __package__ = "scripts"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import time
+import uuid
 import torch
 import warnings
 import uvicorn
@@ -103,6 +104,18 @@ def parse_response(text):
 
 
 def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False):
+    response_id = f"chatcmpl-{uuid.uuid4().hex}"
+    created = int(time.time())
+
+    def make_chunk(delta, finish_reason=None):
+        return json.dumps({
+            "id": response_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": "minimind",
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]
+        }, ensure_ascii=False)
+
     try:
         new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)
         inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
@@ -128,6 +141,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
                 queue.put(None)
 
         Thread(target=_generate).start()
+        yield make_chunk({"role": "assistant"})
 
         full_text = ""
         emitted = 0
@@ -139,7 +153,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
                 break
             if isinstance(text, dict):
                 yield json.dumps(text, ensure_ascii=False)
-                continue
+                return
             full_text += text
 
             if not thinking_ended:
@@ -148,28 +162,28 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
                     thinking_ended = True
                     new_r = full_text[emitted:pos]
                     if new_r:
-                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                        yield make_chunk({"reasoning_content": new_r})
                     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)
+                        yield make_chunk({"content": after})
                         emitted = len(full_text)
                 else:
                     new_r = full_text[emitted:]
                     if new_r:
-                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                        yield make_chunk({"reasoning_content": new_r})
                         emitted = len(full_text)
             else:
                 new_c = full_text[emitted:]
                 if new_c:
-                    yield json.dumps({"choices": [{"delta": {"content": new_c}}]}, ensure_ascii=False)
+                    yield make_chunk({"content": new_c})
                     emitted = len(full_text)
 
         _, _, 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)
+            yield make_chunk({"tool_calls": [{"index": i, **call} for i, call in enumerate(tool_calls)]})
+        yield make_chunk({}, "tool_calls" if tool_calls else "stop")
 
     except Exception as e:
         yield json.dumps({"error": str(e)})

tests/test_serve_openai_api.py

@@ -0,0 +1,153 @@
+"""Run with: python -m unittest discover -s tests -p 'test_serve_openai_api.py' -v.
+
+Uses the repository tokenizer, real TextStreamer, FastAPI's HTTP test transport,
+and the OpenAI SDK. No downloaded weights, API credentials or GPU are needed.
+"""
+import json
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import torch
+from fastapi.testclient import TestClient
+from openai import APIError, OpenAI
+from openai.types.chat import ChatCompletionChunk
+from transformers import AutoTokenizer
+
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from scripts import serve_openai_api as server
+
+
+class ScriptedModel:
+    """Send deterministic token IDs through the actual TextStreamer."""
+
+    def __init__(self, tokenizer, text="Hello world!", error=False):
+        self.tokenizer, self.text, self.error = tokenizer, text, error
+
+    def generate(self, input_ids, streamer=None, **kwargs):
+        tokens = self.tokenizer.encode(self.text, add_special_tokens=False)
+        if streamer is not None:
+            streamer.put(input_ids.cpu())
+            for token in tokens:
+                streamer.put(torch.tensor([token]))
+            if self.error:
+                raise RuntimeError("synthetic generation failure")
+            streamer.end()
+        return torch.cat([input_ids, input_ids.new_tensor([tokens])], dim=1)
+
+
+class StreamingAPITests(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(self.tokenizer), create=True
+        )
+        globals_patch.start()
+        self.addCleanup(globals_patch.stop)
+        self.http = TestClient(server.app)
+        self.sdk = OpenAI(
+            api_key="local-test-key", base_url="http://testserver/v1",
+            http_client=self.http, max_retries=0, _strict_response_validation=True
+        )
+        self.addCleanup(self.sdk.close)
+        self.request = {"model": "minimind", "messages": [{"role": "user", "content": "Hi"}]}
+
+    def frames(self, **extra):
+        response = self.http.post("/v1/chat/completions", json={**self.request, "stream": True, **extra})
+        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: ") and line[6:] != "[DONE]"]
+
+    def sdk_chunks(self, **extra):
+        return list(self.sdk.chat.completions.create(**self.request, stream=True, **extra))
+
+    def test_http_chunks_validate_against_sdk_schema(self):
+        frames = self.frames()
+        for frame in frames:
+            chunk = ChatCompletionChunk.model_validate(frame)
+            self.assertEqual(chunk.object, "chat.completion.chunk")
+            self.assertEqual(chunk.model, "minimind")
+            self.assertEqual(chunk.choices[0].index, 0)
+        self.assertEqual(len({frame["id"] for frame in frames}), 1)
+        self.assertEqual(len({frame["created"] for frame in frames}), 1)
+        self.assertNotEqual(frames[0]["id"], self.frames()[0]["id"])
+
+    def test_sdk_reads_plain_content(self):
+        chunks = self.sdk_chunks()
+        self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks), "Hello world!")
+        self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+    def test_sdk_accumulates_assistant_message(self):
+        with self.sdk.beta.chat.completions.stream(**self.request) as stream:
+            for _ in stream:
+                pass
+            message = stream.get_final_completion().choices[0].message
+        self.assertEqual(message.role, "assistant")
+        self.assertEqual(message.content, "Hello world!")
+
+    def test_reasoning_and_content_are_preserved(self):
+        server.model = ScriptedModel(self.tokenizer, "正在思考\n</think>\n答案")
+        chunks = self.sdk_chunks(extra_body={"open_thinking": True})
+        self.assertEqual("".join(getattr(c.choices[0].delta, "reasoning_content", "") or "" for c in chunks), "正在思考\n")
+        self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks), "答案")
+
+    def test_sdk_accumulates_multiple_tool_calls(self):
+        server.model = ScriptedModel(self.tokenizer,
+            '<tool_call>{"name":"weather","arguments":{"city":"杭州"}}</tool_call>\n'
+            '<tool_call>{"name":"clock","arguments":{}}</tool_call>')
+        chunks = self.sdk_chunks()
+        calls = [call for chunk in chunks for call in chunk.choices[0].delta.tool_calls or []]
+        self.assertEqual([call.index for call in calls], [0, 1])
+        self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls")
+        with self.sdk.beta.chat.completions.stream(**self.request) as stream:
+            for _ in stream:
+                pass
+            calls = stream.get_final_completion().choices[0].message.tool_calls
+        self.assertEqual([call.function.name for call in calls], ["weather", "clock"])
+        self.assertEqual(json.loads(calls[0].function.arguments), {"city": "杭州"})
+        self.assertEqual(json.loads(calls[1].function.arguments), {})
+
+    def test_generation_error_never_reports_success(self):
+        for text in ("", "partial output "):
+            with self.subTest(text=text):
+                server.model = ScriptedModel(self.tokenizer, text, error=True)
+                frames = self.frames()
+                self.assertEqual(frames[-1], {"error": "synthetic generation failure"})
+                self.assertFalse(any(c.get("finish_reason") for frame in frames for c in frame.get("choices", [])))
+                with self.assertRaises(APIError):
+                    self.sdk_chunks()
+
+    def test_empty_completion_still_has_valid_terminal_chunk(self):
+        server.model = ScriptedModel(self.tokenizer, "")
+        chunks = self.sdk_chunks()
+        self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+        self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks), "")
+
+    def test_nonstream_response_is_unchanged(self):
+        response = self.sdk.chat.completions.create(**self.request, stream=False)
+        self.assertEqual(response.choices[0].message.role, "assistant")
+        self.assertEqual(response.choices[0].message.content, "Hello world!")
+        self.assertEqual(response.choices[0].finish_reason, "stop")
+
+    def test_tiny_minimind_inference_through_http(self):
+        torch.manual_seed(42)
+        config = 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
+        )
+        server.model = MiniMindForCausalLM(config).eval()
+        chunks = self.sdk_chunks(max_tokens=2)
+        self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+        self.assertTrue(all(c.id == chunks[0].id for c in chunks))
+
+
+if __name__ == "__main__":
+    unittest.main()