1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
// src/llm/providers/llama_cpp/engine/inference.rs
use super::session::InferenceSession;
use crate::error::{AmbiError, Result};
use crate::types::config::LlamaEngineConfig;
use llama_cpp_2::context::LlamaContext;
use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::{AddBos, LlamaModel};
use llama_cpp_2::sampling::LlamaSampler;
use llama_cpp_2::token::LlamaToken;
use log::{debug, info, warn};
impl InferenceSession {
/// Execute a full completion loop for the given prompt.
///
/// # Parameters
/// - `prompt`: the formatted prompt to be tokenized and processed.
/// - `cfg`: static engine configuration (temperature, penalties, etc.).
/// - `callback`: invoked for every successfully decoded UTF-8 string.
/// Returning `false` signals early termination (e.g. receiver dropped).
///
/// # Returns
/// `Ok(())` on normal completion (end-of-generation or callback
/// termination), otherwise an `AmbiError::EngineError` describing the
/// failure point.
pub(crate) fn run_inference<F>(
prompt: &str,
model: &LlamaModel,
context: &mut LlamaContext,
mut batch: &mut LlamaBatch,
session: &mut InferenceSession,
cfg: &LlamaEngineConfig,
mut callback: F,
) -> Result<()>
where
F: FnMut(String) -> bool,
{
// ----- 1. Tokenize prompt --------------------------------------------------
debug!("\n{}\n========================================", prompt);
let tokens_list = model
.str_to_token(prompt, AddBos::Always)
.map_err(|e| AmbiError::EngineError(format!("Tokenize failed: {}", e)))?;
let current_tokens: Vec<LlamaToken> = tokens_list.to_vec();
// ----- 2. Validate prompt size vs context window --------------------------
if current_tokens.len() >= cfg.n_ctx as usize {
return Err(AmbiError::EngineError(format!(
"Prompt size ({} tokens) exceeds or equals n_ctx limit ({})",
current_tokens.len(),
cfg.n_ctx
)));
}
let dynamic_max_tokens = std::cmp::min(
cfg.max_tokens as usize,
(cfg.n_ctx as usize).saturating_sub(current_tokens.len()),
);
if dynamic_max_tokens < 32 {
return Err(AmbiError::EngineError(format!(
"Insufficient token space left for generation (only {} tokens). \
Increase n_ctx or reduce prompt length.",
dynamic_max_tokens
)));
}
// ----- 3. KV‑cache eviction optimisation ----------------------------------
// Compare the new prompt with the previously cached history to avoid
// re‑evaluating the common prefix.
let mut match_len = 0;
for (t1, t2) in session.history_tokens.iter().zip(current_tokens.iter()) {
if t1 == t2 {
match_len += 1;
} else {
break;
}
}
// If the common prefix is shorter than our full history, we must
// remove the divergent suffix from the KV cache and shift remaining
// entries.
if match_len < session.history_tokens.len() {
let evicted_len = session.history_tokens.len() - match_len;
info!(
"Evicting {} tokens, applying KV‑cache shift to save evaluation cost.",
evicted_len
);
let p0 = match_len as u32; // start of the range to remove
let p1 = (match_len + evicted_len) as u32; // exclusive end
if let Err(e) = context.clear_kv_cache_seq(Some(0), Some(p0), Some(p1)) {
warn!(
"Failed to cleanly remove KV‑cache sequence: {}. Falling back to full reset.",
e
);
context.clear_kv_cache();
match_len = 0;
} else if let Err(e) =
context.kv_cache_seq_add(0, Some(p1), None, -(evicted_len as i32))
{
warn!(
"Failed to shift KV‑cache sequence: {}. Falling back to full reset.",
e
);
context.clear_kv_cache();
match_len = 0;
} else {
// Shift successful; history is now the common prefix.
session.history_tokens.truncate(match_len);
// match_len stays as the new effective length.
}
}
// ----- 4. Evaluate new tokens (the suffix that is not cached) -------------
session.pos = match_len as i32;
let new_tokens = ¤t_tokens[match_len..];
let chunk_size = cfg.n_tokens;
let total_new = new_tokens.len();
let mut processed = 0;
for chunk in new_tokens.chunks(chunk_size) {
batch.clear();
for &t in chunk.iter() {
processed += 1;
let is_last = processed == total_new;
batch
.add(t, session.pos, &[0], is_last)
.map_err(|e| AmbiError::EngineError(format!("Batch add failed: {}", e)))?;
session.pos += 1;
}
if !chunk.is_empty() {
context
.decode(&mut batch)
.map_err(|e| AmbiError::EngineError(format!("Decoding failed: {}", e)))?;
}
}
// Now our history fully reflects the current prompt.
session.history_tokens = current_tokens;
// ----- 5. Configure sampler chain -----------------------------------------
let mut sampler = LlamaSampler::chain_simple([
LlamaSampler::penalties(
cfg.penalty_last_n,
cfg.penalty_repeat,
cfg.penalty_freq,
cfg.penalty_present,
),
LlamaSampler::top_p(cfg.top_p, cfg.min_keep),
LlamaSampler::temp(cfg.temp),
LlamaSampler::dist(cfg.seed),
]);
// ----- 6. Generation loop -------------------------------------------------
let mut decoded_count = 0;
loop {
let next_token = sampler.sample(&context, batch.n_tokens() - 1);
sampler.accept(next_token);
// Check for end-of-generation or length limit.
if model.is_eog_token(next_token) || decoded_count >= dynamic_max_tokens {
break;
}
// Record token for future KV‑cache matching.
session.history_tokens.push(next_token);
// Decode token bytes and assemble valid UTF-8 strings.
if let Ok(bytes) = model.token_to_piece_bytes(next_token, cfg.buffer_size, true, None) {
session.utf8_buffer.extend_from_slice(&bytes);
match std::str::from_utf8(&session.utf8_buffer) {
Ok(valid_str) => {
if !callback(valid_str.to_string()) {
// Callback requested early stop (e.g. stream dropped).
break;
}
session.utf8_buffer.clear();
}
Err(e) => {
let valid_up_to = e.valid_up_to();
if valid_up_to > 0 {
// Emit the valid prefix even if the buffer is incomplete.
let valid_str = unsafe {
std::str::from_utf8_unchecked(&session.utf8_buffer[..valid_up_to])
};
if !callback(valid_str.to_string()) {
break;
}
session.utf8_buffer.drain(..valid_up_to);
}
// Otherwise keep the incomplete bytes for the next token.
}
}
} // If token_to_piece_bytes fails we simply skip this token (rare edge case).
// Prepare the batch for the next iteration.
batch.clear();
batch
.add(next_token, session.pos, &[0], true)
.map_err(|e| AmbiError::EngineError(format!("Batch add failed: {}", e)))?;
context
.decode(&mut batch)
.map_err(|e| AmbiError::EngineError(format!("Decoding failed: {}", e)))?;
session.pos += 1;
decoded_count += 1;
}
debug!("Generation finished after {} new tokens.", decoded_count);
Ok(())
}
}