154 lines
7.1 KiB
Python
154 lines
7.1 KiB
Python
"""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()
|