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>
This commit is contained in:
parent
a0988508f0
commit
b87c4a1e13
10 changed files with 354 additions and 144 deletions
|
@ -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()
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue