polyc_llm/

model_info.rs

//! Bundled model catalog: per-model context-window facts and the longest-prefix
//! lookup that drives compaction.
//!
//! The catalog is a compile-checked Rust table (no JSON file, no runtime
//! parsing). It is the single home for "how big is this model's window". The
//! control plane (compaction trigger, via `model_select`) is the consumer today;
//! it lives in the shared `polyc-llm` crate so a future harness-side reader gets
//! the same numbers rather than a divergent copy. An unknown slug never refuses —
//! it falls back to a conservative window with a single `tracing::warn`, taking
//! a "warn, don't fail" stance.

/// Per-model context facts, looked up by slug.
///
/// `Eq` is intentionally NOT derived: `force_temperature` is an `f32`, for which
/// total equality is undefined. `PartialEq` is enough for the table's tests.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ModelInfo {
    /// Maximum total tokens (prompt + generation) the model accepts.
    pub context_window: u32,
    /// Whole-percent fraction of `context_window` that is considered usable
    /// (default 95). usable = `context_window` * pct / 100.
    pub effective_context_window_percent: u16,
    /// Optional forced sampling temperature for this model, overriding whatever
    /// the caller requested. `Some` for models that require a fixed temperature
    /// (e.g. GLM-4.6/4.7 want `1.0`); `None` leaves the caller's value alone.
    ///
    /// Lives here — keyed by slug, with the same longest-prefix discipline as
    /// the context window — so the pin is resolved from the *per-request*
    /// model, never baked from a provider's `default_model`.
    pub force_temperature: Option<f32>,
}

impl ModelInfo {
    /// Default usable fraction of a model's window — reserved headroom so the
    /// prompt never runs flush against the hard limit.
    pub const DEFAULT_EFFECTIVE_PERCENT: u16 = 95;

    /// A catalog entry with the default effective percentage and no temperature
    /// pin.
    #[must_use]
    pub const fn new(context_window: u32) -> Self {
        Self {
            context_window,
            effective_context_window_percent: Self::DEFAULT_EFFECTIVE_PERCENT,
            force_temperature: None,
        }
    }

    /// Builder: set a forced sampling temperature (see
    /// [`force_temperature`](Self::force_temperature)).
    #[must_use]
    pub const fn with_forced_temperature(mut self, temperature: f32) -> Self {
        self.force_temperature = Some(temperature);
        self
    }

    /// floor(`context_window` * `effective_context_window_percent` / 100): the
    /// usable token budget after reserving headroom.
    #[must_use]
    pub const fn effective_window(self) -> u32 {
        let usable = (self.context_window as u64 * self.effective_context_window_percent as u64)
            / Self::PERCENT_DENOM;
        // `usable <= context_window` (the percent is a fraction of 100), so it
        // always fits a u32; clamp defensively rather than wrap on the cast.
        // `u32::try_from` is not yet const-stable, so the checked cast is
        // expressed manually.
        if usable > u32::MAX as u64 {
            u32::MAX
        } else {
            // SAFETY (value): guarded above to fit in u32, so this never truncates.
            #[allow(clippy::cast_possible_truncation)]
            {
                usable as u32
            }
        }
    }

    /// Whole-percent denominator for [`effective_window`](Self::effective_window).
    const PERCENT_DENOM: u64 = 100;
}

/// Conservative fallback window for an unknown slug. `128_000` chosen so an
/// unknown model compacts early rather than overflowing a real (possibly
/// smaller) window.
pub const FALLBACK_CONTEXT_WINDOW: u32 = 128_000;

/// The bundled catalog. Keys are matched by longest prefix (see
/// [`lookup_model`]), so dated/preview suffixes resolve to their base entry and
/// more-specific slugs (`gpt-4o`) beat shorter ones (`gpt-4`).
static CATALOG: &[(&str, ModelInfo)] = &[
    // Gemini 3+ only — pre-3 (2.5 and earlier) is deprecated and unsupported.
    // The `gemini-3` catch-all covers 3 / 3.1 / flash / pro / flash-lite (all 1M).
    ("gemini-3-flash", ModelInfo::new(1_048_576)),
    ("gemini-3.1-pro", ModelInfo::new(1_048_576)),
    ("gemini-3-pro", ModelInfo::new(1_048_576)),
    ("gemini-3", ModelInfo::new(1_048_576)),
    ("gpt-4o", ModelInfo::new(128_000)),
    ("gpt-4-turbo", ModelInfo::new(128_000)),
    ("gpt-4", ModelInfo::new(8_192)),
    ("llama3.2", ModelInfo::new(131_072)),
    ("llama3", ModelInfo::new(131_072)),
    // Z.AI GLM family (served OpenAI-compatible). Only GLM 5.2 is 1M; the rest
    // are 128k–204.8k. Longest-prefix keeps `glm-5.1`/`glm-5.2` ahead of `glm-5`.
    ("glm-5.2", ModelInfo::new(1_048_576)),
    ("glm-5.1", ModelInfo::new(200_000)),
    ("glm-5", ModelInfo::new(204_800)),
    // GLM-4.6/4.7 want a fixed temperature of 1.0 (matches opencode); the 5.x
    // family and 4.5 must NOT be pinned.
    (
        "glm-4.7",
        ModelInfo::new(204_800).with_forced_temperature(1.0),
    ),
    (
        "glm-4.6",
        ModelInfo::new(204_800).with_forced_temperature(1.0),
    ),
    ("glm-4.5", ModelInfo::new(131_072)),
];

