diff --git a/crates/brightstaff/src/router/router_model_v1.rs b/crates/brightstaff/src/router/router_model_v1.rs index 2dccd61f..3117fb07 100644 --- a/crates/brightstaff/src/router/router_model_v1.rs +++ b/crates/brightstaff/src/router/router_model_v1.rs @@ -73,7 +73,7 @@ impl RouterModel for RouterModelV1 { fn generate_request( &self, messages: &[Message], - usage_preferences: &Option>, + usage_preferences_from_request: &Option>, ) -> ChatCompletionsRequest { // remove system prompt, tool calls, tool call response and messages without content // if content is empty its likely a tool call @@ -150,36 +150,17 @@ impl RouterModel for RouterModelV1 { }) .collect::>(); - let llm_route_json = usage_preferences - .as_ref() - .map(|prefs| { - let llm_route: Vec = prefs - .iter() - .flat_map(|pref| { - let routing_preferences = pref.routing_preferences.clone(); - routing_preferences - .into_iter() - .map(|routing_pref| RoutingPreference { - name: routing_pref.name, - description: routing_pref.description, - }) - }) - .collect(); - serde_json::to_string(&llm_route).unwrap_or_default() - }) - .unwrap_or_else(|| self.llm_route_json_str.clone()); - - let messages_content = ARCH_ROUTER_V1_SYSTEM_PROMPT - .replace("{routes}", &llm_route_json) - .replace( - "{conversation}", - &serde_json::to_string(&selected_conversation_list).unwrap_or_default(), - ); + // Generate the router request message based on the usage preferences. + // If preferences are passed in request then we use them otherwise we use the default routing model preferences. + let router_message = match convert_to_router_preferences(usage_preferences_from_request) { + Some(prefs) => generate_router_message(&prefs, &selected_conversation_list), + None => generate_router_message(&self.llm_route_json_str, &selected_conversation_list), + }; ChatCompletionsRequest { model: self.routing_model.clone(), messages: vec![Message { - content: Some(ContentType::Text(messages_content)), + content: Some(ContentType::Text(router_message)), role: USER_ROLE.to_string(), }], temperature: Some(0.01), @@ -245,6 +226,37 @@ impl RouterModel for RouterModelV1 { } } +fn generate_router_message(prefs: &str, selected_conversation_list: &Vec) -> String { + ARCH_ROUTER_V1_SYSTEM_PROMPT + .replace("{routes}", prefs) + .replace( + "{conversation}", + &serde_json::to_string(&selected_conversation_list).unwrap_or_default(), + ) +} + +fn convert_to_router_preferences( + prefs_from_request: &Option>, +) -> Option { + if let Some(usage_preferences) = prefs_from_request { + let routing_preferences = usage_preferences + .iter() + .flat_map(|pref| { + pref.routing_preferences + .iter() + .map(|routing_pref| RoutingPreference { + name: routing_pref.name.clone(), + description: routing_pref.description.clone(), + }) + }) + .collect::>(); + + return Some(serde_json::to_string(&routing_preferences).unwrap_or_default()); + } + + None +} + fn fix_json_response(body: &str) -> String { let mut updated_body = body.to_string();