Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions engine/baml-lib/baml-types/src/tracing/events.rs
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,7 @@ pub struct HTTPRequest {
#[serde(deserialize_with = "deserialize_headers")]
headers: HashMap<String, String>,
pub body: HTTPBody,
pub client_details: std::sync::Arc<ClientDetails>,
}

impl HTTPRequest {
Expand All @@ -398,13 +399,15 @@ impl HTTPRequest {
method: String,
headers: HashMap<String, String>,
body: HTTPBody,
client_details: ClientDetails,
) -> Self {
Self {
id,
url,
method,
headers,
body,
client_details: std::sync::Arc::new(client_details),
}
}

Expand Down Expand Up @@ -614,6 +617,11 @@ mod tests {
"POST".to_string(),
headers.clone(),
HTTPBody::new(b"test body".to_vec()),
ClientDetails {
name: "test-client".to_string(),
provider: "test-provider".to_string(),
options: IndexMap::new(),
},
);

// Test that .headers() returns original headers
Expand Down
2 changes: 1 addition & 1 deletion engine/baml-rpc/src/auth/api_key.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ impl ApiKeyEnvironment {
Self::DEVELOPMENT => "dev",
Self::STAGING => "stg",
Self::PRODUCTION => "prod",
_ => &self.value[..3.min(self.value.len())],
_ => &self.value[..self.value.len().min(4)],
}
}
}
3 changes: 3 additions & 0 deletions engine/baml-rpc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,5 +30,8 @@ pub use ui::{
ListProjectsRequest, ListProjectsResponse, Project, UpdateProject, UpdateProjectRequest,
UpdateProjectResponse,
},
ui_function_call_http_calls::{
GetFunctionCallHttpCalls, GetFunctionCallHttpCallsRequest, GetFunctionCallHttpCallsResponse,
},
ui_function_calls::{ListFunctionCalls, ListFunctionCallsRequest, ListFunctionCallsResponse},
};
1 change: 1 addition & 0 deletions engine/baml-rpc/src/runtime_api/trace_event.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@ pub enum IntermediateData<'a> {
url: String,
method: String,
headers: HashMap<String, String>,
client_details: RpcClientDetails,
body: HTTPBody<'a>,
},
RawLLMResponse {
Expand Down
2 changes: 2 additions & 0 deletions engine/baml-rpc/src/ui/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ pub mod ui_baml_src;
pub mod ui_control_plane_orgs;
pub mod ui_control_plane_projects;
pub mod ui_dashboard;
pub mod ui_dashboard_cost;
pub mod ui_function_call_http_calls;
pub mod ui_function_calls;
pub mod ui_traces;
pub mod ui_types;
35 changes: 35 additions & 0 deletions engine/baml-rpc/src/ui/ui_dashboard.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,13 +151,48 @@ pub struct StatusCountOverTime {
pub error_count: u64,
}

#[derive(Debug, Serialize, Deserialize, TS)]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New time series and metric types (TimeSeriesFloatPoint, TimeSeriesIntPoint, MetricSeriesFloat, MetricSeriesInt) are defined here. These are duplicated in ui_dashboard_cost. Consider extracting shared types to a common module to reduce code duplication.

#[ts(export)]
pub struct TimeSeriesFloatPoint {
#[ts(type = "number")]
pub interval_start: EpochMsTimestamp,
pub value: f64,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct TimeSeriesIntPoint {
#[ts(type = "number")]
pub interval_start: EpochMsTimestamp,
pub value: u64,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct MetricSeriesFloat {
pub total: f64,
pub series: Vec<TimeSeriesFloatPoint>,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct MetricSeriesInt {
pub total: u64,
pub series: Vec<TimeSeriesIntPoint>,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GetDashboardDataResponse {
pub status_counts_by_function: Vec<StatusCountByFunction>,
pub parsing_errors_by_client: Vec<ParsingErrorCountByClient>,
pub parsing_errors_by_function: Vec<ParsingErrorCountByFunction>,
pub status_counts: Vec<StatusCountOverTime>,
pub latency_p75_ms: MetricSeriesFloat,
pub latency_p95_ms: MetricSeriesFloat,
pub latency_avg_ms: MetricSeriesFloat,
pub total_llm_calls: MetricSeriesInt,
pub total_traces: MetricSeriesInt,
}

pub struct GetDashboardData;
Expand Down
65 changes: 65 additions & 0 deletions engine/baml-rpc/src/ui/ui_dashboard_cost.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
use baml_ids::ProjectId;
use serde::{Deserialize, Serialize};
use ts_rs::TS;

use crate::{base::EpochMsTimestamp, rpc::ApiEndpoint};

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GetCostStatsRequest {
#[ts(type = "string")]
pub project_id: ProjectId,
#[ts(type = "FilterExpression<number>", optional)]
pub start_time:
Option<super::ui_function_calls::FilterExpression<crate::base::EpochMsTimestamp>>,
#[ts(type = "FilterExpression<number>", optional)]
pub end_time: Option<super::ui_function_calls::FilterExpression<crate::base::EpochMsTimestamp>>,
#[ts(optional)]
pub relative_time: Option<super::ui_function_calls::RelativeTime>,
#[ts(optional)]
pub function_name: Option<super::ui_function_calls::FilterExpression<String>>,
#[ts(optional)]
pub status: Option<
super::ui_function_calls::FilterExpression<super::ui_function_calls::FunctionCallStatus>,
>,
#[ts(optional)]
pub tags: Option<Vec<super::ui_function_calls::TagFilter>>,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct TimeSeriesFloatPoint {
#[ts(type = "number")]
pub interval_start: EpochMsTimestamp,
pub value: f64,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct MetricSeriesFloat {
pub total: f64,
pub series: Vec<TimeSeriesFloatPoint>,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct CostByClientBreakdownItem {
pub client_name: String,
pub total_cost: f64,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GetCostStatsResponse {
pub total_cost: MetricSeriesFloat,
pub cost_by_client: Vec<CostByClientBreakdownItem>,
}

pub struct GetCostStats;

impl ApiEndpoint for GetCostStats {
type Request<'a> = GetCostStatsRequest;
type Response<'a> = GetCostStatsResponse;

const PATH: &'static str = "/v1/dashboard/cost";
}
29 changes: 29 additions & 0 deletions engine/baml-rpc/src/ui/ui_function_call_http_calls.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
use serde::{Deserialize, Serialize};
use ts_rs::TS;

use super::ui_types::UiHttpCall;
use crate::{rpc::ApiEndpoint, FunctionCallId, ProjectId};

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GetFunctionCallHttpCallsRequest {
#[ts(type = "string")]
pub project_id: ProjectId,
#[ts(type = "string")]
pub function_call_id: FunctionCallId,
}

#[derive(Debug, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct GetFunctionCallHttpCallsResponse {
pub http_calls: Vec<UiHttpCall>,
}

pub struct GetFunctionCallHttpCalls;

impl ApiEndpoint for GetFunctionCallHttpCalls {
type Request<'a> = GetFunctionCallHttpCallsRequest;
type Response<'a> = GetFunctionCallHttpCallsResponse;

const PATH: &'static str = "/v1/function-calls/http-calls";
}
4 changes: 4 additions & 0 deletions engine/baml-rpc/src/ui/ui_traces.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,9 @@ pub struct ListTracesRequest {
/// Search term to filter across function_call_id, function_name, tags, error, input (args), and output
#[ts(optional)]
pub search: Option<String>,
/// Filter to only show LLM function calls (function_type = 'baml_llm')
#[ts(optional)]
pub llm_only: Option<FilterExpression<bool>>,
}

impl Default for ListTracesRequest {
Expand Down Expand Up @@ -156,6 +159,7 @@ impl Default for ListTracesRequest {
streamed: None,
relative_time: None,
search: None,
llm_only: None,
}
}
}
Expand Down
12 changes: 7 additions & 5 deletions engine/baml-rpc/src/ui/ui_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,11 +154,13 @@ pub struct UiUsageEstimate {
pub input_tokens: Option<u64>,
#[ts(type = "number | null")]
pub output_tokens: Option<u64>,
// TODO: add cost estimate data here
// This is tricky to do because we need provider & model to effectively
// resolve the token costs. Even restricting to just openai is non-straightforward,
// and frankly I'm skeptical that restricting to just openai is a sufficiently common
// implementation use case.
// Cost estimates calculated from provider-specific pricing
#[ts(type = "number | null")]
pub input_cost: Option<f64>,
#[ts(type = "number | null")]
pub output_cost: Option<f64>,
#[ts(type = "number | null")]
pub total_cost: Option<f64>,
}

#[derive(Debug, Serialize, Deserialize, TS)]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@ impl aws_smithy_runtime_api::client::interceptors::Intercept for CollectorInterc
request.method().to_string(),
headers,
HTTPBody::new(request.body().bytes().unwrap_or_default().to_vec()),
self.client_details.clone(),
);
let call_stack = self.call_stack.clone();
let request = Arc::new(request);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,11 @@ pub(crate) async fn build_and_log_outbound_request(
.unwrap_or_default()
.into(),
),
ClientDetails {
name: client.context().name.clone(),
provider: client.context().provider.clone(),
options: client.context().options.clone(),
},
)),
);
BAML_TRACER.lock().unwrap().put(Arc::new(event));
Expand Down
7 changes: 6 additions & 1 deletion engine/baml-runtime/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use anyhow::{Context, Result};
use baml_ids::{FunctionCallId, HttpRequestId};
use baml_types::{
expr::{Expr, ExprMetadata},
tracing::events::{HTTPBody, HTTPRequest, TraceEvent},
tracing::events::{ClientDetails, HTTPBody, HTTPRequest, TraceEvent},
BamlMap, BamlValue, BamlValueWithMeta, Completion, Constraint,
};
use cfg_if::cfg_if;
Expand Down Expand Up @@ -1064,6 +1064,11 @@ impl BamlRuntime {
.unwrap_or_default()
.into(),
),
ClientDetails {
name: "unknown".to_string(),
provider: "unknown".to_string(),
options: IndexMap::new(),
},
))
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ mod tests {
runtime_api::baml_value::{
BamlValue, Media, MediaValue, TypeIndex, ValueContent, ValueMetadata,
},
RpcClientDetails,
};
use indexmap::IndexMap;

Expand Down Expand Up @@ -614,6 +615,11 @@ mod tests {
body: HTTPBody {
raw: Cow::Borrowed(original_body.as_bytes()),
},
client_details: RpcClientDetails {
name: "test-client".to_string(),
provider: "openai".to_string(),
options: IndexMap::new(),
},
});

// Process the trace data (simulating the actual pipeline)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,11 @@ impl<'a> IntoRpcEvent<'a, baml_rpc::runtime_api::IntermediateData<'a>>
url: self.url().to_string(),
method: self.method().to_string(),
headers: redact_headers(self.headers().clone()),
client_details: RpcClientDetails {
name: self.client_details.name.clone(),
provider: self.client_details.provider.clone(),
options: self.client_details.options.clone(),
},
body: self.body().to_rpc_event(lookup),
}
}
Expand Down
Loading
Loading