diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index ffd717280..2160e883c 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -60,6 +60,27 @@ class Message(BaseModel): tool_calls: list[ToolCall] | None = None +class TokenizeRequest(BaseModel): + """POST /v1/tokenize: count tokens for a prompt without generating. + + Two input shapes, both accepted: + - raw text: ``{"input": "..."}`` (llama.cpp /tokenize-compatible) + - messages: ``{"messages": [...]}`` — rendered through the SAME chat + template a generation would use, so the count matches the + ``usage.prompt_tokens`` a real request would report. + """ + model_config = ConfigDict(extra="allow") + + input: str | list[str] | None = None + messages: list[Message] | None = None + add_special_tokens: bool = True + + +class TokenizeResponse(BaseModel): + tokens: int + token_ids: list[int] = Field(default_factory=list) + + class ChatCompletionRequest(BaseModel): model_config = ConfigDict(extra="allow") diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index dc2f73a97..a0abf2a68 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -18,6 +18,8 @@ CompletionRequest, ModelCard, ModelList, + TokenizeRequest, + TokenizeResponse, ToolChoiceObject, ) from .function_call_parser import ToolCallItem @@ -32,6 +34,7 @@ ToolCallArgsDelta, ToolCallsDelta, ToolCallStart, + count_prompt_tokens, generate_events, generate_full, prerender_error, @@ -123,6 +126,48 @@ async def v1_chat_completions(req: ChatCompletionRequest, request: Request): return gate return await handle_chat_completion(req, request, state, get_model_sampling()) + @app.post("/v1/tokenize") + async def v1_tokenize(req: TokenizeRequest): + """Count tokens without generating. Raw text or chat messages. + + llama.cpp exposes the same endpoint (POST /tokenize) so OpenAI-compatible + clients can pre-validate prompt size against max_model_len before sending. + The messages path renders through the SAME chat template a generation + would use (via count_prompt_tokens), so the count matches the + usage.prompt_tokens a real request would report. + """ + state = get_state() + if req.input is None and req.messages is None: + return create_error_response( + "either 'input' (raw text or list of texts) or 'messages' is required" + ) + try: + if req.messages is not None: + # Messages path: render via the chat template, matching a real generation. + if not req.messages: + return create_error_response("messages: at least one message is required") + spec_msgs = [m.model_dump(exclude_none=True) for m in req.messages] + n_tokens = await count_prompt_tokens( + render_messages(spec_msgs), None, {}, state + ) + return TokenizeResponse(tokens=n_tokens) + # Raw text path: direct tokenizer encode, no chat template. + tokenizer = await asyncio.to_thread(state.frontend_tokenizer) + texts = req.input if isinstance(req.input, list) else [req.input] + token_ids: list[int] = [] + for text in texts: + ids = await asyncio.to_thread( + tokenizer.tokenizer.encode, + text, + add_special_tokens=req.add_special_tokens, + ) + token_ids.extend(ids) + return TokenizeResponse(tokens=len(token_ids), token_ids=token_ids) + except GenerationError as exc: + return create_error_response(str(exc), code=exc.code) + except Exception as exc: # noqa: BLE001 -- tokenizer init/load failure -> server error + return create_error_response(f"tokenization failed: {exc}", status_code=500) + @app.post("/v1/completions") async def v1_completions(req: CompletionRequest, request: Request): log_request("/v1/completions", req, request) diff --git a/tests/server/test_tokenize_endpoint.py b/tests/server/test_tokenize_endpoint.py new file mode 100644 index 000000000..fad7b04bf --- /dev/null +++ b/tests/server/test_tokenize_endpoint.py @@ -0,0 +1,69 @@ +"""Unit tests for the POST /v1/tokenize endpoint.""" +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from freetoken.server.api_models import TokenizeRequest +from freetoken.server.openai_api import register_openai_routes + + +class FakeTokenizer: + """Stands in for the frontend TokenizeManager: raw encode only.""" + + class _Tok: + @staticmethod + def encode(text, add_special_tokens=True, **_): + # Deterministic stand-in: one token per 4 chars, ids are positions. + n = max(1, len(text) // 4) + return list(range(100, 100 + n)) + + tokenizer = _Tok() + + +class FakeState: + def __init__(self) -> None: + from types import SimpleNamespace + + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + default_thinking_mode="auto", + maintenance_state="serving", + ) + + def frontend_tokenizer(self): + return FakeTokenizer() + + +@pytest.fixture() +def client(): + app = FastAPI() + state = FakeState() + register_openai_routes(app, get_state=lambda: state, get_model_sampling=lambda: {}) + return TestClient(app) + + +def test_tokenize_raw_text(client): + resp = client.post("/v1/tokenize", json={"input": "hello world, this is a test"}) + assert resp.status_code == 200 + body = resp.json() + assert body["tokens"] > 0 + assert isinstance(body["token_ids"], list) + assert len(body["token_ids"]) == body["tokens"] + + +def test_tokenize_requires_input_or_messages(client): + resp = client.post("/v1/tokenize", json={}) + assert resp.status_code == 400 + assert "required" in resp.json()["error"]["message"] + + +def test_tokenize_response_model_defaults(): + req = TokenizeRequest(input="hi") + assert req.input == "hi" + assert req.messages is None + assert req.add_special_tokens is True + req2 = TokenizeRequest(messages=[{"role": "user", "content": "hi"}]) + assert req2.input is None \ No newline at end of file