Track token usage per message in DB

[jankrepl]

Why

Tracking token usage will be useful for internal analytics but also potentially we can expose this to the user.

Description

No estimates - we use the response bodies of openrouter/openai.

Create a new table called token_usage and it should have the following keys (feel free to change the key names - I am not claiming they are optimal).

• id: UUID - primary key
• message_id: UUID - foreign key linking to messages
• type: Enum["input-cached", "input-notcached", "completion"]
• task: Enum["chat-completion", "tool-selection"]
• count: int

There is some inner dependency in a sense that

• task="tool-selection" will be linked to messages of entity USER and task="chat-completion" will be linked to messages of entity AI_MESSAGE and AI_TOOL

Note that the idea is to make it fine-grained - we can do aggregations in SQL.

feel free to make modifications and challenge the above schema

Example

messages table

message_id	entity	...
ab2546b2-a803-4316-89ee-4e76527df7a0	USER
8e0339ac-810e-4949-95c2-a651894b0ba8	AI_MESSAGE
3498b267-71fe-4b8e-a562-092dc42a129c	USER
f815e2c8-94ff-430c-8a78-ad1cc5ea872a	AI_TOOL
8fc14491-0966-46bc-b16b-1ea693cb99ac	TOOL
793d54ba-5c2e-42d9-93a0-4dcd2ad51e65	AI_MESSAGE

token_usage table

id	message_id	type	task	count
8a20fd68-37ab-456a-a5fd-98f0f0dc7921	ab2546b2-a803-4316-89ee-4e76527df7a0	input-cached	tool-selection	33
26804c92-2e22-4cfd-a86c-59d98ee2ee80	ab2546b2-a803-4316-89ee-4e76527df7a0	input-notcached	tool-selection	120
b6932504-ced9-4701-bfd0-dd95b88a9942	ab2546b2-a803-4316-89ee-4e76527df7a0	completion	tool-selection	259
7440508a-a5a2-4f92-8911-f74cafea41c7	8e0339ac-810e-4949-95c2-a651894b0ba8	input-cached	chat-completion	510
86d0f066-7799-4d51-9d40-622e760b4b16	8e0339ac-810e-4949-95c2-a651894b0ba8	input-notcached	chat-completion	772
1b2edfc7-eee9-4e9c-a016-fd1812d83e7c	8e0339ac-810e-4949-95c2-a651894b0ba8	completion	chat-completion	1200
432e20d5-d8dd-4c0e-afce-6a7d9706340b	3498b267-71fe-4b8e-a562-092dc42a129c	input-cached	tool-selection	35
9295b68b-d891-46b6-b707-947424b63f6c	3498b267-71fe-4b8e-a562-092dc42a129c	input-notcached	tool-selection	172
e50bdedf-d9ac-4509-9921-fd9e9ad15d19	3498b267-71fe-4b8e-a562-092dc42a129c	completion	tool-selection	351
9475eb84-44c5-45e7-862d-b5c8b3525d0a	f815e2c8-94ff-430c-8a78-ad1cc5ea872a	input-cached	chat-completion	22
536fa3d9-4281-46f9-b133-0e2dbb2bdd55	f815e2c8-94ff-430c-8a78-ad1cc5ea872a	input-notcached	chat-completion	518
bc284cf7-0b4e-45d3-8b35-0359a8e236e7	f815e2c8-94ff-430c-8a78-ad1cc5ea872a	completion	chat-completion	932
2670766a-0cdb-4ab1-b4f7-dc04293266bc	793d54ba-5c2e-42d9-93a0-4dcd2ad51e65	input-cached	chat-completion	523
aae73226-722b-47e2-994f-c1fb5282b7de	793d54ba-5c2e-42d9-93a0-4dcd2ad51e65	input-notcached	chat-completion	513
2d97320c-dc40-4524-bc78-c54328389442	793d54ba-5c2e-42d9-93a0-4dcd2ad51e65	completion	chat-completion	1652

Comments

• We have more endpoints (e.g. question suggestions, title generation ...), however, for now I would not include them.
• I would not track cost since the pricing of models changes over time - the token counts should be way more stable