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