ai/
model.rs

1use std::default::Default;
2use std::fmt::{self, Display};
3use std::str::FromStr;
4use std::sync::OnceLock;
5
6use anyhow::{bail, Result};
7use serde::{Deserialize, Serialize};
8use tiktoken_rs::CoreBPE;
9use tiktoken_rs::model::get_context_size;
10
11use crate::profile;
12
13// Cached tokenizer for performance
14static TOKENIZER: OnceLock<CoreBPE> = OnceLock::new();
15
16// Model identifiers - using screaming case for constants
17const MODEL_GPT4_1: &str = "gpt-4.1";
18const MODEL_GPT4_1_MINI: &str = "gpt-4.1-mini";
19const MODEL_GPT4_1_NANO: &str = "gpt-4.1-nano";
20const MODEL_GPT4_5: &str = "gpt-4.5";
21
22/// Represents the available AI models for commit message generation.
23/// Each model has different capabilities and token limits.
24///
25/// Known variants exist so we can specialize tokenizer/context-size handling for
26/// them, but any other model string is carried through verbatim via `Other` so
27/// users can point git-ai at arbitrary models (e.g. local ollama endpoints).
28#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize, Default)]
29pub enum Model {
30  /// GPT-4.1 - highest quality of the GPT-4.1 family
31  GPT41,
32  /// Mini version of GPT-4.1: faster/cheaper, the **default** for this per-commit tool
33  #[default]
34  GPT41Mini,
35  /// Nano version of GPT-4.1 for very fast processing
36  GPT41Nano,
37  /// GPT-4.5 model for advanced capabilities
38  GPT45,
39  /// Any other model string, carried through verbatim (e.g. local/ollama models).
40  Other(String)
41}
42
43impl Model {
44  /// Counts the number of tokens in the given text for the current model.
45  /// This is used to ensure we stay within the model's token limits.
46  ///
47  /// # Arguments
48  /// * `text` - The text to count tokens for
49  ///
50  /// # Returns
51  /// * `Result<usize>` - The number of tokens or an error
52  pub fn count_tokens(&self, text: &str) -> Result<usize> {
53    profile!("Count tokens");
54
55    // Fast path for empty text
56    if text.is_empty() {
57      return Ok(0);
58    }
59
60    // Always use the proper tokenizer for accurate counts
61    // We cannot afford to underestimate tokens as it may cause API failures
62    let tokenizer = TOKENIZER.get_or_init(|| get_tokenizer(self.as_ref()));
63
64    // Use direct tokenization for accurate token count
65    let tokens = tokenizer.encode_ordinary(text);
66    Ok(tokens.len())
67  }
68
69  /// Gets the maximum context size for the current model.
70  ///
71  /// # Returns
72  /// * `usize` - The maximum number of tokens the model can process
73  pub fn context_size(&self) -> usize {
74    profile!("Get context size");
75    // tiktoken-rs 0.12 returns Option; fall back to 4096 (the historical default
76    // returned by tiktoken-rs 0.7 for unrecognized models) when the model is unknown.
77    get_context_size(self.as_ref()).unwrap_or(4096)
78  }
79
80  /// Truncates the given text to fit within the specified token limit.
81  ///
82  /// # Arguments
83  /// * `text` - The text to truncate
84  /// * `max_tokens` - The maximum number of tokens allowed
85  ///
86  /// # Returns
87  /// * `Result<String>` - The truncated text or an error
88  pub(crate) fn truncate(&self, text: &str, max_tokens: usize) -> Result<String> {
89    profile!("Truncate text");
90    self.walk_truncate(text, max_tokens, usize::MAX)
91  }
92
93  /// Truncates text to fit within a token limit using a single tokenization pass.
94  ///
95  /// The previous implementation re-`join`ed words and re-tokenized the full candidate
96  /// string on every binary-search iteration (O(log n) full tokenizations + re-joins),
97  /// which was the dominant cost on large diffs. This version tokenizes the text exactly
98  /// once, keeps the first `max_tokens` tokens, and decodes them back to a string. Because
99  /// a single character can span multiple tokens, slicing the token vector can land on an
100  /// invalid UTF-8 boundary; in that case we drop trailing tokens until the decode succeeds.
101  /// Dropping tokens only ever reduces the count, so the result is guaranteed to re-encode
102  /// to `<= max_tokens` while always being valid UTF-8.
103  ///
104  /// The `_within` parameter is retained for signature compatibility but is no longer used:
105  /// the result is exact rather than an approximation within a tolerance.
106  ///
107  /// # Arguments
108  /// * `text` - The text to truncate
109  /// * `max_tokens` - The maximum number of tokens allowed
110  /// * `_within` - Unused; kept for backward-compatible call sites
111  ///
112  /// # Returns
113  /// * `Result<String>` - The truncated text or an error
114  pub(crate) fn walk_truncate(&self, text: &str, max_tokens: usize, _within: usize) -> Result<String> {
115    profile!("Walk truncate");
116    log::debug!("max_tokens: {max_tokens}");
117
118    // Nothing to keep.
119    if max_tokens == 0 || text.is_empty() {
120      return Ok(String::new());
121    }
122
123    let tokenizer = TOKENIZER.get_or_init(|| get_tokenizer(self.as_ref()));
124
125    // Single tokenization pass.
126    let tokens = tokenizer.encode_ordinary(text);
127    if tokens.len() <= max_tokens {
128      return Ok(text.to_string());
129    }
130
131    // Keep the first `max_tokens` tokens, then back off until the slice decodes to
132    // valid UTF-8 (the slice boundary may fall inside a multi-byte character).
133    let mut end = max_tokens;
134    loop {
135      match tokenizer.decode(&tokens[..end]) {
136        Ok(decoded) => return Ok(decoded),
137        Err(_) if end > 0 => end -= 1,
138        Err(e) => return Err(e)
139      }
140    }
141  }
142}
143
144impl AsRef<str> for Model {
145  fn as_ref(&self) -> &str {
146    match self {
147      Model::GPT41 => MODEL_GPT4_1,
148      Model::GPT41Mini => MODEL_GPT4_1_MINI,
149      Model::GPT41Nano => MODEL_GPT4_1_NANO,
150      Model::GPT45 => MODEL_GPT4_5,
151      Model::Other(name) => name.as_str()
152    }
153  }
154}
155
156// Keep conversion to String for cases that need owned strings
157impl From<&Model> for String {
158  fn from(model: &Model) -> Self {
159    model.as_ref().to_string()
160  }
161}
162
163// Keep the old impl for backwards compatibility where possible
164impl Model {
165  pub fn as_str(&self) -> &str {
166    self.as_ref()
167  }
168}
169
170impl FromStr for Model {
171  type Err = anyhow::Error;
172
173  fn from_str(s: &str) -> Result<Self> {
174    let trimmed = s.trim();
175    let normalized = trimmed.to_lowercase();
176    match normalized.as_str() {
177      "" => bail!("Model name cannot be empty"),
178      "gpt-4.1" => Ok(Model::GPT41),
179      "gpt-4.1-mini" => Ok(Model::GPT41Mini),
180      "gpt-4.1-nano" => Ok(Model::GPT41Nano),
181      "gpt-4.5" => Ok(Model::GPT45),
182      // Backward compatibility for deprecated models - map to closest GPT-4.1 equivalent
183      "gpt-4" | "gpt-4o" => {
184        log::warn!(
185          "Model '{}' is deprecated. Mapping to 'gpt-4.1'. \
186          Please update your configuration with: git ai config set model gpt-4.1",
187          s
188        );
189        Ok(Model::GPT41)
190      }
191      "gpt-4o-mini" | "gpt-3.5-turbo" => {
192        log::warn!(
193          "Model '{}' is deprecated. Mapping to 'gpt-4.1-mini'. \
194          Please update your configuration with: git ai config set model gpt-4.1-mini",
195          s
196        );
197        Ok(Model::GPT41Mini)
198      }
199      // Any other model string is accepted and carried through verbatim (original
200      // case preserved, since local/ollama model names can be case-sensitive).
201      _ => Ok(Model::Other(trimmed.to_string()))
202    }
203  }
204}
205
206impl Display for Model {
207  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
208    write!(f, "{}", self.as_ref())
209  }
210}
211
212// Implement conversion from string types to Model with fallback to default
213impl From<&str> for Model {
214  fn from(s: &str) -> Self {
215    s.parse().unwrap_or_else(|e| {
216      log::error!("Failed to parse model '{}': {}. Falling back to default model 'gpt-4.1'.", s, e);
217      Model::default()
218    })
219  }
220}
221
222impl From<String> for Model {
223  fn from(s: String) -> Self {
224    s.as_str().into()
225  }
226}
227
228/// Returns true if `name` is a built-in/known model or a recognized deprecated
229/// alias. Such names are always considered valid and need no endpoint verification.
230pub fn is_known_or_deprecated(name: &str) -> bool {
231  matches!(
232    name.trim().to_lowercase().as_str(),
233    "gpt-4.1" | "gpt-4.1-mini" | "gpt-4.1-nano" | "gpt-4.5" | "gpt-4" | "gpt-4o" | "gpt-4o-mini" | "gpt-3.5-turbo"
234  )
235}
236
237fn get_tokenizer(_model_str: &str) -> CoreBPE {
238  // TODO: This should be based on the model string, but for now we'll just use cl100k_base
239  // which is used by gpt-3.5-turbo and gpt-4
240  tiktoken_rs::cl100k_base().expect("Failed to create tokenizer")
241}
242
243#[cfg(test)]
244mod tests {
245  use super::*;
246
247  /// C3: A large synthetic text must truncate to <= max_tokens, stay on a valid UTF-8
248  /// boundary, and re-encode to <= max_tokens.
249  #[test]
250  fn test_truncate_large_text_is_exact_and_utf8_safe() {
251    let model = Model::GPT41;
252    // Large input that comfortably exceeds the limit, with multi-byte UTF-8 characters
253    // (é, 世界, emoji) so we exercise the token-boundary back-off.
254    let text = "The quick brown fox café 世界 🚀 jumps over the lazy dog. ".repeat(500);
255    let max_tokens = 100;
256
257    let truncated = model.truncate(&text, max_tokens).unwrap();
258
259    // Valid UTF-8 by construction (it's a Rust String), and re-encodes to <= max_tokens.
260    let recount = model.count_tokens(&truncated).unwrap();
261    assert!(recount <= max_tokens, "re-encoded token count {recount} exceeds max {max_tokens}");
262
263    // It actually truncated (input was far larger than the limit).
264    assert!(truncated.len() < text.len(), "expected truncation to shorten the text");
265    assert!(!truncated.is_empty(), "truncation of large text should not be empty");
266  }
267
268  /// Text already within the limit is returned unchanged.
269  #[test]
270  fn test_truncate_passthrough_when_within_limit() {
271    let model = Model::GPT41;
272    let text = "small bit of text";
273    let truncated = model.truncate(text, 1000).unwrap();
274    assert_eq!(truncated, text);
275  }
276
277  /// max_tokens == 0 yields an empty string (and never panics).
278  #[test]
279  fn test_truncate_zero_tokens() {
280    let model = Model::GPT41;
281    let truncated = model.truncate("anything at all here", 0).unwrap();
282    assert_eq!(truncated, "");
283  }
284
285  /// Truncating multi-byte content to a tiny budget stays on a char boundary.
286  #[test]
287  fn test_truncate_multibyte_small_budget() {
288    let model = Model::GPT41;
289    let text = "日本語のテキストをトークン化してから切り詰めます".repeat(20);
290    let truncated = model.truncate(&text, 5).unwrap();
291    let recount = model.count_tokens(&truncated).unwrap();
292    assert!(recount <= 5, "re-encoded token count {recount} exceeds 5");
293    // Valid UTF-8 (String) — assert it is a prefix-ish valid slice by checking it round-trips.
294    assert!(truncated.is_char_boundary(truncated.len()));
295  }
296}
ai/model.rs

ai/
model.rs
