feat: specify reasoning effort in model ID, or use a default

src/openai/chat/completion.rs

@@ -1,4 +1,6 @@
 use serde::{Deserialize, Serialize};
+use serde_with::{DeserializeFromStr, SerializeDisplay};
+use strum::{EnumString, VariantArray};
 use std::collections::HashMap;
 
 // ============================================================================
@@ -916,8 +918,8 @@ pub enum PromptCacheRetention {
 }
 
 /// Constrains reasoning effort for reasoning models.
-#[derive(Debug, Clone, Deserialize, Serialize)]
-#[serde(rename_all = "lowercase")]
+#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)]
+#[strum(serialize_all = "lowercase")]
 pub enum ReasoningEffort {
     /// Disable reasoning.
     None,
@@ -930,7 +932,8 @@ pub enum ReasoningEffort {
     /// High reasoning effort.
     High,
     /// Extra high reasoning effort.
-    Xhigh,
+    #[strum(serialize = "xhigh")]
+    ExtraHigh,
     /// Maximum reasoning effort.
     Max,
 }

src/server.rs

@@ -10,9 +10,9 @@ use axum::response::Sse;
 use axum_extra::routing::RouterExt;
 use axum_extra::routing::TypedPath;
 use ecow::EcoString;
-use enumset::{EnumSet, EnumSetType};
 use headers::Header;
 use headers::authorization::Bearer;
+use itertools::iproduct;
 use mitsein::vec1::Vec1;
 use rand::Rng;
 use rand::RngExt;
@@ -21,7 +21,14 @@ use secrecy::ExposeSecret;
 use secrecy::SecretString;
 use serde::Deserialize;
 use serde::Serialize;
+use serde_with::DeserializeFromStr;
+use serde_with::SerializeDisplay;
+use strum::EnumString;
+use strum::VariantArray;
 use std::collections::VecDeque;
+use std::convert::Infallible;
+use std::fmt::Display;
+use std::str::FromStr;
 use std::sync::Arc;
 use typed_builder::TypedBuilder;
 use url::Url;
@@ -51,8 +58,8 @@ enum Provider {
     OpenRouter,
 }
 
-#[derive(Debug, Serialize, Deserialize, EnumSetType)]
-#[serde(rename_all = "lowercase")]
+#[derive(Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray)]
+#[strum(serialize_all = "lowercase")]
 enum Shortcut {
     XS,
     S,
@@ -61,6 +68,43 @@ enum Shortcut {
     XL,
 }
 
+
+#[derive(Debug, Clone, DeserializeFromStr, SerializeDisplay)]
+struct ShortcutAndMaybeReasoning(Shortcut, Option<openai::chat::completion::ReasoningEffort>);
+
+impl FromStr for ShortcutAndMaybeReasoning {
+    type Err = Infallible; // TODO
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some((model, reasoning)) = s.split_once(":") {
+            let model = model.parse().expect("TODO");
+            let reasoning = reasoning.parse().expect("TODO");
+
+            Ok(Self(model, Some(reasoning)))
+        } else {
+            let model = s.parse().expect("TODO");
+            let reasoning = None;
+
+            Ok(Self(model, reasoning))
+        }
+    }
+}
+
+impl Display for ShortcutAndMaybeReasoning {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Self(model, reasoning) = self;
+
+        write!(f, "{model}")?;
+        if let Some(reasoning) = reasoning {
+            write!(f, ":")?;
+            write!(f, "{reasoning}")?;
+        }
+
+        Ok(())
+    }
+}
+
+
 impl Shortcut {
     fn suggestions(&self) -> Vec<(Provider, EcoString)> {
         use Provider as P;
@@ -133,6 +177,19 @@ impl Shortcut {
             ],
         }
     }
+
+    fn default_reasoning(&self) -> openai::chat::completion::ReasoningEffort {
+        use Shortcut as S;
+        use openai::chat::completion::ReasoningEffort as R;
+
+        match self {
+            S::XS => R::None,
+            S::S => R::Minimal,
+            S::M => R::Low,
+            S::L => R::Medium,
+            S::XL => R::High,
+        }
+    }
 }
 
 fn take_random_by_swapping<T>(mut vec1: Vec1<T>, mut rng: impl Rng) -> (T, Option<Vec1<T>>) {
@@ -202,12 +259,12 @@ impl IntoResponse for ChatCompletionError {
 async fn create_chat_completion(
     _: V1ChatCompletionsPath,
     State(state): State<AppState>,
-    Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<Shortcut>>,
+    Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<ShortcutAndMaybeReasoning>>,
 ) -> Result<ChatCompletionResponse, ChatCompletionError> {
-    let suggestions = request.model.suggestions();
+    let ShortcutAndMaybeReasoning(shortcut, reasoning_in_id) = request.model.clone();
     let mut model_upstreams = VecDeque::new();
 
-    for (provider, model) in suggestions {
+    for (provider, model) in shortcut.suggestions() {
         let upstreams = match provider {
             Provider::CloudflareWorkersAi => vec![], // TODO
             Provider::GoogleAiStudio => state
@@ -329,6 +386,8 @@ async fn create_chat_completion(
             web_search_options,
         } = request.clone();
 
+        let reasoning_effort = Some(reasoning_effort.or(reasoning_in_id.clone()).unwrap_or(shortcut.default_reasoning()));
+
         let request = openai::chat::completion::CreateChatCompletionRequest {
             model,
             messages,
@@ -455,11 +514,14 @@ async fn create_chat_completion(
 struct V1ModelsPath;
 
 #[tracing::instrument]
-async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<Shortcut>> {
-    let models = EnumSet::all()
+async fn list_models(_: V1ModelsPath) -> Json<openai::models::ListModelsResponse<ShortcutAndMaybeReasoning>> {
+    let shortcuts = Shortcut::VARIANTS.into_iter().cloned();
+    let reasoning_efforts = openai::chat::completion::ReasoningEffort::VARIANTS.into_iter().cloned().map(Some).chain(None);
+
+    let models = iproduct!(shortcuts, reasoning_efforts)
         .into_iter()
-        .map(|shortcut| openai::models::Model {
-            id: shortcut,
+        .map(|(shortcut, reasoning)| openai::models::Model {
+            id: ShortcutAndMaybeReasoning(shortcut, reasoning),
             created: 0,
             object: openai::models::ModelObject::Model,
             owned_by: "nobody".into(),