1mod authorization;
2pub mod db;
3mod telemetry;
4mod token;
5
6use crate::api::events::SnowflakeRow;
7use crate::api::CloudflareIpCountryHeader;
8use crate::build_kinesis_client;
9use crate::{
10 build_clickhouse_client, db::UserId, executor::Executor, Cents, Config, Error, Result,
11};
12use anyhow::{anyhow, Context as _};
13use authorization::authorize_access_to_language_model;
14use axum::routing::get;
15use axum::{
16 body::Body,
17 http::{self, HeaderName, HeaderValue, Request, StatusCode},
18 middleware::{self, Next},
19 response::{IntoResponse, Response},
20 routing::post,
21 Extension, Json, Router, TypedHeader,
22};
23use chrono::{DateTime, Duration, Utc};
24use collections::HashMap;
25use db::TokenUsage;
26use db::{usage_measure::UsageMeasure, ActiveUserCount, LlmDatabase};
27use futures::{Stream, StreamExt as _};
28use reqwest_client::ReqwestClient;
29use rpc::{
30 proto::Plan, LanguageModelProvider, PerformCompletionParams, EXPIRED_LLM_TOKEN_HEADER_NAME,
31};
32use rpc::{ListModelsResponse, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME};
33use serde_json::json;
34use std::{
35 pin::Pin,
36 sync::Arc,
37 task::{Context, Poll},
38};
39use strum::IntoEnumIterator;
40use telemetry::{report_llm_rate_limit, report_llm_usage, LlmRateLimitEventRow, LlmUsageEventRow};
41use tokio::sync::RwLock;
42use util::ResultExt;
43
44pub use token::*;
45
46pub struct LlmState {
47 pub config: Config,
48 pub executor: Executor,
49 pub db: Arc<LlmDatabase>,
50 pub http_client: ReqwestClient,
51 pub kinesis_client: Option<aws_sdk_kinesis::Client>,
52 pub clickhouse_client: Option<clickhouse::Client>,
53 active_user_count_by_model:
54 RwLock<HashMap<(LanguageModelProvider, String), (DateTime<Utc>, ActiveUserCount)>>,
55}
56
57const ACTIVE_USER_COUNT_CACHE_DURATION: Duration = Duration::seconds(30);
58
59impl LlmState {
60 pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {
61 let database_url = config
62 .llm_database_url
63 .as_ref()
64 .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
65 let max_connections = config
66 .llm_database_max_connections
67 .ok_or_else(|| anyhow!("missing LLM_DATABASE_MAX_CONNECTIONS"))?;
68
69 let mut db_options = db::ConnectOptions::new(database_url);
70 db_options.max_connections(max_connections);
71 let mut db = LlmDatabase::new(db_options, executor.clone()).await?;
72 db.initialize().await?;
73
74 let db = Arc::new(db);
75
76 let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
77 let http_client =
78 ReqwestClient::user_agent(&user_agent).context("failed to construct http client")?;
79
80 let this = Self {
81 executor,
82 db,
83 http_client,
84 kinesis_client: if config.kinesis_access_key.is_some() {
85 build_kinesis_client(&config).await.log_err()
86 } else {
87 None
88 },
89 clickhouse_client: config
90 .clickhouse_url
91 .as_ref()
92 .and_then(|_| build_clickhouse_client(&config).log_err()),
93 active_user_count_by_model: RwLock::new(HashMap::default()),
94 config,
95 };
96
97 Ok(Arc::new(this))
98 }
99
100 pub async fn get_active_user_count(
101 &self,
102 provider: LanguageModelProvider,
103 model: &str,
104 ) -> Result<ActiveUserCount> {
105 let now = Utc::now();
106
107 {
108 let active_user_count_by_model = self.active_user_count_by_model.read().await;
109 if let Some((last_updated, count)) =
110 active_user_count_by_model.get(&(provider, model.to_string()))
111 {
112 if now - *last_updated < ACTIVE_USER_COUNT_CACHE_DURATION {
113 return Ok(*count);
114 }
115 }
116 }
117
118 let mut cache = self.active_user_count_by_model.write().await;
119 let new_count = self.db.get_active_user_count(provider, model, now).await?;
120 cache.insert((provider, model.to_string()), (now, new_count));
121 Ok(new_count)
122 }
123}
124
125pub fn routes() -> Router<(), Body> {
126 Router::new()
127 .route("/models", get(list_models))
128 .route("/completion", post(perform_completion))
129 .layer(middleware::from_fn(validate_api_token))
130}
131
132async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoResponse {
133 let token = req
134 .headers()
135 .get(http::header::AUTHORIZATION)
136 .and_then(|header| header.to_str().ok())
137 .ok_or_else(|| {
138 Error::http(
139 StatusCode::BAD_REQUEST,
140 "missing authorization header".to_string(),
141 )
142 })?
143 .strip_prefix("Bearer ")
144 .ok_or_else(|| {
145 Error::http(
146 StatusCode::BAD_REQUEST,
147 "invalid authorization header".to_string(),
148 )
149 })?;
150
151 let state = req.extensions().get::<Arc<LlmState>>().unwrap();
152 match LlmTokenClaims::validate(token, &state.config) {
153 Ok(claims) => {
154 if state.db.is_access_token_revoked(&claims.jti).await? {
155 return Err(Error::http(
156 StatusCode::UNAUTHORIZED,
157 "unauthorized".to_string(),
158 ));
159 }
160
161 tracing::Span::current()
162 .record("user_id", claims.user_id)
163 .record("login", claims.github_user_login.clone())
164 .record("authn.jti", &claims.jti)
165 .record("is_staff", claims.is_staff);
166
167 req.extensions_mut().insert(claims);
168 Ok::<_, Error>(next.run(req).await.into_response())
169 }
170 Err(ValidateLlmTokenError::Expired) => Err(Error::Http(
171 StatusCode::UNAUTHORIZED,
172 "unauthorized".to_string(),
173 [(
174 HeaderName::from_static(EXPIRED_LLM_TOKEN_HEADER_NAME),
175 HeaderValue::from_static("true"),
176 )]
177 .into_iter()
178 .collect(),
179 )),
180 Err(_err) => Err(Error::http(
181 StatusCode::UNAUTHORIZED,
182 "unauthorized".to_string(),
183 )),
184 }
185}
186
187async fn list_models(
188 Extension(state): Extension<Arc<LlmState>>,
189 Extension(claims): Extension<LlmTokenClaims>,
190 country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
191) -> Result<Json<ListModelsResponse>> {
192 let country_code = country_code_header.map(|header| header.to_string());
193
194 let mut accessible_models = Vec::new();
195
196 for (provider, model) in state.db.all_models() {
197 let authorize_result = authorize_access_to_language_model(
198 &state.config,
199 &claims,
200 country_code.as_deref(),
201 provider,
202 &model.name,
203 );
204
205 if authorize_result.is_ok() {
206 accessible_models.push(rpc::LanguageModel {
207 provider,
208 name: model.name,
209 });
210 }
211 }
212
213 Ok(Json(ListModelsResponse {
214 models: accessible_models,
215 }))
216}
217
218async fn perform_completion(
219 Extension(state): Extension<Arc<LlmState>>,
220 Extension(claims): Extension<LlmTokenClaims>,
221 country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
222 Json(params): Json<PerformCompletionParams>,
223) -> Result<impl IntoResponse> {
224 let model = normalize_model_name(
225 state.db.model_names_for_provider(params.provider),
226 params.model,
227 );
228
229 authorize_access_to_language_model(
230 &state.config,
231 &claims,
232 country_code_header
233 .map(|header| header.to_string())
234 .as_deref(),
235 params.provider,
236 &model,
237 )?;
238
239 check_usage_limit(&state, params.provider, &model, &claims).await?;
240
241 let stream = match params.provider {
242 LanguageModelProvider::Anthropic => {
243 let api_key = if claims.is_staff {
244 state
245 .config
246 .anthropic_staff_api_key
247 .as_ref()
248 .context("no Anthropic AI staff API key configured on the server")?
249 } else {
250 state
251 .config
252 .anthropic_api_key
253 .as_ref()
254 .context("no Anthropic AI API key configured on the server")?
255 };
256
257 let mut request: anthropic::Request =
258 serde_json::from_str(params.provider_request.get())?;
259
260 // Override the model on the request with the latest version of the model that is
261 // known to the server.
262 //
263 // Right now, we use the version that's defined in `model.id()`, but we will likely
264 // want to change this code once a new version of an Anthropic model is released,
265 // so that users can use the new version, without having to update Zed.
266 request.model = match model.as_str() {
267 "claude-3-5-sonnet" => anthropic::Model::Claude3_5Sonnet.id().to_string(),
268 "claude-3-opus" => anthropic::Model::Claude3Opus.id().to_string(),
269 "claude-3-haiku" => anthropic::Model::Claude3Haiku.id().to_string(),
270 "claude-3-sonnet" => anthropic::Model::Claude3Sonnet.id().to_string(),
271 _ => request.model,
272 };
273
274 let (chunks, rate_limit_info) = anthropic::stream_completion_with_rate_limit_info(
275 &state.http_client,
276 anthropic::ANTHROPIC_API_URL,
277 api_key,
278 request,
279 )
280 .await
281 .map_err(|err| match err {
282 anthropic::AnthropicError::ApiError(ref api_error) => match api_error.code() {
283 Some(anthropic::ApiErrorCode::RateLimitError) => {
284 tracing::info!(
285 target: "upstream rate limit exceeded",
286 user_id = claims.user_id,
287 login = claims.github_user_login,
288 authn.jti = claims.jti,
289 is_staff = claims.is_staff,
290 provider = params.provider.to_string(),
291 model = model
292 );
293
294 Error::http(
295 StatusCode::TOO_MANY_REQUESTS,
296 "Upstream Anthropic rate limit exceeded.".to_string(),
297 )
298 }
299 Some(anthropic::ApiErrorCode::InvalidRequestError) => {
300 Error::http(StatusCode::BAD_REQUEST, api_error.message.clone())
301 }
302 Some(anthropic::ApiErrorCode::OverloadedError) => {
303 Error::http(StatusCode::SERVICE_UNAVAILABLE, api_error.message.clone())
304 }
305 Some(_) => {
306 Error::http(StatusCode::INTERNAL_SERVER_ERROR, api_error.message.clone())
307 }
308 None => Error::Internal(anyhow!(err)),
309 },
310 anthropic::AnthropicError::Other(err) => Error::Internal(err),
311 })?;
312
313 if let Some(rate_limit_info) = rate_limit_info {
314 tracing::info!(
315 target: "upstream rate limit",
316 is_staff = claims.is_staff,
317 provider = params.provider.to_string(),
318 model = model,
319 tokens_remaining = rate_limit_info.tokens_remaining,
320 requests_remaining = rate_limit_info.requests_remaining,
321 requests_reset = ?rate_limit_info.requests_reset,
322 tokens_reset = ?rate_limit_info.tokens_reset,
323 );
324 }
325
326 chunks
327 .map(move |event| {
328 let chunk = event?;
329 let (
330 input_tokens,
331 output_tokens,
332 cache_creation_input_tokens,
333 cache_read_input_tokens,
334 ) = match &chunk {
335 anthropic::Event::MessageStart {
336 message: anthropic::Response { usage, .. },
337 }
338 | anthropic::Event::MessageDelta { usage, .. } => (
339 usage.input_tokens.unwrap_or(0) as usize,
340 usage.output_tokens.unwrap_or(0) as usize,
341 usage.cache_creation_input_tokens.unwrap_or(0) as usize,
342 usage.cache_read_input_tokens.unwrap_or(0) as usize,
343 ),
344 _ => (0, 0, 0, 0),
345 };
346
347 anyhow::Ok(CompletionChunk {
348 bytes: serde_json::to_vec(&chunk).unwrap(),
349 input_tokens,
350 output_tokens,
351 cache_creation_input_tokens,
352 cache_read_input_tokens,
353 })
354 })
355 .boxed()
356 }
357 LanguageModelProvider::OpenAi => {
358 let api_key = state
359 .config
360 .openai_api_key
361 .as_ref()
362 .context("no OpenAI API key configured on the server")?;
363 let chunks = open_ai::stream_completion(
364 &state.http_client,
365 open_ai::OPEN_AI_API_URL,
366 api_key,
367 serde_json::from_str(params.provider_request.get())?,
368 )
369 .await?;
370
371 chunks
372 .map(|event| {
373 event.map(|chunk| {
374 let input_tokens =
375 chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;
376 let output_tokens =
377 chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;
378 CompletionChunk {
379 bytes: serde_json::to_vec(&chunk).unwrap(),
380 input_tokens,
381 output_tokens,
382 cache_creation_input_tokens: 0,
383 cache_read_input_tokens: 0,
384 }
385 })
386 })
387 .boxed()
388 }
389 LanguageModelProvider::Google => {
390 let api_key = state
391 .config
392 .google_ai_api_key
393 .as_ref()
394 .context("no Google AI API key configured on the server")?;
395 let chunks = google_ai::stream_generate_content(
396 &state.http_client,
397 google_ai::API_URL,
398 api_key,
399 serde_json::from_str(params.provider_request.get())?,
400 )
401 .await?;
402
403 chunks
404 .map(|event| {
405 event.map(|chunk| {
406 // TODO - implement token counting for Google AI
407 CompletionChunk {
408 bytes: serde_json::to_vec(&chunk).unwrap(),
409 input_tokens: 0,
410 output_tokens: 0,
411 cache_creation_input_tokens: 0,
412 cache_read_input_tokens: 0,
413 }
414 })
415 })
416 .boxed()
417 }
418 };
419
420 Ok(Response::new(Body::wrap_stream(TokenCountingStream {
421 state,
422 claims,
423 provider: params.provider,
424 model,
425 tokens: TokenUsage::default(),
426 inner_stream: stream,
427 })))
428}
429
430fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
431 if let Some(known_model_name) = known_models
432 .iter()
433 .filter(|known_model_name| name.starts_with(known_model_name.as_str()))
434 .max_by_key(|known_model_name| known_model_name.len())
435 {
436 known_model_name.to_string()
437 } else {
438 name
439 }
440}
441
442/// The maximum monthly spending an individual user can reach on the free tier
443/// before they have to pay.
444pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10);
445
446/// The default value to use for maximum spend per month if the user did not
447/// explicitly set a maximum spend.
448///
449/// Used to prevent surprise bills.
450pub const DEFAULT_MAX_MONTHLY_SPEND: Cents = Cents::from_dollars(10);
451
452async fn check_usage_limit(
453 state: &Arc<LlmState>,
454 provider: LanguageModelProvider,
455 model_name: &str,
456 claims: &LlmTokenClaims,
457) -> Result<()> {
458 if claims.is_staff {
459 return Ok(());
460 }
461
462 let model = state.db.model(provider, model_name)?;
463 let usage = state
464 .db
465 .get_usage(
466 UserId::from_proto(claims.user_id),
467 provider,
468 model_name,
469 Utc::now(),
470 )
471 .await?;
472 let free_tier = claims.free_tier_monthly_spending_limit();
473
474 if usage.spending_this_month >= free_tier {
475 if !claims.has_llm_subscription {
476 return Err(Error::http(
477 StatusCode::PAYMENT_REQUIRED,
478 "Maximum spending limit reached for this month.".to_string(),
479 ));
480 }
481
482 if (usage.spending_this_month - free_tier) >= Cents(claims.max_monthly_spend_in_cents) {
483 return Err(Error::Http(
484 StatusCode::FORBIDDEN,
485 "Maximum spending limit reached for this month.".to_string(),
486 [(
487 HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME),
488 HeaderValue::from_static("true"),
489 )]
490 .into_iter()
491 .collect(),
492 ));
493 }
494 }
495
496 let active_users = state.get_active_user_count(provider, model_name).await?;
497
498 let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);
499 let users_in_recent_days = active_users.users_in_recent_days.max(1);
500
501 let per_user_max_requests_per_minute =
502 model.max_requests_per_minute as usize / users_in_recent_minutes;
503 let per_user_max_tokens_per_minute =
504 model.max_tokens_per_minute as usize / users_in_recent_minutes;
505 let per_user_max_tokens_per_day = model.max_tokens_per_day as usize / users_in_recent_days;
506
507 let checks = [
508 (
509 usage.requests_this_minute,
510 per_user_max_requests_per_minute,
511 UsageMeasure::RequestsPerMinute,
512 ),
513 (
514 usage.tokens_this_minute,
515 per_user_max_tokens_per_minute,
516 UsageMeasure::TokensPerMinute,
517 ),
518 (
519 usage.tokens_this_day,
520 per_user_max_tokens_per_day,
521 UsageMeasure::TokensPerDay,
522 ),
523 ];
524
525 for (used, limit, usage_measure) in checks {
526 if used > limit {
527 let resource = match usage_measure {
528 UsageMeasure::RequestsPerMinute => "requests_per_minute",
529 UsageMeasure::TokensPerMinute => "tokens_per_minute",
530 UsageMeasure::TokensPerDay => "tokens_per_day",
531 };
532
533 tracing::info!(
534 target: "user rate limit",
535 user_id = claims.user_id,
536 login = claims.github_user_login,
537 authn.jti = claims.jti,
538 is_staff = claims.is_staff,
539 provider = provider.to_string(),
540 model = model.name,
541 requests_this_minute = usage.requests_this_minute,
542 tokens_this_minute = usage.tokens_this_minute,
543 tokens_this_day = usage.tokens_this_day,
544 users_in_recent_minutes = users_in_recent_minutes,
545 users_in_recent_days = users_in_recent_days,
546 max_requests_per_minute = per_user_max_requests_per_minute,
547 max_tokens_per_minute = per_user_max_tokens_per_minute,
548 max_tokens_per_day = per_user_max_tokens_per_day,
549 );
550
551 SnowflakeRow::new(
552 "Language Model Rate Limited",
553 claims.metrics_id,
554 claims.is_staff,
555 claims.system_id.clone(),
556 json!({
557 "usage": usage,
558 "users_in_recent_minutes": users_in_recent_minutes,
559 "users_in_recent_days": users_in_recent_days,
560 "max_requests_per_minute": per_user_max_requests_per_minute,
561 "max_tokens_per_minute": per_user_max_tokens_per_minute,
562 "max_tokens_per_day": per_user_max_tokens_per_day,
563 "plan": match claims.plan {
564 Plan::Free => "free".to_string(),
565 Plan::ZedPro => "zed_pro".to_string(),
566 },
567 "model": model.name.clone(),
568 "provider": provider.to_string(),
569 "usage_measure": resource.to_string(),
570 }),
571 )
572 .write(&state.kinesis_client, &state.config.kinesis_stream)
573 .await
574 .log_err();
575
576 if let Some(client) = state.clickhouse_client.as_ref() {
577 report_llm_rate_limit(
578 client,
579 LlmRateLimitEventRow {
580 time: Utc::now().timestamp_millis(),
581 user_id: claims.user_id as i32,
582 is_staff: claims.is_staff,
583 plan: match claims.plan {
584 Plan::Free => "free".to_string(),
585 Plan::ZedPro => "zed_pro".to_string(),
586 },
587 model: model.name.clone(),
588 provider: provider.to_string(),
589 usage_measure: resource.to_string(),
590 requests_this_minute: usage.requests_this_minute as u64,
591 tokens_this_minute: usage.tokens_this_minute as u64,
592 tokens_this_day: usage.tokens_this_day as u64,
593 users_in_recent_minutes: users_in_recent_minutes as u64,
594 users_in_recent_days: users_in_recent_days as u64,
595 max_requests_per_minute: per_user_max_requests_per_minute as u64,
596 max_tokens_per_minute: per_user_max_tokens_per_minute as u64,
597 max_tokens_per_day: per_user_max_tokens_per_day as u64,
598 },
599 )
600 .await
601 .log_err();
602 }
603
604 return Err(Error::http(
605 StatusCode::TOO_MANY_REQUESTS,
606 format!("Rate limit exceeded. Maximum {} reached.", resource),
607 ));
608 }
609 }
610
611 Ok(())
612}
613
614struct CompletionChunk {
615 bytes: Vec<u8>,
616 input_tokens: usize,
617 output_tokens: usize,
618 cache_creation_input_tokens: usize,
619 cache_read_input_tokens: usize,
620}
621
622struct TokenCountingStream<S> {
623 state: Arc<LlmState>,
624 claims: LlmTokenClaims,
625 provider: LanguageModelProvider,
626 model: String,
627 tokens: TokenUsage,
628 inner_stream: S,
629}
630
631impl<S> Stream for TokenCountingStream<S>
632where
633 S: Stream<Item = Result<CompletionChunk, anyhow::Error>> + Unpin,
634{
635 type Item = Result<Vec<u8>, anyhow::Error>;
636
637 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
638 match Pin::new(&mut self.inner_stream).poll_next(cx) {
639 Poll::Ready(Some(Ok(mut chunk))) => {
640 chunk.bytes.push(b'\n');
641 self.tokens.input += chunk.input_tokens;
642 self.tokens.output += chunk.output_tokens;
643 self.tokens.input_cache_creation += chunk.cache_creation_input_tokens;
644 self.tokens.input_cache_read += chunk.cache_read_input_tokens;
645 Poll::Ready(Some(Ok(chunk.bytes)))
646 }
647 Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
648 Poll::Ready(None) => Poll::Ready(None),
649 Poll::Pending => Poll::Pending,
650 }
651 }
652}
653
654impl<S> Drop for TokenCountingStream<S> {
655 fn drop(&mut self) {
656 let state = self.state.clone();
657 let claims = self.claims.clone();
658 let provider = self.provider;
659 let model = std::mem::take(&mut self.model);
660 let tokens = self.tokens;
661 self.state.executor.spawn_detached(async move {
662 let usage = state
663 .db
664 .record_usage(
665 UserId::from_proto(claims.user_id),
666 claims.is_staff,
667 provider,
668 &model,
669 tokens,
670 claims.has_llm_subscription,
671 Cents(claims.max_monthly_spend_in_cents),
672 claims.free_tier_monthly_spending_limit(),
673 Utc::now(),
674 )
675 .await
676 .log_err();
677
678 if let Some(usage) = usage {
679 tracing::info!(
680 target: "user usage",
681 user_id = claims.user_id,
682 login = claims.github_user_login,
683 authn.jti = claims.jti,
684 is_staff = claims.is_staff,
685 requests_this_minute = usage.requests_this_minute,
686 tokens_this_minute = usage.tokens_this_minute,
687 );
688
689 let properties = json!({
690 "plan": match claims.plan {
691 Plan::Free => "free".to_string(),
692 Plan::ZedPro => "zed_pro".to_string(),
693 },
694 "model": model,
695 "provider": provider,
696 "usage": usage,
697 "tokens": tokens
698 });
699 SnowflakeRow::new(
700 "Language Model Used",
701 claims.metrics_id,
702 claims.is_staff,
703 claims.system_id.clone(),
704 properties,
705 )
706 .write(&state.kinesis_client, &state.config.kinesis_stream)
707 .await
708 .log_err();
709
710 if let Some(clickhouse_client) = state.clickhouse_client.as_ref() {
711 report_llm_usage(
712 clickhouse_client,
713 LlmUsageEventRow {
714 time: Utc::now().timestamp_millis(),
715 user_id: claims.user_id as i32,
716 is_staff: claims.is_staff,
717 plan: match claims.plan {
718 Plan::Free => "free".to_string(),
719 Plan::ZedPro => "zed_pro".to_string(),
720 },
721 model,
722 provider: provider.to_string(),
723 input_token_count: tokens.input as u64,
724 cache_creation_input_token_count: tokens.input_cache_creation as u64,
725 cache_read_input_token_count: tokens.input_cache_read as u64,
726 output_token_count: tokens.output as u64,
727 requests_this_minute: usage.requests_this_minute as u64,
728 tokens_this_minute: usage.tokens_this_minute as u64,
729 tokens_this_day: usage.tokens_this_day as u64,
730 input_tokens_this_month: usage.tokens_this_month.input as u64,
731 cache_creation_input_tokens_this_month: usage
732 .tokens_this_month
733 .input_cache_creation
734 as u64,
735 cache_read_input_tokens_this_month: usage
736 .tokens_this_month
737 .input_cache_read
738 as u64,
739 output_tokens_this_month: usage.tokens_this_month.output as u64,
740 spending_this_month: usage.spending_this_month.0 as u64,
741 lifetime_spending: usage.lifetime_spending.0 as u64,
742 },
743 )
744 .await
745 .log_err();
746 }
747 }
748 })
749 }
750}
751
752pub fn log_usage_periodically(state: Arc<LlmState>) {
753 state.executor.clone().spawn_detached(async move {
754 loop {
755 state
756 .executor
757 .sleep(std::time::Duration::from_secs(30))
758 .await;
759
760 for provider in LanguageModelProvider::iter() {
761 for model in state.db.model_names_for_provider(provider) {
762 if let Some(active_user_count) = state
763 .get_active_user_count(provider, &model)
764 .await
765 .log_err()
766 {
767 tracing::info!(
768 target: "active user counts",
769 provider = provider.to_string(),
770 model = model,
771 users_in_recent_minutes = active_user_count.users_in_recent_minutes,
772 users_in_recent_days = active_user_count.users_in_recent_days,
773 );
774 }
775 }
776 }
777
778 if let Some(usages) = state
779 .db
780 .get_application_wide_usages_by_model(Utc::now())
781 .await
782 .log_err()
783 {
784 for usage in usages {
785 tracing::info!(
786 target: "computed usage",
787 provider = usage.provider.to_string(),
788 model = usage.model,
789 requests_this_minute = usage.requests_this_minute,
790 tokens_this_minute = usage.tokens_this_minute,
791 );
792 }
793 }
794 }
795 })
796}