/// Longest-prefix match.
///
/// Iterate the catalog and keep the entry whose key is a prefix of `slug`
/// (`slug.starts_with(key)`) with the GREATEST `key.len()`. An exact match is
/// naturally the longest possible prefix. Returns `None` if no key is a prefix.
#[must_use]
pub fn lookup_model(slug: &str) -> Option<ModelInfo> {
    CATALOG
        .iter()
        .filter(|(k, _)| slug.starts_with(k))
        .max_by_key(|(k, _)| k.len())
        .map(|(_, info)| *info)
}

/// [`lookup_model`] with the fallback applied plus ONE `tracing::warn` naming
/// the slug and the fallback window. NEVER refuses.
#[must_use]
pub fn lookup_model_or_fallback(slug: &str) -> ModelInfo {
    lookup_model(slug).unwrap_or_else(|| {
        tracing::warn!(
            model = %slug,
            fallback_context_window = FALLBACK_CONTEXT_WINDOW,
            "unknown model slug; using conservative fallback context window"
        );
        ModelInfo::new(FALLBACK_CONTEXT_WINDOW)
    })
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use super::*;

    #[test]
    fn exact_match_resolves() {
        assert_eq!(lookup_model("gpt-4").unwrap().context_window, 8_192);
        assert_eq!(lookup_model("gpt-4o").unwrap().context_window, 128_000);
    }

    #[test]
    fn longest_prefix_beats_shorter_keys() {
        // A dated gpt-4o slug must resolve to gpt-4o (128k), NOT gpt-4 (8k).
        assert_eq!(
            lookup_model("gpt-4o-2024-08-06").unwrap().context_window,
            128_000
        );
        // gpt-4-turbo beats gpt-4.
        assert_eq!(
            lookup_model("gpt-4-turbo-2024").unwrap().context_window,
            128_000
        );
    }

    #[test]
    fn preview_and_tag_suffixes_resolve_to_base() {
        assert_eq!(
            lookup_model("gemini-3.1-pro-preview")
                .unwrap()
                .context_window,
            1_048_576
        );
        assert_eq!(
            lookup_model("gemini-3-flash-preview-09-2025")
                .unwrap()
                .context_window,
            1_048_576
        );
        assert_eq!(
            lookup_model("llama3.2:latest").unwrap().context_window,
            131_072
        );
    }

    #[test]
    fn glm_family_windows_resolve() {
        // Only GLM 5.2 is 1M; siblings are smaller.
        assert_eq!(lookup_model("glm-5.2").unwrap().context_window, 1_048_576);
        assert_eq!(lookup_model("glm-5.1").unwrap().context_window, 200_000);
        assert_eq!(lookup_model("glm-5").unwrap().context_window, 204_800);
        assert_eq!(lookup_model("glm-4.6").unwrap().context_window, 204_800);
        assert_eq!(lookup_model("glm-4.5").unwrap().context_window, 131_072);
    }

    #[test]
    fn glm_force_temperature_only_for_4_6_and_4_7() {
        assert_eq!(
            lookup_model("glm-4.6").unwrap().force_temperature,
            Some(1.0)
        );
        assert_eq!(
            lookup_model("glm-4.7-flash").unwrap().force_temperature,
            Some(1.0)
        );
        // Catalog default and the 5.x family carry no pin.
        assert_eq!(lookup_model("glm-4.5").unwrap().force_temperature, None);
        assert_eq!(lookup_model("glm-5.2").unwrap().force_temperature, None);
        assert_eq!(lookup_model("gpt-4o").unwrap().force_temperature, None);
    }

    #[test]
    fn glm_longest_prefix_and_suffixes() {
        // A dated/suffixed glm-5.2 slug must resolve to glm-5.2 (1M), NOT glm-5.
        assert_eq!(
            lookup_model("glm-5.2-0712").unwrap().context_window,
            1_048_576
        );
        // glm-5.1 must not be shadowed by glm-5.
        assert_eq!(lookup_model("glm-5.1-air").unwrap().context_window, 200_000);
    }

    #[test]
    fn unknown_slug_has_no_catalog_entry_but_falls_back() {
        assert!(lookup_model("foo-model").is_none());
        assert_eq!(
            lookup_model_or_fallback("foo-model").context_window,
            FALLBACK_CONTEXT_WINDOW
        );
    }

    #[test]
    fn effective_window_math() {
        assert_eq!(ModelInfo::new(1_048_576).effective_window(), 996_147);
        assert_eq!(ModelInfo::new(128_000).effective_window(), 121_600);
    }
}