1239 lines
48 KiB
Rust
1239 lines
48 KiB
Rust
//! Turn orchestration: reacts to EOU events from the STT worker, schedules the
|
|
//! latency filler, runs the reasoning LLM, fires the TTS request, and handles
|
|
//! barge-in flush.
|
|
//!
|
|
//! Two pipeline shapes are supported:
|
|
//! - **With LLM** (`LLM_ENABLED=true`): transcript → LLM → sentence chunks →
|
|
//! TTS. The assistant text is also forwarded to the client so a chat box can
|
|
//! display it.
|
|
//! - **Without LLM**: transcript → TTS (the original pass-through behaviour).
|
|
|
|
use crate::audio::{AudioMixer, SampleRingBuffer};
|
|
use crate::audit::SessionAudit;
|
|
use crate::config::Config;
|
|
use crate::lang::{self, LangBanks, LanguageState};
|
|
use crate::llm::{ChatMessage, LlmClient};
|
|
use crate::prompt::SessionPrompts;
|
|
use crate::stt::{Tokenizer, Transcript};
|
|
use crate::tools::{ChatMarkdownStripper, ToolCallStripper};
|
|
use crate::tts::TtsEngine;
|
|
use crate::{assets, orchestrator::backchannel::DEFAULT_CATEGORY};
|
|
use serde_json::json;
|
|
use std::collections::HashMap;
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::sync::mpsc;
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::{debug, error, info, warn};
|
|
|
|
/// Capacity (samples) of the per-request staging buffer used by the TTS
|
|
/// pipeline — 30 s at the mixer sample rate. Large enough to hold a whole
|
|
/// chunk even when a later request finishes synthesizing while an earlier one
|
|
/// is still playing out.
|
|
const TTS_SLOT_CAPACITY_SAMPLES: usize = crate::config::OUTPUT_SAMPLE_RATE as usize * 30;
|
|
|
|
/// One in-flight synthesis request in the TTS pipeline.
|
|
struct TtsSlot {
|
|
/// Staging buffer the request streams its PCM into. The relay forwards it
|
|
/// to the mixer queue in submission order.
|
|
buffer: SampleRingBuffer,
|
|
/// Task running the synthesis; finished means the chunk is fully staged.
|
|
handle: tokio::task::JoinHandle<()>,
|
|
}
|
|
|
|
/// Move every sample currently staged in `src` into `dst`, respecting `dst`'s
|
|
/// free space so already-queued (i.e. currently playing) audio is never
|
|
/// dropped. Returns the number of samples moved.
|
|
fn pump(src: &SampleRingBuffer, dst: &SampleRingBuffer) -> usize {
|
|
let mut tmp = [0i16; 8192];
|
|
let mut moved = 0;
|
|
loop {
|
|
let want = dst.remaining_capacity().min(src.len()).min(tmp.len());
|
|
if want == 0 {
|
|
break;
|
|
}
|
|
let n = src.pop_into(&mut tmp[..want]);
|
|
if n == 0 {
|
|
break;
|
|
}
|
|
dst.push(&tmp[..n]);
|
|
moved += n;
|
|
}
|
|
moved
|
|
}
|
|
|
|
/// Preloaded PCM samples keyed by name (e.g. "uhm"). Used for the filler bank
|
|
/// injected at end-of-utterance (played round-robin), for the backchannel clips
|
|
/// played while the user speaks, and for the start-filler bank (both played in
|
|
/// random order).
|
|
pub struct FillerBank {
|
|
samples: HashMap<String, Vec<i16>>,
|
|
/// Names in a stable order, for round-robin selection.
|
|
names: Vec<String>,
|
|
/// xorshift64 state for [`Self::random`] (seeded from the clock at load).
|
|
rng: AtomicU64,
|
|
}
|
|
|
|
impl FillerBank {
|
|
/// Load every `*.raw` file in `dir` as raw 16-bit LE mono PCM. `label` is
|
|
/// only used in log messages (e.g. "filler", "backchannel").
|
|
pub fn load(dir: &std::path::Path, label: &str) -> Self {
|
|
let mut samples = HashMap::new();
|
|
if let Ok(entries) = std::fs::read_dir(dir) {
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
if path.extension().and_then(|e| e.to_str()) != Some("raw") {
|
|
continue;
|
|
}
|
|
let name = path
|
|
.file_stem()
|
|
.and_then(|n| n.to_str())
|
|
.unwrap_or("unknown")
|
|
.to_string();
|
|
match std::fs::read(&path) {
|
|
Ok(bytes) => {
|
|
let pcm = bytes
|
|
.chunks_exact(2)
|
|
.map(|c| i16::from_le_bytes([c[0], c[1]]))
|
|
.collect::<Vec<i16>>();
|
|
info!(label, name, samples = pcm.len(), "loaded {label} clip");
|
|
samples.insert(name, pcm);
|
|
}
|
|
Err(e) => {
|
|
warn!(label, path = %path.display(), error = %e, "failed to read {label} clip")
|
|
}
|
|
}
|
|
}
|
|
if samples.is_empty() {
|
|
warn!(label, dir = %dir.display(), "no {label} clips found");
|
|
}
|
|
} else {
|
|
warn!(label, dir = %dir.display(), "{label} directory not found");
|
|
}
|
|
|
|
let mut names: Vec<String> = samples.keys().cloned().collect();
|
|
names.sort();
|
|
|
|
// Seed the xorshift64 state from the wall clock (OR 1 so it is never
|
|
// zero, which would make xorshift degenerate).
|
|
let seed = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.map(|d| d.as_nanos() as u64)
|
|
.unwrap_or(0x9E37_79B9_7F4A_7C15)
|
|
| 1;
|
|
|
|
Self {
|
|
samples,
|
|
names,
|
|
rng: AtomicU64::new(seed),
|
|
}
|
|
}
|
|
|
|
/// Get the samples for a clip by name. Used by the backchannel to play a
|
|
/// specific (category-selected) clip instead of a uniform random one.
|
|
pub fn get(&self, name: &str) -> Option<&Vec<i16>> {
|
|
self.samples.get(name)
|
|
}
|
|
|
|
/// Clip names in the stable (sorted) order used by [`Self::nth`]. The
|
|
/// backchannel uses it to map the phrase list (and its categories) onto the
|
|
/// synthesized clips, since clips are named `{prefix}{index}` in list order.
|
|
pub fn names(&self) -> &[String] {
|
|
&self.names
|
|
}
|
|
|
|
/// Clip at position `index` (wrapping), with its name. Used for the
|
|
/// round-robin end-of-utterance filler.
|
|
pub fn nth(&self, index: usize) -> Option<(&str, &Vec<i16>)> {
|
|
if self.names.is_empty() {
|
|
return None;
|
|
}
|
|
let name = &self.names[index % self.names.len()];
|
|
self.samples.get(name).map(|s| (name.as_str(), s))
|
|
}
|
|
|
|
/// Clip chosen uniformly at random, with its name. Used by the start filler
|
|
/// and the backchannel so the cue varies from turn to turn.
|
|
pub fn random(&self) -> Option<(&str, &Vec<i16>)> {
|
|
if self.names.is_empty() {
|
|
return None;
|
|
}
|
|
// xorshift64 step over the shared state.
|
|
let mut x = self.rng.load(Ordering::Relaxed);
|
|
x ^= x << 13;
|
|
x ^= x >> 7;
|
|
x ^= x << 17;
|
|
self.rng.store(x, Ordering::Relaxed);
|
|
|
|
let name = &self.names[(x % self.names.len() as u64) as usize];
|
|
self.samples.get(name).map(|s| (name.as_str(), s))
|
|
}
|
|
|
|
/// Whether the bank holds no clips.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.names.is_empty()
|
|
}
|
|
}
|
|
|
|
/// Orchestrates a single conversation turn.
|
|
pub struct TurnManager {
|
|
mixer: AudioMixer,
|
|
tts: Arc<dyn TtsEngine>,
|
|
/// Active-language clip banks, shared with the mic task (backchannel). The
|
|
/// end-of-turn and start fillers are read from here so all three cues
|
|
/// follow the same language — see [`crate::lang`].
|
|
langs: Arc<LanguageState>,
|
|
/// Optional tokenizer to decode STT token IDs into text. If `None`, a
|
|
/// placeholder string is sent to the TTS (useful for testing the pipeline
|
|
/// without a real tokenizer).
|
|
tokenizer: Option<Tokenizer>,
|
|
/// Cancellation token for the currently running TTS request.
|
|
cancel: Arc<parking_lot::Mutex<Option<CancellationToken>>>,
|
|
/// Round-robin cursor over the filler bank (end-of-utterance fillers).
|
|
filler_cursor: AtomicUsize,
|
|
/// Reasoning LLM client. `None` => direct STT → TTS pipeline.
|
|
llm: Option<Arc<LlmClient>>,
|
|
/// Server configuration (LLM prompt / streaming / history window).
|
|
config: Arc<Config>,
|
|
/// Rolling conversation history (user + assistant turns).
|
|
history: parking_lot::Mutex<Vec<ChatMessage>>,
|
|
/// Outbound text-event channel (JSON strings) consumed by the session
|
|
/// sender task and forwarded to the browser (chat box).
|
|
out_tx: Option<mpsc::UnboundedSender<String>>,
|
|
/// User turns seen so far, used to schedule the periodic language check.
|
|
lang_turns: AtomicU64,
|
|
/// Whether a language-detection call is currently in flight (one at a time).
|
|
lang_detecting: Arc<AtomicBool>,
|
|
/// Per-session audit sink: the language switches are recorded here (every
|
|
/// other event reaches the log through the outbound tap in the session
|
|
/// sender — see [`crate::audit`]).
|
|
audit: Arc<SessionAudit>,
|
|
/// Per-session system-prompt registry: the effective prompt can be
|
|
/// overridden at runtime (see [`crate::prompt_api`]).
|
|
prompts: Arc<SessionPrompts>,
|
|
/// This session's id, the key into [`Self::prompts`].
|
|
session_id: String,
|
|
}
|
|
|
|
impl TurnManager {
|
|
pub fn new(
|
|
mixer: AudioMixer,
|
|
tts: Arc<dyn TtsEngine>,
|
|
langs: Arc<LanguageState>,
|
|
tokenizer: Option<Tokenizer>,
|
|
llm: Option<Arc<LlmClient>>,
|
|
config: Arc<Config>,
|
|
out_tx: Option<mpsc::UnboundedSender<String>>,
|
|
audit: Arc<SessionAudit>,
|
|
prompts: Arc<SessionPrompts>,
|
|
session_id: String,
|
|
) -> Self {
|
|
Self {
|
|
mixer,
|
|
tts,
|
|
langs,
|
|
tokenizer,
|
|
cancel: Arc::new(parking_lot::Mutex::new(None)),
|
|
filler_cursor: AtomicUsize::new(0),
|
|
llm,
|
|
config,
|
|
history: parking_lot::Mutex::new(Vec::new()),
|
|
out_tx,
|
|
lang_turns: AtomicU64::new(0),
|
|
lang_detecting: Arc::new(AtomicBool::new(false)),
|
|
audit,
|
|
prompts,
|
|
session_id,
|
|
}
|
|
}
|
|
|
|
/// Pick the next filler clip round-robin from the bank of the active
|
|
/// language.
|
|
fn next_filler(&self) -> Option<(String, Vec<i16>)> {
|
|
let banks = self.langs.current();
|
|
let i = self.filler_cursor.fetch_add(1, Ordering::Relaxed);
|
|
banks
|
|
.filler
|
|
.nth(i)
|
|
.map(|(name, s)| (name.to_string(), s.clone()))
|
|
}
|
|
|
|
/// Schedule the *start filler* acknowledgment on the mixer's dedicated
|
|
/// start-filler channel.
|
|
///
|
|
/// Like [`Self::schedule_filler`] this is a latency mask, but with its own
|
|
/// (shorter) delay, `START_FILLER_GO_DELAY`: the cue is only queued when the
|
|
/// first response audio is still **late** after that window — if the first
|
|
/// audio chunk has already reached the mixer (even if it is still inside the
|
|
/// pre-buffer and playback has not started), the agent is already talking
|
|
/// and the cue is dropped. One phrase is picked **at random** from the bank
|
|
/// (built from `START_FILLER_TEXTS`), so the cue varies turn to turn.
|
|
/// Disabled by `START_FILLER_ENABLED=false`; a barge-in (`cancel`) or a new
|
|
/// turn drops the pending cue, and `mixer.flush` clears it once queued.
|
|
fn schedule_start_filler(&self, cancel: &CancellationToken) {
|
|
if !self.config.start_filler_enabled {
|
|
return;
|
|
}
|
|
// Own the clip: the spawned task must be `'static` (same as the
|
|
// round-robin filler above).
|
|
let banks = self.langs.current();
|
|
let Some((name, samples)) = banks
|
|
.start_filler
|
|
.random()
|
|
.map(|(n, s)| (n.to_string(), s.clone()))
|
|
else {
|
|
debug!("start filler enabled but no clip loaded — skipping");
|
|
return;
|
|
};
|
|
|
|
let delay_ms = self.config.start_filler_go_delay_ms;
|
|
let mixer = self.mixer.clone();
|
|
let cancel = cancel.clone();
|
|
let out_tx = self.out_tx.clone();
|
|
tokio::spawn(async move {
|
|
tokio::select! {
|
|
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
|
|
// Barge-in / new turn: never play a stale acknowledgment.
|
|
_ = cancel.cancelled() => return,
|
|
}
|
|
if mixer.tts_audio_arrived() {
|
|
debug!(
|
|
filler = %name,
|
|
delay_ms, "first response audio already arrived — start filler skipped"
|
|
);
|
|
return;
|
|
}
|
|
let queued = mixer.start_filler_queue.push(&samples);
|
|
debug!(filler = %name, samples = queued, delay_ms, "queued start filler");
|
|
if let Some(tx) = &out_tx {
|
|
let _ = tx.send(
|
|
json!({ "type": "cue", "kind": "start_filler", "name": name }).to_string(),
|
|
);
|
|
}
|
|
});
|
|
}
|
|
|
|
/// Schedule the single end-of-turn filler acknowledgment.
|
|
///
|
|
/// The filler is a latency mask: it buys time while the reply is being
|
|
/// generated, so it is only useful when the response audio is *late*. It is
|
|
/// therefore not played unconditionally — if the reply starts playing
|
|
/// within `FILLER_DELAY_MS`, the filler is dropped instead of speaking over
|
|
/// the beginning of the answer. At most one clip is emitted per turn and a
|
|
/// barge-in (`cancel`) drops the pending one. Disabled entirely by
|
|
/// `FILLER_ENABLED=false`.
|
|
fn schedule_filler(&self, cancel: &CancellationToken) {
|
|
if !self.config.filler_enabled {
|
|
return;
|
|
}
|
|
let Some((name, samples)) = self.next_filler() else {
|
|
debug!("filler enabled but the bank is empty — skipping");
|
|
return;
|
|
};
|
|
|
|
let delay_ms = self.config.filler_delay_ms;
|
|
let mixer = self.mixer.clone();
|
|
let cancel = cancel.clone();
|
|
let out_tx = self.out_tx.clone();
|
|
let samples_len = samples.len();
|
|
tokio::spawn(async move {
|
|
tokio::select! {
|
|
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
|
|
// Barge-in / new turn: never play a stale acknowledgment.
|
|
_ = cancel.cancelled() => return,
|
|
}
|
|
if mixer.tts_playback_started() {
|
|
debug!(
|
|
filler = %name,
|
|
delay_ms, "response audio already playing — filler skipped"
|
|
);
|
|
return;
|
|
}
|
|
mixer.filler_queue.push(&samples);
|
|
debug!(filler = %name, samples = samples_len, delay_ms, "pushed filler (delayed)");
|
|
if let Some(tx) = &out_tx {
|
|
let _ =
|
|
tx.send(json!({ "type": "cue", "kind": "filler", "name": name }).to_string());
|
|
}
|
|
});
|
|
}
|
|
|
|
/// Send a JSON event to the client, if a channel is attached.
|
|
fn emit(&self, value: serde_json::Value) {
|
|
if let Some(tx) = &self.out_tx {
|
|
let _ = tx.send(value.to_string());
|
|
}
|
|
}
|
|
|
|
/// React to an EOU / transcript event from the STT worker.
|
|
pub async fn on_eou(&self, transcript: &Transcript) {
|
|
info!(
|
|
tokens = transcript.token_ids.len(),
|
|
duration_ms = transcript.duration_ms,
|
|
text = %transcript.text,
|
|
"EOU detected"
|
|
);
|
|
|
|
// The Parakeet pipeline decodes text internally; fall back to the
|
|
// external tokenizer only if no text was produced.
|
|
let text = if !transcript.text.trim().is_empty() {
|
|
transcript.text.trim().to_string()
|
|
} else if transcript.token_ids.is_empty() {
|
|
// No tokens at all: nothing was said (a VAD false-trigger that
|
|
// slipped past the gate). Skip the turn instead of forwarding the
|
|
// "<transcript 0 tokens>" placeholder to the LLM.
|
|
debug!("no tokens decoded, skipping turn");
|
|
return;
|
|
} else {
|
|
match &self.tokenizer {
|
|
Some(tok) => tok.decode(&transcript.token_ids),
|
|
None => format!("<transcript {} tokens>", transcript.token_ids.len()),
|
|
}
|
|
};
|
|
if text.trim().is_empty() {
|
|
debug!("empty transcript after decode, skipping turn");
|
|
return;
|
|
}
|
|
self.handle_user_turn(text).await;
|
|
}
|
|
|
|
/// Handle a user turn typed in the chat box (no STT involved).
|
|
pub async fn on_user_text(&self, text: &str) {
|
|
let text = text.trim();
|
|
if text.is_empty() {
|
|
return;
|
|
}
|
|
self.handle_user_turn(text.to_string()).await;
|
|
}
|
|
|
|
/// Shared entry point for every user turn, whether spoken or typed.
|
|
///
|
|
/// 1. Forward the recognized turn to the client chat box.
|
|
/// 2. Cancel any in-flight reply (barge-in within the turn).
|
|
/// 3. Run the LLM (if configured) or send the text straight to the TTS.
|
|
///
|
|
/// The latency masks are not queued here: both the start filler and the
|
|
/// end-of-turn filler are *scheduled* once the response is dispatched, so
|
|
/// they only play if the audio is late — see [`Self::schedule_filler`] and
|
|
/// [`Self::schedule_start_filler`].
|
|
async fn handle_user_turn(&self, text: String) {
|
|
self.emit(json!({ "type": "user", "text": text }));
|
|
|
|
self.maybe_detect_language(&text);
|
|
self.cancel_current_tts();
|
|
|
|
match &self.llm {
|
|
Some(llm) => self.run_llm_turn(llm.clone(), text).await,
|
|
None => self.dispatch_tts(&text),
|
|
}
|
|
}
|
|
|
|
/// Kick off a language check for `text` when the detection schedule allows
|
|
/// it. Runs in the background so the reply is never delayed: the switch (if
|
|
/// any) takes effect from the next turn. See [`crate::lang`].
|
|
///
|
|
/// The first turn is always checked; afterwards a sample is sent only every
|
|
/// `LANG_REDETECT_EVERY` turns, so a language change mid-conversation is
|
|
/// picked up without paying for a detection call on every turn.
|
|
fn maybe_detect_language(&self, text: &str) {
|
|
if !self.config.lang_detect_enabled {
|
|
return;
|
|
}
|
|
let Some(llm) = self.llm.clone() else {
|
|
return;
|
|
};
|
|
// Only the backchannel, the end-of-turn filler and the start filler are
|
|
// language-aware; with all of them disabled there is nothing to switch.
|
|
if !self.config.backchannel.enabled
|
|
&& !self.config.filler_enabled
|
|
&& !self.config.start_filler_enabled
|
|
{
|
|
return;
|
|
}
|
|
if text.chars().count() < self.config.lang_min_chars {
|
|
return;
|
|
}
|
|
let turn = self.lang_turns.fetch_add(1, Ordering::Relaxed);
|
|
if turn % self.config.lang_redetect_every != 0 {
|
|
return;
|
|
}
|
|
// One detection at a time: a slow LLM call must not pile up.
|
|
if self.lang_detecting.swap(true, Ordering::AcqRel) {
|
|
return;
|
|
}
|
|
let langs = self.langs.clone();
|
|
let tts = self.tts.clone();
|
|
let config = self.config.clone();
|
|
let detecting = self.lang_detecting.clone();
|
|
let audit = self.audit.clone();
|
|
let text = text.to_string();
|
|
tokio::spawn(async move {
|
|
switch_language_if_needed(langs, llm, tts, config, &text, audit).await;
|
|
detecting.store(false, Ordering::Release);
|
|
});
|
|
}
|
|
|
|
/// Run one full LLM turn: stream the reply, chunk it into sentences and
|
|
/// synthesize each sentence as soon as it is complete.
|
|
async fn run_llm_turn(&self, llm: Arc<LlmClient>, user_text: String) {
|
|
// Build the message list from the rolling history. It grows across
|
|
// tool rounds (assistant call + tool results) but only the *spoken*
|
|
// text is written back to the rolling history.
|
|
let mut messages = {
|
|
let mut hist = self.history.lock();
|
|
hist.push(ChatMessage::user(&user_text));
|
|
let max = self.config.llm_history_messages.max(2);
|
|
if hist.len() > max {
|
|
let drop = hist.len() - max;
|
|
hist.drain(..drop);
|
|
}
|
|
let mut msgs = Vec::with_capacity(hist.len() + 1);
|
|
// Effective prompt: per-session override, else process-wide override,
|
|
// else the configured default (`LLM_SYSTEM_PROMPT`).
|
|
let effective = self.prompts.resolve(
|
|
Some(&self.session_id),
|
|
&self.config.llm_system_prompt,
|
|
self.config.llm_system_prompt_from_env,
|
|
);
|
|
let system = effective.prompt.trim();
|
|
if !system.is_empty() {
|
|
msgs.push(ChatMessage::system(system));
|
|
}
|
|
msgs.extend(hist.iter().cloned());
|
|
msgs
|
|
};
|
|
|
|
let cancel = CancellationToken::new();
|
|
*self.cancel.lock() = Some(cancel.clone());
|
|
self.mixer.mark_stream_start();
|
|
self.schedule_filler(&cancel);
|
|
self.schedule_start_filler(&cancel);
|
|
debug!(
|
|
endpoint = llm.endpoint(),
|
|
model = llm.model(),
|
|
messages = messages.len(),
|
|
stream = self.config.llm_stream,
|
|
"dispatching LLM turn"
|
|
);
|
|
|
|
// TTS sink: a *pipelined* synthesizer. Each chunk is requested as soon
|
|
// as the LLM produces it instead of waiting for the previous chunk's
|
|
// stream to end, so the per-request start-up latency overlaps the
|
|
// previous chunk's playback. Requests write into their own staging
|
|
// buffer; a relay forwards those buffers to the mixer queue strictly in
|
|
// submission order, which keeps playback in order and gapless.
|
|
let (sentence_tx, mut sentence_rx) = mpsc::unbounded_channel::<String>();
|
|
let tts = self.tts.clone();
|
|
let tts_queue = self.mixer.tts_queue.clone();
|
|
let cancel_tts = cancel.clone();
|
|
let mixer_for_tts = self.mixer.clone();
|
|
let pipeline_depth = self.config.tts_pipeline_depth;
|
|
let tts_task = tokio::spawn(async move {
|
|
let (slot_tx, mut slot_rx) = mpsc::unbounded_channel::<TtsSlot>();
|
|
let permits = Arc::new(tokio::sync::Semaphore::new(pipeline_depth));
|
|
|
|
// Collector: start one synthesis task per chunk immediately, so its
|
|
// start-up latency runs concurrently with the chunks already
|
|
// playing. `permits` caps how many requests are in flight at once.
|
|
let collector = {
|
|
let tts = tts.clone();
|
|
let cancel = cancel_tts.clone();
|
|
let permits = permits.clone();
|
|
tokio::spawn(async move {
|
|
while let Some(sentence) = sentence_rx.recv().await {
|
|
if cancel.is_cancelled() {
|
|
break;
|
|
}
|
|
let buffer = SampleRingBuffer::new(TTS_SLOT_CAPACITY_SAMPLES);
|
|
let buf_task = buffer.clone();
|
|
let permits = permits.clone();
|
|
let cancel_task = cancel.clone();
|
|
let tts_task_engine = tts.clone();
|
|
let handle = tokio::spawn(async move {
|
|
let _permit = tokio::select! {
|
|
p = permits.acquire_owned() => match p {
|
|
Ok(p) => p,
|
|
Err(_) => return,
|
|
},
|
|
_ = cancel_task.cancelled() => return,
|
|
};
|
|
if cancel_task.is_cancelled() {
|
|
return;
|
|
}
|
|
if let Err(e) = tts_task_engine
|
|
.synthesize(&sentence, &buf_task, &cancel_task)
|
|
.await
|
|
{
|
|
error!(error = %e, "TTS synthesis failed (pipelined)");
|
|
}
|
|
});
|
|
if slot_tx.send(TtsSlot { buffer, handle }).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
})
|
|
};
|
|
|
|
// Relay: forward each chunk's staged audio to the mixer queue in
|
|
// order, waiting for the current chunk to finish before advancing.
|
|
while let Some(slot) = slot_rx.recv().await {
|
|
let started = std::time::Instant::now();
|
|
debug!(
|
|
staged = slot.buffer.len(),
|
|
finished = slot.handle.is_finished(),
|
|
"TTS pipeline: relaying chunk to mixer"
|
|
);
|
|
let mut relayed = 0usize;
|
|
loop {
|
|
relayed += pump(&slot.buffer, &tts_queue);
|
|
if cancel_tts.is_cancelled() {
|
|
break;
|
|
}
|
|
if slot.handle.is_finished() && slot.buffer.is_empty() {
|
|
debug!(
|
|
relayed,
|
|
elapsed_ms = started.elapsed().as_millis(),
|
|
queue_len = tts_queue.len(),
|
|
"TTS pipeline chunk fully relayed to the mixer"
|
|
);
|
|
break;
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(2)).await;
|
|
}
|
|
if cancel_tts.is_cancelled() {
|
|
break;
|
|
}
|
|
}
|
|
let _ = collector.await;
|
|
if !cancel_tts.is_cancelled() {
|
|
mixer_for_tts.mark_stream_end();
|
|
}
|
|
});
|
|
|
|
// Tool-calling loop. Each round streams one reply; tool call blocks are
|
|
// stripped from the spoken/displayed text (the model emits them as
|
|
// text — see `crate::tools`), the calls run over HTTP, and the results
|
|
// are handed back for the next round. A round that streams back empty
|
|
// is re-asked without streaming, where the server returns the call it
|
|
// dropped from the SSE path. Bounded by `LLM_TOOL_MAX_ITERATIONS`. With
|
|
// no tools configured the loop runs exactly once.
|
|
// Re-read per turn: a runtime override (`POST /api/tools`) is picked up
|
|
// by the next turn without a restart.
|
|
let registry = llm.tools();
|
|
let max_rounds = self.config.llm_tool_max_iterations.max(1);
|
|
let mut shown = String::new();
|
|
let mut round = 0usize;
|
|
|
|
loop {
|
|
// LLM producer: streams clean text deltas into a channel.
|
|
let (delta_tx, mut delta_rx) = mpsc::unbounded_channel::<String>();
|
|
let stream_mode = self.config.llm_stream;
|
|
let llm_cancel = cancel.clone();
|
|
let llm_round = llm.clone();
|
|
let round_messages = messages.clone();
|
|
let llm_task = tokio::spawn(async move {
|
|
if stream_mode {
|
|
llm_round
|
|
.chat_stream(&round_messages, &llm_cancel, delta_tx)
|
|
.await
|
|
} else {
|
|
let res = llm_round.chat(&round_messages, &llm_cancel).await;
|
|
if let Ok(text) = &res {
|
|
let _ = delta_tx.send(text.clone());
|
|
}
|
|
res
|
|
}
|
|
});
|
|
|
|
// Consume deltas: emit `assistant_delta` events and feed the
|
|
// splitter, so the TTS starts as soon as the first few words are
|
|
// available. The stripper holds back tool-call blocks so they are
|
|
// never spoken nor shown.
|
|
let mut splitter = SentenceSplitter::new(
|
|
self.config.llm_stream_first_chunk_chars,
|
|
self.config.llm_stream_chunk_chars,
|
|
);
|
|
let mut stripper = ToolCallStripper::new();
|
|
let mut markdown = ChatMarkdownStripper::new();
|
|
let mut raw = String::new();
|
|
while let Some(delta) = delta_rx.recv().await {
|
|
if cancel.is_cancelled() {
|
|
break;
|
|
}
|
|
raw.push_str(&delta);
|
|
let visible = stripper.push(&delta);
|
|
let (spoken_delta, markdown_docs) = markdown.push(&visible);
|
|
for doc in markdown_docs {
|
|
if !shown.trim().is_empty() {
|
|
shown.push_str("\n\n");
|
|
}
|
|
shown.push_str(&doc);
|
|
self.emit(json!({ "type": "assistant_markdown", "text": doc }));
|
|
}
|
|
if spoken_delta.is_empty() {
|
|
continue;
|
|
}
|
|
self.emit(json!({ "type": "assistant_delta", "text": spoken_delta }));
|
|
shown.push_str(&spoken_delta);
|
|
for sentence in splitter.push(&spoken_delta) {
|
|
if sentence_tx.send(sentence).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if !cancel.is_cancelled() {
|
|
let tail = stripper.finish();
|
|
let (spoken_tail, markdown_docs) = markdown.push(&tail);
|
|
for doc in markdown_docs {
|
|
if !shown.trim().is_empty() {
|
|
shown.push_str("\n\n");
|
|
}
|
|
shown.push_str(&doc);
|
|
self.emit(json!({ "type": "assistant_markdown", "text": doc }));
|
|
}
|
|
let (spoken_rest, markdown_rest) = markdown.finish();
|
|
for doc in markdown_rest {
|
|
if !shown.trim().is_empty() {
|
|
shown.push_str("\n\n");
|
|
}
|
|
shown.push_str(&doc);
|
|
self.emit(json!({ "type": "assistant_markdown", "text": doc }));
|
|
}
|
|
let tail = format!("{spoken_tail}{spoken_rest}");
|
|
if !tail.is_empty() {
|
|
self.emit(json!({ "type": "assistant_delta", "text": tail }));
|
|
for sentence in splitter.push(&tail) {
|
|
if sentence_tx.send(sentence).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
shown.push_str(&tail);
|
|
}
|
|
if let Some(rest) = splitter.finish() {
|
|
let _ = sentence_tx.send(rest);
|
|
}
|
|
}
|
|
|
|
match llm_task.await {
|
|
Ok(Ok(_)) => {}
|
|
Ok(Err(e)) => {
|
|
error!(error = %e, "LLM request failed");
|
|
self.emit(json!({ "type": "error", "text": e.to_string() }));
|
|
}
|
|
Err(e) => {
|
|
error!(error = %e, "LLM task join failed");
|
|
self.emit(json!({ "type": "error", "text": format!("LLM task failed: {e}") }));
|
|
}
|
|
}
|
|
|
|
if cancel.is_cancelled() {
|
|
break;
|
|
}
|
|
|
|
let mut calls = stripper.take_calls();
|
|
|
|
// Streaming rescue. Some OpenAI-compatible servers swallow the
|
|
// model's native tool call on the SSE path — vLLM with a
|
|
// `--tool-call-parser` that does not match the model's format ends
|
|
// the stream with no content at all — while the same request
|
|
// *without* streaming returns the call intact. When a round
|
|
// produced nothing and tools are configured, re-ask once,
|
|
// non-streaming, and take the call from there.
|
|
if calls.is_empty()
|
|
&& registry.is_some()
|
|
&& !cancel.is_cancelled()
|
|
&& shown.trim().is_empty()
|
|
&& round + 1 < max_rounds
|
|
{
|
|
match llm.chat_reply(&messages, &cancel).await {
|
|
Ok(reply) => {
|
|
calls = reply.tool_calls;
|
|
if calls.is_empty() && !reply.text.trim().is_empty() {
|
|
// No call after all: the blocking reply holds the
|
|
// actual answer. Split out any chat-only markdown
|
|
// so long documents are shown in chat but not read aloud.
|
|
let mut markdown = ChatMarkdownStripper::new();
|
|
let (spoken_reply, markdown_docs) = markdown.push(&reply.text);
|
|
let (spoken_tail, markdown_rest) = markdown.finish();
|
|
for doc in markdown_docs.into_iter().chain(markdown_rest.into_iter()) {
|
|
if !shown.trim().is_empty() {
|
|
shown.push_str("\n\n");
|
|
}
|
|
shown.push_str(&doc);
|
|
self.emit(json!({ "type": "assistant_markdown", "text": doc }));
|
|
}
|
|
let spoken_reply = format!("{spoken_reply}{spoken_tail}");
|
|
if !spoken_reply.trim().is_empty() {
|
|
self.emit(
|
|
json!({ "type": "assistant_delta", "text": &spoken_reply }),
|
|
);
|
|
shown.push_str(&spoken_reply);
|
|
let _ = sentence_tx.send(spoken_reply.trim().to_string());
|
|
}
|
|
}
|
|
}
|
|
Err(e) => error!(error = %e, "tool-call fallback request failed"),
|
|
}
|
|
}
|
|
|
|
if calls.is_empty() {
|
|
break;
|
|
}
|
|
let Some(registry) = registry.as_ref() else {
|
|
warn!("model emitted a tool call but no tools are configured — ignoring");
|
|
break;
|
|
};
|
|
if round + 1 >= max_rounds {
|
|
warn!(max_rounds, "tool-call iteration limit reached — stopping");
|
|
break;
|
|
}
|
|
round += 1;
|
|
|
|
// Record the model's own call followed by each result, so the next
|
|
// round sees the exchange (standard OpenAI tool round-trip).
|
|
messages.push(ChatMessage::assistant(raw.clone()));
|
|
for call in &calls {
|
|
info!(tool = %call.name, round, "executing tool call");
|
|
self.emit(json!({
|
|
"type": "tool_call",
|
|
"name": &call.name,
|
|
"arguments": &call.arguments,
|
|
}));
|
|
let result = registry.execute(call).await;
|
|
messages.push(ChatMessage::tool(result));
|
|
}
|
|
}
|
|
|
|
drop(sentence_tx);
|
|
let _ = tts_task.await;
|
|
|
|
if !cancel.is_cancelled() && !shown.trim().is_empty() {
|
|
self.history
|
|
.lock()
|
|
.push(ChatMessage::assistant(shown.clone()));
|
|
self.emit(json!({ "type": "assistant_end" }));
|
|
}
|
|
|
|
let mut guard = self.cancel.lock();
|
|
if guard.as_ref().map(|c| c.is_cancelled()).unwrap_or(false) {
|
|
*guard = None;
|
|
}
|
|
}
|
|
|
|
/// Handle an explicit client-side interrupt (barge-in).
|
|
/// Flushes all queued audio and cancels the in-flight LLM/TTS request.
|
|
pub fn interrupt(&self) {
|
|
info!("interrupt received, flushing audio");
|
|
self.cancel_current_tts();
|
|
self.mixer.flush();
|
|
}
|
|
|
|
/// Whether the agent is currently speaking (TTS in flight or queued audio).
|
|
/// Used by the backchannel gate to avoid cuing on top of the agent's own
|
|
/// TTS — see [`Self::is_audible`] for the barge-in gate.
|
|
pub fn is_speaking(&self) -> bool {
|
|
self.mixer.is_active()
|
|
}
|
|
|
|
/// Whether the agent has *audible* output (reply playing out, or a
|
|
/// latency-mask cue queued). Unlike [`Self::is_speaking`] this is `false`
|
|
/// while the reply is still being synthesised or sits in the pre-buffer.
|
|
/// Used by the VAD-driven barge-in: a transient must not cancel a reply
|
|
/// before the user has heard any of it.
|
|
pub fn is_audible(&self) -> bool {
|
|
self.mixer.is_audible()
|
|
}
|
|
|
|
/// Update the voice-cloning reference at runtime. Pass `None` for
|
|
/// `audio_path` to clear the reference.
|
|
pub fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
|
|
self.tts.set_reference(audio_path, text);
|
|
}
|
|
|
|
/// Set the gain of the mixer's external audio channel (URL / radio /
|
|
/// YouTube background source). See [`crate::media`].
|
|
pub fn set_external_gain(&self, gain: f32) {
|
|
self.mixer.set_external_gain(gain);
|
|
}
|
|
|
|
/// Synthesize `text` directly with the TTS engine, bypassing the LLM.
|
|
/// Used by the playground's "Speak" box to test the TTS pipeline.
|
|
pub async fn speak(&self, text: &str) {
|
|
let text = text.trim();
|
|
if text.is_empty() {
|
|
return;
|
|
}
|
|
info!(text_len = text.len(), "speak request (direct TTS)");
|
|
|
|
// Cancel any in-flight TTS and clear the mixer queues so the new
|
|
// utterance starts cleanly (no overlap with the previous one).
|
|
self.cancel_current_tts();
|
|
self.mixer.flush();
|
|
|
|
self.dispatch_tts(text);
|
|
}
|
|
|
|
/// Fire one TTS request for `text` into the mixer's TTS queue.
|
|
fn dispatch_tts(&self, text: &str) {
|
|
let cancel = CancellationToken::new();
|
|
*self.cancel.lock() = Some(cancel.clone());
|
|
|
|
let tts = self.tts.clone();
|
|
let queue = self.mixer.tts_queue.clone();
|
|
let cancel_holder = self.cancel.clone();
|
|
let mixer = self.mixer.clone();
|
|
let text_owned = text.to_string();
|
|
self.mixer.mark_stream_start();
|
|
self.schedule_filler(&cancel);
|
|
self.schedule_start_filler(&cancel);
|
|
debug!(
|
|
engine = tts.name(),
|
|
sample_rate = tts.sample_rate(),
|
|
text_len = text.len(),
|
|
"dispatching TTS synthesis"
|
|
);
|
|
|
|
tokio::spawn(async move {
|
|
if let Err(e) = tts.synthesize(&text_owned, &queue, &cancel).await {
|
|
error!(error = %e, "TTS synthesis failed");
|
|
}
|
|
// Signal completion (unless cancelled by a barge-in) so the mixer
|
|
// drains short utterances instead of waiting for the pre-buffer.
|
|
if !cancel.is_cancelled() {
|
|
mixer.mark_stream_end();
|
|
debug!(queue_len = queue.len(), "TTS stream end signalled to mixer");
|
|
}
|
|
let mut guard = cancel_holder.lock();
|
|
if guard.as_ref().map(|c| c.is_cancelled()).unwrap_or(false) {
|
|
*guard = None;
|
|
}
|
|
});
|
|
}
|
|
|
|
fn cancel_current_tts(&self) {
|
|
if let Some(token) = self.cancel.lock().take() {
|
|
token.cancel();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Detect the language of `text` and, when it differs from the active one,
|
|
/// build and publish the clip banks for it.
|
|
///
|
|
/// The phrase list comes from the built-in tables ([`crate::lang`]) for the
|
|
/// common languages, or from the LLM when `LANG_LLM_FALLBACK_ENABLED` is set.
|
|
/// Clips are synthesized only when missing — [`assets::ensure_lang_backchannel_clips`]
|
|
/// keeps a per-language index — and live in their own subdirectory, so no other
|
|
/// language is touched. Switching back to the default language restores the
|
|
/// base banks and needs no synthesis at all.
|
|
async fn switch_language_if_needed(
|
|
langs: Arc<LanguageState>,
|
|
llm: Arc<LlmClient>,
|
|
tts: Arc<dyn TtsEngine>,
|
|
config: Arc<Config>,
|
|
text: &str,
|
|
audit: Arc<SessionAudit>,
|
|
) {
|
|
let Some(lang) = lang::detect_language(&llm, text).await else {
|
|
return;
|
|
};
|
|
let current = langs.current();
|
|
if lang == current.lang {
|
|
return;
|
|
}
|
|
if lang == config.lang_default {
|
|
info!(lang = %lang, "user language back to the default — restoring the base clip banks");
|
|
audit.event(
|
|
"language_switch",
|
|
json!({ "from": current.lang, "to": lang, "generated": false }),
|
|
);
|
|
langs.set((*langs.base()).clone());
|
|
return;
|
|
}
|
|
|
|
// Phrase lists: built-in table first, LLM fallback otherwise.
|
|
let backchannel = match lang::backchannel_texts(&lang) {
|
|
Some(pairs) => pairs,
|
|
None if config.lang_llm_fallback_enabled => {
|
|
match lang::generate_backchannel_texts(&llm, &lang).await {
|
|
Some(phrases) => phrases
|
|
.into_iter()
|
|
.map(|p| (DEFAULT_CATEGORY.to_string(), p))
|
|
.collect(),
|
|
None => return,
|
|
}
|
|
}
|
|
None => return,
|
|
};
|
|
let filler = match lang::filler_texts(&lang) {
|
|
Some(texts) => texts,
|
|
None if config.lang_llm_fallback_enabled => {
|
|
match lang::generate_filler_texts(&llm, &lang).await {
|
|
Some(texts) => texts,
|
|
None => return,
|
|
}
|
|
}
|
|
None => return,
|
|
};
|
|
// Start filler is optional: when the language has no phrase list (and the
|
|
// LLM fallback is unavailable) the cue simply stays silent for it, without
|
|
// preventing the switch of the banks above.
|
|
let start_filler = if !config.start_filler_enabled {
|
|
None
|
|
} else {
|
|
match lang::start_filler_texts(&lang) {
|
|
Some(texts) => Some(texts),
|
|
None if config.lang_llm_fallback_enabled => {
|
|
lang::generate_start_filler_texts(&llm, &lang).await
|
|
}
|
|
None => None,
|
|
}
|
|
};
|
|
|
|
info!(
|
|
from = %current.lang,
|
|
to = %lang,
|
|
backchannel = backchannel.len(),
|
|
filler = filler.len(),
|
|
start_filler = start_filler.as_ref().map(|t| t.len()).unwrap_or(0),
|
|
"user language changed — generating the clip banks for it"
|
|
);
|
|
|
|
let phrases: Vec<String> = backchannel.iter().map(|(_, p)| p.clone()).collect();
|
|
let categories: Vec<String> = backchannel.iter().map(|(c, _)| c.clone()).collect();
|
|
|
|
// Idempotent: an unchanged phrase list is skipped via the per-directory
|
|
// `texts.txt` index, so an already-generated language is just reloaded.
|
|
let voice_profile = config.cue_voice_profile();
|
|
assets::migrate_legacy_cue_bank(
|
|
&config.backchannel.dir,
|
|
&lang,
|
|
&voice_profile,
|
|
"backchannel",
|
|
);
|
|
assets::migrate_legacy_cue_bank(&config.fillers_dir, &lang, &voice_profile, "filler");
|
|
if start_filler.is_some() {
|
|
assets::migrate_legacy_cue_bank(
|
|
&config.start_filler_dir,
|
|
&lang,
|
|
&voice_profile,
|
|
"start-filler",
|
|
);
|
|
}
|
|
assets::ensure_lang_backchannel_clips(
|
|
&config.backchannel.dir,
|
|
&lang,
|
|
&voice_profile,
|
|
&phrases,
|
|
tts.as_ref(),
|
|
)
|
|
.await;
|
|
assets::ensure_lang_filler_clips(
|
|
&config.fillers_dir,
|
|
&lang,
|
|
&voice_profile,
|
|
&filler,
|
|
tts.as_ref(),
|
|
)
|
|
.await;
|
|
if let Some(texts) = &start_filler {
|
|
assets::ensure_lang_start_filler_clips(
|
|
&config.start_filler_dir,
|
|
&lang,
|
|
&voice_profile,
|
|
texts,
|
|
tts.as_ref(),
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let backchannel_bank = Arc::new(FillerBank::load(
|
|
&assets::cue_bank_dir(&config.backchannel.dir, &lang, &voice_profile),
|
|
"backchannel",
|
|
));
|
|
let filler_bank = Arc::new(FillerBank::load(
|
|
&assets::cue_bank_dir(&config.fillers_dir, &lang, &voice_profile),
|
|
"filler",
|
|
));
|
|
let start_filler_bank = Arc::new(FillerBank::load(
|
|
&assets::cue_bank_dir(&config.start_filler_dir, &lang, &voice_profile),
|
|
"start-filler",
|
|
));
|
|
|
|
langs.set(LangBanks {
|
|
lang,
|
|
backchannel: backchannel_bank,
|
|
backchannel_texts: phrases,
|
|
backchannel_categories: categories,
|
|
filler: filler_bank,
|
|
start_filler: start_filler_bank,
|
|
});
|
|
audit.event(
|
|
"language_switch",
|
|
json!({ "from": current.lang, "to": langs.current().lang, "generated": true }),
|
|
);
|
|
}
|
|
|
|
/// Spawn an orchestrator loop that consumes transcripts from `rx` and drives
|
|
/// the turn manager.
|
|
pub fn spawn_orchestrator(manager: Arc<TurnManager>, mut rx: mpsc::UnboundedReceiver<Transcript>) {
|
|
tokio::spawn(async move {
|
|
while let Some(transcript) = rx.recv().await {
|
|
manager.on_eou(&transcript).await;
|
|
}
|
|
});
|
|
}
|
|
|
|
/// Splits a streaming reply into speech-sized chunks so the TTS can start
|
|
/// speaking while the model is still generating the rest.
|
|
///
|
|
/// A chunk always ends on `.`, `!`, `?`, `;` or a newline. On top of that, a
|
|
/// word-boundary flush keeps latency bounded when the model runs on without
|
|
/// punctuation: the *first* chunk is cut once [`first_chunk_chars`] characters
|
|
/// are buffered (kept small, so the agent starts on the first few words), and
|
|
/// every following chunk at [`chunk_chars`] (larger, for better prosody). A
|
|
/// threshold of `0` disables that flush.
|
|
///
|
|
/// [`first_chunk_chars`]: SentenceSplitter::first_chunk_chars
|
|
/// [`chunk_chars`]: SentenceSplitter::chunk_chars
|
|
struct SentenceSplitter {
|
|
buf: String,
|
|
/// Bytes buffered before the first chunk is flushed on a word boundary.
|
|
first_chunk_chars: usize,
|
|
/// Same, for every chunk after the first.
|
|
chunk_chars: usize,
|
|
/// How many chunks were handed out so far (drives the threshold above).
|
|
emitted: usize,
|
|
}
|
|
|
|
impl SentenceSplitter {
|
|
fn new(first_chunk_chars: usize, chunk_chars: usize) -> Self {
|
|
Self {
|
|
buf: String::new(),
|
|
first_chunk_chars,
|
|
chunk_chars,
|
|
emitted: 0,
|
|
}
|
|
}
|
|
|
|
/// Feed a delta and return every complete chunk it produced.
|
|
fn push(&mut self, delta: &str) -> Vec<String> {
|
|
self.buf.push_str(delta);
|
|
let mut out = Vec::new();
|
|
|
|
let mut start = 0usize;
|
|
for (i, b) in self.buf.bytes().enumerate() {
|
|
// ASCII terminators only: safe to slice at `i + 1`.
|
|
if matches!(b, b'.' | b'!' | b'?' | b';' | b'\n') {
|
|
let segment = self.buf[start..=i].trim();
|
|
if !segment.is_empty() {
|
|
out.push(segment.to_string());
|
|
self.emitted += 1;
|
|
}
|
|
start = i + 1;
|
|
}
|
|
}
|
|
if start > 0 {
|
|
self.buf.drain(..start);
|
|
}
|
|
|
|
// Word-boundary flush: cut at the first whitespace at or past the
|
|
// threshold, so the chunk holds at least that many bytes.
|
|
let threshold = if self.emitted == 0 {
|
|
self.first_chunk_chars
|
|
} else {
|
|
self.chunk_chars
|
|
};
|
|
if threshold > 0 && self.buf.len() >= threshold {
|
|
let safe = floor_char_boundary(&self.buf, threshold);
|
|
if let Some(rel) = self.buf[safe..].find(' ') {
|
|
let end = safe + rel;
|
|
let segment = self.buf[..end].trim().to_string();
|
|
if !segment.is_empty() {
|
|
out.push(segment);
|
|
self.emitted += 1;
|
|
}
|
|
self.buf.drain(..end);
|
|
}
|
|
}
|
|
|
|
out
|
|
}
|
|
|
|
/// Flush whatever is left once the stream has ended.
|
|
fn finish(&mut self) -> Option<String> {
|
|
let rest = self.buf.trim().to_string();
|
|
self.buf.clear();
|
|
if rest.is_empty() {
|
|
None
|
|
} else {
|
|
Some(rest)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Largest char boundary `<= i`.
|
|
fn floor_char_boundary(s: &str, mut i: usize) -> usize {
|
|
if i >= s.len() {
|
|
return s.len();
|
|
}
|
|
while i > 0 && !s.is_char_boundary(i) {
|
|
i -= 1;
|
|
}
|
|
i
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn splits_on_sentence_terminators() {
|
|
let mut s = SentenceSplitter::new(20, 120);
|
|
assert_eq!(
|
|
s.push("Hola. ¿Cómo estás? Bien"),
|
|
vec!["Hola.".to_string(), "¿Cómo estás?".to_string()]
|
|
);
|
|
assert_eq!(s.finish(), Some("Bien".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn splits_across_deltas() {
|
|
let mut s = SentenceSplitter::new(20, 120);
|
|
assert!(s.push("Hola").is_empty());
|
|
assert_eq!(s.push(" mundo!"), vec!["Hola mundo!".to_string()]);
|
|
}
|
|
|
|
#[test]
|
|
fn first_chunk_flushes_early_on_word_boundary() {
|
|
let mut s = SentenceSplitter::new(20, 120);
|
|
let out = s.push("Ciao come stai oggi tutto bene");
|
|
assert_eq!(out.len(), 1);
|
|
assert!(out[0].len() >= 20, "chunk too short: {:?}", out[0]);
|
|
}
|
|
|
|
#[test]
|
|
fn later_chunks_use_the_larger_threshold() {
|
|
let mut s = SentenceSplitter::new(20, 120);
|
|
// First chunk is small and punctuation-free.
|
|
let first = s.push("Ciao come stai oggi tutto bene");
|
|
assert_eq!(first.len(), 1);
|
|
// The next run must reach the larger threshold before flushing.
|
|
let second = s.push(&"palabra ".repeat(20)); // 160 bytes
|
|
assert_eq!(second.len(), 1);
|
|
assert!(
|
|
second[0].len() >= 120,
|
|
"chunk too short: {}",
|
|
second[0].len()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn zero_threshold_disables_word_boundary_flush() {
|
|
let mut s = SentenceSplitter::new(0, 0);
|
|
assert!(s.push(&"palabra ".repeat(40)).is_empty());
|
|
assert!(!s.finish().unwrap_or_default().is_empty());
|
|
}
|
|
}
|