emelex 1.1.1

Apple Silicon local inference toolkit powered by MLX
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//! The streaming bridge: engine per-token callback → bounded channel →
//! rig's raw streaming choices.
//!
//! Cancellation is cooperative: when the consumer drops the stream, the
//! channel receiver goes away; the engine observes that during media
//! preprocessing, between evaluated prefill chunks, or at the next token.

use rig_core::{
	completion::CompletionError,
	streaming::{RawStreamingChoice, RawStreamingToolCall, StreamingResult},
};

use super::response::StreamingResponse;
use crate::{
	client::StreamTextReconciler,
	convert::{self, EngineRequest},
	engine::{
		generate::{GenerateReply, Session},
		streaming::TokenKind,
	},
};

type Item = Result<RawStreamingChoice<StreamingResponse>, CompletionError>;

/// Run one generation on the client's dedicated inference thread,
/// bridging tokens into a stream.
pub(super) fn spawn(
	inner: &crate::client::Inner,
	engine_request: EngineRequest,
) -> StreamingResult<StreamingResponse> {
	let (tx, rx) = tokio::sync::mpsc::channel::<Item>(64);
	let job_tx = tx.clone();
	let submitted = inner.submit(Box::new(move |session| {
		run_job(session, &job_tx, engine_request);
	}));
	if let Err(reason) = submitted {
		// Inference thread is gone: surface it as the stream's only item.
		let _ = tx.try_send(Err(CompletionError::ProviderError(reason.to_string())));
	}
	drop(tx);
	Box::pin(futures::stream::unfold(rx, |mut rx| async move {
		rx.recv().await.map(|item| (item, rx))
	}))
}

fn run_job(session: &Session, tx: &tokio::sync::mpsc::Sender<Item>, engine_request: EngineRequest) {
	if tx.is_closed() {
		return;
	}
	let EngineRequest {
		messages,
		tools,
		options,
	} = engine_request;
	let mut visible_text = StreamTextReconciler::default();
	// The worker loop also catches panics, but catching here lets the panic
	// surface as an explicit stream error instead of silent truncation.
	let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
		let request_cancelled = || tx.is_closed();
		session.generate_cached_cancellable(
			&messages,
			tools.as_deref(),
			options,
			&request_cancelled,
			|_| true,
			|token| {
				let choice = match token.kind {
					TokenKind::Text => {
						let Some(text) = visible_text.push_text(token.text) else {
							return !tx.is_closed();
						};
						if text.is_empty() {
							return !tx.is_closed();
						}
						RawStreamingChoice::Message(text)
					}
					TokenKind::Reasoning => RawStreamingChoice::ReasoningDelta {
						id: None,
						reasoning: token.text,
					},
					// Raw markup is emitted structurally only after terminal
					// validation. Keep cancellation observable meanwhile.
					TokenKind::ToolCall => {
						visible_text.observe_tool_span();
						return !tx.is_closed();
					}
				};
				if matches!(
					&choice,
					RawStreamingChoice::ReasoningDelta { reasoning, .. } if reasoning.is_empty()
				) {
					return !tx.is_closed();
				}
				tx.blocking_send(Ok(choice)).is_ok()
			},
		)
	}));
	match outcome {
		Ok(Ok(reply)) => send_completed(tx, &visible_text, &reply),
		Ok(Err(error)) => {
			let _ = tx.blocking_send(Err(CompletionError::from(crate::error::from_engine(error))));
		}
		Err(_) => {
			let _ = tx.blocking_send(Err(CompletionError::ProviderError(
				"generation panicked; the panic message was printed to stderr".to_string(),
			)));
		}
	}
}

fn send_completed(
	tx: &tokio::sync::mpsc::Sender<Item>,
	visible_text: &StreamTextReconciler,
	reply: &GenerateReply,
) {
	tracing::debug!(
		text_chars = reply.text.len(),
		tool_calls = reply.tool_calls.len(),
		finish = super::response::finish_reason_label(reply.finish_reason),
		"generation finished"
	);
	let suffix = match visible_text.terminal_suffix(&reply.text) {
		Ok(suffix) => suffix,
		Err(error) => {
			let _ = tx.blocking_send(Err(CompletionError::from(error)));
			return;
		}
	};
	if !suffix.is_empty()
		&& tx
			.blocking_send(Ok(RawStreamingChoice::Message(suffix.to_string())))
			.is_err()
	{
		return;
	}
	for call in &reply.tool_calls {
		// Ignore send failures from here on: the receiver being gone just
		// means nobody is listening anymore.
		let _ = tx.blocking_send(Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall::new(
			call.id.clone(),
			call.name.clone(),
			call.arguments.clone(),
		))));
	}
	let response = StreamingResponse {
		usage: convert::usage_data(reply.usage),
		finish_reason: super::response::finish_reason_label(reply.finish_reason).to_string(),
		speculation: convert::speculation_data(reply.speculation.as_ref()),
	};
	let _ = tx.blocking_send(Ok(RawStreamingChoice::FinalResponse(response)));
}