Add Anthropic count_tokens endpoint

README.md

@@ -1009,6 +1009,7 @@ Supported endpoints:
 - `POST /v1/responses`
 - `POST /v1/completions`
 - `POST /v1/messages`
+- `POST /v1/messages/count_tokens`
 
 The Flash and PRO model endpoints are compatibility aliases. They both report
 the model currently loaded from the GGUF passed with `-m`; the endpoint name does
@@ -1031,6 +1032,11 @@ clients. It accepts `system`, `messages`, `tools`, `tool_choice`, `max_tokens`,
 `temperature`, `top_p`, `top_k`, `stream`, `stop_sequences`, and thinking
 controls. Tool uses are returned as Anthropic `tool_use` blocks.
 
+`/v1/messages/count_tokens` accepts the same Anthropic message fields and
+returns `{"input_tokens":N}` without queueing a generation request. The count
+uses the same rendered prompt and tokenizer path as `/v1/messages`, including
+system text and tool definitions.
+
 Default sampled API generation uses `temperature=1`, `top_p=1`, and
 `min_p=0.05`, so the default filter is relative probability rather than
 nucleus mass. In thinking mode DwarfStar applies those fixed sampling defaults

ds4_server.c

@@ -12361,6 +12361,16 @@ static bool send_models(server *s, int fd) {
     return ok;
 }
 
+static bool send_anthropic_count_tokens(int fd, bool enable_cors, int input_tokens) {
+    buf b = {0};
+    buf_puts(&b, "{\"input_tokens\":");
+    buf_printf(&b, "%d", input_tokens);
+    buf_puts(&b, "}\n");
+    bool ok = http_response(fd, enable_cors, 200, "application/json", b.ptr);
+    buf_free(&b);
+    return ok;
+}
+
 static void client_done(server *s) {
     pthread_mutex_lock(&s->mu);
     if (s->clients > 0) s->clients--;
@@ -12408,7 +12418,18 @@ static void *client_main(void *arg) {
     char err[160];
     bool ok = false;
     const int ctx_size = s->ctx_size;
-    if (!strcmp(hr.method, "POST") && !strcmp(hr.path, "/v1/messages")) {
+    if (!strcmp(hr.method, "POST") && !strcmp(hr.path, "/v1/messages/count_tokens")) {
+        ok = parse_anthropic_request(s->engine, s, hr.body, s->default_tokens,
+                                     ctx_size, &req, err, sizeof(err));
+        http_request_free(&hr);
+        if (!ok) {
+            http_error(fd, s->enable_cors, 400, err);
+            goto done;
+        }
+        send_anthropic_count_tokens(fd, s->enable_cors, req.prompt.len);
+        request_free(&req);
+        goto done;
+    } else if (!strcmp(hr.method, "POST") && !strcmp(hr.path, "/v1/messages")) {
         ok = parse_anthropic_request(s->engine, s, hr.body, s->default_tokens,
                                      ctx_size, &req, err, sizeof(err));
     } else if (!strcmp(hr.method, "POST") && !strcmp(hr.path, "/v1/chat/completions")) {
@@ -13595,6 +13616,24 @@ static void test_cors_preflight_response_is_no_content(void) {
     close(sv[1]);
 }
 
+static void test_anthropic_count_tokens_response_shape(void) {
+    int sv[2];
+    TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv) == 0);
+    if (sv[0] < 0 || sv[1] < 0) return;
+
+    TEST_ASSERT(send_anthropic_count_tokens(sv[0], true, 123));
+    shutdown(sv[0], SHUT_WR);
+    char *out = read_socket_text(sv[1]);
+    TEST_ASSERT(strstr(out, "HTTP/1.1 200 OK") != NULL);
+    TEST_ASSERT(strstr(out, "Content-Type: application/json") != NULL);
+    TEST_ASSERT(strstr(out, "Access-Control-Allow-Origin: *") != NULL);
+    TEST_ASSERT(strstr(out, "{\"input_tokens\":123}\n") != NULL);
+    TEST_ASSERT(strstr(out, "\"usage\"") == NULL);
+    free(out);
+    close(sv[0]);
+    close(sv[1]);
+}
+
 static void test_cors_sse_headers(void) {
     int sv[2];
     TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv) == 0);
@@ -17441,6 +17480,7 @@ static void ds4_server_unit_tests_run(void) {
     test_context_length_error_uses_protocol_standard_shape();
     test_cors_headers_are_opt_in();
     test_cors_preflight_response_is_no_content();
+    test_anthropic_count_tokens_response_shape();
     test_cors_sse_headers();
     test_anthropic_live_stream_sends_incremental_blocks();
     test_anthropic_usage_reports_cache_details();

tests/test_anthropic_count_tokens.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+"""Check that Anthropic count_tokens matches /v1/messages input usage."""
+
+import argparse
+import json
+import sys
+import urllib.error
+import urllib.request
+
+
+def post_json(base_url, path, payload, timeout):
+    body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
+    request = urllib.request.Request(
+        base_url.rstrip("/") + path,
+        data=body,
+        headers={
+            "Content-Type": "application/json",
+            "anthropic-version": "2023-06-01",
+        },
+        method="POST",
+    )
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            return json.loads(response.read())
+    except urllib.error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError("%s returned HTTP %d: %s" %
+                           (path, exc.code, detail)) from exc
+
+
+def require_nonnegative_int(value, field):
+    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+        raise RuntimeError("%s must be a non-negative integer" % field)
+    return value
+
+
+def message_input_tokens(response):
+    usage = response.get("usage")
+    if not isinstance(usage, dict):
+        raise RuntimeError("/v1/messages response is missing usage")
+    return sum(
+        require_nonnegative_int(usage.get(field, 0), "usage.%s" % field)
+        for field in (
+            "input_tokens",
+            "cache_read_input_tokens",
+            "cache_creation_input_tokens",
+        )
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compare /v1/messages/count_tokens with /v1/messages usage"
+    )
+    parser.add_argument("--url", default="http://127.0.0.1:8000")
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--timeout", type=float, default=1800.0)
+    args = parser.parse_args()
+    if args.timeout <= 0:
+        parser.error("--timeout must be positive")
+
+    payload = {
+        "model": args.model,
+        "system": "You are a concise assistant.",
+        "messages": [
+            {
+                "role": "user",
+                "content": "Explain token counting briefly.",
+            }
+        ],
+        "tools": [
+            {
+                "name": "lookup",
+                "description": "Look something up",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string"},
+                    },
+                    "required": ["query"],
+                },
+            }
+        ],
+        "max_tokens": 1,
+        "temperature": 0,
+    }
+
+    count_response = post_json(
+        args.url, "/v1/messages/count_tokens", payload, args.timeout
+    )
+    count = require_nonnegative_int(
+        count_response.get("input_tokens"), "count_tokens.input_tokens"
+    )
+
+    message_response = post_json(
+        args.url, "/v1/messages", payload, args.timeout
+    )
+    message_count = message_input_tokens(message_response)
+    if count != message_count:
+        print(
+            "FAIL model=%s count_tokens=%d messages_input_tokens=%d" %
+            (args.model, count, message_count),
+            file=sys.stderr,
+        )
+        return 1
+
+    print(
+        "PASS model=%s input_tokens=%d" % (args.model, count)
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())