superfluid_daemon/
inference.rs

1//! Shared inference requests, session construction, sampling and execution.
2
3pub(crate) use crate::generation_output::*;
4use crate::wal::{channel, role, EventBody};
5use crate::ModelSamplingDefaults;
6use crate::{Daemon, DaemonError, GenParams};
7use superfluid_abi::finish;
8use serde::Deserialize;
9use std::sync::Arc;
10
11pub(crate) const DEFAULT_MAX_TOKENS: u32 = 2048;
12
13#[derive(Clone, Copy)]
14pub struct TokenLimits {
15    pub default_max_tokens: Option<u32>,
16    pub max_context: u32,
17}
18
19impl TokenLimits {
20    pub fn resolve(&self, explicit: Option<u32>, prompt_tokens: impl FnOnce() -> u32) -> u32 {
21        explicit
22            .or(self.default_max_tokens)
23            .unwrap_or_else(|| self.max_context.saturating_sub(prompt_tokens()).max(1))
24    }
25}
26pub(crate) const DEFAULT_MAX_CONTEXT: u32 = 8192;
27
28#[derive(Clone, Copy)]
29pub(crate) struct RequestPolicy {
30    pub(crate) limits: TokenLimits,
31    pub(crate) sampling: SamplingOverrides,
32    pub(crate) qos: RequestQos,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub struct RequestQos {
37    pub class: u8,
38    pub batch_invariant: bool,
39}
40
41impl RequestQos {
42    pub const AGENT: RequestQos = RequestQos {
43        class: crate::qos::FOREGROUND_AGENT,
44        batch_invariant: false,
45    };
46
47    pub(crate) fn apply(self, daemon: &Daemon, session: u64) -> Result<(), DaemonError> {
48        // A fork inherits its parent's class, so a continued session is
49        // compared with what it holds, not with the default.
50        let held = {
51            let store = daemon.store();
52            let store = store.lock().expect("store");
53            let s = store.session(session)?;
54            RequestQos { class: s.qos_class, batch_invariant: s.batch_invariant }
55        };
56        if self != held {
57            daemon.set_qos(session, self.class, self.batch_invariant)?;
58        }
59        Ok(())
60    }
61}
62
63impl Default for RequestQos {
64    fn default() -> RequestQos {
65        RequestQos::AGENT
66    }
67}
68
69#[derive(Debug, Clone, Copy, Default, PartialEq)]
70pub struct SamplingOverrides {
71    pub temperature: Option<f32>,
72    pub top_p: Option<f32>,
73    pub top_k: Option<u32>,
74    pub min_p: Option<f32>,
75    pub repeat_penalty: Option<f32>,
76    pub spec_max_temperature: Option<f32>,
77    /// Optional policy for an embedding whose omitted request parameters
78    /// predate model-published defaults. None preserves superfluid's policy.
79    pub fallback: Option<SamplingFallback>,
80}
81
82/// An embedding's defaults, below request and operator values: temperature and repetition
83/// replace the model's; truncation applies only when the model publishes none.
84#[derive(Debug, Clone, Copy, PartialEq)]
85pub struct SamplingFallback {
86    pub temperature: f32,
87    pub repeat_penalty: f32,
88    pub top_p: f32,
89    pub top_k: u32,
90    pub min_p: f32,
91    /// Interpret a lone frequency_penalty as multiplicative (1 + value).
92    /// An explicit presence_penalty or repeat_penalty retains additive
93    /// frequency semantics; presence_penalty alone disables fallback repeat.
94    pub frequency_penalty_as_repeat: bool,
95}
96
97impl SamplingOverrides {
98    pub fn for_daemon(&self, daemon: &Daemon) -> Self {
99        let mut o = *self;
100        if o.spec_max_temperature.is_some() && daemon.speculation().is_none() {
101            o.spec_max_temperature = None;
102        }
103        o
104    }
105
106    pub fn truncates(&self) -> bool {
107        self.top_p.is_some() || self.top_k.is_some() || self.min_p.is_some()
108    }
109}
110
111#[derive(Deserialize, Default)]
112pub(crate) struct ChatRequest {
113    pub(crate) model: Option<String>,
114    #[serde(default)]
115    pub(crate) messages: Vec<ChatMessage>,
116    #[serde(default)]
117    pub(crate) tools: Vec<serde_json::Value>,
118    #[serde(default)]
119    pub(crate) tool_choice: Option<serde_json::Value>,
120    #[serde(default)]
121    pub(crate) stream: bool,
122    #[serde(default)]
123    pub(crate) stream_options: Option<StreamOptions>,
124    pub(crate) max_completion_tokens: Option<u32>,
125    pub(crate) max_tokens: Option<u32>,
126    pub(crate) temperature: Option<f32>,
127    pub(crate) top_p: Option<f32>,
128    pub(crate) top_k: Option<u32>,
129    pub(crate) seed: Option<u64>,
130    pub(crate) presence_penalty: Option<f32>,
131    pub(crate) frequency_penalty: Option<f32>,
132    #[serde(alias = "repetition_penalty")]
133    pub(crate) repeat_penalty: Option<f32>,
134    pub(crate) min_p: Option<f32>,
135    pub(crate) n: Option<u32>,
136    #[serde(default)]
137    pub(crate) ignore_eos: bool,
138    #[serde(default)]
139    pub(crate) stop: StopField,
140    pub(crate) response_format: Option<serde_json::Value>,
141    pub(crate) logit_bias: Option<std::collections::HashMap<String, f32>>,
142    pub(crate) logprobs: Option<bool>,
143    pub(crate) top_logprobs: Option<u32>,
144    pub(crate) chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
145    pub(crate) enable_thinking: Option<bool>,
146    pub(crate) reasoning_effort: Option<String>,
147    /// Set by `/v1/responses` only: the request continues a stored session
148    /// instead of starting one, and `messages` are the turn appended to it.
149    #[serde(skip)]
150    pub(crate) continues: Option<Continuation>,
151    /// The route the access log names, when another dialect runs the chat
152    /// path; chat completions when unset.
153    #[serde(skip)]
154    pub(crate) route: Option<&'static str>,
155}
156
157/// A stored session forked at `at_event`, the event after its last.
158#[derive(Debug, Clone, Copy, PartialEq, Eq)]
159pub(crate) struct Continuation {
160    pub(crate) session: u64,
161    pub(crate) at_event: u64,
162}
163
164impl ChatRequest {
165    pub(crate) fn route(&self) -> &'static str {
166        self.route.unwrap_or("POST /v1/chat/completions")
167    }
168
169    pub(crate) fn thinking(&self) -> Option<bool> {
170        self.chat_template_kwargs
171            .as_ref()
172            .and_then(|m| m.get("enable_thinking"))
173            .and_then(|v| v.as_bool())
174            .or(self.enable_thinking)
175    }
176
177    pub(crate) fn template_kwargs(&self) -> serde_json::Map<String, serde_json::Value> {
178        let mut m = self.chat_template_kwargs.clone().unwrap_or_default();
179        if let Some(t) = self.enable_thinking {
180            m.entry("enable_thinking".to_string())
181                .or_insert(serde_json::Value::Bool(t));
182        }
183        if let Some(e) = &self.reasoning_effort {
184            m.entry("reasoning_effort".to_string())
185                .or_insert(serde_json::Value::String(e.clone()));
186        }
187        m
188    }
189}
190
191pub(crate) fn schema_of(rf: &Option<serde_json::Value>) -> Option<String> {
192    let rf = rf.as_ref()?;
193    match rf.get("type").and_then(|t| t.as_str()) {
194        Some("json_object") => Some(r#"{"type":"object"}"#.to_string()),
195        Some("json_schema") => rf
196            .get("json_schema")
197            .and_then(|js| js.get("schema").or(Some(js)))
198            .map(|s| s.to_string()),
199        _ => None,
200    }
201}
202
203#[derive(Debug, Clone, PartialEq, Eq, Default)]
204pub(crate) enum ToolChoice {
205    #[default]
206    Auto,
207    None,
208    Required,
209    Function(String),
210}
211
212impl ToolChoice {
213    pub(crate) fn parse(v: &Option<serde_json::Value>) -> ToolChoice {
214        let Some(v) = v else { return ToolChoice::Auto };
215        if let Some(s) = v.as_str() {
216            return match s {
217                "none" => ToolChoice::None,
218                "required" => ToolChoice::Required,
219                _ => ToolChoice::Auto,
220            };
221        }
222        let name = v
223            .get("function")
224            .and_then(|f| f.get("name"))
225            .or_else(|| v.get("name"))
226            .and_then(|n| n.as_str());
227        match name {
228            Some(n) if !n.is_empty() => ToolChoice::Function(n.to_string()),
229            _ => ToolChoice::Auto,
230        }
231    }
232
233    pub(crate) fn forces_call(&self) -> bool {
234        matches!(self, ToolChoice::Required | ToolChoice::Function(_))
235    }
236
237    fn declares_tools(&self) -> bool {
238        !matches!(self, ToolChoice::None)
239    }
240
241    fn constrain_to<'a>(
242        &self,
243        tools: &'a [serde_json::Value],
244    ) -> Option<Vec<&'a serde_json::Value>> {
245        match self {
246            ToolChoice::Auto | ToolChoice::None => None,
247            ToolChoice::Required => Some(tools.iter().collect()),
248            ToolChoice::Function(name) => {
249                let picked: Vec<&serde_json::Value> = tools
250                    .iter()
251                    .filter(|t| {
252                        let f = t.get("function").unwrap_or(*t);
253                        f.get("name").and_then(|n| n.as_str()) == Some(name.as_str())
254                    })
255                    .collect();
256                (!picked.is_empty()).then_some(picked)
257            }
258        }
259    }
260}
261
262pub(crate) struct RequestHandles<'a> {
263    daemon: &'a Daemon,
264    grammar: Option<u32>,
265    bias: Option<u32>,
266}
267
268impl<'a> RequestHandles<'a> {
269    pub(crate) fn new(daemon: &'a Daemon, grammar: Option<u32>, bias: Option<u32>) -> Self {
270        RequestHandles {
271            daemon,
272            grammar,
273            bias,
274        }
275    }
276
277    pub(crate) fn apply(&self, extras: &mut crate::scheduler::GenExtras) {
278        if let Some(h) = self.grammar {
279            extras.grammar_handle = h;
280        }
281        if let Some(h) = self.bias {
282            extras.logit_bias_handle = h;
283        }
284    }
285}
286
287impl Drop for RequestHandles<'_> {
288    fn drop(&mut self) {
289        if let Some(h) = self.grammar.take() {
290            self.daemon.grammar_free(h);
291        }
292        if let Some(h) = self.bias.take() {
293            self.daemon.logit_bias_free(h);
294        }
295    }
296}
297
298pub(crate) fn grammar_handle_for(
299    daemon: &Daemon,
300    req: &ChatRequest,
301    choice: &ToolChoice,
302) -> Result<Option<u32>, DaemonError> {
303    if let ToolChoice::Function(name) = choice {
304        let declared = req.tools.iter().any(|t| {
305            let f = t.get("function").unwrap_or(t);
306            f.get("name").and_then(|n| n.as_str()) == Some(name.as_str())
307        });
308        if !declared {
309            return Err(DaemonError::Constraint(format!(
310                "tool_choice names function {name:?}, which is not among the declared tools"
311            )));
312        }
313    }
314    if !req.tools.is_empty() && *choice != ToolChoice::None {
315        let jsons: Vec<String> = match choice.constrain_to(&req.tools) {
316            Some(tools) => tools.iter().map(|t| t.to_string()).collect(),
317            None => req.tools.iter().map(|t| t.to_string()).collect(),
318        };
319        if let Some(tag) = daemon.tool_structural_tag(&jsons, choice.forces_call()) {
320            if let Ok(h) = daemon.grammar_create_structural(&tag) {
321                if h != 0 {
322                    return Ok(Some(h));
323                }
324            }
325        }
326    }
327    if let Some(tools) = choice.constrain_to(&req.tools) {
328        if tools.is_empty() {
329            return Err(DaemonError::Constraint(
330                "tool_choice demands a call but the request declares no tools".to_string(),
331            ));
332        }
333        let jsons: Vec<String> = tools.iter().map(|t| t.to_string()).collect();
334        if let Some(schema) = daemon.tool_call_grammar(&jsons) {
335            if let Ok(h) = daemon.grammar_create(&schema) {
336                if h != 0 {
337                    return Ok(Some(h));
338                }
339            }
340        }
341        return Err(DaemonError::Constraint(
342            "tool_choice demands a call, but no tool-call grammar could be compiled for this \
343             model's dialect on this engine"
344                .to_string(),
345        ));
346    }
347    match schema_of(&req.response_format) {
348        None if req
349            .response_format
350            .as_ref()
351            .and_then(|rf| rf.get("type"))
352            .and_then(|t| t.as_str())
353            == Some("json_schema") =>
354        {
355            Err(DaemonError::Constraint(
356                "response_format type json_schema requires a json_schema.schema block".to_string(),
357            ))
358        }
359        None => Ok(None),
360        Some(s) => {
361            if let Some(tag) = daemon.response_format_tag(&s) {
362                return match daemon.grammar_create_structural(&tag) {
363                    Ok(h) if h != 0 => Ok(Some(h)),
364                    _ => Err(DaemonError::Constraint(
365                        "response_format grammar did not compile in this model's answer frame on this engine"
366                            .to_string(),
367                    )),
368                };
369            }
370            match daemon.grammar_create(&s) {
371                Ok(h) if h != 0 => Ok(Some(h)),
372                Ok(_) => Err(DaemonError::Constraint(
373                    "response_format grammar did not compile on this engine (schema rejected or \
374                     constrained decoding unavailable in this build)"
375                        .to_string(),
376                )),
377                Err(e) => Err(DaemonError::Constraint(format!(
378                    "response_format grammar could not be created: {e}"
379                ))),
380            }
381        }
382    }
383}
384
385#[derive(Deserialize, Default, Clone)]
386#[serde(untagged)]
387pub(crate) enum StopField {
388    #[default]
389    None,
390    One(String),
391    Many(Vec<String>),
392}
393
394impl StopField {
395    pub(crate) fn list(&self) -> Vec<String> {
396        match self {
397            StopField::None => Vec::new(),
398            StopField::One(s) => vec![s.clone()],
399            StopField::Many(v) => v.iter().filter(|s| !s.is_empty()).cloned().collect(),
400        }
401    }
402}
403
404fn stop_prefix_holdback(text: &str, stops: &[String]) -> usize {
405    let mut hold = 0usize;
406    for s in stops {
407        if s.is_empty() {
408            continue;
409        }
410        let max = s.len().saturating_sub(1).min(text.len());
411        for k in (1..=max).rev() {
412            if !text.is_char_boundary(text.len() - k) || !s.is_char_boundary(k) {
413                continue;
414            }
415            if text.ends_with(&s[..k]) {
416                hold = hold.max(k);
417                break;
418            }
419        }
420    }
421    hold
422}
423
424pub(crate) fn truncate_at_stop(text: &str, stops: &[String]) -> (String, bool) {
425    let mut cut = None;
426    for s in stops {
427        if let Some(i) = text.find(s.as_str()) {
428            cut = Some(cut.map_or(i, |c: usize| c.min(i)));
429        }
430    }
431    match cut {
432        Some(i) => (text[..i].to_string(), true),
433        None => (text.to_string(), false),
434    }
435}
436
437#[cfg(test)]
438mod penalty_alias_tests {
439    use super::{
440        advertised_generation_defaults, draw_seed_if_unseeded, params_of_for, repeat_penalty_for,
441        ChatRequest, CompletionRequest, SamplingOverrides, DEFAULT_REPEAT_PENALTY, DEFAULT_TOP_K,
442    };
443    use crate::GenParams;
444    use crate::ModelSamplingDefaults;
445
446    #[test]
447    fn unseeded_sampled_requests_draw_distinct_seeds() {
448        let m = ModelSamplingDefaults::default();
449        let op = SamplingOverrides::default();
450        let mut a = params_of_for(Some(0.7), None, None, None, None, &m, &op);
451        let mut b = params_of_for(Some(0.7), None, None, None, None, &m, &op);
452        assert_eq!((a.seed, b.seed), (0, 0), "params_of_for itself never draws");
453        assert!(draw_seed_if_unseeded(&mut a, None));
454        assert!(draw_seed_if_unseeded(&mut b, None));
455        assert_ne!(a.seed, 0);
456        assert_ne!(a.seed, b.seed, "two requests, two streams");
457        let mut greedy = params_of_for(Some(0.0), None, None, None, None, &m, &op);
458        assert!(!draw_seed_if_unseeded(&mut greedy, None));
459        assert_eq!(greedy.seed, 0, "greedy has no use for a seed");
460        let mut explicit = params_of_for(Some(0.7), None, None, None, Some(42), &m, &op);
461        assert!(!draw_seed_if_unseeded(&mut explicit, Some(42)));
462        assert_eq!(explicit.seed, 42, "the client's seed is honoured as is");
463    }
464
465    #[test]
466    fn repetition_penalty_is_an_alias_for_repeat_penalty() {
467        let canonical: ChatRequest =
468            serde_json::from_str(r#"{"messages":[],"repeat_penalty":1.3}"#)
469                .expect("canonical spelling");
470        let alias: ChatRequest =
471            serde_json::from_str(r#"{"messages":[],"repetition_penalty":1.3}"#)
472                .expect("alias spelling");
473        assert_eq!(canonical.repeat_penalty, Some(1.3));
474        assert_eq!(
475            alias.repeat_penalty,
476            Some(1.3),
477            "repetition_penalty was dropped"
478        );
479
480        let c_canon: CompletionRequest =
481            serde_json::from_str(r#"{"prompt":"hi","repeat_penalty":1.3}"#).expect("canonical");
482        let c_alias: CompletionRequest =
483            serde_json::from_str(r#"{"prompt":"hi","repetition_penalty":1.3}"#).expect("alias");
484        assert_eq!(c_canon.repeat_penalty, Some(1.3));
485        assert_eq!(
486            c_alias.repeat_penalty,
487            Some(1.3),
488            "repetition_penalty was dropped on /v1/completions"
489        );
490    }
491
492    fn runtime_only(
493        temperature: Option<f32>,
494        top_p: Option<f32>,
495        top_k: Option<u32>,
496        seed: Option<u64>,
497    ) -> GenParams {
498        params_of_for(
499            temperature,
500            top_p,
501            top_k,
502            None,
503            seed,
504            &ModelSamplingDefaults::default(),
505            &SamplingOverrides::default(),
506        )
507    }
508
509    #[test]
510    fn unspecified_truncation_defaults_to_top_k() {
511        let p = runtime_only(None, None, None, None);
512        assert_eq!(
513            p.top_k, DEFAULT_TOP_K,
514            "neither knob given -> default top_k"
515        );
516        assert_eq!(p.top_p, 1.0, "top_p still reports its nominal default");
517        assert_eq!(p.temperature, 1.0);
518    }
519
520    #[test]
521    fn explicit_sampling_knobs_are_never_overridden() {
522        let explicit_top_p = runtime_only(None, Some(1.0), None, None);
523        assert_eq!(
524            explicit_top_p.top_k, 0,
525            "explicit top_p disables the default"
526        );
527        assert_eq!(runtime_only(None, Some(0.9), None, None).top_k, 0);
528        assert_eq!(
529            runtime_only(None, None, Some(7), None).top_k,
530            7,
531            "explicit top_k wins"
532        );
533        assert_eq!(
534            runtime_only(None, None, Some(0), None).top_k,
535            0,
536            "explicit 0 means disabled"
537        );
538    }
539
540    #[test]
541    fn greedy_gets_no_truncation_default() {
542        assert_eq!(runtime_only(Some(0.0), None, None, None).top_k, 0);
543    }
544
545    #[test]
546    fn advertised_defaults_match_the_effective_ones() {
547        let adv = advertised_generation_defaults(
548            &ModelSamplingDefaults::default(),
549            &SamplingOverrides::default(),
550        );
551        let effective = runtime_only(None, None, None, None);
552        assert_eq!(
553            adv.top_k, effective.top_k,
554            "advertised top_k must be the effective one"
555        );
556        assert_eq!(adv.top_p, effective.top_p);
557        assert_eq!(adv.temperature, effective.temperature);
558        assert_eq!(adv.top_k, DEFAULT_TOP_K);
559    }
560
561    #[test]
562    fn round_tripping_advertised_defaults_is_stable() {
563        let adv = advertised_generation_defaults(
564            &ModelSamplingDefaults::default(),
565            &SamplingOverrides::default(),
566        );
567        let echoed = runtime_only(
568            Some(adv.temperature),
569            Some(adv.top_p),
570            Some(adv.top_k),
571            None,
572        );
573        assert_eq!(echoed.top_k, adv.top_k, "echoed top_k survives");
574        assert_eq!(echoed.top_p, adv.top_p);
575        assert_eq!(echoed.temperature, adv.temperature);
576    }
577
578    #[test]
579    fn model_defaults_beat_the_runtime_fallback() {
580        let qwen38 = ModelSamplingDefaults {
581            temperature: Some(1.0),
582            top_p: Some(0.95),
583            top_k: Some(20),
584            min_p: Some(0.0),
585            do_sample: Some(true),
586            repetition_penalty: None,
587        };
588        let p = params_of_for(
589            None,
590            None,
591            None,
592            None,
593            None,
594            &qwen38,
595            &SamplingOverrides::default(),
596        );
597        assert_eq!(p.top_k, 20, "the model asked for 20, not the fallback 40");
598        assert_eq!(p.top_p, 0.95);
599        assert_eq!(p.temperature, 1.0);
600    }
601
602    #[test]
603    fn caller_beats_the_model_defaults() {
604        let m = ModelSamplingDefaults {
605            top_k: Some(20),
606            top_p: Some(0.95),
607            ..Default::default()
608        };
609        assert_eq!(
610            params_of_for(
611                None,
612                None,
613                Some(5),
614                None,
615                None,
616                &m,
617                &SamplingOverrides::default()
618            )
619            .top_k,
620            5
621        );
622        assert_eq!(
623            params_of_for(
624                None,
625                Some(0.5),
626                None,
627                None,
628                None,
629                &m,
630                &SamplingOverrides::default()
631            )
632            .top_p,
633            0.5
634        );
635        assert_eq!(
636            params_of_for(
637                Some(0.3),
638                None,
639                None,
640                None,
641                None,
642                &m,
643                &SamplingOverrides::default()
644            )
645            .temperature,
646            0.3
647        );
648    }
649
650    #[test]
651    fn empty_model_defaults_keep_the_fallback() {
652        let p = runtime_only(None, None, None, None);
653        assert_eq!(p.top_k, DEFAULT_TOP_K);
654    }
655
656    #[test]
657    fn spec_max_temperature_caps_only_the_bundle_default() {
658        let m = ModelSamplingDefaults {
659            temperature: Some(1.0),
660            ..Default::default()
661        };
662        let op = SamplingOverrides {
663            spec_max_temperature: Some(0.6),
664            ..Default::default()
665        };
666        assert_eq!(
667            params_of_for(None, None, None, None, None, &m, &op).temperature,
668            0.6
669        );
670        assert_eq!(
671            params_of_for(Some(0.9), None, None, None, None, &m, &op).temperature,
672            0.9,
673            "caller wins"
674        );
675        let op_t = SamplingOverrides {
676            temperature: Some(0.8),
677            ..op
678        };
679        assert_eq!(
680            params_of_for(None, None, None, None, None, &m, &op_t).temperature,
681            0.8,
682            "--temperature wins"
683        );
684        let cool = ModelSamplingDefaults {
685            temperature: Some(0.3),
686            ..Default::default()
687        };
688        assert_eq!(
689            params_of_for(None, None, None, None, None, &cool, &op).temperature,
690            0.3
691        );
692        let greedy = ModelSamplingDefaults {
693            do_sample: Some(false),
694            ..Default::default()
695        };
696        assert_eq!(
697            params_of_for(None, None, None, None, None, &greedy, &op).temperature,
698            0.0
699        );
700        assert_eq!(
701            params_of_for(
702                None,
703                None,
704                None,
705                None,
706                None,
707                &ModelSamplingDefaults::default(),
708                &op
709            )
710            .temperature,
711            0.6
712        );
713        assert_eq!(advertised_generation_defaults(&m, &op).temperature, 0.6);
714    }
715
716    #[test]
717    fn do_sample_false_means_greedy() {
718        let m = ModelSamplingDefaults {
719            do_sample: Some(false),
720            ..Default::default()
721        };
722        assert_eq!(
723            params_of_for(
724                None,
725                None,
726                None,
727                None,
728                None,
729                &m,
730                &SamplingOverrides::default()
731            )
732            .temperature,
733            0.0
734        );
735        assert_eq!(
736            params_of_for(
737                Some(0.8),
738                None,
739                None,
740                None,
741                None,
742                &m,
743                &SamplingOverrides::default()
744            )
745            .temperature,
746            0.8
747        );
748    }
749
750    #[test]
751    fn model_top_p_alone_suppresses_the_fallback_top_k() {
752        let m = ModelSamplingDefaults {
753            top_p: Some(0.9),
754            ..Default::default()
755        };
756        let p = params_of_for(
757            None,
758            None,
759            None,
760            None,
761            None,
762            &m,
763            &SamplingOverrides::default(),
764        );
765        assert_eq!(p.top_p, 0.9);
766        assert_eq!(
767            p.top_k, 0,
768            "no fallback top_k when the model already truncates"
769        );
770    }
771
772    #[test]
773    fn explicit_top_p_suppresses_the_model_top_k() {
774        let qwen38 = ModelSamplingDefaults {
775            top_p: Some(0.95),
776            top_k: Some(20),
777            ..Default::default()
778        };
779        let p = params_of_for(
780            None,
781            Some(1.0),
782            None,
783            None,
784            None,
785            &qwen38,
786            &SamplingOverrides::default(),
787        );
788        assert_eq!(p.top_p, 1.0, "caller asked for the full distribution");
789        assert_eq!(p.top_k, 0, "and must not inherit the model's top_k");
790    }
791
792    #[test]
793    fn explicit_top_k_suppresses_the_model_top_p() {
794        let qwen38 = ModelSamplingDefaults {
795            top_p: Some(0.95),
796            top_k: Some(20),
797            ..Default::default()
798        };
799        let p = params_of_for(
800            None,
801            None,
802            Some(5),
803            None,
804            None,
805            &qwen38,
806            &SamplingOverrides::default(),
807        );
808        assert_eq!(p.top_k, 5);
809        assert_eq!(
810            p.top_p, 1.0,
811            "model top_p is set aside once the caller truncates"
812        );
813    }
814
815    #[test]
816    fn model_min_p_alone_is_a_truncation_claim() {
817        let m = ModelSamplingDefaults {
818            min_p: Some(0.05),
819            ..Default::default()
820        };
821        let p = params_of_for(
822            None,
823            None,
824            None,
825            None,
826            None,
827            &m,
828            &SamplingOverrides::default(),
829        );
830        assert_eq!(p.min_p, 0.05);
831        assert_eq!(
832            p.top_k, 0,
833            "no fallback top_k stacked on the model's policy"
834        );
835    }
836
837    #[test]
838    fn advertised_defaults_follow_the_model() {
839        let qwen38 = ModelSamplingDefaults {
840            top_p: Some(0.95),
841            top_k: Some(20),
842            ..Default::default()
843        };
844        let adv = advertised_generation_defaults(&qwen38, &SamplingOverrides::default());
845        assert_eq!(adv.top_k, 20);
846        assert_eq!(adv.top_p, 0.95);
847        let echoed = params_of_for(
848            Some(adv.temperature),
849            Some(adv.top_p),
850            Some(adv.top_k),
851            None,
852            None,
853            &qwen38,
854            &SamplingOverrides::default(),
855        );
856        assert_eq!(
857            echoed.top_k, 20,
858            "echoing the advertised values changes nothing"
859        );
860        assert_eq!(echoed.top_p, 0.95);
861    }
862
863    #[test]
864    fn operator_flags_sit_between_caller_and_bundle() {
865        let qwen38 = ModelSamplingDefaults {
866            temperature: Some(0.6),
867            top_p: Some(0.95),
868            top_k: Some(20),
869            min_p: Some(0.0),
870            do_sample: Some(true),
871            repetition_penalty: None,
872        };
873        let op = SamplingOverrides {
874            top_k: Some(64),
875            ..Default::default()
876        };
877
878        let p = params_of_for(None, None, None, None, None, &qwen38, &op);
879        assert_eq!(p.top_k, 64, "operator top_k overrides the bundle's");
880        assert_eq!(
881            p.top_p, 1.0,
882            "and takes the whole level: the bundle's top_p is set aside"
883        );
884        assert_eq!(
885            p.temperature, 0.6,
886            "temperature is a separate rung and still comes from the bundle"
887        );
888
889        let p = params_of_for(None, None, Some(7), None, None, &qwen38, &op);
890        assert_eq!(p.top_k, 7, "an explicit request value wins over the flag");
891
892        let p = params_of_for(
893            None,
894            None,
895            None,
896            None,
897            None,
898            &ModelSamplingDefaults::default(),
899            &op,
900        );
901        assert_eq!(p.top_k, 64);
902    }
903
904    #[test]
905    fn operator_temperature_precedence() {
906        let greedy_model = ModelSamplingDefaults {
907            do_sample: Some(false),
908            ..Default::default()
909        };
910        let op = SamplingOverrides {
911            temperature: Some(0.9),
912            ..Default::default()
913        };
914        assert_eq!(
915            params_of_for(None, None, None, None, None, &greedy_model, &op).temperature,
916            0.9,
917            "the operator outranks the checkpoint"
918        );
919        assert_eq!(
920            params_of_for(Some(0.0), None, None, None, None, &greedy_model, &op).temperature,
921            0.0,
922            "an explicit request temperature still wins"
923        );
924    }
925
926    #[test]
927    fn an_unknown_penalty_is_omitted_but_a_pinned_one_is_not() {
928        let advertise = |resolved: Option<ModelSamplingDefaults>, op: SamplingOverrides| {
929            (resolved.is_some() || op.repeat_penalty.is_some())
930                .then(|| repeat_penalty_for(None, &op, &resolved.clone().unwrap_or_default()))
931        };
932        let none = SamplingOverrides::default();
933
934        assert_eq!(advertise(None, none), None);
935
936        let op = SamplingOverrides {
937            repeat_penalty: Some(1.2),
938            ..Default::default()
939        };
940        assert_eq!(advertise(None, op), Some(1.2));
941
942        let quiet = ModelSamplingDefaults::default();
943        assert_eq!(advertise(Some(quiet), none), Some(DEFAULT_REPEAT_PENALTY));
944        let opted_out = ModelSamplingDefaults {
945            repetition_penalty: Some(1.0),
946            ..Default::default()
947        };
948        assert_eq!(advertise(Some(opted_out), none), Some(1.0));
949    }
950
951    #[test]
952    fn operator_owned_settings_survive_a_failed_bundle_lookup() {
953        let advertise = |resolved: Option<ModelSamplingDefaults>, op: SamplingOverrides| {
954            let eff = advertised_generation_defaults(&resolved.clone().unwrap_or_default(), &op);
955            let temp = (resolved.is_some() || op.temperature.is_some()).then_some(eff.temperature);
956            let trunc = (resolved.is_some() || op.truncates()).then_some((eff.top_p, eff.top_k));
957            (temp, trunc)
958        };
959
960        assert_eq!(advertise(None, SamplingOverrides::default()), (None, None));
961
962        let op = SamplingOverrides {
963            top_k: Some(64),
964            ..Default::default()
965        };
966        assert_eq!(advertise(None, op), (None, Some((1.0, 64))));
967
968        let op = SamplingOverrides {
969            temperature: Some(0.4),
970            ..Default::default()
971        };
972        assert_eq!(advertise(None, op), (Some(0.4), None));
973
974        let qwen38 = ModelSamplingDefaults {
975            temperature: Some(0.6),
976            top_p: Some(0.95),
977            top_k: Some(20),
978            min_p: Some(0.0),
979            do_sample: Some(true),
980            repetition_penalty: None,
981        };
982        let op = SamplingOverrides {
983            top_k: Some(64),
984            ..Default::default()
985        };
986        assert_eq!(advertise(Some(qwen38), op), (Some(0.6), Some((1.0, 64))));
987    }
988
989    #[test]
990    fn repeat_penalty_resolves_caller_then_operator_then_bundle_then_default() {
991        let none = SamplingOverrides::default();
992        let no_model = ModelSamplingDefaults::default();
993
994        assert_eq!(
995            repeat_penalty_for(None, &none, &no_model),
996            DEFAULT_REPEAT_PENALTY
997        );
998        assert_eq!(repeat_penalty_for(Some(1.2), &none, &no_model), 1.2);
999
1000        let model = ModelSamplingDefaults {
1001            repetition_penalty: Some(1.1),
1002            ..Default::default()
1003        };
1004        assert_eq!(
1005            repeat_penalty_for(None, &none, &model),
1006            1.1,
1007            "bundle beats the runtime default"
1008        );
1009
1010        let opted_out = ModelSamplingDefaults {
1011            repetition_penalty: Some(1.0),
1012            ..Default::default()
1013        };
1014        assert_eq!(repeat_penalty_for(None, &none, &opted_out), 1.0);
1015
1016        let op = SamplingOverrides {
1017            repeat_penalty: Some(1.2),
1018            ..Default::default()
1019        };
1020        assert_eq!(
1021            repeat_penalty_for(None, &op, &model),
1022            1.2,
1023            "flag beats the bundle"
1024        );
1025        assert_eq!(
1026            repeat_penalty_for(Some(1.0), &op, &model),
1027            1.0,
1028            "caller's 1.0 opts out"
1029        );
1030    }
1031
1032    #[test]
1033    fn echoing_advertised_min_p_is_stable() {
1034        let model = ModelSamplingDefaults {
1035            top_p: Some(0.95),
1036            top_k: Some(20),
1037            min_p: Some(0.05),
1038            ..Default::default()
1039        };
1040        let adv = advertised_generation_defaults(&model, &SamplingOverrides::default());
1041        assert_eq!(adv.min_p, 0.05, "advertised min_p comes from the model");
1042        let echoed = params_of_for(
1043            Some(adv.temperature),
1044            Some(adv.top_p),
1045            Some(adv.top_k),
1046            Some(adv.min_p),
1047            None,
1048            &model,
1049            &SamplingOverrides::default(),
1050        );
1051        assert_eq!(echoed.min_p, 0.05, "echoed min_p survives the round trip");
1052        assert_eq!(echoed.top_p, 0.95);
1053        assert_eq!(echoed.top_k, 20);
1054    }
1055}
1056
1057#[cfg(test)]
1058mod stop_tests {
1059    use super::{stop_prefix_holdback, truncate_at_stop};
1060
1061    fn stops(s: &[&str]) -> Vec<String> {
1062        s.iter().map(|x| x.to_string()).collect()
1063    }
1064
1065    #[test]
1066    fn holdback_is_the_longest_stop_prefix_at_the_tail() {
1067        let st = stops(&["END", "<|eot|>"]);
1068        assert_eq!(stop_prefix_holdback("hello", &st), 0);
1069        assert_eq!(stop_prefix_holdback("hello E", &st), 1);
1070        assert_eq!(stop_prefix_holdback("hello EN", &st), 2);
1071        assert_eq!(stop_prefix_holdback("hello END", &st), 0);
1072        assert_eq!(stop_prefix_holdback("x <|eot", &st), 5);
1073        assert_eq!(stop_prefix_holdback("a <|e", &stops(&["<|eot|>", "e"])), 3);
1074        assert_eq!(
1075            stop_prefix_holdback("héllo —", &stops(&["— fin"])),
1076            "—".len()
1077        );
1078        assert_eq!(stop_prefix_holdback("abc", &stops(&[""])), 0);
1079    }
1080
1081    #[test]
1082    fn delta_stream_never_emits_a_stop_prefix_that_the_next_delta_completes() {
1083        let st = stops(&["END"]);
1084        let mut acc = String::new();
1085        let mut sent = 0usize;
1086        let mut emitted = String::new();
1087        for delta in ["hello E", "ND world"] {
1088            acc.push_str(delta);
1089            let (mut t, hit) = truncate_at_stop(&acc, &st);
1090            if !hit {
1091                let keep = t.len() - stop_prefix_holdback(&t, &st);
1092                t.truncate(keep);
1093            }
1094            emitted.push_str(t.get(sent..).unwrap_or(""));
1095            sent = sent.max(t.len());
1096            if hit {
1097                break;
1098            }
1099        }
1100        assert_eq!(emitted, "hello ");
1101    }
1102}
1103
1104#[derive(Clone, Copy, Default)]
1105pub(crate) struct DecodeSpan {
1106    first: Option<std::time::Instant>,
1107    last: Option<std::time::Instant>,
1108    first_count: u64,
1109    last_count: u64,
1110    tokens: u64,
1111    elapsed: std::time::Duration,
1112    measured_tokens: u64,
1113    measured_elapsed: std::time::Duration,
1114    /// Whether the scheduler timed the generation: its figure then stands,
1115    /// none included, and the events are not consulted.
1116    scheduler_timed: bool,
1117}
1118
1119impl DecodeSpan {
1120    pub(crate) fn stamp(&mut self, produced: u64) {
1121        self.stamp_at(std::time::Instant::now(), produced);
1122    }
1123
1124    fn stamp_at(&mut self, now: std::time::Instant, produced: u64) {
1125        if self.first.is_none() {
1126            self.first = Some(now);
1127            self.first_count = produced;
1128        }
1129        self.last = Some(now);
1130        self.last_count = produced.max(self.last_count);
1131    }
1132
1133    pub(crate) fn first_token(&self) -> Option<std::time::Instant> {
1134        self.first
1135    }
1136
1137    pub(crate) fn finish(&mut self) {
1138        if let (Some(f), Some(l)) = (self.first, self.last) {
1139            self.tokens += self.last_count.saturating_sub(self.first_count);
1140            self.elapsed += l.saturating_duration_since(f);
1141        }
1142        self.first = None;
1143        self.last = None;
1144        self.first_count = 0;
1145        self.last_count = 0;
1146    }
1147
1148    /// Adds a generation's decode as the scheduler timed it, tick end to tick
1149    /// end. Events alone cannot: a tick hands over all of its tokens at once,
1150    /// and a channel change splits one tick into several events, so the rest
1151    /// of the first tick counted against the next one's time. A generation
1152    /// that ended in its first tick adds `(0, 0.0)`: it has no decode to time,
1153    /// whatever the events of that tick say.
1154    pub(crate) fn add_measured(&mut self, (tokens, secs): (u64, f64)) {
1155        self.scheduler_timed = true;
1156        if tokens > 0 && secs > 0.0 {
1157            self.measured_tokens += tokens;
1158            self.measured_elapsed += std::time::Duration::from_secs_f64(secs);
1159        }
1160    }
1161
1162    fn measured(&self) -> (u64, Option<std::time::Duration>) {
1163        if self.scheduler_timed {
1164            match self.measured_tokens {
1165                0 => (0, None),
1166                n => (n, Some(self.measured_elapsed)),
1167            }
1168        } else if self.elapsed == std::time::Duration::ZERO {
1169            (0, None)
1170        } else {
1171            (self.tokens, Some(self.elapsed))
1172        }
1173    }
1174}
1175
1176#[derive(Clone, Copy)]
1177pub(crate) struct GenTiming {
1178    pub t0: std::time::Instant,
1179    pub first: Option<std::time::Instant>,
1180    pub decode: DecodeSpan,
1181}
1182
1183fn fim_log_cached(c: &crate::Completion) -> u64 {
1184    if c.cached {
1185        c.prompt_tokens as u64
1186    } else {
1187        0
1188    }
1189}
1190
1191fn prefill_rate(prompt: u64, cached: u64, ttft: Option<std::time::Duration>) -> String {
1192    let prefilled = prompt - cached.min(prompt);
1193    ttft.filter(|d| d.as_secs_f64() > 0.0 && prefilled > 0)
1194        .map(|d| format!("{:.0}", prefilled as f64 / d.as_secs_f64()))
1195        .unwrap_or_else(|| "-".into())
1196}
1197
1198fn decode_rate(tokens: u64, span: Option<std::time::Duration>) -> String {
1199    span.filter(|d| tokens > 0 && d.as_millis() as u64 >= tokens)
1200        .map(|d| format!("{:.1}", tokens as f64 / d.as_secs_f64()))
1201        .unwrap_or_else(|| "-".into())
1202}
1203
1204pub(crate) fn access_log(
1205    route: &str,
1206    model: &str,
1207    prompt: u64,
1208    cached: u64,
1209    completion: u64,
1210    timing: GenTiming,
1211    finish: &str,
1212) {
1213    let GenTiming { t0, first, decode } = timing;
1214    let total_ms = t0.elapsed().as_millis();
1215    let ttft = first.map(|f| f.saturating_duration_since(t0));
1216    let (decode_tokens, decode_span) = decode.measured();
1217    let prefill = prefill_rate(prompt, cached, ttft);
1218    let decode = decode_rate(decode_tokens, decode_span);
1219    tracing::info!(
1220        "<-- 200 {route} model={model} prompt={prompt} cached={cached} completion={completion} ttft={}ms prefill={prefill} tok/s decode={decode} tok/s total={total_ms}ms finish={finish}",
1221        ttft.map(|d| d.as_millis()).unwrap_or(0),
1222    );
1223}
1224
1225pub(crate) fn extras_of(
1226    presence: Option<f32>,
1227    frequency: Option<f32>,
1228    repeat: Option<f32>,
1229    op: &SamplingOverrides,
1230    model: &ModelSamplingDefaults,
1231) -> crate::scheduler::GenExtras {
1232    let mut frequency_penalty = frequency.unwrap_or(0.0);
1233    let mut repeat_penalty = repeat_penalty_for(repeat, op, model);
1234    if op.fallback.is_some_and(|f| f.frequency_penalty_as_repeat)
1235        && repeat.is_none()
1236        && op.repeat_penalty.is_none()
1237    {
1238        if presence.is_some() {
1239            repeat_penalty = 1.0;
1240        } else if let Some(f) = frequency {
1241            repeat_penalty = 1.0 + f;
1242            frequency_penalty = 0.0;
1243        }
1244    }
1245    crate::scheduler::GenExtras {
1246        ephemeral: true,
1247        presence_penalty: presence.unwrap_or(0.0),
1248        frequency_penalty,
1249        repeat_penalty,
1250        grammar_handle: 0,
1251        logit_bias_handle: 0,
1252        want_logprobs: false,
1253        top_logprobs: 0,
1254        wants_deltas: false,
1255        open_channel: 0,
1256        thinking: None,
1257        ignore_eos: false,
1258        seed_drawn: false,
1259        tool_schemas: None,
1260    }
1261}
1262
1263pub(crate) fn logit_bias_handle_of(
1264    daemon: &Daemon,
1265    map: &Option<std::collections::HashMap<String, f32>>,
1266) -> Option<u32> {
1267    let map = map.as_ref()?;
1268    if map.is_empty() {
1269        return None;
1270    }
1271    let mut tokens = Vec::with_capacity(map.len());
1272    let mut values = Vec::with_capacity(map.len());
1273    for (k, v) in map {
1274        if let Ok(t) = k.parse::<i32>() {
1275            tokens.push(t);
1276            values.push(*v);
1277        }
1278    }
1279    if tokens.is_empty() {
1280        return None;
1281    }
1282    daemon.logit_bias_create(&tokens, &values).ok()
1283}
1284
1285#[derive(Debug, Default, Deserialize)]
1286pub(crate) struct StreamOptions {
1287    #[serde(default)]
1288    #[allow(dead_code)]
1289    pub(crate) include_usage: bool,
1290    #[serde(default)]
1291    pub(crate) continuous_usage_stats: bool,
1292}
1293
1294#[derive(Deserialize)]
1295pub(crate) struct CompletionRequest {
1296    pub(crate) model: Option<String>,
1297    #[serde(default)]
1298    pub(crate) ignore_eos: bool,
1299    #[serde(default)]
1300    pub(crate) prompt: String,
1301    #[serde(default)]
1302    pub(crate) stream: bool,
1303    #[serde(default)]
1304    pub(crate) echo: bool,
1305    pub(crate) max_tokens: Option<u32>,
1306    pub(crate) temperature: Option<f32>,
1307    pub(crate) top_p: Option<f32>,
1308    pub(crate) top_k: Option<u32>,
1309    pub(crate) seed: Option<u64>,
1310    pub(crate) presence_penalty: Option<f32>,
1311    pub(crate) frequency_penalty: Option<f32>,
1312    #[serde(alias = "repetition_penalty")]
1313    pub(crate) repeat_penalty: Option<f32>,
1314    pub(crate) min_p: Option<f32>,
1315    pub(crate) logprobs: Option<u32>,
1316    pub(crate) suffix: Option<String>,
1317    pub(crate) fim_mode: Option<String>,
1318}
1319
1320pub(crate) const DEFAULT_TOP_K: u32 = 40;
1321
1322pub(crate) const DEFAULT_REPEAT_PENALTY: f32 = 1.0;
1323
1324#[cfg(test)]
1325mod embedding_sampling_tests {
1326    use super::*;
1327
1328    fn policy() -> SamplingOverrides {
1329        SamplingOverrides {
1330            fallback: Some(SamplingFallback {
1331                temperature: 0.0, repeat_penalty: 1.07,
1332                top_p: 0.82, top_k: 17, min_p: 0.03,
1333                frequency_penalty_as_repeat: true,
1334            }),
1335            ..Default::default()
1336        }
1337    }
1338
1339    #[test]
1340    fn omitted_parameters_and_props_share_embedding_defaults() {
1341        let m = ModelSamplingDefaults::default();
1342        let op = policy();
1343        let p = advertised_generation_defaults(&m, &op);
1344        assert_eq!((p.temperature, p.top_p, p.top_k, p.min_p), (0.0, 0.82, 17, 0.03));
1345        assert_eq!(repeat_penalty_for(None, &op, &m), 1.07);
1346        let original = advertised_generation_defaults(&m, &SamplingOverrides::default());
1347        assert_eq!((original.temperature, original.top_p, original.top_k, original.min_p), (1.0, 1.0, 40, 0.0));
1348        assert_eq!(repeat_penalty_for(None, &SamplingOverrides::default(), &m), 1.0);
1349    }
1350
1351    #[test]
1352    fn model_truncation_wins_as_a_set_but_not_model_temperature_or_repeat() {
1353        for model in [
1354            ModelSamplingDefaults { top_p: Some(0.95), ..Default::default() },
1355            ModelSamplingDefaults { top_k: Some(9), ..Default::default() },
1356            ModelSamplingDefaults { min_p: Some(0.1), ..Default::default() },
1357        ] {
1358            let model = ModelSamplingDefaults { temperature: Some(0.7), repetition_penalty: Some(1.2), ..model };
1359            let p = advertised_generation_defaults(&model, &policy());
1360            assert_eq!(p.temperature, 0.0);
1361            assert_eq!((p.top_p, p.top_k, p.min_p), (model.top_p.unwrap_or(1.0), model.top_k.unwrap_or(0), model.min_p.unwrap_or(0.0)));
1362            assert_eq!(repeat_penalty_for(None, &policy(), &model), 1.07);
1363        }
1364    }
1365
1366    #[test]
1367    fn caller_and_operator_values_override_embedding_without_merging_truncation() {
1368        let m = ModelSamplingDefaults { top_k: Some(90), ..Default::default() };
1369        let op = SamplingOverrides { temperature: Some(0.6), top_k: Some(12), repeat_penalty: Some(1.3), ..policy() };
1370        let p = advertised_generation_defaults(&m, &op);
1371        assert_eq!((p.temperature, p.top_p, p.top_k, p.min_p), (0.6, 1.0, 12, 0.0));
1372        assert_eq!(repeat_penalty_for(None, &op, &m), 1.3);
1373        let p = params_of_for(Some(0.2), Some(1.0), None, None, Some(42), &m, &op);
1374        assert_eq!((p.temperature, p.top_p, p.top_k, p.min_p, p.seed), (0.2, 1.0, 0, 0.0, 42));
1375        assert_eq!(repeat_penalty_for(Some(1.0), &op, &m), 1.0);
1376    }
1377
1378    #[test]
1379    fn penalty_mapping_is_opt_in_and_respects_explicit_parameters() {
1380        let m = ModelSamplingDefaults::default();
1381        for (presence, frequency, repeat, expected) in [
1382            (None, None, None, (0.0, 0.0, 1.07)),
1383            (None, Some(0.2), None, (0.0, 0.0, 1.2)),
1384            (None, Some(0.0), None, (0.0, 0.0, 1.0)),
1385            (Some(0.0), Some(0.2), None, (0.0, 0.2, 1.0)),
1386            (Some(0.3), None, None, (0.3, 0.0, 1.0)),
1387            (None, Some(0.2), Some(1.1), (0.0, 0.2, 1.1)),
1388        ] {
1389            let e = extras_of(presence, frequency, repeat, &policy(), &m);
1390            assert_eq!((e.presence_penalty, e.frequency_penalty, e.repeat_penalty), expected);
1391        }
1392        let e = extras_of(None, Some(0.2), None, &SamplingOverrides::default(), &m);
1393        assert_eq!((e.frequency_penalty, e.repeat_penalty), (0.2, 1.0));
1394    }
1395}
1396
1397pub(crate) fn advertised_generation_defaults(
1398    model: &ModelSamplingDefaults,
1399    op: &SamplingOverrides,
1400) -> GenParams {
1401    params_of_for(None, None, None, None, None, model, op)
1402}
1403
1404pub(crate) fn repeat_penalty_for(
1405    caller: Option<f32>,
1406    op: &SamplingOverrides,
1407    model: &ModelSamplingDefaults,
1408) -> f32 {
1409    caller
1410        .or(op.repeat_penalty)
1411        .or(op.fallback.map(|f| f.repeat_penalty))
1412        .or(model.repetition_penalty)
1413        .unwrap_or(DEFAULT_REPEAT_PENALTY)
1414}
1415
1416pub(crate) fn draw_seed_if_unseeded(params: &mut GenParams, client_seed: Option<u64>) -> bool {
1417    if client_seed.is_some() || params.temperature <= 0.0 {
1418        return false;
1419    }
1420    params.seed = fresh_seed();
1421    true
1422}
1423
1424fn fresh_seed() -> u64 {
1425    use std::hash::{BuildHasher, Hasher};
1426    let h = std::collections::hash_map::RandomState::new()
1427        .build_hasher()
1428        .finish();
1429    if h == 0 {
1430        1
1431    } else {
1432        h
1433    }
1434}
1435
1436pub(crate) fn params_of_for(
1437    temperature: Option<f32>,
1438    top_p: Option<f32>,
1439    top_k: Option<u32>,
1440    min_p: Option<f32>,
1441    seed: Option<u64>,
1442    model: &ModelSamplingDefaults,
1443    op: &SamplingOverrides,
1444) -> GenParams {
1445    let model_temperature = match model.do_sample {
1446        Some(false) => Some(0.0),
1447        _ => model.temperature,
1448    };
1449    let chosen = temperature.is_some() || op.temperature.is_some();
1450    let temperature = temperature
1451        .or(op.temperature)
1452        .or(op.fallback.map(|f| f.temperature))
1453        .or(model_temperature)
1454        .unwrap_or(1.0);
1455    let temperature = match op.spec_max_temperature {
1456        Some(cap) if !chosen => temperature.min(cap),
1457        _ => temperature,
1458    };
1459
1460    let caller_truncates = top_p.is_some() || top_k.is_some() || min_p.is_some();
1461    let model_truncates = model.top_p.is_some() || model.top_k.is_some() || model.min_p.is_some();
1462    let (eff_top_p, eff_top_k, eff_min_p) = if caller_truncates {
1463        (top_p, top_k, min_p)
1464    } else if op.truncates() {
1465        (op.top_p, op.top_k, op.min_p)
1466    } else if model_truncates {
1467        (model.top_p, model.top_k, model.min_p)
1468    } else if let Some(f) = op.fallback {
1469        (Some(f.top_p), Some(f.top_k), Some(f.min_p))
1470    } else if temperature > 0.0 {
1471        (None, Some(DEFAULT_TOP_K), None)
1472    } else {
1473        (None, None, None)
1474    };
1475    GenParams {
1476        temperature,
1477        top_p: eff_top_p.unwrap_or(1.0),
1478        min_p: eff_min_p.unwrap_or(0.0),
1479        top_k: eff_top_k.unwrap_or(0),
1480        seed: seed.unwrap_or(0),
1481    }
1482}
1483
1484#[derive(Deserialize, Default)]
1485pub(crate) struct ChatMessage {
1486    pub(crate) role: String,
1487    #[serde(default)]
1488    pub(crate) content: Option<serde_json::Value>,
1489    #[serde(default)]
1490    pub(crate) tool_calls: Vec<ReqToolCall>,
1491    #[serde(default)]
1492    pub(crate) tool_call_id: Option<String>,
1493    #[serde(default)]
1494    pub(crate) name: Option<String>,
1495    #[serde(default)]
1496    pub(crate) reasoning_content: Option<String>,
1497    #[serde(default)]
1498    pub(crate) reasoning: Option<String>,
1499}
1500
1501pub(crate) fn reasoning_of(m: &ChatMessage, role: u32) -> Option<String> {
1502    if role != role::ASSISTANT {
1503        return None;
1504    }
1505    m.reasoning_content
1506        .as_deref()
1507        .or(m.reasoning.as_deref())
1508        .map(str::trim)
1509        .filter(|s| !s.is_empty())
1510        .map(str::to_string)
1511}
1512
1513#[derive(Clone, Copy, PartialEq, Eq)]
1514pub(crate) enum MediaKind {
1515    Image,
1516    Audio,
1517}
1518
1519fn require_media(daemon: &Daemon, kind: MediaKind, param: &str) -> Result<(), DaemonError> {
1520    match kind {
1521        MediaKind::Image => daemon.require("modalities", "image_encode", "unsupported_input", Some(param), "does not accept images"),
1522        MediaKind::Audio => daemon.require("modalities", "gemma_audio_encode", "unsupported_input", Some(param), "does not accept audio"),
1523    }
1524}
1525
1526fn content_param(messages: &[ChatMessage], m: &ChatMessage) -> String {
1527    match messages.iter().position(|x| std::ptr::eq(x, m)) {
1528        Some(i) => format!("messages[{i}].content"),
1529        None => "messages".to_string(),
1530    }
1531}
1532
1533pub(crate) struct SplitContent {
1534    pub(crate) text: String,
1535    pub(crate) images: Vec<(String, Vec<u8>, MediaKind)>,
1536}
1537
1538pub(crate) fn absorbable_system_text(m: &ChatMessage) -> Result<Option<String>, DaemonError> {
1539    if m.role != "system" {
1540        return Ok(None);
1541    }
1542    let split = split_content(&m.content)?;
1543    Ok(split.images.is_empty().then_some(split.text))
1544}
1545
1546pub(crate) fn split_content(v: &Option<serde_json::Value>) -> Result<SplitContent, DaemonError> {
1547    let mut out = SplitContent {
1548        text: String::new(),
1549        images: Vec::new(),
1550    };
1551    match v {
1552        None => {}
1553        Some(serde_json::Value::String(s)) => out.text = s.clone(),
1554        Some(serde_json::Value::Array(parts)) => {
1555            for p in parts {
1556                match p.get("type").and_then(|t| t.as_str()) {
1557                    Some("text") => {
1558                        if let Some(t) = p.get("text").and_then(|t| t.as_str()) {
1559                            out.text.push_str(t);
1560                        }
1561                    }
1562                    Some("image_url") => {
1563                        let url = p
1564                            .get("image_url")
1565                            .and_then(|u| u.get("url"))
1566                            .and_then(|u| u.as_str())
1567                            .unwrap_or("");
1568                        let Some(rest) = url.strip_prefix("data:") else {
1569                            return Err(DaemonError::Protocol(
1570                                "image_url must be a data: URL (remote images are not fetched)",
1571                            ));
1572                        };
1573                        let Some((_, b64)) = rest.split_once(";base64,") else {
1574                            return Err(DaemonError::Protocol("image_url data URL must be base64"));
1575                        };
1576                        let bytes = base64_decode(b64)?;
1577                        out.images
1578                            .push((std::mem::take(&mut out.text), bytes, MediaKind::Image));
1579                    }
1580                    Some("input_audio") => {
1581                        let ia = p.get("input_audio");
1582                        let fmt = ia
1583                            .and_then(|a| a.get("format"))
1584                            .and_then(|f| f.as_str())
1585                            .unwrap_or("wav");
1586                        if fmt != "wav" {
1587                            return Err(DaemonError::Protocol(
1588                                "input_audio format must be wav (16-bit PCM or float32)",
1589                            ));
1590                        }
1591                        let b64 = ia
1592                            .and_then(|a| a.get("data"))
1593                            .and_then(|d| d.as_str())
1594                            .ok_or(DaemonError::Protocol("input_audio requires base64 data"))?;
1595                        let bytes = base64_decode(b64)?;
1596                        out.images
1597                            .push((std::mem::take(&mut out.text), bytes, MediaKind::Audio));
1598                    }
1599                    _ => {}
1600                }
1601            }
1602        }
1603        Some(_) => {
1604            return Err(DaemonError::Protocol(
1605                "message content must be a string or parts",
1606            ))
1607        }
1608    }
1609    Ok(out)
1610}
1611
1612#[derive(Deserialize)]
1613pub(crate) struct ReqToolCall {
1614    #[serde(default)]
1615    pub(crate) id: Option<String>,
1616    pub(crate) function: Option<ReqFunction>,
1617}
1618
1619#[derive(Deserialize)]
1620pub(crate) struct ReqFunction {
1621    pub(crate) name: String,
1622    #[serde(default)]
1623    pub(crate) arguments: serde_json::Value,
1624}
1625
1626pub(crate) fn build_session(
1627    daemon: &Daemon,
1628    req: &ChatRequest,
1629    op: &SamplingOverrides,
1630    model: &ModelSamplingDefaults,
1631) -> Result<(u64, bool), DaemonError> {
1632    if let Err(why) = daemon.validate_template_kwargs(&req.template_kwargs()) {
1633        return Err(DaemonError::Constraint(why));
1634    }
1635    let mut params = params_of_for(
1636        req.temperature,
1637        req.top_p,
1638        req.top_k,
1639        req.min_p,
1640        req.seed,
1641        model,
1642        &op.for_daemon(daemon),
1643    );
1644    let seed_drawn = draw_seed_if_unseeded(&mut params, req.seed);
1645
1646    if let Some(c) = req.continues {
1647        // Only a codec that renders each message alone renders a turn
1648        // appended to a stored conversation as it would render it fresh.
1649        if !daemon.renders_per_message() {
1650            return Err(DaemonError::Protocol(
1651                "this model's chat template renders whole conversations; a stored one is not continued by appending",
1652            ));
1653        }
1654        let session = daemon.fork(c.session, c.at_event, Some(params))?;
1655        // The stored conversation ends after a query: a turn here follows
1656        // none of this request's user messages, or comes after the last one.
1657        let last_query = req.messages.iter().rposition(|m| m.role == "user");
1658        for (i, m) in req.messages.iter().enumerate() {
1659            append_chat_message(daemon, session, req, m, last_query.is_none_or(|q| i > q))?;
1660        }
1661        return Ok((session, seed_drawn));
1662    }
1663
1664    let session = daemon.create(None, params)?;
1665
1666    if !daemon.renders_per_message() {
1667        let mut msgs: Vec<crate::codec::ChatMessage> = Vec::with_capacity(req.messages.len());
1668        for m in &req.messages {
1669            let split = split_content(&m.content)?;
1670            let role = match m.role.as_str() {
1671                "system" => role::SYSTEM,
1672                "assistant" => role::ASSISTANT,
1673                "tool" => role::TOOL,
1674                _ => role::USER,
1675            };
1676            if !split.images.is_empty() {
1677                let mut parts: Vec<crate::codec::ContentPart> = Vec::new();
1678                let trailing = split.text;
1679                for (pre, bytes, kind) in split.images {
1680                    require_media(daemon, kind, &content_param(&req.messages, m))?;
1681                    if !pre.is_empty() {
1682                        parts.push(crate::codec::ContentPart::Text(pre));
1683                    }
1684                    let hash = daemon.put_media(&bytes)?;
1685                    parts.push(match kind {
1686                        MediaKind::Image => crate::codec::ContentPart::Image { blob: hash },
1687                        MediaKind::Audio => crate::codec::ContentPart::Audio { blob: hash },
1688                    });
1689                }
1690                if !trailing.is_empty() {
1691                    parts.push(crate::codec::ContentPart::Text(trailing));
1692                }
1693                let mut msg = crate::codec::ChatMessage::new(role, String::new());
1694                msg.parts = parts;
1695                msg.reasoning = reasoning_of(m, role);
1696                msgs.push(msg);
1697                continue;
1698            }
1699            let mut msg = crate::codec::ChatMessage::new(role, split.text);
1700            for c in &m.tool_calls {
1701                let Some(f) = &c.function else { continue };
1702                let args = match &f.arguments {
1703                    serde_json::Value::String(s) => s.clone(),
1704                    other => other.to_string(),
1705                };
1706                msg.tool_calls.push(crate::codec::ToolCallMsg {
1707                    id: c.id.clone().unwrap_or_default(),
1708                    name: f.name.clone(),
1709                    arguments: args,
1710                });
1711            }
1712            msg.reasoning = reasoning_of(m, role);
1713            if role == role::TOOL {
1714                msg.tool_call_id = m.tool_call_id.clone();
1715                msg.name = m.name.clone();
1716            }
1717            msgs.push(msg);
1718        }
1719        let tools: Vec<String> = req.tools.iter().map(|t| t.to_string()).collect();
1720        daemon.append_conversation(session, &msgs, &tools, &req.template_kwargs())?;
1721        return Ok((session, seed_drawn));
1722    }
1723
1724    let (system, tool_jsons) = system_head(req)?;
1725    let rest = if system.is_some() { &req.messages[1..] } else { &req.messages[..] };
1726    daemon.append_system_full(session, system, tool_jsons, &req.template_kwargs())?;
1727    // A turn after the last user query is a step of a tool loop.
1728    let last_query = req.messages.iter().rposition(|m| m.role == "user");
1729    let skipped = req.messages.len() - rest.len();
1730    for (i, m) in rest.iter().enumerate() {
1731        let after_query = last_query.is_some_and(|q| skipped + i > q);
1732        append_chat_message(daemon, session, req, m, after_query)?;
1733    }
1734    Ok((session, seed_drawn))
1735}
1736
1737/// What a per-message codec renders first: the leading system message when
1738/// it is text only, and the tools when the request may call them. Two
1739/// requests with the same head render the same first turn.
1740pub(crate) fn system_head(req: &ChatRequest) -> Result<(Option<String>, Vec<String>), DaemonError> {
1741    let system = req.messages.first().map(absorbable_system_text).transpose()?.flatten();
1742    let tools = if !req.tools.is_empty() && ToolChoice::parse(&req.tool_choice).declares_tools() {
1743        req.tools.iter().map(|t| t.to_string()).collect()
1744    } else {
1745        Vec::new()
1746    };
1747    Ok((system, tools))
1748}
1749
1750fn append_chat_message(
1751    daemon: &Daemon,
1752    session: u64,
1753    req: &ChatRequest,
1754    m: &ChatMessage,
1755    after_query: bool,
1756) -> Result<(), DaemonError> {
1757    let split = split_content(&m.content)?;
1758    if !split.images.is_empty() {
1759        let r = match m.role.as_str() {
1760            "assistant" => role::ASSISTANT,
1761            "system" => role::SYSTEM,
1762            _ => role::USER,
1763        };
1764        let n = split.images.len();
1765        let mut trailing = split.text;
1766        for (i, (pre, bytes, kind)) in split.images.into_iter().enumerate() {
1767            require_media(daemon, kind, &content_param(&req.messages, m))?;
1768            let hash = daemon.put_media(&bytes)?;
1769            let post = if i + 1 == n {
1770                std::mem::take(&mut trailing)
1771            } else {
1772                String::new()
1773            };
1774            daemon.append_image(session, r, &hash, &pre, &post)?;
1775        }
1776        return Ok(());
1777    }
1778    let content = split.text;
1779    match m.role.as_str() {
1780        "assistant" if !m.tool_calls.is_empty() => {
1781            let calls: Vec<(String, String)> = m
1782                .tool_calls
1783                .iter()
1784                .filter_map(|c| c.function.as_ref())
1785                .map(|f| {
1786                    let args = match &f.arguments {
1787                        serde_json::Value::String(s) => s.clone(),
1788                        v => v.to_string(),
1789                    };
1790                    (f.name.clone(), args)
1791                })
1792                .collect();
1793            daemon.append_assistant_with_tool_calls(session, content, calls, after_query)?;
1794        }
1795        "assistant" => {
1796            daemon.append_message(session, role::ASSISTANT, content)?;
1797        }
1798        "system" => {
1799            daemon.append_message(session, role::SYSTEM, content)?;
1800        }
1801        "tool" => {
1802            daemon.append_tool_history(session, content)?;
1803        }
1804        _ => {
1805            daemon.append_message(session, role::USER, content)?;
1806        }
1807    }
1808    Ok(())
1809}
1810
1811pub(crate) struct Collected {
1812    pub(crate) content: String,
1813    pub(crate) reasoning: String,
1814    pub(crate) tool_calls: Vec<(u64, String, String)>,
1815    pub(crate) unparsed_tool: String,
1816    pub(crate) broken_calls: Vec<(u64, String, String)>,
1817}
1818
1819pub(crate) fn promote_forced_tool_call(
1820    daemon: &Daemon,
1821    c: &mut Collected,
1822    session: u64,
1823    schemas: Option<&crate::codec::ToolSchemas>,
1824) {
1825    if !c.tool_calls.is_empty() {
1826        return;
1827    }
1828    let raw = c.content.trim().to_string();
1829    if raw.is_empty() {
1830        return;
1831    }
1832    if let Some((name, args)) = daemon.parse_tool_call_with(&raw, schemas) {
1833        c.tool_calls.push((session, name, args));
1834        c.content.clear();
1835    }
1836}
1837
1838pub(crate) fn claim_named_calls(daemon: &Daemon, c: &mut Collected, events: &[crate::CommittedEvent], choice: &ToolChoice) {
1839    if *choice == ToolChoice::None {
1840        return;
1841    }
1842    let mut left = String::new();
1843    for e in events {
1844        if let EventBody::ToolParseFailure { raw } = &e.body {
1845            match daemon.named_tool_call(raw) {
1846                Some((name, args)) => c.broken_calls.push((e.event_id, name, args)),
1847                None => left.push_str(raw),
1848            }
1849        }
1850    }
1851    c.unparsed_tool = left;
1852}
1853
1854pub(crate) fn collect(events: &[crate::CommittedEvent]) -> Collected {
1855    let mut c = Collected {
1856        content: String::new(),
1857        reasoning: String::new(),
1858        tool_calls: Vec::new(),
1859        unparsed_tool: String::new(),
1860        broken_calls: Vec::new(),
1861    };
1862    for e in events {
1863        match &e.body {
1864            EventBody::Generated {
1865                text, channel: ch, ..
1866            } => match *ch {
1867                channel::REASONING => c.reasoning.push_str(text),
1868                channel::TOOL_CALL => {}
1869                _ => c.content.push_str(text),
1870            },
1871            EventBody::ToolUse { name, arguments } => {
1872                c.tool_calls
1873                    .push((e.event_id, name.clone(), arguments.clone()));
1874            }
1875            EventBody::ToolParseFailure { raw } => c.unparsed_tool.push_str(raw),
1876            _ => {}
1877        }
1878    }
1879    c
1880}
1881
1882pub(crate) fn finish_reason(
1883    fin: u32,
1884    produced: u32,
1885    max_tokens: u32,
1886    had_tool_calls: bool,
1887) -> &'static str {
1888    let cut = match fin {
1889        finish::LENGTH => true,
1890        finish::EOS | finish::CANCELLED => false,
1891        _ => produced >= max_tokens,
1892    };
1893    if cut {
1894        "length"
1895    } else if had_tool_calls {
1896        "tool_calls"
1897    } else {
1898        "stop"
1899    }
1900}
1901
1902pub(crate) fn prompt_len(daemon: &Daemon, session: u64) -> u32 {
1903    daemon
1904        .store()
1905        .lock()
1906        .expect("store")
1907        .session(session)
1908        .map(|s| s.tokens.len().min(u32::MAX as usize) as u32)
1909        .unwrap_or(0)
1910}
1911
1912pub(crate) fn next_id() -> String {
1913    use std::sync::atomic::{AtomicU64, Ordering};
1914    static N: AtomicU64 = AtomicU64::new(1);
1915    format!("chatcmpl-superfluid-{}", N.fetch_add(1, Ordering::Relaxed))
1916}
1917
1918pub(crate) fn unix_now() -> u64 {
1919    std::time::SystemTime::now()
1920        .duration_since(std::time::UNIX_EPOCH)
1921        .map(|d| d.as_secs())
1922        .unwrap_or(0)
1923}
1924
1925#[cfg(test)]
1926mod access_log_rate_tests {
1927    use super::*;
1928    use std::time::Duration;
1929
1930    #[test]
1931    fn prefill_counts_only_what_was_prefilled() {
1932        let cases = [
1933            (12961u64, 0u64, 10444u64, "1241"),
1934            (14481, 12800, 5716, "294"),
1935            (15627, 14900, 5057, "144"),
1936            (16045, 15900, 2138, "68"),
1937        ];
1938        for (prompt, cached, ttft_ms, want) in cases {
1939            let got = prefill_rate(prompt, cached, Some(Duration::from_millis(ttft_ms)));
1940            assert_eq!(got, want, "prompt={prompt} cached={cached}");
1941        }
1942    }
1943
1944    #[test]
1945    fn a_fully_warm_prompt_prefilled_nothing() {
1946        assert_eq!(
1947            prefill_rate(4096, 4096, Some(Duration::from_millis(80))),
1948            "-"
1949        );
1950        assert_eq!(
1951            prefill_rate(4096, 9000, Some(Duration::from_millis(80))),
1952            "-"
1953        );
1954    }
1955
1956    #[test]
1957    fn a_fim_cache_hit_logs_its_prompt_as_cached() {
1958        let c = |cached| crate::Completion {
1959            text: String::new(),
1960            tokens: vec![1; 16],
1961            cached,
1962            expired: false,
1963            finish: superfluid_abi::finish::NONE,
1964            prompt_tokens: 21,
1965        };
1966        assert_eq!(fim_log_cached(&c(true)), 21);
1967        assert_eq!(
1968            prefill_rate(
1969                21,
1970                fim_log_cached(&c(true)),
1971                Some(Duration::from_micros(60))
1972            ),
1973            "-"
1974        );
1975        assert_eq!(fim_log_cached(&c(false)), 0);
1976    }
1977
1978    #[test]
1979    fn prefill_needs_a_ttft() {
1980        assert_eq!(prefill_rate(1024, 0, None), "-");
1981        assert_eq!(prefill_rate(1024, 0, Some(Duration::ZERO)), "-");
1982    }
1983
1984    fn run(d: &mut DecodeSpan, base: std::time::Instant, stamps: &[(u64, u64)]) {
1985        for (ms, produced) in stamps {
1986            d.stamp_at(base + Duration::from_millis(*ms), *produced);
1987        }
1988        d.finish();
1989    }
1990
1991    #[test]
1992    fn decode_is_the_inter_token_rate() {
1993        let base = std::time::Instant::now();
1994        let mut d = DecodeSpan::default();
1995        run(&mut d, base, &[(0, 1), (781, 200)]);
1996        let (tokens, span) = d.measured();
1997        assert_eq!(tokens, 199);
1998        assert_eq!(decode_rate(tokens, span), "254.8");
1999    }
2000
2001    #[test]
2002    fn decode_excludes_the_whole_first_commit() {
2003        let base = std::time::Instant::now();
2004        let mut d = DecodeSpan::default();
2005        run(&mut d, base, &[(0, 5), (1000, 100)]);
2006        assert_eq!(d.measured().0, 95, "the first commit's 5 predate the span");
2007        assert_eq!(decode_rate(d.measured().0, d.measured().1), "95.0");
2008    }
2009
2010    #[test]
2011    fn decode_spans_do_not_swallow_the_gaps_between_choices() {
2012        let base = std::time::Instant::now();
2013        let mut d = DecodeSpan::default();
2014        run(&mut d, base, &[(0, 1), (500, 51)]);
2015        run(&mut d, base, &[(4500, 1), (5000, 51)]);
2016        let (tokens, span) = d.measured();
2017        assert_eq!(tokens, 100);
2018        assert_eq!(span, Some(Duration::from_millis(1000)), "gap excluded");
2019        assert_eq!(decode_rate(tokens, span), "100.0");
2020    }
2021
2022    #[test]
2023    fn decode_counts_only_what_was_delivered() {
2024        let base = std::time::Instant::now();
2025        let mut d = DecodeSpan::default();
2026        run(&mut d, base, &[(0, 1), (500, 50)]);
2027        let (tokens, span) = d.measured();
2028        assert_eq!(tokens, 49, "the undelivered tail is usage's business");
2029        assert_eq!(decode_rate(tokens, span), "98.0");
2030    }
2031
2032    #[test]
2033    fn decode_declines_to_guess() {
2034        let base = std::time::Instant::now();
2035        let mut d = DecodeSpan::default();
2036        d.finish();
2037        assert_eq!(d.measured(), (0, None));
2038        assert_eq!(decode_rate(0, None), "-");
2039
2040        let mut one = DecodeSpan::default();
2041        run(&mut one, base, &[(0, 40)]);
2042        assert_eq!(decode_rate(one.measured().0, one.measured().1), "-");
2043
2044        assert_eq!(decode_rate(200, Some(Duration::from_millis(150))), "-");
2045        assert_eq!(decode_rate(100, None), "-");
2046    }
2047
2048    #[test]
2049    fn a_tick_split_by_a_channel_change_is_not_counted_against_the_next_tick() {
2050        // Qwen3.6-35B-A3B, 128 tokens after a 4096-token prompt: the tick that
2051        // ended the prompt also decoded and committed 19 text tokens, then 79
2052        // reasoning tokens 9 ms later; the next tick added 30 in 0.79 s.
2053        let base = std::time::Instant::now();
2054        let mut d = DecodeSpan::default();
2055        run(&mut d, base, &[(0, 19), (9, 98), (801, 128)]);
2056        assert_eq!(decode_rate(d.measured().0, d.measured().1), "136.1", "events alone overstate it");
2057        d.add_measured((30, 0.792));
2058        assert_eq!(decode_rate(d.measured().0, d.measured().1), "37.9");
2059    }
2060
2061    #[test]
2062    fn a_generation_delivered_in_one_tick_has_no_decode_rate() {
2063        let base = std::time::Instant::now();
2064        let mut d = DecodeSpan::default();
2065        run(&mut d, base, &[(0, 128)]);
2066        d.add_measured((0, 0.0));
2067        assert_eq!(decode_rate(d.measured().0, d.measured().1), "-");
2068    }
2069
2070    #[test]
2071    fn a_one_tick_generation_split_by_a_channel_change_has_no_decode_rate() {
2072        // One tick's tokens committed as two events 9 ms apart: the events
2073        // alone make a rate of the gap between them.
2074        let base = std::time::Instant::now();
2075        let mut d = DecodeSpan::default();
2076        run(&mut d, base, &[(0, 19), (9, 20)]);
2077        assert_eq!(decode_rate(d.measured().0, d.measured().1), "111.1", "events alone");
2078        d.add_measured((0, 0.0));
2079        assert_eq!(d.measured(), (0, None), "the scheduler saw no tick after the first");
2080    }
2081
2082    #[test]
2083    fn measured_decode_adds_up_over_choices_and_ignores_empty_ones() {
2084        let mut d = DecodeSpan::default();
2085        d.add_measured((63, 0.5));
2086        d.add_measured((0, 0.0));
2087        d.add_measured((63, 0.5));
2088        assert_eq!(d.measured(), (126, Some(Duration::from_secs(1))));
2089        assert_eq!(decode_rate(d.measured().0, d.measured().1), "126.0");
2090    }
2091}
2092
2093#[cfg(test)]
2094mod key_order_tests {
2095    use super::ChatRequest;
2096
2097    #[test]
2098    fn a_tool_keeps_the_key_order_it_was_sent_in() {
2099        let req: ChatRequest = serde_json::from_str(
2100            r#"{"messages":[],"tools":[{"type":"function","function":{"name":"replace_text",
2101                "parameters":{"type":"object","properties":{"path":{"type":"string"},
2102                "old_text":{"type":"string"},"new_text":{"type":"string"}}}}}]}"#,
2103        )
2104        .expect("request");
2105        let t = req.tools[0].to_string();
2106        let at = |k: &str| t.find(k).unwrap_or_else(|| panic!("{k} missing: {t}"));
2107        assert!(at("\"type\"") < at("\"function\""), "{t}");
2108        assert!(
2109            at("\"path\"") < at("\"old_text\"") && at("\"old_text\"") < at("\"new_text\""),
2110            "{t}"
2111        );
2112    }
2113}
2114
2115#[cfg(test)]
2116mod system_absorption_tests {
2117    use super::*;
2118
2119    #[test]
2120    fn assistant_reasoning_is_carried_and_other_roles_are_not() {
2121        let mut m = msg("assistant", serde_json::Value::Null);
2122        m.reasoning_content = Some("  plan  ".into());
2123        assert_eq!(reasoning_of(&m, role::ASSISTANT), Some("plan".into()));
2124        m.reasoning_content = None;
2125        m.reasoning = Some("alt".into());
2126        assert_eq!(
2127            reasoning_of(&m, role::ASSISTANT),
2128            Some("alt".into()),
2129            "the other spelling"
2130        );
2131        m.reasoning = Some("   ".into());
2132        assert_eq!(
2133            reasoning_of(&m, role::ASSISTANT),
2134            None,
2135            "blank reasoning is none"
2136        );
2137        let mut u = msg("user", serde_json::json!("hi"));
2138        u.reasoning_content = Some("x".into());
2139        assert_eq!(
2140            reasoning_of(&u, role::USER),
2141            None,
2142            "only assistant turns carry reasoning"
2143        );
2144    }
2145
2146    fn msg(role: &str, content: serde_json::Value) -> ChatMessage {
2147        ChatMessage {
2148            role: role.to_string(),
2149            content: Some(content),
2150            tool_calls: Vec::new(),
2151            tool_call_id: None,
2152            name: None,
2153            reasoning_content: None,
2154            reasoning: None,
2155        }
2156    }
2157
2158    fn image_part() -> serde_json::Value {
2159        serde_json::json!({
2160            "type": "image_url",
2161            "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="}
2162        })
2163    }
2164
2165    #[test]
2166    fn text_only_system_message_is_absorbed() {
2167        let m = msg("system", serde_json::json!("you are helpful"));
2168        assert_eq!(
2169            absorbable_system_text(&m).expect("split"),
2170            Some("you are helpful".to_string()),
2171            "a text-only system message folds into the scaffold"
2172        );
2173    }
2174
2175    #[test]
2176    fn multimodal_system_message_is_left_for_the_media_loop() {
2177        let m = msg(
2178            "system",
2179            serde_json::json!([{"type": "text", "text": "look at this"}, image_part()]),
2180        );
2181        let split = split_content(&m.content).expect("split");
2182        assert_eq!(split.images.len(), 1, "fixture must carry one image part");
2183
2184        assert_eq!(
2185            absorbable_system_text(&m).expect("split"),
2186            None,
2187            "a system message with media must NOT be absorbed — its images \
2188             would be dropped on the floor"
2189        );
2190    }
2191
2192    #[test]
2193    fn audio_system_message_is_left_for_the_media_loop() {
2194        let m = msg(
2195            "system",
2196            serde_json::json!([
2197                {"type": "text", "text": "listen"},
2198                {"type": "input_audio", "input_audio": {"data": "UklGRiQAAABXQVZF", "format": "wav"}}
2199            ]),
2200        );
2201        let split = split_content(&m.content).expect("split");
2202        assert_eq!(split.images.len(), 1, "fixture must carry one audio part");
2203        assert_eq!(absorbable_system_text(&m).expect("split"), None);
2204    }
2205
2206    #[test]
2207    fn non_system_roles_are_never_absorbed() {
2208        for role in ["user", "assistant", "tool", "developer"] {
2209            let m = msg(role, serde_json::json!("hello"));
2210            assert_eq!(
2211                absorbable_system_text(&m).expect("split"),
2212                None,
2213                "{role} must not be absorbed into the system scaffold"
2214            );
2215        }
2216    }
2217}
2218
2219pub(crate) enum RunError {
2220    Daemon(DaemonError),
2221    Worker(String),
2222}
2223impl From<DaemonError> for RunError {
2224    fn from(error: DaemonError) -> Self {
2225        Self::Daemon(error)
2226    }
2227}
2228
2229pub(crate) async fn completions(
2230    daemon: Arc<Daemon>,
2231    model_name: String,
2232    req: CompletionRequest,
2233    policy: RequestPolicy,
2234    hold: Option<crate::keypolicy::ClientKey>,
2235    hooks: crate::keepalive::Hooks,
2236) -> Result<Generation, RunError> {
2237    let qos = policy.qos;
2238    let max_tokens = policy
2239        .limits
2240        .resolve(req.max_tokens, || daemon.tokenize(&req.prompt).len() as u32);
2241    let id = format!(
2242        "cmpl-superfluid-{}",
2243        next_id().rsplit('-').next().unwrap_or("0")
2244    );
2245    let created = unix_now();
2246    let model_defaults = {
2247        let d = daemon.clone();
2248        tokio::task::spawn_blocking(move || {
2249            crate::with_sync_class(Some(qos.class), || d.model_sampling_defaults())
2250        })
2251        .await
2252        .unwrap_or_default()
2253    };
2254    let mut params = params_of_for(
2255        req.temperature,
2256        req.top_p,
2257        req.top_k,
2258        req.min_p,
2259        req.seed,
2260        &model_defaults,
2261        &policy.sampling.for_daemon(&daemon),
2262    );
2263    let seed_drawn = draw_seed_if_unseeded(&mut params, req.seed);
2264    let mut extras = extras_of(
2265        req.presence_penalty,
2266        req.frequency_penalty,
2267        req.repeat_penalty,
2268        &policy.sampling,
2269        &model_defaults,
2270    );
2271    extras.seed_drawn = seed_drawn;
2272    let want_logprobs = req.logprobs.is_some();
2273    extras.want_logprobs = want_logprobs;
2274    extras.top_logprobs = req.logprobs.unwrap_or(0).min(20) as u8;
2275    extras.ignore_eos = req.ignore_eos;
2276    let echo_offset = if req.echo { req.prompt.len() } else { 0 };
2277
2278    if req.stream {
2279        let body = completion_stream(daemon, req, max_tokens, params, extras, qos, hold).await?;
2280        return Ok(Generation {
2281            kind: Kind::Completion,
2282            id,
2283            model: model_name,
2284            created,
2285            warm_header: None,
2286            body,
2287        });
2288    }
2289
2290    let daemon_lp = Arc::clone(&daemon);
2291    let model_log = model_name.clone();
2292    let out = blocking(qos, hold, move || {
2293        let mut hooks = hooks;
2294        let t0 = std::time::Instant::now();
2295        let mut decode = DecodeSpan::default();
2296        let mut produced = 0u64;
2297        let session = daemon.create(None, params)?;
2298        qos.apply(&daemon, session)?;
2299        daemon.append(session, Some(req.prompt.clone()), Vec::new())?;
2300        if hooks.is_on() {
2301            daemon.session_fits(session, None)?;
2302        }
2303        hooks.admitted(&daemon, session);
2304        hooks.check()?;
2305        let out = daemon.generate_streaming_ex(session, max_tokens, extras, |ev| {
2306            hooks.cancel_if_gone();
2307            if let EventBody::Generated { span, .. } = &ev.body {
2308                if !span.is_empty() {
2309                    produced += span.len() as u64;
2310                    decode.stamp(produced);
2311                }
2312            }
2313            Ok(())
2314        })?;
2315        let first_tok = decode.first_token();
2316        decode.finish();
2317        decode.add_measured(out.decode);
2318        let c = collect(&out.events);
2319        let usage = usage(&daemon, session, out.tokens_generated, out.warm_prefix);
2320        access_log(
2321            "POST /v1/completions",
2322            &model_log,
2323            usage.prompt.unwrap_or(0),
2324            out.warm_prefix,
2325            out.tokens_generated as u64,
2326            GenTiming {
2327                t0,
2328                first: first_tok,
2329                decode,
2330            },
2331            finish_reason(out.finish, out.tokens_generated, max_tokens, false),
2332        );
2333        let mut text = String::new();
2334        if req.echo {
2335            text.push_str(&req.prompt);
2336        }
2337        text.push_str(&c.content);
2338        text.push_str(&c.reasoning);
2339        text.push_str(&c.unparsed_tool);
2340        Ok::<_, DaemonError>((out, text, usage))
2341    })
2342    .await;
2343    let (out, text, usage) = out?;
2344    let logprobs = want_logprobs.then(|| decode_logprobs(&daemon_lp, &out.logprobs, echo_offset));
2345    let mut frame = Frame::text(
2346        text,
2347        logprobs,
2348        Some(finish_reason(
2349            out.finish,
2350            out.tokens_generated,
2351            max_tokens,
2352            false,
2353        )),
2354    );
2355    frame.usage = Some(usage);
2356    Ok(Generation {
2357        kind: Kind::Completion,
2358        id,
2359        model: model_name,
2360        created,
2361        warm_header: None,
2362        body: Body::Complete(frame),
2363    })
2364}
2365
2366#[allow(clippy::too_many_arguments)]
2367async fn completion_stream(
2368    daemon: Arc<Daemon>,
2369    req: CompletionRequest,
2370    max_tokens: u32,
2371    params: GenParams,
2372    extras: crate::scheduler::GenExtras,
2373    qos: RequestQos,
2374    hold: Option<crate::keypolicy::ClientKey>,
2375) -> Result<Body, RunError> {
2376    let want_logprobs = req.logprobs.is_some();
2377    let admitted_daemon = Arc::clone(&daemon);
2378    let (req, session) = blocking(qos, hold.clone(), move || {
2379        let session = admitted_daemon.create(None, params)?;
2380        qos.apply(&admitted_daemon, session)?;
2381        admitted_daemon.append(session, Some(req.prompt.clone()), Vec::new())?;
2382        admitted_daemon.session_fits(session, None)?;
2383        Ok((req, session))
2384    })
2385    .await?;
2386    let rx = spawn_stream(qos, hold, move |tx| {
2387        let _closed = crate::keepalive::cancel_when_closed(&daemon, session, tx);
2388        let send = |v: Frame| {
2389            tx.blocking_send(Ok(v))
2390                .map_err(|_| DaemonError::Protocol("generation consumer gone"))
2391        };
2392        let lp_q: std::rc::Rc<
2393            std::cell::RefCell<std::collections::VecDeque<crate::scheduler::TokenLogprob>>,
2394        > = std::rc::Rc::new(std::cell::RefCell::new(std::collections::VecDeque::new()));
2395        let lp_fill = std::rc::Rc::clone(&lp_q);
2396        let mut offset: usize = 0;
2397        let run = || -> Result<(), DaemonError> {
2398            if req.echo {
2399                offset = req.prompt.len();
2400                send(Frame::text(&req.prompt, None, None))?;
2401            }
2402            let out = daemon.generate_streaming_ex_lp(
2403                session,
2404                max_tokens,
2405                extras,
2406                |ev| {
2407                    if let EventBody::Generated {
2408                        text,
2409                        channel: ch,
2410                        span,
2411                        ..
2412                    } = &ev.body
2413                    {
2414                        let these: Vec<crate::scheduler::TokenLogprob> = if want_logprobs {
2415                            let mut q = lp_q.borrow_mut();
2416                            (0..span.len()).filter_map(|_| q.pop_front()).collect()
2417                        } else {
2418                            Vec::new()
2419                        };
2420                        if !text.is_empty() && *ch != channel::TOOL_CALL {
2421                            let lp = if want_logprobs {
2422                                let v = Some(decode_logprobs(&daemon, &these, offset));
2423                                offset += text.len();
2424                                v
2425                            } else {
2426                                None
2427                            };
2428                            return send(Frame::text(text, lp, None));
2429                        }
2430                    }
2431                    Ok(())
2432                },
2433                |items| {
2434                    lp_fill.borrow_mut().extend(items.iter().cloned());
2435                },
2436            )?;
2437            let fin = finish_reason(out.finish, out.tokens_generated, max_tokens, false);
2438            let mut terminal = Frame::text("", None, Some(fin));
2439            terminal.usage = Some(usage(
2440                &daemon,
2441                session,
2442                out.tokens_generated,
2443                out.warm_prefix,
2444            ));
2445            send(terminal)
2446        };
2447        run()
2448    });
2449    Ok(Body::Stream(rx))
2450}
2451pub(crate) async fn chat(
2452    daemon: Arc<Daemon>,
2453    model_name: String,
2454    req: ChatRequest,
2455    policy: RequestPolicy,
2456    hold: Option<crate::keypolicy::ClientKey>,
2457    hooks: crate::keepalive::Hooks,
2458) -> Result<Generation, RunError> {
2459    let qos = policy.qos;
2460    let explicit_max = req.max_completion_tokens.or(req.max_tokens);
2461    let id = next_id();
2462    let created = unix_now();
2463
2464    if req.stream {
2465        let body = chat_stream(daemon, model_name.clone(), req, policy, explicit_max, hold).await?;
2466        return Ok(Generation {
2467            kind: Kind::Chat,
2468            id,
2469            model: model_name,
2470            created,
2471            warm_header: None,
2472            body,
2473        });
2474    }
2475
2476    let n = req.n.unwrap_or(1).max(1);
2477    let route = req.route();
2478    let stops = req.stop.list();
2479    let model_log = model_name.clone();
2480    let out = blocking(qos, hold, move || {
2481        let mut hooks = hooks;
2482        let t0 = std::time::Instant::now();
2483        let mut decode = DecodeSpan::default();
2484        let mut ttft_first: Option<std::time::Instant> = None;
2485        let mut last_finish = "-";
2486        let model_defaults = daemon.model_sampling_defaults();
2487        let mut extras = extras_of(
2488            req.presence_penalty,
2489            req.frequency_penalty,
2490            req.repeat_penalty,
2491            &policy.sampling,
2492            &model_defaults,
2493        );
2494        let choice = ToolChoice::parse(&req.tool_choice);
2495        if choice.forces_call() {
2496            extras.open_channel = crate::wal::channel::TOOL_CALL;
2497        }
2498        let tool_schemas = crate::codec::ToolSchemas::of_values(&req.tools);
2499        extras.tool_schemas = tool_schemas.clone();
2500        let handles = RequestHandles::new(
2501            &daemon,
2502            grammar_handle_for(&daemon, &req, &choice)?,
2503            logit_bias_handle_of(&daemon, &req.logit_bias),
2504        );
2505        handles.apply(&mut extras);
2506        let want_logprobs = req.logprobs.unwrap_or(false);
2507        extras.want_logprobs = want_logprobs;
2508        extras.top_logprobs = req.top_logprobs.unwrap_or(0).min(20) as u8;
2509        extras.thinking = req.thinking();
2510        extras.ignore_eos = req.ignore_eos;
2511        let cancels = daemon.cancel_registry();
2512        let RequestPolicy {
2513            limits,
2514            sampling: op,
2515            qos,
2516        } = policy;
2517        let mut choices = Vec::with_capacity(n as usize);
2518        let mut completion_tokens = 0u64;
2519        let mut prompt_tokens = 0u64;
2520        let mut warm = 0u64;
2521        let mut warm_first: Option<u64> = None;
2522        let mut spec = (0u64, 0u64);
2523        let mut last_session = None;
2524        for i in 0..n {
2525            if i > 0 && hooks.gone() {
2526                break;
2527            }
2528            let (session, seed_drawn) = build_session(&daemon, &req, &op, &model_defaults)?;
2529            last_session = Some(session);
2530            let mut extras = extras.clone();
2531            extras.seed_drawn = seed_drawn;
2532            qos.apply(&daemon, session)?;
2533            if hooks.is_on() {
2534                daemon.session_fits(session, req.thinking())?;
2535            }
2536            hooks.admitted(&daemon, session);
2537            hooks.check()?;
2538            let max_tokens = limits.resolve(explicit_max, || prompt_len(&daemon, session));
2539            let mut acc = String::new();
2540            let mut produced = 0u64;
2541            let stops = stops.clone();
2542            let cancels = std::sync::Arc::clone(&cancels);
2543            let mut stopped = false;
2544            let out = daemon.generate_streaming_ex(session, max_tokens, extras, |ev| {
2545                hooks.cancel_if_gone();
2546                if let EventBody::Generated {
2547                    text,
2548                    channel: ch,
2549                    span,
2550                    ..
2551                } = &ev.body
2552                {
2553                    if !span.is_empty() {
2554                        produced += span.len() as u64;
2555                        decode.stamp(produced);
2556                    }
2557                    if *ch == channel::TEXT && !stops.is_empty() && !stopped {
2558                        acc.push_str(text);
2559                        if stops.iter().any(|s| acc.contains(s.as_str())) {
2560                            cancels.cancel(session);
2561                            stopped = true;
2562                        }
2563                    }
2564                }
2565                Ok(())
2566            })?;
2567            ttft_first = ttft_first.or_else(|| decode.first_token());
2568            decode.finish();
2569            decode.add_measured(out.decode);
2570            let mut c = collect(&out.events);
2571            claim_named_calls(&daemon, &mut c, &out.events, &choice);
2572            if choice.forces_call() {
2573                promote_forced_tool_call(&daemon, &mut c, session, tool_schemas.as_deref());
2574            }
2575            let (content, hit_stop) = truncate_at_stop(&c.content, &stops);
2576            c.content = content;
2577            let usage = usage(&daemon, session, out.tokens_generated, out.warm_prefix);
2578            completion_tokens += usage.completion;
2579            spec.0 += out.spec.0;
2580            spec.1 += out.spec.1;
2581            prompt_tokens = usage.prompt.unwrap_or(prompt_tokens);
2582            warm = warm.max(out.warm_prefix);
2583            warm_first.get_or_insert(out.warm_prefix);
2584            let finish = if !c.broken_calls.is_empty() {
2585                "length"
2586            } else if hit_stop {
2587                "stop"
2588            } else {
2589                finish_reason(
2590                    out.finish,
2591                    out.tokens_generated,
2592                    max_tokens,
2593                    !c.tool_calls.is_empty() && choice != ToolChoice::None,
2594                )
2595            };
2596            choices.push(Choice {
2597                index: i,
2598                message: message(&c, &choice),
2599                finish: Some(finish),
2600                logprobs: want_logprobs.then(|| decode_logprobs(&daemon, &out.logprobs, 0)),
2601                ..Choice::default()
2602            });
2603            last_finish = finish;
2604        }
2605        access_log(
2606            route,
2607            &model_log,
2608            prompt_tokens,
2609            warm_first.unwrap_or(warm),
2610            completion_tokens,
2611            GenTiming {
2612                t0,
2613                first: ttft_first,
2614                decode,
2615            },
2616            last_finish,
2617        );
2618        Ok::<_, DaemonError>((choices, prompt_tokens, completion_tokens, warm, spec, last_session))
2619    })
2620    .await;
2621    let (choices, prompt_tokens, completion_tokens, warm, spec, session) = out?;
2622    let frame = Frame {
2623        choices,
2624        usage: Some(Usage::complete(
2625            prompt_tokens,
2626            completion_tokens,
2627            Some(warm.min(prompt_tokens)),
2628        )),
2629        spec: (spec.0 > 0).then_some(spec),
2630        session,
2631        ..Frame::default()
2632    };
2633    Ok(Generation {
2634        kind: Kind::Chat,
2635        id,
2636        model: model_name,
2637        created,
2638        warm_header: Some(warm),
2639        body: Body::Complete(frame),
2640    })
2641}
2642
2643async fn chat_stream(
2644    daemon: Arc<Daemon>,
2645    model_name: String,
2646    req: ChatRequest,
2647    policy: RequestPolicy,
2648    explicit_max: Option<u32>,
2649    hold: Option<crate::keypolicy::ClientKey>,
2650) -> Result<Body, RunError> {
2651    let RequestPolicy {
2652        limits,
2653        sampling: op,
2654        qos,
2655    } = policy;
2656    let t0 = std::time::Instant::now();
2657    let admitted_daemon = Arc::clone(&daemon);
2658    let (req, session, seed_drawn, model_defaults) = blocking(qos, hold.clone(), move || {
2659        let model_defaults = admitted_daemon.model_sampling_defaults();
2660        let (session, seed_drawn) = build_session(&admitted_daemon, &req, &op, &model_defaults)?;
2661        qos.apply(&admitted_daemon, session)?;
2662        admitted_daemon.session_fits(session, req.thinking())?;
2663        Ok((req, session, seed_drawn, model_defaults))
2664    })
2665    .await?;
2666    let rx = spawn_stream(qos, hold, move |tx| {
2667        let _closed = crate::keepalive::cancel_when_closed(&daemon, session, tx);
2668        let send = |v: Frame| -> Result<(), DaemonError> {
2669            tx.blocking_send(Ok(v))
2670                .map_err(|_| DaemonError::Protocol("generation consumer gone"))
2671        };
2672        let run = || -> Result<(), DaemonError> {
2673            send(Frame::chat(
2674                Message {
2675                    assistant: true,
2676                    ..Message::default()
2677                },
2678                None,
2679            ))?;
2680            let mut tool_index = 0usize;
2681            let mut orphaned_calls = 0usize;
2682            let mut extras = extras_of(
2683                req.presence_penalty,
2684                req.frequency_penalty,
2685                req.repeat_penalty,
2686                &op,
2687                &model_defaults,
2688            );
2689            extras.seed_drawn = seed_drawn;
2690            let choice = ToolChoice::parse(&req.tool_choice);
2691            if choice.forces_call() {
2692                extras.open_channel = crate::wal::channel::TOOL_CALL;
2693            }
2694            let tool_schemas = crate::codec::ToolSchemas::of_values(&req.tools);
2695            extras.tool_schemas = tool_schemas.clone();
2696            let mut frag =
2697                crate::tool_fragment::ToolCallStream::new(if choice == ToolChoice::None {
2698                    crate::tool_fragment::ToolFragmentMode::None
2699                } else {
2700                    daemon.tool_fragment_mode()
2701                })
2702                .with_schemas(tool_schemas.clone());
2703            let mut open_call_id = String::new();
2704            let handles = RequestHandles::new(
2705                &daemon,
2706                grammar_handle_for(&daemon, &req, &choice)?,
2707                logit_bias_handle_of(&daemon, &req.logit_bias),
2708            );
2709            handles.apply(&mut extras);
2710            let want_logprobs = req.logprobs.unwrap_or(false);
2711            extras.want_logprobs = want_logprobs;
2712            extras.top_logprobs = req.top_logprobs.unwrap_or(0).min(20) as u8;
2713            extras.thinking = req.thinking();
2714            extras.ignore_eos = req.ignore_eos;
2715            let lp_q: std::rc::Rc<
2716                std::cell::RefCell<std::collections::VecDeque<crate::scheduler::TokenLogprob>>,
2717            > = std::rc::Rc::new(std::cell::RefCell::new(std::collections::VecDeque::new()));
2718            let lp_fill = std::rc::Rc::clone(&lp_q);
2719            let max_tokens = limits.resolve(explicit_max, || prompt_len(&daemon, session));
2720            let stops = req.stop.list();
2721            let cancels = daemon.cancel_registry();
2722            let delta_mode = !want_logprobs;
2723            let continuous_usage = req
2724                .stream_options
2725                .as_ref()
2726                .is_some_and(|o| o.continuous_usage_stats);
2727            struct StreamText {
2728                usage_now: Option<u32>,
2729                committed: u32,
2730                decode: DecodeSpan,
2731                content_acc: String,
2732                content_sent: usize,
2733                stopped: bool,
2734                dead: bool,
2735            }
2736            let st = std::rc::Rc::new(std::cell::RefCell::new(StreamText {
2737                usage_now: None,
2738                committed: 0,
2739                decode: DecodeSpan::default(),
2740                content_acc: String::new(),
2741                content_sent: 0,
2742                stopped: false,
2743                dead: false,
2744            }));
2745            let push_text =
2746                |s: &mut StreamText, text: &str, lp: Option<Logprobs>| -> Result<(), DaemonError> {
2747                    if s.stopped {
2748                        return Ok(());
2749                    }
2750                    if choice.forces_call() {
2751                        s.content_acc.push_str(text);
2752                        return Ok(());
2753                    }
2754                    let start = s.content_sent;
2755                    s.content_acc.push_str(text);
2756                    let tail = &s.content_acc[start..];
2757                    let mut cut: Option<usize> = None;
2758                    for stop in &stops {
2759                        if let Some(i) = tail.find(stop.as_str()) {
2760                            cut = Some(cut.map_or(i, |c| c.min(i)));
2761                        }
2762                    }
2763                    let hit = cut.is_some();
2764                    let end = match cut {
2765                        Some(i) => start + i,
2766                        None => start + tail.len() - stop_prefix_holdback(tail, &stops),
2767                    };
2768                    let piece = s.content_acc.get(start..end).unwrap_or("").to_string();
2769                    s.content_sent = s.content_sent.max(end);
2770                    if !piece.is_empty() {
2771                        let mut ck = Frame::chat(Message::content(piece), None);
2772                        if let Some(lp) = lp {
2773                            ck.choices[0].logprobs = Some(lp);
2774                        }
2775                        if let Some(n) = s.usage_now {
2776                            ck.usage = Some(Usage::partial(n));
2777                        }
2778                        send(ck)?;
2779                    }
2780                    if hit {
2781                        cancels.cancel(session);
2782                        s.stopped = true;
2783                    }
2784                    Ok(())
2785                };
2786            let st_delta = std::rc::Rc::clone(&st);
2787            let on_delta = |ch: u32, text: &str, produced: u32| {
2788                if !delta_mode || text.is_empty() {
2789                    return;
2790                }
2791                let mut s = st_delta.borrow_mut();
2792                if continuous_usage {
2793                    s.usage_now = Some(produced);
2794                }
2795                if s.dead {
2796                    return;
2797                }
2798                s.decode.stamp(produced as u64);
2799                let r = if ch == channel::REASONING {
2800                    let mut ck = Frame::chat(Message::reasoning(text), None);
2801                    if let Some(n) = s.usage_now {
2802                        ck.usage = Some(Usage::partial(n));
2803                    }
2804                    send(ck)
2805                } else {
2806                    push_text(&mut s, text, None)
2807                };
2808                if r.is_err() {
2809                    cancels.cancel(session);
2810                    s.dead = true;
2811                }
2812            };
2813            let st_ev = std::rc::Rc::clone(&st);
2814            let on_event = |ev: &crate::CommittedEvent| {
2815                match &ev.body {
2816                    EventBody::Generated {
2817                        text,
2818                        channel: ch,
2819                        span,
2820                        ..
2821                    } => {
2822                        let mut s = st_ev.borrow_mut();
2823                        let covered_by_delta = delta_mode
2824                            && (*ch == channel::TEXT || *ch == channel::REASONING)
2825                            && !text.is_empty();
2826                        s.committed += span.len() as u32;
2827                        let committed = s.committed as u64;
2828                        if !covered_by_delta && !span.is_empty() {
2829                            s.decode.stamp(committed);
2830                        }
2831                        if continuous_usage && !delta_mode {
2832                            s.usage_now = Some(s.committed);
2833                        }
2834                        let these: Vec<crate::scheduler::TokenLogprob> = if want_logprobs {
2835                            let mut q = lp_q.borrow_mut();
2836                            (0..span.len()).filter_map(|_| q.pop_front()).collect()
2837                        } else {
2838                            Vec::new()
2839                        };
2840                        if *ch == channel::TOOL_CALL {
2841                            for d in frag.feed(text) {
2842                                if d.opening {
2843                                    open_call_id = format!("call_{}", ev.event_id);
2844                                }
2845                                let json = Message::tool(&d, tool_index, &open_call_id);
2846                                send(Frame::chat(json, None))?;
2847                            }
2848                            return Ok(());
2849                        }
2850                        if text.is_empty() {
2851                            return Ok(());
2852                        }
2853                        if delta_mode {
2854                            return Ok(());
2855                        }
2856                        if *ch == channel::REASONING {
2857                            let mut ck = Frame::chat(Message::reasoning(text), None);
2858                            if let Some(n) = s.usage_now {
2859                                ck.usage = Some(Usage::partial(n));
2860                            }
2861                            return send(ck);
2862                        }
2863                        let lp = want_logprobs.then(|| decode_logprobs(&daemon, &these, 0));
2864                        push_text(&mut s, text, lp)
2865                    }
2866                    EventBody::ToolParseFailure { raw } => {
2867                        if choice != ToolChoice::None && !frag.opened() {
2868                            if let Some((name, arguments)) = daemon.named_tool_call(raw) {
2869                                let delta = Message::complete_tool(
2870                                    tool_index,
2871                                    format!("call_{}", ev.event_id),
2872                                    name,
2873                                    arguments,
2874                                );
2875                                tool_index += 1;
2876                                orphaned_calls += 1;
2877                                return send(Frame::chat(delta, None));
2878                            }
2879                        }
2880                        if frag.opened() {
2881                            orphaned_calls += 1;
2882                            tool_index += 1;
2883                        }
2884                        frag.reset();
2885                        let mut ck = Frame::chat(Message::content(raw), None);
2886                        if continuous_usage {
2887                            ck.usage = Some(Usage::partial(st_ev.borrow().committed));
2888                        }
2889                        send(ck)
2890                    }
2891                    EventBody::ToolUse { name, arguments } => {
2892                        if choice == ToolChoice::None {
2893                            let raw = format!(
2894                                "{{\"name\":{},\"arguments\":{}}}",
2895                                serde_json::Value::String(name.clone()),
2896                                arguments
2897                            );
2898                            let mut ck = Frame::chat(Message::content(raw), None);
2899                            if continuous_usage {
2900                                ck.usage = Some(Usage::partial(st_ev.borrow().committed));
2901                            }
2902                            return send(ck);
2903                        }
2904                        if frag.opened() {
2905                            for d in frag.finish(arguments) {
2906                                send(Frame::chat(
2907                                    Message::tool(&d, tool_index, &open_call_id),
2908                                    None,
2909                                ))?;
2910                            }
2911                            frag.reset();
2912                            tool_index += 1;
2913                            return Ok(());
2914                        }
2915                        let delta = Message::complete_tool(
2916                            tool_index,
2917                            format!("call_{}", ev.event_id),
2918                            name.clone(),
2919                            arguments.clone(),
2920                        );
2921                        tool_index += 1;
2922                        send(Frame::chat(delta, None))
2923                    }
2924                    _ => Ok(()),
2925                }
2926            };
2927            let on_logprobs = |items: &[crate::scheduler::TokenLogprob]| {
2928                lp_fill.borrow_mut().extend(items.iter().cloned());
2929            };
2930            let out = if delta_mode {
2931                daemon.generate_streaming_delta(
2932                    session,
2933                    max_tokens,
2934                    extras,
2935                    on_event,
2936                    on_logprobs,
2937                    on_delta,
2938                )?
2939            } else {
2940                daemon.generate_streaming_ex_lp(
2941                    session,
2942                    max_tokens,
2943                    extras,
2944                    on_event,
2945                    on_logprobs,
2946                )?
2947            };
2948            let st = st.borrow();
2949            let (mut decode, stopped) = (st.decode, st.stopped);
2950            let first_tok = decode.first_token();
2951            decode.finish();
2952            decode.add_measured(out.decode);
2953            let (content_acc, content_sent) = (st.content_acc.clone(), st.content_sent);
2954            drop(st);
2955            let mut c = collect(&out.events);
2956            if !stopped && !choice.forces_call() && content_sent < content_acc.len() {
2957                let rest = content_acc[content_sent..].to_string();
2958                let mut ck = Frame::chat(Message::content(rest), None);
2959                if continuous_usage {
2960                    ck.usage = Some(Usage::partial(out.tokens_generated));
2961                }
2962                send(ck)?;
2963            }
2964            if choice.forces_call() {
2965                c.content = content_acc.clone();
2966                promote_forced_tool_call(&daemon, &mut c, session, tool_schemas.as_deref());
2967                if tool_index > 0 {
2968                } else if let Some((id, name, args)) = c.tool_calls.first() {
2969                    send(Frame::chat(
2970                        Message::complete_tool(
2971                            tool_index,
2972                            format!("call_{id}"),
2973                            name.clone(),
2974                            args.clone(),
2975                        ),
2976                        None,
2977                    ))?;
2978                } else if !c.content.is_empty() {
2979                    let mut ck = Frame::chat(Message::content(c.content), None);
2980                    if continuous_usage {
2981                        ck.usage = Some(Usage::partial(out.tokens_generated));
2982                    }
2983                    send(ck)?;
2984                }
2985            }
2986            let fin = if stopped {
2987                "stop"
2988            } else {
2989                finish_reason(
2990                    out.finish,
2991                    out.tokens_generated,
2992                    max_tokens,
2993                    !c.tool_calls.is_empty() && choice != ToolChoice::None,
2994                )
2995            };
2996            let fin = if orphaned_calls > 0 { "length" } else { fin };
2997            let mut terminal = Frame::chat(Message::default(), Some(fin));
2998            terminal.usage = Some(usage(
2999                &daemon,
3000                session,
3001                out.tokens_generated,
3002                out.warm_prefix,
3003            ));
3004            access_log(
3005                &format!("{} (stream)", req.route()),
3006                &model_name,
3007                terminal.usage.as_ref().and_then(|u| u.prompt).unwrap_or(0),
3008                out.warm_prefix,
3009                out.tokens_generated as u64,
3010                GenTiming {
3011                    t0,
3012                    first: first_tok,
3013                    decode,
3014                },
3015                fin,
3016            );
3017            terminal.warm = Some(out.warm_prefix);
3018            terminal.spec = (out.spec.0 > 0).then_some(out.spec);
3019            terminal.session = Some(session);
3020            send(terminal)
3021        };
3022        run()
3023    });
3024
3025    Ok(Body::Stream(rx))
3026}
3027const FIM_DEFAULT_MAX_TOKENS: u32 = 128;
3028
3029#[allow(clippy::too_many_arguments)]
3030pub(crate) async fn fim_completions(
3031    daemon: Arc<Daemon>,
3032    model_name: String,
3033    req: CompletionRequest,
3034    limits: TokenLimits,
3035    sampling: SamplingOverrides,
3036    mode: u8,
3037    key: Option<crate::keypolicy::ClientKey>,
3038) -> Result<Generation, RunError> {
3039    if !daemon.fim_supported() {
3040        return Err(DaemonError::NoFimDialect.into());
3041    }
3042    let hold = key;
3043    let max_tokens = req
3044        .max_tokens
3045        .or(limits.default_max_tokens)
3046        .unwrap_or(FIM_DEFAULT_MAX_TOKENS)
3047        .max(1);
3048    let serving = daemon
3049        .fim_satellite()
3050        .unwrap_or_else(|| Arc::clone(&daemon));
3051    let model_defaults = {
3052        let d = Arc::clone(&serving);
3053        tokio::task::spawn_blocking(move || {
3054            crate::with_sync_class(Some(crate::qos::INLINE_COMPLETION), || {
3055                d.model_sampling_defaults()
3056            })
3057        })
3058        .await
3059        .unwrap_or_default()
3060    };
3061    let op = sampling.for_daemon(&serving);
3062    let mut params = params_of_for(
3063        req.temperature,
3064        req.top_p,
3065        req.top_k,
3066        req.min_p,
3067        req.seed,
3068        &model_defaults,
3069        &op,
3070    );
3071    let seed_drawn = draw_seed_if_unseeded(&mut params, req.seed);
3072    let mut extras = extras_of(
3073        req.presence_penalty,
3074        req.frequency_penalty,
3075        req.repeat_penalty,
3076        &op,
3077        &model_defaults,
3078    );
3079    extras.seed_drawn = seed_drawn;
3080    extras.ignore_eos = req.ignore_eos;
3081    let fim = crate::FimRequest {
3082        prefix: req.prompt,
3083        suffix: req.suffix.unwrap_or_default(),
3084        mode,
3085        max_tokens,
3086        params,
3087        extras,
3088    };
3089    let id = format!(
3090        "cmpl-superfluid-{}",
3091        next_id().rsplit('-').next().unwrap_or("0")
3092    );
3093    let created = unix_now();
3094    let served_by = daemon
3095        .fim_satellite_name()
3096        .unwrap_or_else(|| model_name.clone());
3097    let log_model = served_by.clone();
3098    let fin_of = move |c: &crate::Completion| {
3099        if c.expired {
3100            "length"
3101        } else {
3102            finish_reason(c.finish, c.tokens.len() as u32, max_tokens, false)
3103        }
3104    };
3105    let ext = move |c: &crate::Completion| FimMetadata {
3106        cached: c.cached,
3107        expired: c.expired,
3108        served_by: served_by.clone(),
3109    };
3110    if req.stream {
3111        let admitted_daemon = Arc::clone(&daemon);
3112        let fim = blocking(
3113            RequestQos {
3114                class: crate::qos::INLINE_COMPLETION,
3115                batch_invariant: false,
3116            },
3117            hold.clone(),
3118            move || {
3119                admitted_daemon.fim_fits(&fim)?;
3120                Ok(fim)
3121            },
3122        )
3123        .await?;
3124        let (tx, rx) = tokio::sync::mpsc::channel(64);
3125        tokio::task::spawn_blocking(move || {
3126            let _hold = hold;
3127            let t0 = std::time::Instant::now();
3128            let mut decode = DecodeSpan::default();
3129            let mut delivered = 0usize;
3130            let res = daemon.complete_streaming(&fim, |piece, produced| {
3131                decode.stamp(produced as u64);
3132                delivered = delivered.max(produced);
3133                match tx.blocking_send(Ok(Frame::text(piece, None, None))) {
3134                    Ok(()) => std::ops::ControlFlow::Continue(()),
3135                    Err(_) => std::ops::ControlFlow::Break(()),
3136                }
3137            });
3138            match res {
3139                Ok(c) => {
3140                    let fin = fin_of(&c);
3141                    let first = decode.first_token();
3142                    decode.finish();
3143                    access_log(
3144                        "POST /v1/completions (fim)",
3145                        &log_model,
3146                        c.prompt_tokens as u64,
3147                        fim_log_cached(&c),
3148                        c.tokens.len().max(delivered) as u64,
3149                        GenTiming { t0, first, decode },
3150                        fin,
3151                    );
3152                    let mut last = Frame::text("", None, Some(fin));
3153                    last.fim = Some(ext(&c));
3154                    last.usage = Some(Usage::complete(
3155                        c.prompt_tokens as u64,
3156                        c.tokens.len() as u64,
3157                        None,
3158                    ));
3159                    let _ = tx.blocking_send(Ok(last));
3160                }
3161                Err(e) => {
3162                    let client_error = matches!(
3163                        e,
3164                        DaemonError::StreamTooLong { .. } | DaemonError::NoFimDialect
3165                    );
3166                    let _ = tx.blocking_send(Err(StreamFailure::new(e, client_error)));
3167                }
3168            }
3169        });
3170        return Ok(Generation {
3171            kind: Kind::Completion,
3172            id,
3173            model: model_name,
3174            created,
3175            warm_header: None,
3176            body: Body::Stream(rx),
3177        });
3178    }
3179    let out = tokio::task::spawn_blocking(move || {
3180        let _hold = hold;
3181        let t0 = std::time::Instant::now();
3182        let mut decode = DecodeSpan::default();
3183        let c = daemon.complete_streaming(&fim, |_, produced| {
3184            decode.stamp(produced as u64);
3185            std::ops::ControlFlow::Continue(())
3186        })?;
3187        let first = decode.first_token();
3188        decode.finish();
3189        access_log(
3190            "POST /v1/completions (fim)",
3191            &log_model,
3192            c.prompt_tokens as u64,
3193            fim_log_cached(&c),
3194            c.tokens.len() as u64,
3195            GenTiming { t0, first, decode },
3196            fin_of(&c),
3197        );
3198        Ok::<_, DaemonError>(c)
3199    })
3200    .await;
3201    let c = out.map_err(|e| RunError::Worker(e.to_string()))??;
3202    let mut frame = Frame::text(&c.text, None, Some(fin_of(&c)));
3203    frame.usage = Some(Usage::complete(
3204        c.prompt_tokens as u64,
3205        c.tokens.len() as u64,
3206        None,
3207    ));
3208    frame.fim = Some(ext(&c));
3209    Ok(Generation {
3210        kind: Kind::Completion,
3211        id,
3212        model: model_name,
3213        created,
3214        warm_header: None,
3215        body: Body::Complete(frame),
3216    })
3217}
3218
3219pub(crate) fn chat_once(
3220    daemon: &Daemon,
3221    req: &ChatRequest,
3222    policy: RequestPolicy,
3223) -> Result<Frame, DaemonError> {
3224    let RequestPolicy {
3225        limits,
3226        sampling: op,
3227        qos,
3228    } = policy;
3229    let _class = crate::SyncClassScope::enter(Some(qos.class));
3230    let model_defaults = daemon.model_sampling_defaults();
3231    let mut extras = extras_of(
3232        req.presence_penalty,
3233        req.frequency_penalty,
3234        req.repeat_penalty,
3235        &op,
3236        &model_defaults,
3237    );
3238    let choice = ToolChoice::parse(&req.tool_choice);
3239    if choice.forces_call() {
3240        extras.open_channel = crate::wal::channel::TOOL_CALL;
3241    }
3242    let tool_schemas = crate::codec::ToolSchemas::of_values(&req.tools);
3243    extras.tool_schemas = tool_schemas.clone();
3244    let handles = RequestHandles::new(
3245        daemon,
3246        grammar_handle_for(daemon, req, &choice)?,
3247        logit_bias_handle_of(daemon, &req.logit_bias),
3248    );
3249    handles.apply(&mut extras);
3250    let want_logprobs = req.logprobs.unwrap_or(false);
3251    extras.want_logprobs = want_logprobs;
3252    extras.top_logprobs = req.top_logprobs.unwrap_or(0).min(20) as u8;
3253    extras.thinking = req.thinking();
3254    extras.ignore_eos = req.ignore_eos;
3255    let stops = req.stop.list();
3256    let (session, seed_drawn) = build_session(daemon, req, &op, &model_defaults)?;
3257    extras.seed_drawn = seed_drawn;
3258    qos.apply(daemon, session)?;
3259    let max_tokens = limits.resolve(req.max_completion_tokens.or(req.max_tokens), || {
3260        prompt_len(daemon, session)
3261    });
3262    let cancels = daemon.cancel_registry();
3263    let mut acc = String::new();
3264    let mut stopped = false;
3265    let out = daemon.generate_streaming_ex(session, max_tokens, extras, |ev| {
3266        if let EventBody::Generated {
3267            text, channel: ch, ..
3268        } = &ev.body
3269        {
3270            if *ch == channel::TEXT && !stops.is_empty() && !stopped {
3271                acc.push_str(text);
3272                if stops.iter().any(|s| acc.contains(s.as_str())) {
3273                    cancels.cancel(session);
3274                    stopped = true;
3275                }
3276            }
3277        }
3278        Ok(())
3279    });
3280    let out = out?;
3281    let mut c = collect(&out.events);
3282    claim_named_calls(daemon, &mut c, &out.events, &choice);
3283    if choice.forces_call() {
3284        promote_forced_tool_call(daemon, &mut c, session, tool_schemas.as_deref());
3285    }
3286    let (content, hit_stop) = truncate_at_stop(&c.content, &stops);
3287    c.content = content;
3288    let usage = usage(daemon, session, out.tokens_generated, out.warm_prefix);
3289    let finish = if !c.broken_calls.is_empty() {
3290        "length"
3291    } else if hit_stop {
3292        "stop"
3293    } else {
3294        finish_reason(
3295            out.finish,
3296            out.tokens_generated,
3297            max_tokens,
3298            !c.tool_calls.is_empty() && choice != ToolChoice::None,
3299        )
3300    };
3301    Ok(Frame {
3302        choices: vec![Choice {
3303            message: message(&c, &choice),
3304            finish: Some(finish),
3305            logprobs: want_logprobs.then(|| decode_logprobs(daemon, &out.logprobs, 0)),
3306            ..Choice::default()
3307        }],
3308        usage: Some(usage),
3309        spec: (out.spec.0 > 0).then_some(out.spec),
3310        ..Frame::default()
3311    })
3312}
3313
3314pub(crate) async fn blocking<T: Send + 'static>(
3315    qos: RequestQos,
3316    hold: Option<crate::keypolicy::ClientKey>,
3317    work: impl FnOnce() -> Result<T, DaemonError> + Send + 'static,
3318) -> Result<T, RunError> {
3319    tokio::task::spawn_blocking(move || {
3320        let _hold = hold;
3321        let _class = crate::SyncClassScope::enter(Some(qos.class));
3322        work()
3323    })
3324    .await
3325    .map_err(|e| RunError::Worker(e.to_string()))?
3326    .map_err(RunError::Daemon)
3327}
3328
3329pub(crate) fn base64_decode(data: &str) -> Result<Vec<u8>, DaemonError> {
3330    use base64::Engine;
3331    let cleaned: String = data.chars().filter(|c| !c.is_whitespace()).collect();
3332    base64::engine::general_purpose::STANDARD
3333        .decode(&cleaned)
3334        .or_else(|_| base64::engine::general_purpose::STANDARD_NO_PAD.decode(&cleaned))
3335        .map_err(|_| DaemonError::Protocol("invalid base64 image data"))
3336}
3337
3338pub(crate) fn spawn_stream<T: Send + 'static>(
3339    qos: RequestQos,
3340    hold: Option<crate::keypolicy::ClientKey>,
3341    work: impl FnOnce(&tokio::sync::mpsc::Sender<Result<T, StreamFailure>>) -> Result<(), DaemonError>
3342        + Send
3343        + 'static,
3344) -> tokio::sync::mpsc::Receiver<Result<T, StreamFailure>> {
3345    let (tx, rx) = tokio::sync::mpsc::channel(64);
3346    tokio::task::spawn_blocking(move || {
3347        let _hold = hold;
3348        let _class = crate::SyncClassScope::enter(Some(qos.class));
3349        if let Err(error) = work(&tx) {
3350            let client_error = matches!(
3351                error,
3352                DaemonError::Constraint(_)
3353                    | DaemonError::TemplateRefused(_)
3354                    | DaemonError::Unsupported(_)
3355                    | DaemonError::StreamTooLong { .. }
3356                    | DaemonError::NoFimDialect
3357            );
3358            let _ = tx.blocking_send(Err(StreamFailure::new(error, client_error)));
3359        }
3360    });
3361    rx
3362}
superfluid_daemon/inference.rs

superfluid_daemon/
inference.rs
