Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Track input and output tokens per minute separately (#28097)

Browse source 
This PR adds tracking for input and output tokens per minute separately
from the current aggregate tokens per minute.

We are not yet rate-limiting based on these measures.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-04 11:37:06 -04:00 • committed by GitHub
parent c94b587e1a
commit 5fe86f7e70
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 101 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

14
crates/collab/src/llm/db/tests/usage_tests.rs
View file

@ -83,6 +83,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 2,
tokens_this_minute: 3000,
input_tokens_this_minute: 3000,
output_tokens_this_minute: 0,
tokens_this_day: 3000,
tokens_this_month: TokenUsage {
input: 3000,
@ -102,6 +104,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 1,
tokens_this_minute: 2000,
input_tokens_this_minute: 2000,
output_tokens_this_minute: 0,
tokens_this_day: 3000,
tokens_this_month: TokenUsage {
input: 3000,
@ -140,6 +144,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 2,
tokens_this_minute: 5000,
input_tokens_this_minute: 5000,
output_tokens_this_minute: 0,
tokens_this_day: 6000,
tokens_this_month: TokenUsage {
input: 6000,
@ -160,6 +166,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 0,
tokens_this_minute: 0,
input_tokens_this_minute: 0,
output_tokens_this_minute: 0,
tokens_this_day: 5000,
tokens_this_month: TokenUsage {
input: 6000,
@ -197,6 +205,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 1,
tokens_this_minute: 4000,
input_tokens_this_minute: 4000,
output_tokens_this_minute: 0,
tokens_this_day: 9000,
tokens_this_month: TokenUsage {
input: 10000,
@ -240,6 +250,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 1,
tokens_this_minute: 1500,
input_tokens_this_minute: 1500,
output_tokens_this_minute: 0,
tokens_this_day: 1500,
tokens_this_month: TokenUsage {
input: 1000,
@ -278,6 +290,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
Usage {
requests_this_minute: 2,
tokens_this_minute: 2800,
input_tokens_this_minute: 2500,
output_tokens_this_minute: 0,
tokens_this_day: 2800,
tokens_this_month: TokenUsage {
input: 2000,