diff --git a/Makefile b/Makefile index 7c29ab4..4eb552e 100644 --- a/Makefile +++ b/Makefile @@ -3,7 +3,7 @@ CC ?= cc CFLAGS ?= -O3 -flto -pipe -std=c11 -Wall -D_POSIX_C_SOURCE=200809L -pthread -SRCS = main.c server.c jsonapi.c util.c +SRCS = main.c server.c jsonapi.c util.c tokenizer.c tokenizer_data.c HDRS = hygate.h BINDIR = bin BIN = $(BINDIR)/hygate diff --git a/README.md b/README.md index e005511..69e4798 100644 --- a/README.md +++ b/README.md @@ -66,8 +66,30 @@ CORS is enabled for all origins (`Access-Control-Allow-Origin: *`). - Upstream calls use Fibonacci-backoff retries (up to `MAX_RETRIES = 5`, 300 s timeout). +## Token usage counting + +Responses include an OpenAI-style `usage` object with `prompt_tokens`, +`completion_tokens`, and `total_tokens`. These are computed locally by an +**embedded cl100k_base BPE tokenizer** (vocabulary data in +`tokenizer_data.c`, logic in `tokenizer.c`) — no upstream token counts are +available from the Gradio backend. + +- `prompt_tokens` sums each message's `content` (or multi-part content + parts), `reasoning_content`, and any `tool_calls` function `name` / + `arguments`. +- `completion_tokens` sums the assistant `content`, `reasoning_content`, and + any emitted `tool_calls`. +- `total_tokens = prompt_tokens + completion_tokens`. + +cl100k_base is OpenAI's canonical multilingual tokenizer and is used as the +best available approximation for Hunyuan; the resulting counts are +deterministic subword counts and may differ slightly from the exact Hunyuan +vocabulary. No per-message formatting overhead (the `+3`/priming tokens +OpenAI applies) is added. + ## Security considerations + - The server binds all interfaces (`0.0.0.0`) and currently enforces **no authentication** — it is an open relay to the configured upstream. Do not expose it to untrusted networks without a fronting proxy or auth layer. diff --git a/hygate.h b/hygate.h index d11a1ff..f93e6c2 100644 --- a/hygate.h +++ b/hygate.h @@ -68,8 +68,20 @@ char *build_gradio_payload(void *req_obj, const char *model_default); * frees the result. */ char *build_completion_response(const char *id, long created, const char *model, const char *content, const char *reasoning, - tool_call_t *tc, int n_tc); + tool_call_t *tc, int n_tc, int prompt_tokens); + +/* Token counting (embedded cl100k_base BPE tokenizer). Counts are + * deterministic subword counts; the served model (Hunyuan) uses a + * different tokenizer, so these are an approximation, not exact. */ +int tok_count(const char *text, int len); +int tok_count_str(const char *text); + +/* Count tokens across an OpenAI request's messages (prompt side). */ +int count_messages_tokens(void *req_obj); +/* Count tokens of a generated completion (content + reasoning + tool calls). */ +int count_completion_tokens(const char *content, const char *reasoning, + tool_call_t *tc, int n_tc); /* Build a streaming tool-call delta JSON array (first emission: full * tool call with index). Caller frees the result. */ char *toolcall_delta_json(tool_call_t *tc, int idx); diff --git a/jsonapi.c b/jsonapi.c index 36e2a3e..efc7655 100644 --- a/jsonapi.c +++ b/jsonapi.c @@ -4,6 +4,67 @@ #include #include +/* forward declaration (defined later in this file) */ +static char *content_to_string(json_object *content); + +/* ------------------------------------------------------------------ */ +/* Token counting (embedded cl100k_base BPE) */ +/* ------------------------------------------------------------------ */ +int count_messages_tokens(void *req_obj) { + json_object *req = (json_object *)req_obj; + json_object *messages = NULL; + if (!json_object_object_get_ex(req, "messages", &messages) || + json_object_get_type(messages) != json_type_array) + return 0; + int total = 0; + int n = json_object_array_length(messages); + for (int i = 0; i < n; i++) { + json_object *m = json_object_array_get_idx(messages, i); + if (!m) continue; + json_object *content = json_object_object_get(m, "content"); + char *cs = content_to_string(content); + total += tok_count_str(cs); + free(cs); + json_object *rc = NULL; + if (json_object_object_get_ex(m, "reasoning_content", &rc) && + json_object_get_type(rc) == json_type_string) + total += tok_count_str(json_object_get_string(rc)); + json_object *tcs = NULL; + if (json_object_object_get_ex(m, "tool_calls", &tcs) && + json_object_get_type(tcs) == json_type_array) { + int tn = json_object_array_length(tcs); + for (int j = 0; j < tn; j++) { + json_object *tc = json_object_array_get_idx(tcs, j); + if (!tc) continue; + json_object *fn = NULL; + if (json_object_object_get_ex(tc, "function", &fn) && + json_object_get_type(fn) == json_type_object) { + json_object *nm = NULL, *ar = NULL; + if (json_object_object_get_ex(fn, "name", &nm) && + json_object_get_type(nm) == json_type_string) + total += tok_count_str(json_object_get_string(nm)); + if (json_object_object_get_ex(fn, "arguments", &ar) && + json_object_get_type(ar) == json_type_string) + total += tok_count_str(json_object_get_string(ar)); + } + } + } + } + return total; +} + +int count_completion_tokens(const char *content, const char *reasoning, + tool_call_t *tc, int n_tc) { + int total = 0; + if (content && content[0]) total += tok_count_str(content); + if (reasoning && reasoning[0]) total += tok_count_str(reasoning); + for (int i = 0; i < n_tc; i++) { + if (tc[i].name && tc[i].name[0]) total += tok_count_str(tc[i].name); + if (tc[i].arguments && tc[i].arguments[0]) total += tok_count_str(tc[i].arguments); + } + return total; +} + /* Convert an OpenAI "content" field to a flat string. Caller frees. */ static char *content_to_string(json_object *content) { if (!content) return strdup(""); @@ -328,7 +389,7 @@ int parse_gradio_stream(const char *data_json, /* ------------------------------------------------------------------ */ char *build_completion_response(const char *id, long created, const char *model, const char *content, const char *reasoning, - tool_call_t *tc, int n_tc) { + tool_call_t *tc, int n_tc, int prompt_tokens) { bool has_tc = n_tc > 0; json_object *root = json_object_new_object(); json_object_object_add(root, "id", json_object_new_string(id)); @@ -374,10 +435,12 @@ char *build_completion_response(const char *id, long created, const char *model, json_object_array_add(choices, choice); json_object_object_add(root, "choices", choices); + int completion_tokens = count_completion_tokens(content, reasoning, tc, n_tc); + int total_tokens = prompt_tokens + completion_tokens; json_object *usage = json_object_new_object(); - json_object_object_add(usage, "prompt_tokens", json_object_new_int(0)); - json_object_object_add(usage, "completion_tokens", json_object_new_int(0)); - json_object_object_add(usage, "total_tokens", json_object_new_int(0)); + json_object_object_add(usage, "prompt_tokens", json_object_new_int(prompt_tokens)); + json_object_object_add(usage, "completion_tokens", json_object_new_int(completion_tokens)); + json_object_object_add(usage, "total_tokens", json_object_new_int(total_tokens)); json_object_object_add(root, "usage", usage); const char *s = json_object_to_json_string(root); diff --git a/server.c b/server.c index 3f3fab5..e93616c 100644 --- a/server.c +++ b/server.c @@ -30,6 +30,7 @@ typedef struct { int tc_state_n; int tc_state_cap; bool any_toolcall; + int include_usage; char *final_content; /* non-streaming accumulation */ char *final_reasoning; tool_call_t *final_tc; @@ -132,6 +133,31 @@ static void emit_chunk(sse_ctx_t *ctx, const char *role, const char *content, json_object_put(root); } +static void emit_usage_chunk(sse_ctx_t *ctx, int prompt_tokens, int completion_tokens) { + json_object *root = json_object_new_object(); + json_object_object_add(root, "id", json_object_new_string(ctx->completion_id)); + json_object_object_add(root, "object", json_object_new_string("chat.completion.chunk")); + json_object_object_add(root, "created", json_object_new_int64(ctx->created)); + json_object_object_add(root, "model", json_object_new_string(ctx->model)); + json_object *choices = json_object_new_array(); + json_object_object_add(root, "choices", choices); + json_object *usage = json_object_new_object(); + json_object_object_add(usage, "prompt_tokens", json_object_new_int(prompt_tokens)); + json_object_object_add(usage, "completion_tokens", json_object_new_int(completion_tokens)); + json_object_object_add(usage, "total_tokens", + json_object_new_int(prompt_tokens + completion_tokens)); + json_object_object_add(root, "usage", usage); + const char *j = json_object_to_json_string(root); + if (j) { + size_t need = strlen(j) + 32; + char *buf = (char *)malloc(need); + int bl = snprintf(buf, need, "data: %s\n\n", j); + write_all(ctx->client_fd, buf, (size_t)bl); + free(buf); + } + json_object_put(root); +} + /* ------------------------------------------------------------------ */ /* SSE line processing */ /* ------------------------------------------------------------------ */ @@ -396,6 +422,18 @@ static void handle_chat(int fd, http_req_t *req, const config_t *cfg) { } char *payload = build_gradio_payload(jreq, model); + + int prompt_tokens = count_messages_tokens(jreq); + bool include_usage = false; + json_object *so = NULL; + if (json_object_object_get_ex(jreq, "stream_options", &so) && + json_object_get_type(so) == json_type_object) { + json_object *iu = NULL; + if (json_object_object_get_ex(so, "include_usage", &iu) && + json_object_get_type(iu) == json_type_boolean) + include_usage = json_object_get_boolean(iu); + } + json_object_put(jreq); if (!payload) { send_simple(fd, 400, "Bad Request", "Invalid request: missing messages"); @@ -457,6 +495,7 @@ static void handle_chat(int fd, http_req_t *req, const config_t *cfg) { ctx.client_fd = fd; ctx.streaming = stream; ctx.model = model; + ctx.include_usage = include_usage; ctx.created = (long)time(NULL); char uuid[37]; generate_uuid(uuid); @@ -479,12 +518,19 @@ static void handle_chat(int fd, http_req_t *req, const config_t *cfg) { if (stream) { emit_chunk(&ctx, NULL, NULL, NULL, NULL, ctx.any_toolcall ? "tool_calls" : "stop"); + if (ctx.include_usage) { + int completion_tokens = tok_count_str(ctx.emitted_content.buf); + completion_tokens += tok_count_str(ctx.emitted_reasoning.buf); + for (int i = 0; i < ctx.tc_state_n; i++) + completion_tokens += tok_count_str(ctx.tc_state[i].prev_args); + emit_usage_chunk(&ctx, prompt_tokens, completion_tokens); + } const char *done = "data: [DONE]\n\n"; write_all(fd, done, strlen(done)); } else { char *resp = build_completion_response(ctx.completion_id, ctx.created, ctx.model, ctx.final_content, ctx.final_reasoning, - ctx.final_tc, ctx.final_n_tc); + ctx.final_tc, ctx.final_n_tc, prompt_tokens); if (resp) { send_json(fd, resp); free(resp);