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.as_ref().map(|limits| limits.remaining),
320 input_tokens_remaining = rate_limit_info.input_tokens.as_ref().map(|limits| limits.remaining),
321 output_tokens_remaining = rate_limit_info.output_tokens.as_ref().map(|limits| limits.remaining),
322 requests_remaining = rate_limit_info.requests.as_ref().map(|limits| limits.remaining),
323 requests_reset = ?rate_limit_info.requests.as_ref().map(|limits| limits.reset),
324 tokens_reset = ?rate_limit_info.tokens.as_ref().map(|limits| limits.reset),
325 input_tokens_reset = ?rate_limit_info.input_tokens.as_ref().map(|limits| limits.reset),
326 output_tokens_reset = ?rate_limit_info.output_tokens.as_ref().map(|limits| limits.reset),
327 );
328 }
329
330 chunks
331 .map(move |event| {
332 let chunk = event?;
333 let (
334 input_tokens,
335 output_tokens,
336 cache_creation_input_tokens,
337 cache_read_input_tokens,
338 ) = match &chunk {
339 anthropic::Event::MessageStart {
340 message: anthropic::Response { usage, .. },
341 }
342 | anthropic::Event::MessageDelta { usage, .. } => (
343 usage.input_tokens.unwrap_or(0) as usize,
344 usage.output_tokens.unwrap_or(0) as usize,
345 usage.cache_creation_input_tokens.unwrap_or(0) as usize,
346 usage.cache_read_input_tokens.unwrap_or(0) as usize,
347 ),
348 _ => (0, 0, 0, 0),
349 };
350
351 anyhow::Ok(CompletionChunk {
352 bytes: serde_json::to_vec(&chunk).unwrap(),
353 input_tokens,
354 output_tokens,
355 cache_creation_input_tokens,
356 cache_read_input_tokens,
357 })
358 })
359 .boxed()
360 }
361 LanguageModelProvider::OpenAi => {
362 let api_key = state
363 .config
364 .openai_api_key
365 .as_ref()
366 .context("no OpenAI API key configured on the server")?;
367 let chunks = open_ai::stream_completion(
368 &state.http_client,
369 open_ai::OPEN_AI_API_URL,
370 api_key,
371 serde_json::from_str(params.provider_request.get())?,
372 )
373 .await?;
374
375 chunks
376 .map(|event| {
377 event.map(|chunk| {
378 let input_tokens =
379 chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;
380 let output_tokens =
381 chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;
382 CompletionChunk {
383 bytes: serde_json::to_vec(&chunk).unwrap(),
384 input_tokens,
385 output_tokens,
386 cache_creation_input_tokens: 0,
387 cache_read_input_tokens: 0,
388 }
389 })
390 })
391 .boxed()
392 }
393 LanguageModelProvider::Google => {
394 let api_key = state
395 .config
396 .google_ai_api_key
397 .as_ref()
398 .context("no Google AI API key configured on the server")?;
399 let chunks = google_ai::stream_generate_content(
400 &state.http_client,
401 google_ai::API_URL,
402 api_key,
403 serde_json::from_str(params.provider_request.get())?,
404 )
405 .await?;
406
407 chunks
408 .map(|event| {
409 event.map(|chunk| {
410 // TODO - implement token counting for Google AI
411 CompletionChunk {
412 bytes: serde_json::to_vec(&chunk).unwrap(),
413 input_tokens: 0,
414 output_tokens: 0,
415 cache_creation_input_tokens: 0,
416 cache_read_input_tokens: 0,
417 }
418 })
419 })
420 .boxed()
421 }
422 };
423
424 Ok(Response::new(Body::wrap_stream(TokenCountingStream {
425 state,
426 claims,
427 provider: params.provider,
428 model,
429 tokens: TokenUsage::default(),
430 inner_stream: stream,
431 })))
432}
433
434fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
435 if let Some(known_model_name) = known_models
436 .iter()
437 .filter(|known_model_name| name.starts_with(known_model_name.as_str()))
438 .max_by_key(|known_model_name| known_model_name.len())
439 {
440 known_model_name.to_string()
441 } else {
442 name
443 }
444}
445
446/// The maximum monthly spending an individual user can reach on the free tier
447/// before they have to pay.
448pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10);
449
450/// The default value to use for maximum spend per month if the user did not
451/// explicitly set a maximum spend.
452///
453/// Used to prevent surprise bills.
454pub const DEFAULT_MAX_MONTHLY_SPEND: Cents = Cents::from_dollars(10);
455
456async fn check_usage_limit(
457 state: &Arc<LlmState>,
458 provider: LanguageModelProvider,
459 model_name: &str,
460 claims: &LlmTokenClaims,
461) -> Result<()> {
462 if claims.is_staff {
463 return Ok(());
464 }
465
466 let user_id = UserId::from_proto(claims.user_id);
467 let model = state.db.model(provider, model_name)?;
468 let free_tier = claims.free_tier_monthly_spending_limit();
469
470 let spending_this_month = state
471 .db
472 .get_user_spending_for_month(user_id, Utc::now())
473 .await?;
474 if 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 let monthly_spend = spending_this_month.saturating_sub(free_tier);
483 if monthly_spend >= Cents(claims.max_monthly_spend_in_cents) {
484 return Err(Error::Http(
485 StatusCode::FORBIDDEN,
486 "Maximum spending limit reached for this month.".to_string(),
487 [(
488 HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME),
489 HeaderValue::from_static("true"),
490 )]
491 .into_iter()
492 .collect(),
493 ));
494 }
495 }
496
497 let active_users = state.get_active_user_count(provider, model_name).await?;
498
499 let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);
500 let users_in_recent_days = active_users.users_in_recent_days.max(1);
501
502 let per_user_max_requests_per_minute =
503 model.max_requests_per_minute as usize / users_in_recent_minutes;
504 let per_user_max_tokens_per_minute =
505 model.max_tokens_per_minute as usize / users_in_recent_minutes;
506 let per_user_max_input_tokens_per_minute =
507 model.max_input_tokens_per_minute as usize / users_in_recent_minutes;
508 let per_user_max_output_tokens_per_minute =
509 model.max_output_tokens_per_minute as usize / users_in_recent_minutes;
510 let per_user_max_tokens_per_day = model.max_tokens_per_day as usize / users_in_recent_days;
511
512 let usage = state
513 .db
514 .get_usage(user_id, provider, model_name, Utc::now())
515 .await?;
516
517 let checks = match (provider, model_name) {
518 (LanguageModelProvider::Anthropic, "claude-3-7-sonnet") => vec![
519 (
520 usage.requests_this_minute,
521 per_user_max_requests_per_minute,
522 UsageMeasure::RequestsPerMinute,
523 ),
524 (
525 usage.input_tokens_this_minute,
526 per_user_max_tokens_per_minute,
527 UsageMeasure::InputTokensPerMinute,
528 ),
529 (
530 usage.output_tokens_this_minute,
531 per_user_max_tokens_per_minute,
532 UsageMeasure::OutputTokensPerMinute,
533 ),
534 (
535 usage.tokens_this_day,
536 per_user_max_tokens_per_day,
537 UsageMeasure::TokensPerDay,
538 ),
539 ],
540 _ => vec![
541 (
542 usage.requests_this_minute,
543 per_user_max_requests_per_minute,
544 UsageMeasure::RequestsPerMinute,
545 ),
546 (
547 usage.tokens_this_minute,
548 per_user_max_tokens_per_minute,
549 UsageMeasure::TokensPerMinute,
550 ),
551 (
552 usage.tokens_this_day,
553 per_user_max_tokens_per_day,
554 UsageMeasure::TokensPerDay,
555 ),
556 ],
557 };
558
559 for (used, limit, usage_measure) in checks {
560 if used > limit {
561 let resource = match usage_measure {
562 UsageMeasure::RequestsPerMinute => "requests_per_minute",
563 UsageMeasure::TokensPerMinute => "tokens_per_minute",
564 UsageMeasure::InputTokensPerMinute => "input_tokens_per_minute",
565 UsageMeasure::OutputTokensPerMinute => "output_tokens_per_minute",
566 UsageMeasure::TokensPerDay => "tokens_per_day",
567 };
568
569 tracing::info!(
570 target: "user rate limit",
571 user_id = claims.user_id,
572 login = claims.github_user_login,
573 authn.jti = claims.jti,
574 is_staff = claims.is_staff,
575 provider = provider.to_string(),
576 model = model.name,
577 usage_measure = resource,
578 requests_this_minute = usage.requests_this_minute,
579 tokens_this_minute = usage.tokens_this_minute,
580 input_tokens_this_minute = usage.input_tokens_this_minute,
581 output_tokens_this_minute = usage.output_tokens_this_minute,
582 tokens_this_day = usage.tokens_this_day,
583 users_in_recent_minutes = users_in_recent_minutes,
584 users_in_recent_days = users_in_recent_days,
585 max_requests_per_minute = per_user_max_requests_per_minute,
586 max_tokens_per_minute = per_user_max_tokens_per_minute,
587 max_input_tokens_per_minute = per_user_max_input_tokens_per_minute,
588 max_output_tokens_per_minute = per_user_max_output_tokens_per_minute,
589 max_tokens_per_day = per_user_max_tokens_per_day,
590 );
591
592 SnowflakeRow::new(
593 "Language Model Rate Limited",
594 Some(claims.metrics_id),
595 claims.is_staff,
596 claims.system_id.clone(),
597 json!({
598 "usage": usage,
599 "users_in_recent_minutes": users_in_recent_minutes,
600 "users_in_recent_days": users_in_recent_days,
601 "max_requests_per_minute": per_user_max_requests_per_minute,
602 "max_tokens_per_minute": per_user_max_tokens_per_minute,
603 "max_input_tokens_per_minute": per_user_max_input_tokens_per_minute,
604 "max_output_tokens_per_minute": per_user_max_output_tokens_per_minute,
605 "max_tokens_per_day": per_user_max_tokens_per_day,
606 "plan": match claims.plan {
607 Plan::Free => "free".to_string(),
608 Plan::ZedPro => "zed_pro".to_string(),
609 },
610 "model": model.name.clone(),
611 "provider": provider.to_string(),
612 "usage_measure": resource.to_string(),
613 }),
614 )
615 .write(&state.kinesis_client, &state.config.kinesis_stream)
616 .await
617 .log_err();
618
619 return Err(Error::http(
620 StatusCode::TOO_MANY_REQUESTS,
621 format!("Rate limit exceeded. Maximum {} reached.", resource),
622 ));
623 }
624 }
625
626 Ok(())
627}
628
629struct CompletionChunk {
630 bytes: Vec<u8>,
631 input_tokens: usize,
632 output_tokens: usize,
633 cache_creation_input_tokens: usize,
634 cache_read_input_tokens: usize,
635}
636
637struct TokenCountingStream<S> {
638 state: Arc<LlmState>,
639 claims: LlmTokenClaims,
640 provider: LanguageModelProvider,
641 model: String,
642 tokens: TokenUsage,
643 inner_stream: S,
644}
645
646impl<S> Stream for TokenCountingStream<S>
647where
648 S: Stream<Item = Result<CompletionChunk, anyhow::Error>> + Unpin,
649{
650 type Item = Result<Vec<u8>, anyhow::Error>;
651
652 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
653 match Pin::new(&mut self.inner_stream).poll_next(cx) {
654 Poll::Ready(Some(Ok(mut chunk))) => {
655 chunk.bytes.push(b'\n');
656 self.tokens.input += chunk.input_tokens;
657 self.tokens.output += chunk.output_tokens;
658 self.tokens.input_cache_creation += chunk.cache_creation_input_tokens;
659 self.tokens.input_cache_read += chunk.cache_read_input_tokens;
660 Poll::Ready(Some(Ok(chunk.bytes)))
661 }
662 Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
663 Poll::Ready(None) => Poll::Ready(None),
664 Poll::Pending => Poll::Pending,
665 }
666 }
667}
668
669impl<S> Drop for TokenCountingStream<S> {
670 fn drop(&mut self) {
671 let state = self.state.clone();
672 let claims = self.claims.clone();
673 let provider = self.provider;
674 let model = std::mem::take(&mut self.model);
675 let tokens = self.tokens;
676 self.state.executor.spawn_detached(async move {
677 let usage = state
678 .db
679 .record_usage(
680 UserId::from_proto(claims.user_id),
681 claims.is_staff,
682 provider,
683 &model,
684 tokens,
685 claims.has_llm_subscription,
686 Cents(claims.max_monthly_spend_in_cents),
687 claims.free_tier_monthly_spending_limit(),
688 Utc::now(),
689 )
690 .await
691 .log_err();
692
693 if let Some(usage) = usage {
694 tracing::info!(
695 target: "user usage",
696 user_id = claims.user_id,
697 login = claims.github_user_login,
698 authn.jti = claims.jti,
699 is_staff = claims.is_staff,
700 provider = provider.to_string(),
701 model = model,
702 requests_this_minute = usage.requests_this_minute,
703 tokens_this_minute = usage.tokens_this_minute,
704 input_tokens_this_minute = usage.input_tokens_this_minute,
705 output_tokens_this_minute = usage.output_tokens_this_minute,
706 );
707
708 let properties = json!({
709 "has_llm_subscription": claims.has_llm_subscription,
710 "max_monthly_spend_in_cents": claims.max_monthly_spend_in_cents,
711 "plan": match claims.plan {
712 Plan::Free => "free".to_string(),
713 Plan::ZedPro => "zed_pro".to_string(),
714 },
715 "model": model,
716 "provider": provider,
717 "usage": usage,
718 "tokens": tokens
719 });
720 SnowflakeRow::new(
721 "Language Model Used",
722 Some(claims.metrics_id),
723 claims.is_staff,
724 claims.system_id.clone(),
725 properties,
726 )
727 .write(&state.kinesis_client, &state.config.kinesis_stream)
728 .await
729 .log_err();
730 }
731 })
732 }
733}
734
735pub fn log_usage_periodically(state: Arc<LlmState>) {
736 state.executor.clone().spawn_detached(async move {
737 loop {
738 state
739 .executor
740 .sleep(std::time::Duration::from_secs(30))
741 .await;
742
743 for provider in LanguageModelProvider::iter() {
744 for model in state.db.model_names_for_provider(provider) {
745 if let Some(active_user_count) = state
746 .get_active_user_count(provider, &model)
747 .await
748 .log_err()
749 {
750 tracing::info!(
751 target: "active user counts",
752 provider = provider.to_string(),
753 model = model,
754 users_in_recent_minutes = active_user_count.users_in_recent_minutes,
755 users_in_recent_days = active_user_count.users_in_recent_days,
756 );
757 }
758 }
759 }
760
761 if let Some(usages) = state
762 .db
763 .get_application_wide_usages_by_model(Utc::now())
764 .await
765 .log_err()
766 {
767 for usage in usages {
768 tracing::info!(
769 target: "computed usage",
770 provider = usage.provider.to_string(),
771 model = usage.model,
772 requests_this_minute = usage.requests_this_minute,
773 tokens_this_minute = usage.tokens_this_minute,
774 input_tokens_this_minute = usage.input_tokens_this_minute,
775 output_tokens_this_minute = usage.output_tokens_this_minute,
776 );
777 }
778 }
779 }
780 })
781}