1#![forbid(unsafe_code)]
16
17use candle_core::{Device, Tensor};
18use el_core::{
19 ChatMessage, ChatRequest, ChatResponse, ChatRole, ChatToken, DomainEvent, EdgeError,
20 LlmProvider, Result, SafetyMode, SessionConfig, SessionId, StopReason, Token,
21};
22use el_provenance::LoadPermit;
23use el_runtime::{
24 AnchorGuard, ContrastiveSteerer, ExpertLogits, InferenceEngine, InferenceSession,
25 LightweightFilter, NoSafety, Ports, SafetyModeSelector, SafetySteerer,
26};
27
28pub struct CandleEngine {
30 embed: Tensor,
31 w_out: Tensor,
32 vocab: usize,
33 eos: Token,
34}
35
36impl CandleEngine {
37 pub fn toy(vocab: usize, dim: usize, eos: Token) -> Result<Self> {
41 let device = Device::Cpu;
42
43 let embed_data: Vec<f32> = (0..vocab * dim)
44 .map(|k| {
45 let (i, j) = (k / dim, k % dim);
46 (((i + j) % 7) as f32) * 0.1
47 })
48 .collect();
49 let wout_data: Vec<f32> = (0..dim * vocab)
50 .map(|k| {
51 let (a, b) = (k / vocab, k % vocab);
52 ((((a * 31 + b * 17) % 13) as f32) * 0.1) - 0.6
53 })
54 .collect();
55
56 let embed = Tensor::from_vec(embed_data, (vocab, dim), &device)
57 .map_err(|_| EdgeError::Engine("candle: embed tensor build failed"))?;
58 let w_out = Tensor::from_vec(wout_data, (dim, vocab), &device)
59 .map_err(|_| EdgeError::Engine("candle: w_out tensor build failed"))?;
60
61 Ok(Self {
62 embed,
63 w_out,
64 vocab,
65 eos,
66 })
67 }
68
69 pub fn from_path(path: impl AsRef<std::path::Path>, eos: Token) -> Result<Self> {
79 let file = std::fs::File::open(path.as_ref())
80 .map_err(|_| EdgeError::Engine("model file not found or not readable"))?;
81 Self::load_gguf(&mut std::io::BufReader::new(file), eos)
82 }
83
84 pub fn from_bytes(data: &[u8], eos: Token) -> Result<Self> {
89 Self::load_gguf(&mut std::io::Cursor::new(data), eos)
90 }
91
92 fn load_gguf<R: std::io::Read + std::io::Seek>(reader: &mut R, eos: Token) -> Result<Self> {
93 use candle_core::quantized::gguf_file;
94
95 let content = gguf_file::Content::read(reader)
96 .map_err(|_| EdgeError::Engine("GGUF: invalid or unrecognised file"))?;
97 let device = Device::Cpu;
98
99 let embed = content
100 .tensor(reader, "token_embd.weight", &device)
101 .map_err(|_| EdgeError::Engine("GGUF: missing 'token_embd.weight'"))?
102 .dequantize(&device)
103 .map_err(|_| EdgeError::Engine("GGUF: cannot dequantize embed tensor"))?;
104
105 let (vocab, dim) = match embed.shape().dims() {
106 [v, d] => (*v, *d),
107 _ => return Err(EdgeError::Engine("GGUF: 'token_embd.weight' must be 2-D")),
108 };
109
110 let raw_w_q = match content.tensor(reader, "output.weight", &device) {
111 Ok(t) => t,
112 Err(_) => content
113 .tensor(reader, "lm_head.weight", &device)
114 .map_err(|_| {
115 EdgeError::Engine("GGUF: missing 'output.weight' / 'lm_head.weight'")
116 })?,
117 };
118 let raw_w = raw_w_q
119 .dequantize(&device)
120 .map_err(|_| EdgeError::Engine("GGUF: cannot dequantize output weight"))?;
121
122 let w_out = match raw_w.shape().dims() {
125 [v, _d] if *v == vocab => raw_w
126 .t()
127 .map_err(|_| EdgeError::Engine("GGUF: failed to transpose output weight"))?,
128 _ => raw_w,
129 };
130
131 match w_out.shape().dims() {
134 [d, v] if *d == dim && *v == vocab => {}
135 _ => return Err(EdgeError::Engine(
136 "GGUF: output weight shape incompatible with embed dim — expected [dim, vocab] after transpose",
137 )),
138 }
139
140 Ok(Self {
141 embed,
142 w_out,
143 vocab,
144 eos,
145 })
146 }
147
148 fn forward(&self, last: usize) -> candle_core::Result<Vec<f32>> {
150 let row = self.embed.narrow(0, last, 1)?; let logits = row.matmul(&self.w_out)?; Ok(logits.to_vec2::<f32>()?.remove(0))
153 }
154}
155
156impl InferenceEngine for CandleEngine {
157 fn prefill(&mut self, tokens: &[Token]) -> Result<u32> {
158 Ok(tokens.len() as u32)
159 }
160
161 fn next_logits(&mut self, committed: &[Token]) -> Vec<i32> {
162 let last = committed
163 .last()
164 .copied()
165 .unwrap_or(0)
166 .min(self.vocab as u32 - 1) as usize;
167 match self.forward(last) {
168 Ok(logits) => logits.iter().map(|x| (x * 1000.0).round() as i32).collect(),
169 Err(_) => vec![0; self.vocab],
170 }
171 }
172
173 fn eos_token(&self) -> Token {
174 self.eos
175 }
176
177 fn rollback(&mut self, _keep_committed: u32) -> Result<()> {
180 Ok(())
181 }
182}
183
184pub struct LocalLlmProvider {
190 session: std::sync::Mutex<InferenceSession<CandleEngine>>,
191 vocab: usize,
192}
193
194impl LocalLlmProvider {
195 pub fn from_path(
197 path: impl AsRef<std::path::Path>,
198 eos: Token,
199 permit: LoadPermit,
200 ) -> Result<Self> {
201 let engine = CandleEngine::from_path(path, eos)?;
202 let vocab = engine.vocab;
203 let session = InferenceSession::new(SessionId(1), SessionConfig::default(), engine, permit);
204 Ok(Self {
205 session: std::sync::Mutex::new(session),
206 vocab,
207 })
208 }
209
210 pub fn toy(vocab: usize, dim: usize, eos: Token, permit: LoadPermit) -> Result<Self> {
212 let engine = CandleEngine::toy(vocab, dim, eos)?;
213 let session = InferenceSession::new(SessionId(1), SessionConfig::default(), engine, permit);
214 Ok(Self {
215 session: std::sync::Mutex::new(session),
216 vocab,
217 })
218 }
219
220 fn encode(&self, text: &str) -> Vec<Token> {
221 text.bytes()
222 .map(|b| (b as Token) % self.vocab as Token)
223 .collect()
224 }
225
226 fn decode(tokens: &[Token]) -> String {
227 tokens
228 .iter()
229 .map(|&t| {
230 let b = (t & 0xFF) as u8;
231 if b.is_ascii_graphic() || b == b' ' {
232 b as char
233 } else {
234 '?'
235 }
236 })
237 .collect()
238 }
239
240 fn format_messages(messages: &[ChatMessage]) -> String {
241 messages
242 .iter()
243 .map(|m| {
244 let role = match m.role {
245 ChatRole::System => "system",
246 ChatRole::User => "user",
247 ChatRole::Assistant => "assistant",
248 };
249 format!("{role}: {}", m.content)
250 })
251 .collect::<Vec<_>>()
252 .join("\n")
253 }
254}
255
256impl LlmProvider for LocalLlmProvider {
257 fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
258 let prompt = Self::format_messages(&req.messages);
259 let prompt_tokens = self.encode(&prompt);
260 let prompt_len = prompt_tokens.len() as u32;
261 let max = req.max_tokens.unwrap_or(64);
262
263 let mut session = self.session.lock().unwrap();
264 session.reset();
265 let ports = Ports::permissive();
266 session.load_prompt(&ports, &prompt_tokens)?;
267 session.generate(&ports, max)?;
268
269 let output = session.output().to_vec();
270 let completion_len = output.len() as u32;
271
272 Ok(ChatResponse {
273 content: Self::decode(&output),
274 model: "local/candle".into(),
275 prompt_tokens: prompt_len,
276 completion_tokens: completion_len,
277 })
278 }
279
280 fn chat_stream(&self, req: &ChatRequest, on_token: &mut dyn FnMut(ChatToken)) -> Result<()> {
281 let resp = self.chat(req)?;
282 for ch in resp.content.chars() {
283 on_token(ChatToken {
284 text: ch.to_string(),
285 is_final: false,
286 });
287 }
288 on_token(ChatToken {
289 text: String::new(),
290 is_final: true,
291 });
292 Ok(())
293 }
294}
295
296use candle_transformers::models::quantized_qwen2::ModelWeights as Qwen2Weights;
305use el_core::{ModelId, ModelVersion};
306use el_provenance::{ModelArtifact, SignatureVerifier};
307use tokenizers::Tokenizer;
308
309mod bench {
316 use std::cell::Cell;
317 use std::sync::OnceLock;
318 use std::time::Duration;
319
320 static ENABLED: OnceLock<bool> = OnceLock::new();
321
322 pub fn enabled() -> bool {
324 *ENABLED.get_or_init(|| std::env::var_os("EL_BENCH").is_some())
325 }
326
327 thread_local! {
328 static FWD_TOTAL: Cell<Duration> = const { Cell::new(Duration::ZERO) };
329 static FWD_MODEL: Cell<Duration> = const { Cell::new(Duration::ZERO) };
330 static FWD_CALLS: Cell<u64> = const { Cell::new(0) };
331 }
332
333 pub fn record(total: Duration, model: Duration) {
336 FWD_TOTAL.with(|c| c.set(c.get() + total));
337 FWD_MODEL.with(|c| c.set(c.get() + model));
338 FWD_CALLS.with(|c| c.set(c.get() + 1));
339 }
340
341 pub fn take() -> (Duration, Duration, u64) {
343 (
344 FWD_TOTAL.replace(Duration::ZERO),
345 FWD_MODEL.replace(Duration::ZERO),
346 FWD_CALLS.replace(0),
347 )
348 }
349}
350
351pub struct QwenEngine {
365 model: Qwen2Weights,
366 device: Device,
367 index_pos: usize,
369 fed: usize,
371 prompt: Vec<Token>,
374 last_logits: Vec<i32>,
376 vocab: usize,
377 eos: Token,
378}
379
380impl QwenEngine {
381 pub fn from_path(path: impl AsRef<std::path::Path>, eos: Token) -> Result<Self> {
383 use candle_core::quantized::gguf_file;
384 let mut file = std::fs::File::open(path.as_ref())
385 .map_err(|_| EdgeError::Engine("model file not found or not readable"))?;
386 let content = gguf_file::Content::read(&mut file)
387 .map_err(|_| EdgeError::Engine("GGUF: invalid or unrecognised file"))?;
388 let device = Device::Cpu;
389 let model = Qwen2Weights::from_gguf(content, &mut file, &device)
390 .map_err(|_| EdgeError::Engine("GGUF: failed to load Qwen2 weights"))?;
391 Ok(Self {
392 model,
393 device,
394 index_pos: 0,
395 fed: 0,
396 prompt: Vec::new(),
397 last_logits: Vec::new(),
398 vocab: 0,
399 eos,
400 })
401 }
402
403 fn forward_one(&mut self, token: Token) -> Result<Vec<i32>> {
406 let t_total = bench::enabled().then(std::time::Instant::now);
407
408 let input = Tensor::from_vec(vec![token], (1, 1), &self.device)
409 .map_err(|_| EdgeError::Engine("candle: input tensor build failed"))?;
410
411 let t_model = bench::enabled().then(std::time::Instant::now);
412 let logits = self
413 .model
414 .forward(&input, self.index_pos)
415 .map_err(|_| EdgeError::Engine("candle: Qwen2 forward failed"))?;
416 let model_dur = t_model.map(|t| t.elapsed()).unwrap_or_default();
417
418 self.index_pos += 1;
419 let row = logits
420 .squeeze(0)
421 .map_err(|_| EdgeError::Engine("candle: squeeze logits failed"))?;
422 let floats = row
423 .to_vec1::<f32>()
424 .map_err(|_| EdgeError::Engine("candle: logits to vec failed"))?;
425 let out: Vec<i32> = floats.iter().map(|x| (x * 1000.0).round() as i32).collect();
426
427 if let Some(t) = t_total {
428 bench::record(t.elapsed(), model_dur);
429 }
430 Ok(out)
431 }
432}
433
434impl InferenceEngine for QwenEngine {
435 fn prefill(&mut self, tokens: &[Token]) -> Result<u32> {
436 self.index_pos = 0;
437 self.fed = 0;
438 self.prompt = tokens.to_vec(); for &t in tokens {
440 self.last_logits = self.forward_one(t)?;
441 }
442 self.vocab = self.last_logits.len();
443 Ok(tokens.len() as u32)
444 }
445
446 fn next_logits(&mut self, committed: &[Token]) -> Vec<i32> {
447 while self.fed < committed.len() {
451 let t = committed[self.fed];
452 match self.forward_one(t) {
453 Ok(l) => self.last_logits = l,
454 Err(_) => return vec![0; self.vocab.max(1)],
455 }
456 self.fed += 1;
457 }
458 self.last_logits.clone()
459 }
460
461 fn eos_token(&self) -> Token {
462 self.eos
463 }
464
465 fn rollback(&mut self, _keep_committed: u32) -> Result<()> {
466 self.index_pos = 0;
475 self.fed = 0;
476 for i in 0..self.prompt.len() {
477 let t = self.prompt[i];
478 self.last_logits = self.forward_one(t)?;
479 }
480 Ok(())
481 }
482}
483
484const DEFAULT_UNSAFE_WORDS: &[&str] = &[
500 "bomb",
501 "explosive",
502 "detonator",
503 "methamphetamine",
504 "ricin",
505 "anthrax",
506 "sarin",
507 "nerve agent",
508];
509
510const SAFETY_REFUSAL: &str = "I can't help with that request.";
514
515const CONTRASTIVE_TOP_K: usize = 64;
519
520fn word_to_patterns(tokenizer: &Tokenizer, word: &str) -> Vec<Vec<Token>> {
527 let mut out: Vec<Vec<Token>> = Vec::new();
528 for variant in [format!(" {word}"), word.to_string()] {
529 if let Ok(enc) = tokenizer.encode(variant, false) {
530 let seq = enc.get_ids().to_vec();
531 if !seq.is_empty() && !out.contains(&seq) {
532 out.push(seq);
533 }
534 }
535 }
536 out
537}
538
539#[derive(Debug, Clone)]
543struct SafetyConfig {
544 mode: SafetyMode,
546 banned: Vec<Token>,
550 patterns: Vec<Vec<Token>>,
553 extra_guard_patterns: Vec<Vec<Token>>,
558}
559
560impl SafetyConfig {
561 fn lightweight(tokenizer: &Tokenizer) -> Self {
563 let mut banned = Vec::new();
564 let mut patterns = Vec::new();
565 for &word in DEFAULT_UNSAFE_WORDS {
566 for seq in word_to_patterns(tokenizer, word) {
567 if seq.len() == 1 && !banned.contains(&seq[0]) {
571 banned.push(seq[0]);
572 }
573 if !patterns.contains(&seq) {
574 patterns.push(seq);
575 }
576 }
577 }
578 Self {
579 mode: SafetyMode::Lightweight,
580 banned,
581 patterns,
582 extra_guard_patterns: Vec::new(),
583 }
584 }
585
586 fn ports(&self) -> Ports {
589 let mut ports = Ports::permissive();
590 if matches!(self.mode, SafetyMode::Off) {
591 return ports;
592 }
593 let steerer: Box<dyn SafetySteerer> = if self.banned.is_empty() {
594 Box::new(NoSafety)
595 } else {
596 Box::new(LightweightFilter::new(self.banned.clone()))
597 };
598 ports.safety = steerer;
599 let guard_patterns: Vec<Vec<Token>> = self
602 .patterns
603 .iter()
604 .chain(self.extra_guard_patterns.iter())
605 .cloned()
606 .collect();
607 if !guard_patterns.is_empty() {
608 ports.guard = Some(Box::new(AnchorGuard::hard(guard_patterns)));
609 }
610 if !self.patterns.is_empty() {
613 ports.ingress = Some(Box::new(AnchorGuard::hard(self.patterns.clone())));
614 }
615 ports
616 }
617}
618
619pub struct QwenExpert {
634 engine: std::cell::RefCell<QwenEngine>,
635 fed: std::cell::Cell<usize>,
638 _permit: LoadPermit,
641}
642
643impl QwenExpert {
644 pub fn from_path_primed(
648 path: impl AsRef<std::path::Path>,
649 eos: Token,
650 prompt: &[Token],
651 permit: LoadPermit,
652 ) -> Result<Self> {
653 let mut engine = QwenEngine::from_path(path, eos)?;
654 engine.prefill(prompt)?;
655 Ok(Self {
656 engine: std::cell::RefCell::new(engine),
657 fed: std::cell::Cell::new(0),
658 _permit: permit,
659 })
660 }
661}
662
663impl ExpertLogits for QwenExpert {
664 fn logits(&self, committed: &[Token]) -> Vec<i32> {
665 let mut engine = self.engine.borrow_mut();
666 if committed.len() < self.fed.get() {
672 if engine.rollback(0).is_err() {
673 return Vec::new();
674 }
675 self.fed.set(0);
676 }
677 let out = engine.next_logits(committed);
678 self.fed.set(committed.len());
679 out
680 }
681}
682
683pub struct QwenChatProvider {
695 model_path: std::path::PathBuf,
696 tokenizer: Tokenizer,
697 permit: LoadPermit,
698 eos: Token,
699 default_max_tokens: u32,
700 model_label: String,
701 safety: SafetyConfig,
702 expert_model: Option<std::path::PathBuf>,
705 steer_alpha_milli: i32,
707}
708
709impl QwenChatProvider {
710 pub fn from_paths(
712 model_path: impl AsRef<std::path::Path>,
713 tokenizer_path: impl AsRef<std::path::Path>,
714 ) -> Result<Self> {
715 let model_path = model_path.as_ref().to_path_buf();
716 if !model_path.exists() {
717 return Err(EdgeError::Engine("model file not found"));
718 }
719 let tokenizer = Tokenizer::from_file(tokenizer_path.as_ref())
720 .map_err(|_| EdgeError::Engine("failed to load tokenizer.json"))?;
721
722 let eos = tokenizer.token_to_id("<|im_end|>").unwrap_or(151_645);
724
725 let model_label = model_path
726 .file_stem()
727 .and_then(|s| s.to_str())
728 .map(|s| format!("local/{s}"))
729 .unwrap_or_else(|| "local/qwen2".to_string());
730
731 let safety = SafetyConfig::lightweight(&tokenizer);
732
733 let permit = local_load_permit(&model_path)?;
734 Ok(Self {
735 model_path,
736 tokenizer,
737 permit,
738 eos,
739 default_max_tokens: 512,
740 model_label,
741 safety,
742 expert_model: None,
743 steer_alpha_milli: 1000,
744 })
745 }
746
747 pub fn with_safety(mut self, mode: SafetyMode) -> Self {
753 self.safety.mode = mode;
754 self
755 }
756
757 pub fn with_extra_guard_words<I, S>(mut self, words: I) -> Self
764 where
765 I: IntoIterator<Item = S>,
766 S: AsRef<str>,
767 {
768 for word in words {
769 for seq in word_to_patterns(&self.tokenizer, word.as_ref()) {
770 if !self.safety.extra_guard_patterns.contains(&seq) {
771 self.safety.extra_guard_patterns.push(seq);
772 }
773 }
774 }
775 self
776 }
777
778 pub fn with_expert_model(mut self, path: impl AsRef<std::path::Path>) -> Self {
784 self.expert_model = Some(path.as_ref().to_path_buf());
785 self
786 }
787
788 pub fn with_steer_alpha(mut self, alpha_milli: i32) -> Self {
791 self.steer_alpha_milli = alpha_milli;
792 self
793 }
794
795 fn encode(&self, text: &str) -> Result<Vec<Token>> {
796 let enc = self
797 .tokenizer
798 .encode(text, false)
799 .map_err(|_| EdgeError::Engine("tokenizer encode failed"))?;
800 Ok(enc.get_ids().to_vec())
801 }
802
803 fn decode(&self, ids: &[Token]) -> Result<String> {
804 self.tokenizer
805 .decode(ids, true)
806 .map_err(|_| EdgeError::Engine("tokenizer decode failed"))
807 }
808}
809
810impl LlmProvider for QwenChatProvider {
811 fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
812 let prompt = render_chatml(&req.messages);
813
814 let t_encode = bench::enabled().then(std::time::Instant::now);
815 let prompt_tokens = self.encode(&prompt)?;
816 let d_encode = t_encode.map(|t| t.elapsed()).unwrap_or_default();
817
818 let t_load = bench::enabled().then(std::time::Instant::now);
822 let engine = QwenEngine::from_path(&self.model_path, self.eos)?;
823 let d_load = t_load.map(|t| t.elapsed()).unwrap_or_default();
824
825 let requested = requested_session_safety(self.safety.mode, self.expert_model.is_some());
831 let cfg = SessionConfig {
832 safety: requested,
833 ..SessionConfig::default()
834 };
835 let effective = SafetyModeSelector::resolve(requested, cfg.device);
840 let mut session = InferenceSession::new(SessionId(1), cfg, engine, self.permit);
841 let mut ports = self.safety.ports();
842
843 if matches!(effective, SafetyMode::SecDecoding) {
844 if let Some(expert_path) = &self.expert_model {
845 let expert = QwenExpert::from_path_primed(
846 expert_path,
847 self.eos,
848 &prompt_tokens,
849 local_load_permit(expert_path)?,
850 )?;
851 ports.safety = Box::new(ContrastiveSteerer::new(
852 expert,
853 self.safety.banned.clone(),
854 self.steer_alpha_milli,
855 CONTRASTIVE_TOP_K,
856 effective,
857 ));
858 }
859 }
860
861 let _ = bench::take(); let t_prefill = bench::enabled().then(std::time::Instant::now);
863 session.load_prompt(&ports, &prompt_tokens)?;
864 let d_prefill = t_prefill.map(|t| t.elapsed()).unwrap_or_default();
865 let (pf_total, pf_model, pf_calls) = bench::take();
866
867 let max = req.max_tokens.unwrap_or(self.default_max_tokens);
868 let t_decode = bench::enabled().then(std::time::Instant::now);
869 let stop = session.generate(&ports, max)?;
870 let d_decode = t_decode.map(|t| t.elapsed()).unwrap_or_default();
871 let (dc_total, dc_model, dc_calls) = bench::take();
872
873 let out = session.output().to_vec();
874 let completion_tokens = out.len() as u32;
875
876 let t_detok = bench::enabled().then(std::time::Instant::now);
877 let decoded = self.decode(&out)?.trim().to_string();
878 let d_detok = t_detok.map(|t| t.elapsed()).unwrap_or_default();
879
880 let safety_active = !matches!(self.safety.mode, SafetyMode::Off);
886 let content = if safety_active {
887 let events = session.drain_events();
888 let violations = events
889 .iter()
890 .filter(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))
891 .count();
892 let rollbacks = events
893 .iter()
894 .filter(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. }))
895 .count();
896 let refused = stop == StopReason::Stopped && violations > 0;
897 if violations > 0 || rollbacks > 0 {
898 eprintln!(
899 "[safety] {violations} violation(s), {rollbacks} rollback(s){}",
900 if refused {
901 " → refused (fail-closed)"
902 } else {
903 " → recovered"
904 }
905 );
906 }
907 if refused {
908 SAFETY_REFUSAL.to_string()
909 } else {
910 decoded
911 }
912 } else {
913 decoded
914 };
915
916 if bench::enabled() {
917 report_breakdown(
918 prompt_tokens.len() as u32,
919 completion_tokens,
920 d_load,
921 d_encode,
922 d_prefill,
923 d_decode,
924 d_detok,
925 (pf_total, pf_model, pf_calls),
926 (dc_total, dc_model, dc_calls),
927 );
928 }
929
930 Ok(ChatResponse {
931 content,
932 model: self.model_label.clone(),
933 prompt_tokens: prompt_tokens.len() as u32,
934 completion_tokens,
935 })
936 }
937
938 fn chat_stream(&self, req: &ChatRequest, on_token: &mut dyn FnMut(ChatToken)) -> Result<()> {
939 let resp = self.chat(req)?;
943 for ch in resp.content.chars() {
944 on_token(ChatToken {
945 text: ch.to_string(),
946 is_final: false,
947 });
948 }
949 on_token(ChatToken {
950 text: String::new(),
951 is_final: true,
952 });
953 Ok(())
954 }
955}
956
957#[allow(clippy::too_many_arguments)]
959fn report_breakdown(
960 prompt_tokens: u32,
961 completion_tokens: u32,
962 d_load: std::time::Duration,
963 d_encode: std::time::Duration,
964 d_prefill: std::time::Duration,
965 d_decode: std::time::Duration,
966 d_detok: std::time::Duration,
967 prefill_fwd: (std::time::Duration, std::time::Duration, u64),
968 decode_fwd: (std::time::Duration, std::time::Duration, u64),
969) {
970 let ms = |d: std::time::Duration| d.as_secs_f64() * 1000.0;
971 let total = d_load + d_encode + d_prefill + d_decode + d_detok;
972 let pct = |d: std::time::Duration| {
973 if total.as_secs_f64() > 0.0 {
974 d.as_secs_f64() / total.as_secs_f64() * 100.0
975 } else {
976 0.0
977 }
978 };
979 let tps = |n: u32, d: std::time::Duration| {
980 if d.as_secs_f64() > 0.0 {
981 n as f64 / d.as_secs_f64()
982 } else {
983 0.0
984 }
985 };
986
987 let (pf_total, pf_model, pf_calls) = prefill_fwd;
988 let (dc_total, dc_model, dc_calls) = decode_fwd;
989 let dc_loop = d_decode.saturating_sub(dc_total);
990 let dc_seam = dc_total.saturating_sub(dc_model);
991 let per_tok = |d: std::time::Duration, n: u64| if n > 0 { ms(d) / n as f64 } else { 0.0 };
992
993 eprintln!("\n┌─ EL_BENCH chat() breakdown ───────────────────────────────");
994 eprintln!("│ prompt_tokens={prompt_tokens} completion_tokens={completion_tokens}");
995 eprintln!("│ phase wall(ms) %total throughput");
996 eprintln!(
997 "│ model load {:>9.1} {:>6.1}% (read+dequantize GGUF)",
998 ms(d_load),
999 pct(d_load)
1000 );
1001 eprintln!(
1002 "│ tokenize {:>9.2} {:>6.1}%",
1003 ms(d_encode),
1004 pct(d_encode)
1005 );
1006 eprintln!(
1007 "│ prefill {:>9.1} {:>6.1}% {:>7.1} tok/s",
1008 ms(d_prefill),
1009 pct(d_prefill),
1010 tps(prompt_tokens, d_prefill)
1011 );
1012 eprintln!(
1013 "│ decode {:>9.1} {:>6.1}% {:>7.1} tok/s",
1014 ms(d_decode),
1015 pct(d_decode),
1016 tps(completion_tokens, d_decode)
1017 );
1018 eprintln!(
1019 "│ detokenize {:>9.2} {:>6.1}%",
1020 ms(d_detok),
1021 pct(d_detok)
1022 );
1023 eprintln!("│ TOTAL {:>9.1}", ms(total));
1024 eprintln!("│ ─ forward attribution (where prefill+decode time goes) ─");
1025 eprintln!(
1026 "│ prefill: {} fwd calls, model {:.1}ms, seam {:.1}ms, loop {:.1}ms",
1027 pf_calls,
1028 ms(pf_model),
1029 ms(pf_total.saturating_sub(pf_model)),
1030 ms(d_prefill.saturating_sub(pf_total)),
1031 );
1032 eprintln!(
1033 "│ decode : {} fwd calls, model {:.1}ms, seam {:.1}ms, loop {:.1}ms",
1034 dc_calls,
1035 ms(dc_model),
1036 ms(dc_seam),
1037 ms(dc_loop),
1038 );
1039 eprintln!(
1040 "│ per decoded token: {:.2}ms total = model {:.2} + seam {:.2} + loop {:.2}",
1041 per_tok(d_decode, dc_calls),
1042 per_tok(dc_model, dc_calls),
1043 per_tok(dc_seam, dc_calls),
1044 per_tok(dc_loop, dc_calls),
1045 );
1046 eprintln!("└───────────────────────────────────────────────────────────");
1047}
1048
1049fn render_chatml(messages: &[ChatMessage]) -> String {
1051 let mut s = String::new();
1052 for m in messages {
1053 let role = match m.role {
1054 ChatRole::System => "system",
1055 ChatRole::User => "user",
1056 ChatRole::Assistant => "assistant",
1057 };
1058 s.push_str("<|im_start|>");
1059 s.push_str(role);
1060 s.push('\n');
1061 s.push_str(&m.content);
1062 s.push_str("<|im_end|>\n");
1063 }
1064 s.push_str("<|im_start|>assistant\n");
1065 s
1066}
1067
1068fn requested_session_safety(configured: SafetyMode, has_expert: bool) -> SafetyMode {
1069 match (configured, has_expert) {
1070 (SafetyMode::Off, _) => SafetyMode::Off,
1071 (_, true) => SafetyMode::SecDecoding,
1075 (SafetyMode::SecDecoding | SafetyMode::Csd, false) => SafetyMode::Lightweight,
1079 (mode, false) => mode,
1080 }
1081}
1082
1083fn local_load_permit(path: &std::path::Path) -> Result<LoadPermit> {
1090 struct LocalFileTrust;
1091 impl SignatureVerifier for LocalFileTrust {
1092 fn verify(&self, _bytes: &[u8], _sig: &[u8], _key: u32) -> bool {
1093 true
1094 }
1095 }
1096 let _ = path;
1100 let mut artifact = ModelArtifact::new(
1101 ModelId(1),
1102 ModelVersion::new(0, 1, 0),
1103 el_core::ModelFormat::Gguf,
1104 );
1105 artifact.verify(&LocalFileTrust, b"local-trust", b"", 0);
1106 artifact.ensure_loadable()
1107}
1108
1109#[cfg(test)]
1110mod tests {
1111 use super::*;
1112 use el_runtime::InferenceEngine;
1113
1114 fn ok_permit() -> LoadPermit {
1117 use el_core::{ModelFormat, ModelId, ModelVersion};
1118 use el_provenance::{ModelArtifact, SignatureVerifier};
1119 struct OkV;
1120 impl SignatureVerifier for OkV {
1121 fn verify(&self, _: &[u8], _: &[u8], _: u32) -> bool {
1122 true
1123 }
1124 }
1125 let mut a = ModelArtifact::new(ModelId(1), ModelVersion::new(0, 1, 0), ModelFormat::Gguf);
1126 a.verify(&OkV, b"w", b"s", 0);
1127 a.ensure_loadable().unwrap()
1128 }
1129
1130 fn make_minimal_gguf(vocab: usize, dim: usize) -> Vec<u8> {
1138 let mut w: Vec<u8> = Vec::new();
1139
1140 w.extend_from_slice(b"GGUF");
1142 w.extend_from_slice(&3u32.to_le_bytes()); w.extend_from_slice(&2u64.to_le_bytes()); w.extend_from_slice(&0u64.to_le_bytes()); let tensor_bytes = (vocab * dim * 4) as u64;
1147
1148 let name = b"token_embd.weight";
1150 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1151 w.extend_from_slice(name);
1152 w.extend_from_slice(&2u32.to_le_bytes());
1153 w.extend_from_slice(&(dim as u64).to_le_bytes()); w.extend_from_slice(&(vocab as u64).to_le_bytes()); w.extend_from_slice(&0u32.to_le_bytes()); w.extend_from_slice(&0u64.to_le_bytes()); let name = b"output.weight";
1160 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1161 w.extend_from_slice(name);
1162 w.extend_from_slice(&2u32.to_le_bytes());
1163 w.extend_from_slice(&(dim as u64).to_le_bytes());
1164 w.extend_from_slice(&(vocab as u64).to_le_bytes());
1165 w.extend_from_slice(&0u32.to_le_bytes());
1166 w.extend_from_slice(&tensor_bytes.to_le_bytes()); let pad = (32usize.wrapping_sub(w.len() % 32)) % 32;
1170 w.resize(w.len() + pad, 0u8);
1171
1172 for i in 0..(vocab * dim * 2) {
1174 w.extend_from_slice(&(i as f32 * 0.1f32).to_le_bytes());
1175 }
1176
1177 w
1178 }
1179
1180 #[test]
1183 fn real_candle_forward_is_deterministic_and_right_shape() {
1184 let mut eng = CandleEngine::toy(8, 4, 7).unwrap();
1185 let a = eng.next_logits(&[2]);
1186 let b = eng.next_logits(&[2]);
1187 assert_eq!(a.len(), 8, "logits length == vocab");
1188 assert_eq!(a, b, "fixed weights → deterministic real-tensor forward");
1189 let c = eng.next_logits(&[5]);
1190 assert_ne!(a, c);
1191 }
1192
1193 #[test]
1194 fn drives_the_runtime_end_to_end() {
1195 use el_core::{ModelFormat, ModelId, ModelVersion, SessionConfig, SessionId, StopReason};
1196 use el_provenance::{ModelArtifact, SignatureVerifier};
1197
1198 struct OkVerifier;
1199 impl SignatureVerifier for OkVerifier {
1200 fn verify(&self, _: &[u8], _: &[u8], _: u32) -> bool {
1201 true
1202 }
1203 }
1204 let mut art = ModelArtifact::new(
1205 ModelId(1),
1206 ModelVersion::new(0, 1, 0),
1207 ModelFormat::Safetensors,
1208 );
1209 art.verify(&OkVerifier, b"w", b"s", 1);
1210 let permit = art.ensure_loadable().unwrap();
1211
1212 let eng = CandleEngine::toy(16, 8, 9999).unwrap();
1213 let mut session =
1214 InferenceSession::new(SessionId(1), SessionConfig::default(), eng, permit);
1215 let ports = Ports::permissive();
1216 session.load_prompt(&ports, &[1, 2, 3]).unwrap();
1217
1218 let stop = session.generate(&ports, 4).unwrap();
1219 assert_eq!(stop, StopReason::MaxTokens);
1220 assert_eq!(session.output().len(), 4);
1221 }
1222
1223 #[test]
1226 fn from_bytes_rejects_invalid_magic() {
1227 let r = CandleEngine::from_bytes(b"not a gguf file", 0);
1228 assert!(matches!(r, Err(EdgeError::Engine(_))));
1229 }
1230
1231 #[test]
1232 fn from_bytes_loads_minimal_gguf_and_forward_has_correct_vocab() {
1233 let vocab = 8;
1234 let dim = 4;
1235 let gguf = make_minimal_gguf(vocab, dim);
1236 let mut engine = CandleEngine::from_bytes(&gguf, 7).unwrap();
1237
1238 let logits = engine.next_logits(&[0]);
1239 assert_eq!(logits.len(), vocab, "logit vec width == vocab from GGUF");
1240 assert_eq!(engine.eos_token(), 7);
1241 }
1242
1243 #[test]
1244 fn from_bytes_gguf_forward_is_deterministic() {
1245 let gguf = make_minimal_gguf(8, 4);
1246 let mut eng = CandleEngine::from_bytes(&gguf, 0).unwrap();
1247 assert_eq!(eng.next_logits(&[3]), eng.next_logits(&[3]));
1248 }
1249
1250 fn make_mismatched_gguf(vocab: usize, embed_dim: usize, output_dim: usize) -> Vec<u8> {
1253 let mut w: Vec<u8> = Vec::new();
1254 w.extend_from_slice(b"GGUF");
1255 w.extend_from_slice(&3u32.to_le_bytes());
1256 w.extend_from_slice(&2u64.to_le_bytes());
1257 w.extend_from_slice(&0u64.to_le_bytes());
1258
1259 let embed_bytes = (vocab * embed_dim * 4) as u64;
1260
1261 let name = b"token_embd.weight";
1262 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1263 w.extend_from_slice(name);
1264 w.extend_from_slice(&2u32.to_le_bytes());
1265 w.extend_from_slice(&(embed_dim as u64).to_le_bytes());
1266 w.extend_from_slice(&(vocab as u64).to_le_bytes());
1267 w.extend_from_slice(&0u32.to_le_bytes());
1268 w.extend_from_slice(&0u64.to_le_bytes());
1269
1270 let name = b"output.weight";
1271 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1272 w.extend_from_slice(name);
1273 w.extend_from_slice(&2u32.to_le_bytes());
1274 w.extend_from_slice(&(output_dim as u64).to_le_bytes()); w.extend_from_slice(&(vocab as u64).to_le_bytes());
1276 w.extend_from_slice(&0u32.to_le_bytes());
1277 w.extend_from_slice(&embed_bytes.to_le_bytes());
1278
1279 let pad = (32usize.wrapping_sub(w.len() % 32)) % 32;
1280 w.resize(w.len() + pad, 0u8);
1281
1282 for i in 0..(vocab * embed_dim + vocab * output_dim) {
1283 w.extend_from_slice(&(i as f32 * 0.1f32).to_le_bytes());
1284 }
1285 w
1286 }
1287
1288 #[test]
1289 fn from_path_missing_file_returns_engine_error() {
1290 let r = CandleEngine::from_path(std::path::Path::new("/nonexistent/model.gguf"), 0);
1291 assert!(matches!(r, Err(EdgeError::Engine(_))));
1292 }
1293
1294 #[test]
1295 fn from_bytes_rejects_mismatched_output_dim_at_load_time() {
1296 let gguf = make_mismatched_gguf(8, 4, 7);
1298 let r = CandleEngine::from_bytes(&gguf, 0);
1299 assert!(
1300 matches!(r, Err(EdgeError::Engine(_))),
1301 "mismatched output weight dim must be rejected at load time"
1302 );
1303 }
1304
1305 #[test]
1308 fn local_provider_chat_returns_response() {
1309 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1310 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("hello")])
1311 .with_max_tokens(4);
1312 let resp = p.chat(&req).unwrap();
1313 assert_eq!(resp.model, "local/candle");
1314 assert_eq!(resp.completion_tokens, 4);
1315 assert!(!resp.content.is_empty());
1316 }
1317
1318 #[test]
1319 fn local_provider_stream_ends_with_final_token() {
1320 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1321 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("hi")])
1322 .with_max_tokens(3);
1323 let mut tokens: Vec<el_core::ChatToken> = Vec::new();
1324 p.chat_stream(&req, &mut |t| tokens.push(t)).unwrap();
1325 assert!(tokens.last().unwrap().is_final);
1326 assert!(tokens.len() > 1);
1327 }
1328
1329 #[test]
1330 fn local_provider_session_resets_between_calls() {
1331 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1332 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("a")])
1333 .with_max_tokens(4);
1334 let r1 = p.chat(&req).unwrap();
1335 let r2 = p.chat(&req).unwrap();
1336 assert_eq!(r1.content, r2.content);
1337 }
1338
1339 #[test]
1340 fn local_provider_from_path_missing_file_returns_error() {
1341 let r = LocalLlmProvider::from_path(
1342 std::path::Path::new("/nonexistent/model.gguf"),
1343 0,
1344 ok_permit(),
1345 );
1346 assert!(matches!(r, Err(EdgeError::Engine(_))));
1347 }
1348
1349 #[test]
1352 fn render_chatml_wraps_each_turn_and_opens_assistant() {
1353 let msgs = vec![
1354 ChatMessage::system("be nice"),
1355 ChatMessage::user("hi"),
1356 ChatMessage::assistant("hello"),
1357 ChatMessage::user("bye"),
1358 ];
1359 let got = render_chatml(&msgs);
1360 let want = "<|im_start|>system\nbe nice<|im_end|>\n\
1361 <|im_start|>user\nhi<|im_end|>\n\
1362 <|im_start|>assistant\nhello<|im_end|>\n\
1363 <|im_start|>user\nbye<|im_end|>\n\
1364 <|im_start|>assistant\n";
1365 assert_eq!(got, want);
1366 }
1367
1368 #[test]
1369 fn local_load_permit_passes_the_provenance_gate() {
1370 let permit = local_load_permit(std::path::Path::new("models/qwen.gguf"))
1373 .expect("local permit issued");
1374 assert_eq!(permit.format, el_core::ModelFormat::Gguf);
1375 }
1376
1377 #[test]
1378 fn requested_safety_matches_the_backed_steerer_surface() {
1379 assert_eq!(
1380 requested_session_safety(SafetyMode::Off, true),
1381 SafetyMode::Off,
1382 "Off must stay off even if an expert path is configured"
1383 );
1384 assert_eq!(
1385 requested_session_safety(SafetyMode::Lightweight, true),
1386 SafetyMode::SecDecoding,
1387 "an expert promotes the concrete model-backed path"
1388 );
1389 assert_eq!(
1390 requested_session_safety(SafetyMode::SecDecoding, false),
1391 SafetyMode::Lightweight,
1392 "unbacked SecDecoding must not be reported as active"
1393 );
1394 assert_eq!(
1395 requested_session_safety(SafetyMode::Csd, false),
1396 SafetyMode::Lightweight,
1397 "unbacked Csd must not be reported as active"
1398 );
1399 }
1400
1401 #[test]
1402 fn qwen_provider_from_paths_missing_model_errors() {
1403 let r = QwenChatProvider::from_paths(
1404 std::path::Path::new("/nonexistent/model.gguf"),
1405 std::path::Path::new("/nonexistent/tokenizer.json"),
1406 );
1407 assert!(matches!(r, Err(EdgeError::Engine(_))));
1408 }
1409
1410 #[test]
1413 fn safety_off_wires_no_guard_or_steering() {
1414 let cfg = SafetyConfig {
1417 mode: SafetyMode::Off,
1418 banned: vec![1],
1419 patterns: vec![vec![2]],
1420 extra_guard_patterns: vec![],
1421 };
1422 let ports = cfg.ports();
1423 assert!(ports.guard.is_none(), "Off must not wire the chunk guard");
1424 assert!(ports.ingress.is_none(), "Off must not wire ingress triage");
1425 assert_eq!(
1426 ports.safety.mode(),
1427 SafetyMode::Off,
1428 "Off must keep the no-op steerer"
1429 );
1430 }
1431
1432 #[test]
1433 fn lightweight_wires_guard_and_hard_ban_steerer() {
1434 let cfg = SafetyConfig {
1435 mode: SafetyMode::Lightweight,
1436 banned: vec![1],
1437 patterns: vec![vec![2, 3]],
1438 extra_guard_patterns: vec![],
1439 };
1440 let ports = cfg.ports();
1441 assert!(
1442 ports.guard.is_some(),
1443 "Lightweight must wire the chunk guard"
1444 );
1445 assert!(
1446 ports.ingress.is_some(),
1447 "Lightweight must wire prompt ingress triage (ADR-013)"
1448 );
1449 assert_eq!(
1450 ports.safety.mode(),
1451 SafetyMode::Lightweight,
1452 "a non-empty ban list selects the LightweightFilter steerer"
1453 );
1454 }
1455
1456 #[test]
1457 fn lightweight_without_patterns_has_no_guard_or_ingress() {
1458 let cfg = SafetyConfig {
1462 mode: SafetyMode::Lightweight,
1463 banned: vec![7],
1464 patterns: vec![],
1465 extra_guard_patterns: vec![],
1466 };
1467 let ports = cfg.ports();
1468 assert!(ports.guard.is_none());
1469 assert!(ports.ingress.is_none());
1470 }
1471
1472 #[test]
1473 fn extra_guard_words_drive_guard_but_not_ingress() {
1474 let cfg = SafetyConfig {
1477 mode: SafetyMode::Lightweight,
1478 banned: vec![],
1479 patterns: vec![], extra_guard_patterns: vec![vec![42]], };
1482 let ports = cfg.ports();
1483 assert!(
1484 ports.guard.is_some(),
1485 "extra guard words must drive the output guard"
1486 );
1487 assert!(
1488 ports.ingress.is_none(),
1489 "extra guard words must NOT drive ingress (trajectory demo, not refusal)"
1490 );
1491 }
1492
1493 #[test]
1494 fn qwen_expert_missing_file_errors_and_is_permit_gated() {
1495 let r = QwenExpert::from_path_primed(
1498 std::path::Path::new("/nonexistent/expert.gguf"),
1499 0,
1500 &[1, 2],
1501 ok_permit(),
1502 );
1503 assert!(matches!(r, Err(EdgeError::Engine(_))));
1504 }
1505}