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.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
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,
}

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.
    #[must_use]
    pub const fn new(context_window: u32) -> Self {
        Self {
            context_window,
            effective_context_window_percent: Self::DEFAULT_EFFECTIVE_PERCENT,
        }
    }

    /// 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)),
];

/// 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 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);
    }
}