diff --git a/crates/llm-ls/src/main.rs b/crates/llm-ls/src/main.rs index 054aada..0c3ee8c 100644 --- a/crates/llm-ls/src/main.rs +++ b/crates/llm-ls/src/main.rs @@ -306,7 +306,7 @@ fn build_prompt( ) -> Result { let t = Instant::now(); if fim.enabled { - let mut token_count = context_window; + let mut remaining_token_count = context_window - 3; // account for FIM tokens let mut before_iter = text.lines_at(pos.line as usize + 1).reversed(); let mut after_iter = text.lines_at(pos.line as usize); let mut before_line = before_iter.next(); @@ -332,10 +332,10 @@ fn build_prompt( } else { before_line.len() }; - if tokens > token_count { + if tokens > remaining_token_count { break; } - token_count -= tokens; + remaining_token_count -= tokens; before.push(before_line); } if let Some(after_line) = after_line { @@ -348,10 +348,10 @@ fn build_prompt( } else { after_line.len() }; - if tokens > token_count { + if tokens > remaining_token_count { break; } - token_count -= tokens; + remaining_token_count -= tokens; after.push_str(&after_line); } before_line = before_iter.next(); @@ -369,7 +369,7 @@ fn build_prompt( info!(prompt, build_prompt_ms = time, "built prompt in {time} ms"); Ok(prompt) } else { - let mut token_count = context_window; + let mut remaining_token_count = context_window; let mut before = vec![]; let mut first = true; for mut line in text.lines_at(pos.line as usize + 1).reversed() { @@ -387,10 +387,10 @@ fn build_prompt( } else { line.len() }; - if tokens > token_count { + if tokens > remaining_token_count { break; } - token_count -= tokens; + remaining_token_count -= tokens; before.push(line); } let prompt = before.into_iter().rev().collect::>().join(""); @@ -851,4 +851,3 @@ async fn main() { Server::new(stdin, stdout, socket).serve(service).await; } -