diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 0ded98f27c5..27f9fe347e0 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -357,7 +357,11 @@ impl HostCallResponse { Self::Select(ids) | Self::Filtered(Ok((ids, _))), HostCall::Select { candidates, .. }, ) = (self, call) - && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + && let Some(id) = { + let known: std::collections::HashSet<&str> = + candidates.iter().map(String::as_str).collect(); + ids.iter().find(|id| !known.contains(id.as_str())) + } { return Err(DriveError::UnknownCandidate { id: id.clone() }); } @@ -688,14 +692,22 @@ fn shares_ids(candidates: &[Arc]) -> bool { /// The candidates a host answer names, in its order. An id names every candidate that carries /// it, so deployments sharing an id stay together through a filter, as Python's list keeps them. fn answered_by_id(candidates: &[Arc], ids: &[String]) -> Vec> { - let mut taken = vec![false; candidates.len()]; + let mut positions: std::collections::HashMap<&str, Vec> = + std::collections::HashMap::new(); + for (position, deployment) in candidates.iter().enumerate() { + positions + .entry(deployment.id.as_str()) + .or_default() + .push(position); + } let mut kept = Vec::new(); for id in ids { - for (position, deployment) in candidates.iter().enumerate() { - if !taken[position] && deployment.id == *id { - taken[position] = true; - kept.push(deployment.clone()); - } + if let Some(found) = positions.remove(id.as_str()) { + kept.extend( + found + .into_iter() + .map(|position| candidates[position].clone()), + ); } } kept @@ -1540,6 +1552,26 @@ mod tests { ); } + #[rstest] + fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() { + let candidates: Vec> = ["a", "b", "a", "c"] + .into_iter() + .map(|id| { + let spec: litellm_router_core::DeploymentSpec = + serde_json::from_value(deployment("g", id)).unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + }) + .collect(); + let ids = ["c", "a", "x", "a"].map(String::from); + + let kept = answered_by_id(&candidates, &ids); + + let positions: Vec = kept + .iter() + .map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap()) + .collect(); + assert_eq!(positions, [3, 0, 2]); + } #[rstest] fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() { let mut second = deployment("primary", "a1");