llm.rs

  1mod authorization;
  2pub mod db;
  3mod token;
  4
  5use crate::api::CloudflareIpCountryHeader;
  6use crate::api::events::SnowflakeRow;
  7use crate::build_kinesis_client;
  8use crate::rpc::MIN_ACCOUNT_AGE_FOR_LLM_USE;
  9use crate::{Cents, Config, Error, Result, db::UserId, executor::Executor};
 10use anyhow::{Context as _, anyhow};
 11use authorization::authorize_access_to_language_model;
 12use axum::routing::get;
 13use axum::{
 14    Extension, Json, Router, TypedHeader,
 15    body::Body,
 16    http::{self, HeaderName, HeaderValue, Request, StatusCode},
 17    middleware::{self, Next},
 18    response::{IntoResponse, Response},
 19    routing::post,
 20};
 21use chrono::{DateTime, Duration, Utc};
 22use collections::HashMap;
 23use db::TokenUsage;
 24use db::{ActiveUserCount, LlmDatabase, usage_measure::UsageMeasure};
 25use futures::{Stream, StreamExt as _};
 26use reqwest_client::ReqwestClient;
 27use rpc::{
 28    EXPIRED_LLM_TOKEN_HEADER_NAME, LanguageModelProvider, PerformCompletionParams, proto::Plan,
 29};
 30use rpc::{ListModelsResponse, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME};
 31use serde_json::json;
 32use std::{
 33    pin::Pin,
 34    sync::Arc,
 35    task::{Context, Poll},
 36};
 37use strum::IntoEnumIterator;
 38use tokio::sync::RwLock;
 39use util::ResultExt;
 40
 41pub use token::*;
 42
 43const ACTIVE_USER_COUNT_CACHE_DURATION: Duration = Duration::seconds(30);
 44
 45pub struct LlmState {
 46    pub config: Config,
 47    pub executor: Executor,
 48    pub db: Arc<LlmDatabase>,
 49    pub http_client: ReqwestClient,
 50    pub kinesis_client: Option<aws_sdk_kinesis::Client>,
 51    active_user_count_by_model:
 52        RwLock<HashMap<(LanguageModelProvider, String), (DateTime<Utc>, ActiveUserCount)>>,
 53}
 54
 55impl LlmState {
 56    pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {
 57        let database_url = config
 58            .llm_database_url
 59            .as_ref()
 60            .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
 61        let max_connections = config
 62            .llm_database_max_connections
 63            .ok_or_else(|| anyhow!("missing LLM_DATABASE_MAX_CONNECTIONS"))?;
 64
 65        let mut db_options = db::ConnectOptions::new(database_url);
 66        db_options.max_connections(max_connections);
 67        let mut db = LlmDatabase::new(db_options, executor.clone()).await?;
 68        db.initialize().await?;
 69
 70        let db = Arc::new(db);
 71
 72        let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
 73        let http_client =
 74            ReqwestClient::user_agent(&user_agent).context("failed to construct http client")?;
 75
 76        let this = Self {
 77            executor,
 78            db,
 79            http_client,
 80            kinesis_client: if config.kinesis_access_key.is_some() {
 81                build_kinesis_client(&config).await.log_err()
 82            } else {
 83                None
 84            },
 85            active_user_count_by_model: RwLock::new(HashMap::default()),
 86            config,
 87        };
 88
 89        Ok(Arc::new(this))
 90    }
 91
 92    pub async fn get_active_user_count(
 93        &self,
 94        provider: LanguageModelProvider,
 95        model: &str,
 96    ) -> Result<ActiveUserCount> {
 97        let now = Utc::now();
 98
 99        {
100            let active_user_count_by_model = self.active_user_count_by_model.read().await;
101            if let Some((last_updated, count)) =
102                active_user_count_by_model.get(&(provider, model.to_string()))
103            {
104                if now - *last_updated < ACTIVE_USER_COUNT_CACHE_DURATION {
105                    return Ok(*count);
106                }
107            }
108        }
109
110        let mut cache = self.active_user_count_by_model.write().await;
111        let new_count = self.db.get_active_user_count(provider, model, now).await?;
112        cache.insert((provider, model.to_string()), (now, new_count));
113        Ok(new_count)
114    }
115}
116
117pub fn routes() -> Router<(), Body> {
118    Router::new()
119        .route("/models", get(list_models))
120        .route("/completion", post(perform_completion))
121        .layer(middleware::from_fn(validate_api_token))
122}
123
124async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoResponse {
125    let token = req
126        .headers()
127        .get(http::header::AUTHORIZATION)
128        .and_then(|header| header.to_str().ok())
129        .ok_or_else(|| {
130            Error::http(
131                StatusCode::BAD_REQUEST,
132                "missing authorization header".to_string(),
133            )
134        })?
135        .strip_prefix("Bearer ")
136        .ok_or_else(|| {
137            Error::http(
138                StatusCode::BAD_REQUEST,
139                "invalid authorization header".to_string(),
140            )
141        })?;
142
143    let state = req.extensions().get::<Arc<LlmState>>().unwrap();
144    match LlmTokenClaims::validate(token, &state.config) {
145        Ok(claims) => {
146            if state.db.is_access_token_revoked(&claims.jti).await? {
147                return Err(Error::http(
148                    StatusCode::UNAUTHORIZED,
149                    "unauthorized".to_string(),
150                ));
151            }
152
153            tracing::Span::current()
154                .record("user_id", claims.user_id)
155                .record("login", claims.github_user_login.clone())
156                .record("authn.jti", &claims.jti)
157                .record("is_staff", claims.is_staff);
158
159            req.extensions_mut().insert(claims);
160            Ok::<_, Error>(next.run(req).await.into_response())
161        }
162        Err(ValidateLlmTokenError::Expired) => Err(Error::Http(
163            StatusCode::UNAUTHORIZED,
164            "unauthorized".to_string(),
165            [(
166                HeaderName::from_static(EXPIRED_LLM_TOKEN_HEADER_NAME),
167                HeaderValue::from_static("true"),
168            )]
169            .into_iter()
170            .collect(),
171        )),
172        Err(_err) => Err(Error::http(
173            StatusCode::UNAUTHORIZED,
174            "unauthorized".to_string(),
175        )),
176    }
177}
178
179async fn list_models(
180    Extension(state): Extension<Arc<LlmState>>,
181    Extension(claims): Extension<LlmTokenClaims>,
182    country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
183) -> Result<Json<ListModelsResponse>> {
184    let country_code = country_code_header.map(|header| header.to_string());
185
186    let mut accessible_models = Vec::new();
187
188    for (provider, model) in state.db.all_models() {
189        let authorize_result = authorize_access_to_language_model(
190            &state.config,
191            &claims,
192            country_code.as_deref(),
193            provider,
194            &model.name,
195        );
196
197        if authorize_result.is_ok() {
198            accessible_models.push(rpc::LanguageModel {
199                provider,
200                name: model.name,
201            });
202        }
203    }
204
205    Ok(Json(ListModelsResponse {
206        models: accessible_models,
207    }))
208}
209
210async fn perform_completion(
211    Extension(state): Extension<Arc<LlmState>>,
212    Extension(claims): Extension<LlmTokenClaims>,
213    country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
214    Json(params): Json<PerformCompletionParams>,
215) -> Result<impl IntoResponse> {
216    let model = normalize_model_name(
217        state.db.model_names_for_provider(params.provider),
218        params.model,
219    );
220
221    let bypass_account_age_check = claims.has_llm_subscription || claims.bypass_account_age_check;
222    if !bypass_account_age_check {
223        if Utc::now().naive_utc() - claims.account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE {
224            Err(anyhow!("account too young"))?
225        }
226    }
227
228    authorize_access_to_language_model(
229        &state.config,
230        &claims,
231        country_code_header
232            .map(|header| header.to_string())
233            .as_deref(),
234        params.provider,
235        &model,
236    )?;
237
238    check_usage_limit(&state, params.provider, &model, &claims).await?;
239
240    let stream = match params.provider {
241        LanguageModelProvider::Anthropic => {
242            let api_key = if claims.is_staff {
243                state
244                    .config
245                    .anthropic_staff_api_key
246                    .as_ref()
247                    .context("no Anthropic AI staff API key configured on the server")?
248            } else {
249                state
250                    .config
251                    .anthropic_api_key
252                    .as_ref()
253                    .context("no Anthropic AI API key configured on the server")?
254            };
255
256            let mut request: anthropic::Request =
257                serde_json::from_str(params.provider_request.get())?;
258
259            // Override the model on the request with the latest version of the model that is
260            // known to the server.
261            //
262            // Right now, we use the version that's defined in `model.id()`, but we will likely
263            // want to change this code once a new version of an Anthropic model is released,
264            // so that users can use the new version, without having to update Zed.
265            request.model = match model.as_str() {
266                "claude-3-5-sonnet" => anthropic::Model::Claude3_5Sonnet.id().to_string(),
267                "claude-3-7-sonnet" => anthropic::Model::Claude3_7Sonnet.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 user_id = UserId::from_proto(claims.user_id);
463    let model = state.db.model(provider, model_name)?;
464    let free_tier = claims.free_tier_monthly_spending_limit();
465
466    let spending_this_month = state
467        .db
468        .get_user_spending_for_month(user_id, Utc::now())
469        .await?;
470    if spending_this_month >= free_tier {
471        if !claims.has_llm_subscription {
472            return Err(Error::http(
473                StatusCode::PAYMENT_REQUIRED,
474                "Maximum spending limit reached for this month.".to_string(),
475            ));
476        }
477
478        let monthly_spend = spending_this_month.saturating_sub(free_tier);
479        if monthly_spend >= Cents(claims.max_monthly_spend_in_cents) {
480            return Err(Error::Http(
481                StatusCode::FORBIDDEN,
482                "Maximum spending limit reached for this month.".to_string(),
483                [(
484                    HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME),
485                    HeaderValue::from_static("true"),
486                )]
487                .into_iter()
488                .collect(),
489            ));
490        }
491    }
492
493    let active_users = state.get_active_user_count(provider, model_name).await?;
494
495    let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);
496    let users_in_recent_days = active_users.users_in_recent_days.max(1);
497
498    let per_user_max_requests_per_minute =
499        model.max_requests_per_minute as usize / users_in_recent_minutes;
500    let per_user_max_tokens_per_minute =
501        model.max_tokens_per_minute as usize / users_in_recent_minutes;
502    let per_user_max_input_tokens_per_minute =
503        model.max_input_tokens_per_minute as usize / users_in_recent_minutes;
504    let per_user_max_output_tokens_per_minute =
505        model.max_output_tokens_per_minute as usize / users_in_recent_minutes;
506    let per_user_max_tokens_per_day = model.max_tokens_per_day as usize / users_in_recent_days;
507
508    let usage = state
509        .db
510        .get_usage(user_id, provider, model_name, Utc::now())
511        .await?;
512
513    let checks = [
514        (
515            usage.requests_this_minute,
516            per_user_max_requests_per_minute,
517            UsageMeasure::RequestsPerMinute,
518        ),
519        (
520            usage.tokens_this_minute,
521            per_user_max_tokens_per_minute,
522            UsageMeasure::TokensPerMinute,
523        ),
524        (
525            usage.tokens_this_day,
526            per_user_max_tokens_per_day,
527            UsageMeasure::TokensPerDay,
528        ),
529    ];
530
531    for (used, limit, usage_measure) in checks {
532        if used > limit {
533            let resource = match usage_measure {
534                UsageMeasure::RequestsPerMinute => "requests_per_minute",
535                UsageMeasure::TokensPerMinute => "tokens_per_minute",
536                UsageMeasure::InputTokensPerMinute => "input_tokens_per_minute",
537                UsageMeasure::OutputTokensPerMinute => "output_tokens_per_minute",
538                UsageMeasure::TokensPerDay => "tokens_per_day",
539            };
540
541            tracing::info!(
542                target: "user rate limit",
543                user_id = claims.user_id,
544                login = claims.github_user_login,
545                authn.jti = claims.jti,
546                is_staff = claims.is_staff,
547                provider = provider.to_string(),
548                model = model.name,
549                requests_this_minute = usage.requests_this_minute,
550                tokens_this_minute = usage.tokens_this_minute,
551                input_tokens_this_minute = usage.input_tokens_this_minute,
552                output_tokens_this_minute = usage.output_tokens_this_minute,
553                tokens_this_day = usage.tokens_this_day,
554                users_in_recent_minutes = users_in_recent_minutes,
555                users_in_recent_days = users_in_recent_days,
556                max_requests_per_minute = per_user_max_requests_per_minute,
557                max_tokens_per_minute = per_user_max_tokens_per_minute,
558                max_input_tokens_per_minute = per_user_max_input_tokens_per_minute,
559                max_output_tokens_per_minute = per_user_max_output_tokens_per_minute,
560                max_tokens_per_day = per_user_max_tokens_per_day,
561            );
562
563            SnowflakeRow::new(
564                "Language Model Rate Limited",
565                claims.metrics_id,
566                claims.is_staff,
567                claims.system_id.clone(),
568                json!({
569                    "usage": usage,
570                    "users_in_recent_minutes": users_in_recent_minutes,
571                    "users_in_recent_days": users_in_recent_days,
572                    "max_requests_per_minute": per_user_max_requests_per_minute,
573                    "max_tokens_per_minute": per_user_max_tokens_per_minute,
574                    "max_tokens_per_day": per_user_max_tokens_per_day,
575                    "plan": match claims.plan {
576                        Plan::Free => "free".to_string(),
577                        Plan::ZedPro => "zed_pro".to_string(),
578                    },
579                    "model": model.name.clone(),
580                    "provider": provider.to_string(),
581                    "usage_measure": resource.to_string(),
582                }),
583            )
584            .write(&state.kinesis_client, &state.config.kinesis_stream)
585            .await
586            .log_err();
587
588            return Err(Error::http(
589                StatusCode::TOO_MANY_REQUESTS,
590                format!("Rate limit exceeded. Maximum {} reached.", resource),
591            ));
592        }
593    }
594
595    Ok(())
596}
597
598struct CompletionChunk {
599    bytes: Vec<u8>,
600    input_tokens: usize,
601    output_tokens: usize,
602    cache_creation_input_tokens: usize,
603    cache_read_input_tokens: usize,
604}
605
606struct TokenCountingStream<S> {
607    state: Arc<LlmState>,
608    claims: LlmTokenClaims,
609    provider: LanguageModelProvider,
610    model: String,
611    tokens: TokenUsage,
612    inner_stream: S,
613}
614
615impl<S> Stream for TokenCountingStream<S>
616where
617    S: Stream<Item = Result<CompletionChunk, anyhow::Error>> + Unpin,
618{
619    type Item = Result<Vec<u8>, anyhow::Error>;
620
621    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
622        match Pin::new(&mut self.inner_stream).poll_next(cx) {
623            Poll::Ready(Some(Ok(mut chunk))) => {
624                chunk.bytes.push(b'\n');
625                self.tokens.input += chunk.input_tokens;
626                self.tokens.output += chunk.output_tokens;
627                self.tokens.input_cache_creation += chunk.cache_creation_input_tokens;
628                self.tokens.input_cache_read += chunk.cache_read_input_tokens;
629                Poll::Ready(Some(Ok(chunk.bytes)))
630            }
631            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
632            Poll::Ready(None) => Poll::Ready(None),
633            Poll::Pending => Poll::Pending,
634        }
635    }
636}
637
638impl<S> Drop for TokenCountingStream<S> {
639    fn drop(&mut self) {
640        let state = self.state.clone();
641        let claims = self.claims.clone();
642        let provider = self.provider;
643        let model = std::mem::take(&mut self.model);
644        let tokens = self.tokens;
645        self.state.executor.spawn_detached(async move {
646            let usage = state
647                .db
648                .record_usage(
649                    UserId::from_proto(claims.user_id),
650                    claims.is_staff,
651                    provider,
652                    &model,
653                    tokens,
654                    claims.has_llm_subscription,
655                    Cents(claims.max_monthly_spend_in_cents),
656                    claims.free_tier_monthly_spending_limit(),
657                    Utc::now(),
658                )
659                .await
660                .log_err();
661
662            if let Some(usage) = usage {
663                tracing::info!(
664                    target: "user usage",
665                    user_id = claims.user_id,
666                    login = claims.github_user_login,
667                    authn.jti = claims.jti,
668                    is_staff = claims.is_staff,
669                    requests_this_minute = usage.requests_this_minute,
670                    tokens_this_minute = usage.tokens_this_minute,
671                    input_tokens_this_minute = usage.input_tokens_this_minute,
672                    output_tokens_this_minute = usage.output_tokens_this_minute,
673                );
674
675                let properties = json!({
676                    "has_llm_subscription": claims.has_llm_subscription,
677                    "max_monthly_spend_in_cents": claims.max_monthly_spend_in_cents,
678                    "plan": match claims.plan {
679                        Plan::Free => "free".to_string(),
680                        Plan::ZedPro => "zed_pro".to_string(),
681                    },
682                    "model": model,
683                    "provider": provider,
684                    "usage": usage,
685                    "tokens": tokens
686                });
687                SnowflakeRow::new(
688                    "Language Model Used",
689                    claims.metrics_id,
690                    claims.is_staff,
691                    claims.system_id.clone(),
692                    properties,
693                )
694                .write(&state.kinesis_client, &state.config.kinesis_stream)
695                .await
696                .log_err();
697            }
698        })
699    }
700}
701
702pub fn log_usage_periodically(state: Arc<LlmState>) {
703    state.executor.clone().spawn_detached(async move {
704        loop {
705            state
706                .executor
707                .sleep(std::time::Duration::from_secs(30))
708                .await;
709
710            for provider in LanguageModelProvider::iter() {
711                for model in state.db.model_names_for_provider(provider) {
712                    if let Some(active_user_count) = state
713                        .get_active_user_count(provider, &model)
714                        .await
715                        .log_err()
716                    {
717                        tracing::info!(
718                            target: "active user counts",
719                            provider = provider.to_string(),
720                            model = model,
721                            users_in_recent_minutes = active_user_count.users_in_recent_minutes,
722                            users_in_recent_days = active_user_count.users_in_recent_days,
723                        );
724                    }
725                }
726            }
727
728            if let Some(usages) = state
729                .db
730                .get_application_wide_usages_by_model(Utc::now())
731                .await
732                .log_err()
733            {
734                for usage in usages {
735                    tracing::info!(
736                        target: "computed usage",
737                        provider = usage.provider.to_string(),
738                        model = usage.model,
739                        requests_this_minute = usage.requests_this_minute,
740                        tokens_this_minute = usage.tokens_this_minute,
741                        input_tokens_this_minute = usage.input_tokens_this_minute,
742                        output_tokens_this_minute = usage.output_tokens_this_minute,
743                    );
744                }
745            }
746        }
747    })
748}