assistant: Add health telemetry (#19928)

This PR adds a bit of telemetry for Anthropic models, in order to
understand model health. With this logging, we can monitor and diagnose
dips in performance, for example due to model rollouts.

Release Notes:

- N/A

---------

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>

crates/language_model/Cargo.toml

@@ -46,6 +46,7 @@ serde_json.workspace = true
 settings.workspace = true
 smol.workspace = true
 strum.workspace = true
+telemetry_events.workspace = true
 theme.workspace = true
 thiserror.workspace = true
 tiktoken-rs.workspace = true

crates/language_model/src/language_model.rs

@@ -1,3 +1,4 @@
+pub mod logging;
 mod model;
 pub mod provider;
 mod rate_limiter;
@@ -59,6 +60,7 @@ pub enum LanguageModelCompletionEvent {
     Stop(StopReason),
     Text(String),
     ToolUse(LanguageModelToolUse),
+    StartMessage { message_id: String },
 }
 
 #[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
@@ -76,6 +78,20 @@ pub struct LanguageModelToolUse {
     pub input: serde_json::Value,
 }
 
+pub struct LanguageModelTextStream {
+    pub message_id: Option<String>,
+    pub stream: BoxStream<'static, Result<String>>,
+}
+
+impl Default for LanguageModelTextStream {
+    fn default() -> Self {
+        Self {
+            message_id: None,
+            stream: Box::pin(futures::stream::empty()),
+        }
+    }
+}
+
 pub trait LanguageModel: Send + Sync {
     fn id(&self) -> LanguageModelId;
     fn name(&self) -> LanguageModelName;
@@ -87,6 +103,10 @@ pub trait LanguageModel: Send + Sync {
     fn provider_name(&self) -> LanguageModelProviderName;
     fn telemetry_id(&self) -> String;
 
+    fn api_key(&self, _cx: &AppContext) -> Option<String> {
+        None
+    }
+
     /// Returns the availability of this language model.
     fn availability(&self) -> LanguageModelAvailability {
         LanguageModelAvailability::Public
@@ -113,21 +133,39 @@ pub trait LanguageModel: Send + Sync {
         &self,
         request: LanguageModelRequest,
         cx: &AsyncAppContext,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
+    ) -> BoxFuture<'static, Result<LanguageModelTextStream>> {
         let events = self.stream_completion(request, cx);
 
         async move {
-            Ok(events
-                .await?
-                .filter_map(|result| async move {
+            let mut events = events.await?;
+            let mut message_id = None;
+            let mut first_item_text = None;
+
+            if let Some(first_event) = events.next().await {
+                match first_event {
+                    Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
+                        message_id = Some(id.clone());
+                    }
+                    Ok(LanguageModelCompletionEvent::Text(text)) => {
+                        first_item_text = Some(text);
+                    }
+                    _ => (),
+                }
+            }
+
+            let stream = futures::stream::iter(first_item_text.map(Ok))
+                .chain(events.filter_map(|result| async move {
                     match result {
+                        Ok(LanguageModelCompletionEvent::StartMessage { .. }) => None,
                         Ok(LanguageModelCompletionEvent::Text(text)) => Some(Ok(text)),
                         Ok(LanguageModelCompletionEvent::Stop(_)) => None,
                         Ok(LanguageModelCompletionEvent::ToolUse(_)) => None,
                         Err(err) => Some(Err(err)),
                     }
-                })
-                .boxed())
+                }))
+                .boxed();
+
+            Ok(LanguageModelTextStream { message_id, stream })
         }
         .boxed()
     }

crates/language_model/src/logging.rs

@@ -0,0 +1,90 @@
+use anthropic::{AnthropicError, ANTHROPIC_API_URL};
+use anyhow::{anyhow, Context, Result};
+use client::telemetry::Telemetry;
+use gpui::BackgroundExecutor;
+use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};
+use std::env;
+use std::sync::Arc;
+use telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase};
+use util::ResultExt;
+
+use crate::provider::anthropic::PROVIDER_ID as ANTHROPIC_PROVIDER_ID;
+
+pub fn report_assistant_event(
+    event: AssistantEvent,
+    telemetry: Option<Arc<Telemetry>>,
+    client: Arc<dyn HttpClient>,
+    model_api_key: Option<String>,
+    executor: &BackgroundExecutor,
+) {
+    if let Some(telemetry) = telemetry.as_ref() {
+        telemetry.report_assistant_event(event.clone());
+        if telemetry.metrics_enabled() && event.model_provider == ANTHROPIC_PROVIDER_ID {
+            executor
+                .spawn(async move {
+                    report_anthropic_event(event, client, model_api_key)
+                        .await
+                        .log_err();
+                })
+                .detach();
+        }
+    }
+}
+
+async fn report_anthropic_event(
+    event: AssistantEvent,
+    client: Arc<dyn HttpClient>,
+    model_api_key: Option<String>,
+) -> Result<(), AnthropicError> {
+    let api_key = match model_api_key {
+        Some(key) => key,
+        None => {
+            return Err(AnthropicError::Other(anyhow!(
+                "Anthropic API key is not set"
+            )));
+        }
+    };
+
+    let uri = format!("{ANTHROPIC_API_URL}/v1/log/zed");
+    let request_builder = HttpRequest::builder()
+        .method(Method::POST)
+        .uri(uri)
+        .header("X-Api-Key", api_key)
+        .header("Content-Type", "application/json");
+    let serialized_event: serde_json::Value = serde_json::json!({
+        "completion_type": match event.kind {
+            AssistantKind::Inline => "natural_language_completion_in_editor",
+            AssistantKind::InlineTerminal => "natural_language_completion_in_terminal",
+            AssistantKind::Panel => "conversation_message",
+        },
+        "event": match event.phase {
+            AssistantPhase::Response => "response",
+            AssistantPhase::Invoked => "invoke",
+            AssistantPhase::Accepted => "accept",
+            AssistantPhase::Rejected => "reject",
+        },
+        "metadata": {
+            "language_name": event.language_name,
+            "message_id": event.message_id,
+            "platform": env::consts::OS,
+        }
+    });
+
+    let request = request_builder
+        .body(AsyncBody::from(serialized_event.to_string()))
+        .context("failed to construct request body")?;
+
+    let response = client
+        .send(request)
+        .await
+        .context("failed to send request to Anthropic")?;
+
+    if response.status().is_success() {
+        return Ok(());
+    }
+
+    return Err(AnthropicError::Other(anyhow!(
+        "Failed to log: {}",
+        response.status(),
+    )));
+}

