More improvements to studio2 publishing (#2333)

- Wire through client details on stream data
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Enhance client details handling in HTTP requests, add new metrics and
cost structures, and update protobuf message handling.
> 
>   - **Behavior**:
> - Add `client_details` to `HTTPRequest` in `events.rs` and
`IntermediateData` in `trace_event.rs`.
> - Update `ApiKeyEnvironment::environment()` in `api_key.rs` to handle
longer environment strings.
> - Add new metrics and cost statistics structures in `ui_dashboard.rs`
and `ui_dashboard_cost.rs`.
>   - **Protobuf**:
> - Modify `cffi.pb.go` to include `protoimpl.MessageState` in various
structs for better message handling.
>   - **Misc**:
> - Add `GetFunctionCallHttpCalls` endpoint in
`ui_function_call_http_calls.rs`.
>     - Add `llm_only` filter to `ListTracesRequest` in `ui_traces.rs`.
>     - Add cost estimates to `UiUsageEstimate` in `ui_types.rs`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for df0ed86. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

---------

Co-authored-by: Chris Watts <chris@boundaryml.com>

Author: aaronvg

engine/baml-lib/baml-types/src/tracing/events.rs

@@ -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 {
@@ -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),
         }
     }
 
@@ -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

engine/baml-rpc/src/auth/api_key.rs

@@ -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)],
         }
     }
 }

engine/baml-rpc/src/lib.rs

@@ -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},
 };

engine/baml-rpc/src/runtime_api/trace_event.rs

@@ -120,6 +120,7 @@ pub enum IntermediateData<'a> {
         url: String,
         method: String,
         headers: HashMap<String, String>,
+        client_details: RpcClientDetails,
         body: HTTPBody<'a>,
     },
     RawLLMResponse {

engine/baml-rpc/src/ui/mod.rs

@@ -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;

engine/baml-rpc/src/ui/ui_dashboard.rs

@@ -151,13 +151,48 @@ pub struct StatusCountOverTime {
     pub error_count: u64,
 }
 
+#[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 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;

engine/baml-rpc/src/ui/ui_dashboard_cost.rs

@@ -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";
+}

engine/baml-rpc/src/ui/ui_function_call_http_calls.rs

@@ -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";
+}

engine/baml-rpc/src/ui/ui_traces.rs

@@ -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 {
@@ -156,6 +159,7 @@ impl Default for ListTracesRequest {
             streamed: None,
             relative_time: None,
             search: None,
+            llm_only: None,
         }
     }
 }

engine/baml-rpc/src/ui/ui_types.rs

@@ -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)]