diff --git a/scripts/serve_openai_api.py b/scripts/serve_openai_api.py
index e9cf871..0336742 100644
--- a/scripts/serve_openai_api.py
+++ b/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('')
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)})
diff --git a/tests/test_serve_openai_api.py b/tests/test_serve_openai_api.py
new file mode 100644
index 0000000..baf6a28
--- /dev/null
+++ b/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\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,
+ '{"name":"weather","arguments":{"city":"杭州"}}\n'
+ '{"name":"clock","arguments":{}}')
+ 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()