apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-05-10 00:02:43 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

redesign model_metrics_sources, drop legacy per-provider routing, return ranked model list

Browse source
This commit is contained in:
Adil Hafeez 2026-03-27 12:37:38 -07:00
parent b12bf74e5c
commit 76b1f37052
12 changed files with 639 additions and 429 deletions
Download patch file Download diff file Expand all files Collapse all files

137
crates/brightstaff/src/router/llm.rs
View file

@ -1,11 +1,11 @@
use std::{collections::HashMap, sync::Arc};
use common::{
configuration::{
LlmProvider, ModelUsagePreference, RoutingPreference, TopLevelRoutingPreference,
},
configuration::TopLevelRoutingPreference,
consts::{ARCH_PROVIDER_HINT_HEADER, REQUEST_ID_HEADER, TRACE_PARENT_HEADER},
};
use super::router_model::{ModelUsagePreference, RoutingPreference};
use hermesllm::apis::openai::Message;
use hyper::header;
use thiserror::Error;
@ -22,7 +22,6 @@ pub struct RouterService {
client: reqwest::Client,
router_model: Arc<dyn RouterModel>,
routing_provider_name: String,
llm_usage_defined: bool,
top_level_preferences: HashMap<String, TopLevelRoutingPreference>,
metrics_service: Option<Arc<ModelMetricsService>>,
}
@ -40,60 +39,37 @@ pub type Result<T> = std::result::Result<T, RoutingError>;
impl RouterService {
pub fn new(
providers: Vec<LlmProvider>,
top_level_prefs: Option<Vec<TopLevelRoutingPreference>>,
metrics_service: Option<Arc<ModelMetricsService>>,
router_url: String,
routing_model_name: String,
routing_provider_name: String,
) -> Self {
// Build top-level preference map and sentinel llm_routes when v0.4.0 format is used.
let (top_level_preferences, llm_routes, llm_usage_defined) =
if let Some(top_prefs) = top_level_prefs {
let top_level_map: HashMap<String, TopLevelRoutingPreference> =
top_prefs.into_iter().map(|p| (p.name.clone(), p)).collect();
// Build sentinel routes: route_name → first model (RouterModelV1 needs a model
// mapping, but RouterService overrides the selection via metrics_service).
let sentinel_routes: HashMap<String, Vec<RoutingPreference>> = top_level_map
.iter()
.filter_map(|(name, pref)| {
pref.models.first().map(|first_model| {
(
first_model.clone(),
vec![RoutingPreference {
name: name.clone(),
description: pref.description.clone(),
}],
)
})
})
.collect();
let defined = !top_level_map.is_empty();
(top_level_map, sentinel_routes, defined)
} else {
// Legacy per-provider format.
let providers_with_usage = providers
.iter()
.filter(|provider| provider.routing_preferences.is_some())
.cloned()
.collect::<Vec<LlmProvider>>();
let top_level_preferences: HashMap<String, TopLevelRoutingPreference> = top_level_prefs
.map_or_else(HashMap::new, |prefs| {
prefs.into_iter().map(|p| (p.name.clone(), p)).collect()
});
let routes: HashMap<String, Vec<RoutingPreference>> = providers_with_usage
.iter()
.filter_map(|provider| {
provider
.routing_preferences
.as_ref()
.map(|prefs| (provider.name.clone(), prefs.clone()))
})
.collect();
let defined = !providers_with_usage.is_empty();
(HashMap::new(), routes, defined)
};
// Build sentinel routes for RouterModelV1: route_name → first model.
// RouterModelV1 uses this to build its prompt; RouterService overrides
// the model selection via rank_models() after the route is determined.
let sentinel_routes: HashMap<String, Vec<RoutingPreference>> = top_level_preferences
.iter()
.filter_map(|(name, pref)| {
pref.models.first().map(|first_model| {
(
first_model.clone(),
vec![RoutingPreference {
name: name.clone(),
description: pref.description.clone(),
}],
)
})
})
.collect();
let router_model = Arc::new(router_model_v1::RouterModelV1::new(
llm_routes,
sentinel_routes,
routing_model_name,
router_model_v1::MAX_TOKEN_LEN,
));
@ -103,7 +79,6 @@ impl RouterService {
client: reqwest::Client::new(),
router_model,
routing_provider_name,
llm_usage_defined,
top_level_preferences,
metrics_service,
}
@ -113,10 +88,9 @@ impl RouterService {
&self,
messages: &[Message],
traceparent: &str,
usage_preferences: Option<Vec<ModelUsagePreference>>,
inline_routing_preferences: Option<Vec<TopLevelRoutingPreference>>,
request_id: &str,
) -> Result<Option<(String, String)>> {
) -> Result<Option<(String, Vec<String>)>> {
if messages.is_empty() {
return Ok(None);
}
@ -126,41 +100,27 @@ impl RouterService {
inline_routing_preferences
.map(|prefs| prefs.into_iter().map(|p| (p.name.clone(), p)).collect());
// Determine whether any routing is defined.
let has_top_level = inline_top_map.is_some() || !self.top_level_preferences.is_empty();
if usage_preferences
.as_ref()
.is_none_or(|prefs| prefs.len() < 2)
&& !self.llm_usage_defined
&& !has_top_level
{
// No routing defined — skip the router call entirely.
if inline_top_map.is_none() && self.top_level_preferences.is_empty() {
return Ok(None);
}
// For top-level format, build a synthetic ModelUsagePreference list so RouterModelV1
// For inline overrides, build synthetic ModelUsagePreference list so RouterModelV1
// generates the correct prompt (route name + description pairs).
// For config-level prefs the sentinel routes are already baked into RouterModelV1.
let effective_usage_preferences: Option<Vec<ModelUsagePreference>> =
if let Some(ref inline_map) = inline_top_map {
Some(
inline_map
.values()
.map(|p| ModelUsagePreference {
model: p.models.first().cloned().unwrap_or_default(),
routing_preferences: vec![RoutingPreference {
name: p.name.clone(),
description: p.description.clone(),
}],
})
.collect(),
)
} else if !self.top_level_preferences.is_empty() {
// Config top-level prefs: already encoded as sentinel routes in RouterModelV1,
// pass None so it uses the pre-built llm_route_json_str.
None
} else {
usage_preferences.clone()
};
inline_top_map.as_ref().map(|inline_map| {
inline_map
.values()
.map(|p| ModelUsagePreference {
model: p.models.first().cloned().unwrap_or_default(),
routing_preferences: vec![RoutingPreference {
name: p.name.clone(),
description: p.description.clone(),
}],
})
.collect()
});
let router_request = self
.router_model
@ -209,21 +169,20 @@ impl RouterService {
.router_model
.parse_response(&content, &effective_usage_preferences)?;
let result = if let Some((route_name, _sentinel_model)) = parsed {
// Check if this route belongs to the top-level preference format.
let result = if let Some((route_name, _sentinel)) = parsed {
let top_pref = inline_top_map
.as_ref()
.and_then(|m| m.get(&route_name))
.or_else(|| self.top_level_preferences.get(&route_name));
if let Some(pref) = top_pref {
let selected_model = match &self.metrics_service {
Some(svc) => svc.select_model(&pref.models, &pref.selection_policy).await,
None => pref.models.first().cloned().unwrap_or_default(),
let ranked = match &self.metrics_service {
Some(svc) => svc.rank_models(&pref.models, &pref.selection_policy).await,
None => pref.models.clone(),
};
Some((route_name, selected_model))
Some((route_name, ranked))
} else {
Some((route_name, _sentinel_model))
None
}
} else {
None

296
crates/brightstaff/src/router/model_metrics.rs
View file

@ -2,59 +2,75 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use common::configuration::{ModelMetricsSources, SelectionPolicy, SelectionPreference};
use serde::Deserialize;
use common::configuration::{MetricsSource, SelectionPolicy, SelectionPreference};
use tokio::sync::RwLock;
use tracing::{info, warn};
#[derive(Deserialize)]
struct MetricsResponse {
#[serde(default)]
cost: HashMap<String, f64>,
#[serde(default)]
latency: HashMap<String, f64>,
}
pub struct ModelMetricsService {
cost: Arc<RwLock<HashMap<String, f64>>>,
latency: Arc<RwLock<HashMap<String, f64>>>,
}
impl ModelMetricsService {
pub async fn new(sources: &ModelMetricsSources, client: reqwest::Client) -> Self {
pub async fn new(sources: &[MetricsSource], client: reqwest::Client) -> Self {
let cost_data = Arc::new(RwLock::new(HashMap::new()));
let latency_data = Arc::new(RwLock::new(HashMap::new()));
let metrics = fetch_metrics(&sources.url, &client).await;
info!(
cost_models = metrics.cost.len(),
latency_models = metrics.latency.len(),
url = %sources.url,
"fetched model metrics"
);
*cost_data.write().await = metrics.cost;
*latency_data.write().await = metrics.latency;
for source in sources {
match source {
MetricsSource::CostMetrics {
url,
refresh_interval,
auth,
} => {
let data = fetch_cost_metrics(url, auth.as_ref(), &client).await;
info!(models = data.len(), url = %url, "fetched cost metrics");
*cost_data.write().await = data;
if let Some(interval_secs) = sources.refresh_interval {
let cost_clone = Arc::clone(&cost_data);
let latency_clone = Arc::clone(&latency_data);
let client_clone = client.clone();
let url = sources.url.clone();
tokio::spawn(async move {
let interval = Duration::from_secs(interval_secs);
loop {
tokio::time::sleep(interval).await;
let metrics = fetch_metrics(&url, &client_clone).await;
info!(
cost_models = metrics.cost.len(),
latency_models = metrics.latency.len(),
url = %url,
"refreshed model metrics"
);
*cost_clone.write().await = metrics.cost;
*latency_clone.write().await = metrics.latency;
if let Some(interval_secs) = refresh_interval {
let cost_clone = Arc::clone(&cost_data);
let client_clone = client.clone();
let url = url.clone();
let auth = auth.clone();
let interval = Duration::from_secs(*interval_secs);
tokio::spawn(async move {
loop {
tokio::time::sleep(interval).await;
let data =
fetch_cost_metrics(&url, auth.as_ref(), &client_clone).await;
info!(models = data.len(), url = %url, "refreshed cost metrics");
*cost_clone.write().await = data;
}
});
}
}
});
MetricsSource::PrometheusMetrics {
url,
query,
refresh_interval,
} => {
let data = fetch_prometheus_metrics(url, query, &client).await;
info!(models = data.len(), url = %url, "fetched prometheus latency metrics");
*latency_data.write().await = data;
if let Some(interval_secs) = refresh_interval {
let latency_clone = Arc::clone(&latency_data);
let client_clone = client.clone();
let url = url.clone();
let query = query.clone();
let interval = Duration::from_secs(*interval_secs);
tokio::spawn(async move {
loop {
tokio::time::sleep(interval).await;
let data =
fetch_prometheus_metrics(&url, &query, &client_clone).await;
info!(models = data.len(), url = %url, "refreshed prometheus latency metrics");
*latency_clone.write().await = data;
}
});
}
}
}
}
ModelMetricsService {
@ -63,63 +79,136 @@ impl ModelMetricsService {
}
}
/// Select the best model from `models` according to `policy`.
/// Falls back to `models[0]` if metric data is unavailable for all candidates.
pub async fn select_model(&self, models: &[String], policy: &SelectionPolicy) -> String {
/// Rank `models` by `policy`, returning them in preference order.
/// Models with no metric data are appended at the end in their original order.
pub async fn rank_models(&self, models: &[String], policy: &SelectionPolicy) -> Vec<String> {
match policy.prefer {
SelectionPreference::Cheapest => {
let data = self.cost.read().await;
select_by_ascending_metric(models, &data)
rank_by_ascending_metric(models, &data)
}
SelectionPreference::Fastest => {
let data = self.latency.read().await;
select_by_ascending_metric(models, &data)
}
SelectionPreference::Random => {
let idx = rand_index(models.len());
models[idx].clone()
rank_by_ascending_metric(models, &data)
}
SelectionPreference::Random => shuffle(models),
SelectionPreference::None => models.to_vec(),
}
}
}
fn select_by_ascending_metric(models: &[String], data: &HashMap<String, f64>) -> String {
models
fn rank_by_ascending_metric(models: &[String], data: &HashMap<String, f64>) -> Vec<String> {
let mut with_data: Vec<(&String, f64)> = models
.iter()
.filter_map(|m| data.get(m.as_str()).map(|v| (m, *v)))
.min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(m, _)| m.clone())
.unwrap_or_else(|| models[0].clone())
.collect();
with_data.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
let without_data: Vec<&String> = models
.iter()
.filter(|m| !data.contains_key(m.as_str()))
.collect();
with_data
.iter()
.map(|(m, _)| (*m).clone())
.chain(without_data.iter().map(|m| (*m).clone()))
.collect()
}
/// Simple non-crypto random index using system time nanoseconds.
fn rand_index(len: usize) -> usize {
fn shuffle(models: &[String]) -> Vec<String> {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
let seed = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.subsec_nanos() as usize)
.unwrap_or(0);
nanos % len
let mut result = models.to_vec();
let mut state = seed;
for i in (1..result.len()).rev() {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
let j = state % (i + 1);
result.swap(i, j);
}
result
}
async fn fetch_metrics(url: &str, client: &reqwest::Client) -> MetricsResponse {
match client.get(url).send().await {
Ok(resp) => match resp.json::<MetricsResponse>().await {
async fn fetch_cost_metrics(
url: &str,
auth: Option<&common::configuration::MetricsAuth>,
client: &reqwest::Client,
) -> HashMap<String, f64> {
let mut req = client.get(url);
if let Some(auth) = auth {
if auth.auth_type == "bearer" {
req = req.header("Authorization", format!("Bearer {}", auth.token));
} else {
warn!(auth_type = %auth.auth_type, "unsupported auth type for cost_metrics, skipping auth");
}
}
match req.send().await {
Ok(resp) => match resp.json::<HashMap<String, f64>>().await {
Ok(data) => data,
Err(err) => {
warn!(error = %err, url = %url, "failed to parse metrics response");
MetricsResponse {
cost: HashMap::new(),
latency: HashMap::new(),
}
warn!(error = %err, url = %url, "failed to parse cost metrics response");
HashMap::new()
}
},
Err(err) => {
warn!(error = %err, url = %url, "failed to fetch metrics");
MetricsResponse {
cost: HashMap::new(),
latency: HashMap::new(),
warn!(error = %err, url = %url, "failed to fetch cost metrics");
HashMap::new()
}
}
}
#[derive(serde::Deserialize)]
struct PrometheusResponse {
data: PrometheusData,
}
#[derive(serde::Deserialize)]
struct PrometheusData {
result: Vec<PrometheusResult>,
}
#[derive(serde::Deserialize)]
struct PrometheusResult {
metric: HashMap<String, String>,
value: (f64, String), // (timestamp, value_str)
}
async fn fetch_prometheus_metrics(
url: &str,
query: &str,
client: &reqwest::Client,
) -> HashMap<String, f64> {
let query_url = format!("{}/api/v1/query", url.trim_end_matches('/'));
match client
.get(&query_url)
.query(&[("query", query)])
.send()
.await
{
Ok(resp) => match resp.json::<PrometheusResponse>().await {
Ok(prom) => prom
.data
.result
.into_iter()
.filter_map(|r| {
let model_name = r.metric.get("model_name")?.clone();
let value: f64 = r.value.1.parse().ok()?;
Some((model_name, value))
})
.collect(),
Err(err) => {
warn!(error = %err, url = %query_url, "failed to parse prometheus response");
HashMap::new()
}
},
Err(err) => {
warn!(error = %err, url = %query_url, "failed to fetch prometheus metrics");
HashMap::new()
}
}
}
@ -134,32 +223,35 @@ mod tests {
}
#[test]
fn test_select_by_ascending_metric_picks_lowest() {
fn test_rank_by_ascending_metric_picks_lowest_first() {
let models = vec!["a".to_string(), "b".to_string(), "c".to_string()];
let mut data = HashMap::new();
data.insert("a".to_string(), 0.01);
data.insert("b".to_string(), 0.005);
data.insert("c".to_string(), 0.02);
assert_eq!(select_by_ascending_metric(&models, &data), "b");
assert_eq!(
rank_by_ascending_metric(&models, &data),
vec!["b", "a", "c"]
);
}
#[test]
fn test_select_by_ascending_metric_fallback_to_first() {
fn test_rank_by_ascending_metric_no_data_preserves_order() {
let models = vec!["x".to_string(), "y".to_string()];
let data = HashMap::new();
assert_eq!(select_by_ascending_metric(&models, &data), "x");
assert_eq!(rank_by_ascending_metric(&models, &data), vec!["x", "y"]);
}
#[test]
fn test_select_by_ascending_metric_partial_data() {
fn test_rank_by_ascending_metric_partial_data() {
let models = vec!["a".to_string(), "b".to_string()];
let mut data = HashMap::new();
data.insert("b".to_string(), 100.0);
assert_eq!(select_by_ascending_metric(&models, &data), "b");
assert_eq!(rank_by_ascending_metric(&models, &data), vec!["b", "a"]);
}
#[tokio::test]
async fn test_select_model_cheapest() {
async fn test_rank_models_cheapest() {
let service = ModelMetricsService {
cost: Arc::new(RwLock::new({
let mut m = HashMap::new();
@ -171,13 +263,13 @@ mod tests {
};
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
let result = service
.select_model(&models, &make_policy(SelectionPreference::Cheapest))
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
.await;
assert_eq!(result, "gpt-4o-mini");
assert_eq!(result, vec!["gpt-4o-mini", "gpt-4o"]);
}
#[tokio::test]
async fn test_select_model_fastest() {
async fn test_rank_models_fastest() {
let service = ModelMetricsService {
cost: Arc::new(RwLock::new(HashMap::new())),
latency: Arc::new(RwLock::new({
@ -189,21 +281,57 @@ mod tests {
};
let models = vec!["gpt-4o".to_string(), "claude-sonnet".to_string()];
let result = service
.select_model(&models, &make_policy(SelectionPreference::Fastest))
.rank_models(&models, &make_policy(SelectionPreference::Fastest))
.await;
assert_eq!(result, "claude-sonnet");
assert_eq!(result, vec!["claude-sonnet", "gpt-4o"]);
}
#[tokio::test]
async fn test_select_model_fallback_no_metrics() {
async fn test_rank_models_fallback_no_metrics() {
let service = ModelMetricsService {
cost: Arc::new(RwLock::new(HashMap::new())),
latency: Arc::new(RwLock::new(HashMap::new())),
};
let models = vec!["model-a".to_string(), "model-b".to_string()];
let result = service
.select_model(&models, &make_policy(SelectionPreference::Cheapest))
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
.await;
assert_eq!(result, "model-a");
assert_eq!(result, vec!["model-a", "model-b"]);
}
#[tokio::test]
async fn test_rank_models_partial_data_appended_last() {
let service = ModelMetricsService {
cost: Arc::new(RwLock::new({
let mut m = HashMap::new();
m.insert("gpt-4o".to_string(), 0.005);
m
})),
latency: Arc::new(RwLock::new(HashMap::new())),
};
let models = vec!["gpt-4o-mini".to_string(), "gpt-4o".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
.await;
assert_eq!(result, vec!["gpt-4o", "gpt-4o-mini"]);
}
#[tokio::test]
async fn test_rank_models_none_preserves_order() {
let service = ModelMetricsService {
cost: Arc::new(RwLock::new({
let mut m = HashMap::new();
m.insert("gpt-4o-mini".to_string(), 0.0001);
m.insert("gpt-4o".to_string(), 0.005);
m
})),
latency: Arc::new(RwLock::new(HashMap::new())),
};
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::None))
.await;
// none → original order, despite gpt-4o-mini being cheaper
assert_eq!(result, vec!["gpt-4o", "gpt-4o-mini"]);
}
}

16
crates/brightstaff/src/router/router_model.rs
View file

@ -1,5 +1,5 @@
use common::configuration::ModelUsagePreference;
use hermesllm::apis::openai::{ChatCompletionsRequest, Message};
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Error)]
@ -10,6 +10,20 @@ pub enum RoutingModelError {
pub type Result<T> = std::result::Result<T, RoutingModelError>;
/// Internal route descriptor passed to the router model to build its prompt.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingPreference {
pub name: String,
pub description: String,
}
/// Groups a model with its routing preferences (used internally by RouterModelV1).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelUsagePreference {
pub model: String,
pub routing_preferences: Vec<RoutingPreference>,
}
pub trait RouterModel: Send + Sync {
fn generate_request(
&self,

2
crates/brightstaff/src/router/router_model_v1.rs
View file

@ -1,6 +1,6 @@
use std::collections::HashMap;
use common::configuration::{ModelUsagePreference, RoutingPreference};
use super::router_model::{ModelUsagePreference, RoutingPreference};
use hermesllm::apis::openai::{ChatCompletionsRequest, Message, MessageContent, Role};
use hermesllm::transforms::lib::ExtractText;
use serde::{Deserialize, Serialize};