crates/language_model/src/provider/anthropic.rs

@@ -26,7 +26,7 @@ use theme::ThemeSettings;
 use ui::{prelude::*, Icon, IconName, Tooltip};
 use util::{maybe, ResultExt};
 
-const PROVIDER_ID: &str = "anthropic";
+pub const PROVIDER_ID: &str = "anthropic";
 const PROVIDER_NAME: &str = "Anthropic";
 
 #[derive(Default, Clone, Debug, PartialEq)]
@@ -356,6 +356,10 @@ impl LanguageModel for AnthropicModel {
         format!("anthropic/{}", self.model.id())
     }
 
+    fn api_key(&self, cx: &AppContext) -> Option<String> {
+        self.state.read(cx).api_key.clone()
+    }
+
     fn max_token_count(&self) -> usize {
         self.model.max_token_count()
     }
@@ -520,6 +524,14 @@ pub fn map_to_language_model_completion_events(
                                 ));
                             }
                         }
+                        Event::MessageStart { message } => {
+                            return Some((
+                                Some(Ok(LanguageModelCompletionEvent::StartMessage {
+                                    message_id: message.id,
+                                })),
+                                state,
+                            ))
+                        }
                         Event::MessageDelta { delta, .. } => {
                             if let Some(stop_reason) = delta.stop_reason.as_deref() {
                                 let stop_reason = match stop_reason {