Split Rust code into domain modules
This commit is contained in:
995
src/app.rs
995
src/app.rs
File diff suppressed because it is too large
Load Diff
252
src/app/generation.rs
Normal file
252
src/app/generation.rs
Normal file
@@ -0,0 +1,252 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub(crate) struct ChatMessage {
|
||||||
|
pub(super) id: i32,
|
||||||
|
pub(super) user: bool,
|
||||||
|
pub(super) reasoning: Option<String>,
|
||||||
|
pub(super) reasoning_complete: bool,
|
||||||
|
pub(super) reasoning_open: bool,
|
||||||
|
pub(super) content: String,
|
||||||
|
pub(super) markdown: Vec<markdown::Item>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChatMessage {
|
||||||
|
pub(super) fn append(&mut self, reasoning: bool, chunk: &str) {
|
||||||
|
if reasoning {
|
||||||
|
self.reasoning.get_or_insert_default().push_str(chunk);
|
||||||
|
} else {
|
||||||
|
self.reasoning_complete = true;
|
||||||
|
self.content.push_str(chunk);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn refresh_markdown(&mut self) {
|
||||||
|
if !self.user {
|
||||||
|
let content = if self.reasoning.is_some() {
|
||||||
|
self.content.trim_start()
|
||||||
|
} else {
|
||||||
|
&self.content
|
||||||
|
};
|
||||||
|
self.markdown = markdown::parse(content).collect();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<StoredMessage> for ChatMessage {
|
||||||
|
fn from(message: StoredMessage) -> Self {
|
||||||
|
let mut message = Self {
|
||||||
|
id: message.id,
|
||||||
|
user: message.user,
|
||||||
|
reasoning: message.reasoning,
|
||||||
|
reasoning_complete: message.reasoning_complete,
|
||||||
|
reasoning_open: false,
|
||||||
|
content: message.content,
|
||||||
|
markdown: Vec::new(),
|
||||||
|
};
|
||||||
|
message.refresh_markdown();
|
||||||
|
message
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn start_generation(&mut self) {
|
||||||
|
if self.generating || self.selected_session.is_none() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let prompt = self.composer.trim().to_owned();
|
||||||
|
if prompt.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let model = match ModelChoice::from_id(&self.preferences.selected_model) {
|
||||||
|
Some(model) => model,
|
||||||
|
None => {
|
||||||
|
self.error = Some("The selected model is not supported.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let effective = self.preferences.generation().and_then(|generation| {
|
||||||
|
self.preferences.runtime().and_then(|runtime| {
|
||||||
|
crate::settings::effective_settings(model, &generation, &runtime, &models_path())
|
||||||
|
})
|
||||||
|
});
|
||||||
|
let effective = match effective {
|
||||||
|
Ok(settings) => settings,
|
||||||
|
Err(error) => {
|
||||||
|
self.error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
|
||||||
|
let session_id = self
|
||||||
|
.selected_session
|
||||||
|
.expect("a selected session was checked");
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
let mut messages = self
|
||||||
|
.conversation
|
||||||
|
.iter()
|
||||||
|
.map(|message| ChatTurn {
|
||||||
|
user: message.user,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: message.reasoning.clone(),
|
||||||
|
reasoning_complete: message.reasoning_complete,
|
||||||
|
content: message.content.clone(),
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
messages.push(ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: prompt.clone(),
|
||||||
|
});
|
||||||
|
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
{
|
||||||
|
let Some(service) = &self.generation_service else {
|
||||||
|
self.error = Some("The model runtime is unavailable.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Some(database) = &mut self.database else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let saved = match database.start_chat_turn(session_id, &prompt, assistant_reasoning) {
|
||||||
|
Ok(turn) => turn,
|
||||||
|
Err(error) => {
|
||||||
|
self.error = Some(format!("Could not save the chat turn: {error}"));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let idle_timeout =
|
||||||
|
Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60);
|
||||||
|
self.active_generation = match service.generate(
|
||||||
|
effective.engine,
|
||||||
|
effective.turn,
|
||||||
|
messages,
|
||||||
|
CheckpointTarget::Local(session_checkpoint_path(session_id)),
|
||||||
|
idle_timeout,
|
||||||
|
) {
|
||||||
|
Ok(active) => Some(active),
|
||||||
|
Err(error) => {
|
||||||
|
self.generation_service = None;
|
||||||
|
self.error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let user = ChatMessage::from(saved.0);
|
||||||
|
let mut assistant = ChatMessage::from(saved.1);
|
||||||
|
assistant.reasoning_open = assistant_reasoning;
|
||||||
|
self.composer.clear();
|
||||||
|
self.conversation.push(user);
|
||||||
|
self.conversation.push(assistant);
|
||||||
|
self.generating = true;
|
||||||
|
self.tokens_per_second = None;
|
||||||
|
self.error = None;
|
||||||
|
}
|
||||||
|
#[cfg(not(target_os = "macos"))]
|
||||||
|
{
|
||||||
|
let _ = effective;
|
||||||
|
self.error = Some("Local Metal generation requires macOS.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn poll_generation(&mut self) -> bool {
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
let Some(active) = &mut self.active_generation else {
|
||||||
|
self.generating = false;
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
let mut transcript_changed = false;
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
let mut context_changed = false;
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
loop {
|
||||||
|
match active.events.try_recv() {
|
||||||
|
Ok(GenerationEvent::Loading) => {}
|
||||||
|
Ok(GenerationEvent::Chunk { reasoning, content }) => {
|
||||||
|
if let Some(message) = self.conversation.last_mut()
|
||||||
|
&& !message.user
|
||||||
|
{
|
||||||
|
message.append(reasoning, &content);
|
||||||
|
transcript_changed = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(GenerationEvent::Context {
|
||||||
|
used,
|
||||||
|
limit,
|
||||||
|
tokens_per_second,
|
||||||
|
}) => {
|
||||||
|
self.context_used = used;
|
||||||
|
self.context_limit = limit;
|
||||||
|
self.tokens_per_second = tokens_per_second;
|
||||||
|
context_changed = true;
|
||||||
|
}
|
||||||
|
Ok(GenerationEvent::Finished(result)) => {
|
||||||
|
self.generating = false;
|
||||||
|
if let Err(error) = result {
|
||||||
|
self.error = Some(error);
|
||||||
|
}
|
||||||
|
self.active_generation = None;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Err(TryRecvError::Empty) => break,
|
||||||
|
Err(TryRecvError::Disconnected) => {
|
||||||
|
self.generating = false;
|
||||||
|
self.active_generation = None;
|
||||||
|
self.error = Some("The model runtime stopped unexpectedly.".into());
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
if transcript_changed && let Some(message) = self.conversation.last_mut() {
|
||||||
|
message.refresh_markdown();
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
if transcript_changed
|
||||||
|
&& let Some(message) = self.conversation.last()
|
||||||
|
&& let Some(database) = &mut self.database
|
||||||
|
&& let Err(error) = database.update_message(
|
||||||
|
message.id,
|
||||||
|
message.reasoning.as_deref(),
|
||||||
|
message.reasoning_complete,
|
||||||
|
&message.content,
|
||||||
|
)
|
||||||
|
{
|
||||||
|
if let Some(active) = &self.active_generation {
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
self.error = Some(format!("Could not save generated chat text: {error}"));
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
if context_changed
|
||||||
|
&& let Some(session_id) = self.selected_session
|
||||||
|
&& let Some(database) = &mut self.database
|
||||||
|
{
|
||||||
|
if let Err(error) = database.update_session_context(
|
||||||
|
session_id,
|
||||||
|
self.context_used,
|
||||||
|
self.context_limit,
|
||||||
|
self.tokens_per_second,
|
||||||
|
) {
|
||||||
|
self.error = Some(format!("Could not save context usage: {error}"));
|
||||||
|
} else if let Some(session) = self
|
||||||
|
.projects
|
||||||
|
.iter_mut()
|
||||||
|
.flat_map(|project| &mut project.sessions)
|
||||||
|
.find(|session| session.id == session_id)
|
||||||
|
{
|
||||||
|
session.context_used = self.context_used as i32;
|
||||||
|
session.context_limit = self.context_limit as i32;
|
||||||
|
session.last_tokens_per_second = self.tokens_per_second;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
return transcript_changed;
|
||||||
|
#[cfg(not(target_os = "macos"))]
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
173
src/app/model_manager.rs
Normal file
173
src/app/model_manager.rs
Normal file
@@ -0,0 +1,173 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub(super) enum ModelDownload {
|
||||||
|
Idle,
|
||||||
|
Active(ActiveDownload),
|
||||||
|
Complete(ManagedArtifactId, ModelOperation, DownloadProgress),
|
||||||
|
Failed(ManagedArtifactId, String, DownloadProgress),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||||
|
pub(super) enum ModelOperation {
|
||||||
|
Download,
|
||||||
|
Validate,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub(super) struct ActiveDownload {
|
||||||
|
pub(super) artifact: ManagedArtifactId,
|
||||||
|
pub(super) operation: ModelOperation,
|
||||||
|
pub(super) progress: DownloadProgress,
|
||||||
|
pub(super) sampled_at: Instant,
|
||||||
|
pub(super) sampled_bytes: u64,
|
||||||
|
pub(super) bytes_per_second: f64,
|
||||||
|
pub(super) verified_bytes: Arc<AtomicU64>,
|
||||||
|
pub(super) cancel: Arc<AtomicBool>,
|
||||||
|
pub(super) result: mpsc::Receiver<Result<DownloadOutcome, String>>,
|
||||||
|
pub(super) stopping: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn open_model_manager(&mut self) -> Task<Message> {
|
||||||
|
if let Some(id) = self.model_manager_window {
|
||||||
|
return window::gain_focus(id);
|
||||||
|
}
|
||||||
|
let (id, open) = window::open(window::Settings {
|
||||||
|
size: Size::new(760.0, 560.0),
|
||||||
|
min_size: Some(Size::new(620.0, 420.0)),
|
||||||
|
icon: Some(app_icon()),
|
||||||
|
..Default::default()
|
||||||
|
});
|
||||||
|
self.model_manager_window = Some(id);
|
||||||
|
open.map(Message::ModelManagerOpened)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn start_model_operation(
|
||||||
|
&mut self,
|
||||||
|
artifact: ManagedArtifactId,
|
||||||
|
operation: ModelOperation,
|
||||||
|
) {
|
||||||
|
if matches!(self.model_download, ModelDownload::Active(_)) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let models_path = models_path();
|
||||||
|
let progress = match operation {
|
||||||
|
ModelOperation::Download => model::artifact_download_progress(artifact, &models_path),
|
||||||
|
ModelOperation::Validate => model::artifact_verification_progress(artifact, 0),
|
||||||
|
};
|
||||||
|
let cancel = Arc::new(AtomicBool::new(false));
|
||||||
|
let worker_cancel = Arc::clone(&cancel);
|
||||||
|
let verified_bytes = Arc::new(AtomicU64::new(0));
|
||||||
|
let worker_verified_bytes = Arc::clone(&verified_bytes);
|
||||||
|
let (result_sender, result_receiver) = mpsc::channel();
|
||||||
|
let thread_name = match operation {
|
||||||
|
ModelOperation::Download => "model-download",
|
||||||
|
ModelOperation::Validate => "model-validation",
|
||||||
|
};
|
||||||
|
if let Err(error) = thread::Builder::new()
|
||||||
|
.name(thread_name.to_owned())
|
||||||
|
.spawn(move || {
|
||||||
|
let result = match operation {
|
||||||
|
ModelOperation::Download => model::download_managed_artifact(
|
||||||
|
artifact,
|
||||||
|
&models_path,
|
||||||
|
&worker_cancel,
|
||||||
|
&worker_verified_bytes,
|
||||||
|
),
|
||||||
|
ModelOperation::Validate => model::validate_managed_artifact(
|
||||||
|
artifact,
|
||||||
|
&models_path,
|
||||||
|
&worker_cancel,
|
||||||
|
&worker_verified_bytes,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
let _ = result_sender.send(result);
|
||||||
|
})
|
||||||
|
{
|
||||||
|
self.error = Some(format!("Could not start {thread_name}: {error}"));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.pending_model_delete = None;
|
||||||
|
self.model_download = ModelDownload::Active(ActiveDownload {
|
||||||
|
artifact,
|
||||||
|
operation,
|
||||||
|
sampled_at: Instant::now(),
|
||||||
|
sampled_bytes: progress.completed(),
|
||||||
|
progress,
|
||||||
|
bytes_per_second: 0.0,
|
||||||
|
verified_bytes,
|
||||||
|
cancel,
|
||||||
|
result: result_receiver,
|
||||||
|
stopping: false,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn update_download_progress(&mut self) {
|
||||||
|
let ModelDownload::Active(download) = &mut self.model_download else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let now = Instant::now();
|
||||||
|
let verified = download.verified_bytes.load(Ordering::Relaxed);
|
||||||
|
let mut progress = match download.operation {
|
||||||
|
ModelOperation::Download => {
|
||||||
|
model::artifact_download_progress(download.artifact, &models_path())
|
||||||
|
}
|
||||||
|
ModelOperation::Validate => {
|
||||||
|
model::artifact_verification_progress(download.artifact, verified)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if download.operation == ModelOperation::Download
|
||||||
|
&& let Some(verification) = &mut progress.verification
|
||||||
|
{
|
||||||
|
verification.verified = verified.min(verification.total);
|
||||||
|
}
|
||||||
|
let elapsed = now.duration_since(download.sampled_at).as_secs_f64();
|
||||||
|
let completed = progress.completed();
|
||||||
|
let phase_changed = progress.phase != download.progress.phase;
|
||||||
|
let transferred = completed.saturating_sub(download.sampled_bytes);
|
||||||
|
if phase_changed {
|
||||||
|
download.bytes_per_second = 0.0;
|
||||||
|
} else if transferred > 0 && elapsed > 0.0 {
|
||||||
|
let current = transferred as f64 / elapsed;
|
||||||
|
download.bytes_per_second = if download.bytes_per_second == 0.0 {
|
||||||
|
current
|
||||||
|
} else {
|
||||||
|
download.bytes_per_second * 0.75 + current * 0.25
|
||||||
|
};
|
||||||
|
}
|
||||||
|
download.progress = progress;
|
||||||
|
download.sampled_at = now;
|
||||||
|
download.sampled_bytes = completed;
|
||||||
|
|
||||||
|
let result = match download.result.try_recv() {
|
||||||
|
Ok(result) => Some(result),
|
||||||
|
Err(TryRecvError::Empty) => None,
|
||||||
|
Err(TryRecvError::Disconnected) => {
|
||||||
|
Some(Err("Download worker stopped unexpectedly.".into()))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let artifact = download.artifact;
|
||||||
|
let operation = download.operation;
|
||||||
|
let progress = download.progress.clone();
|
||||||
|
if let Some(result) = result {
|
||||||
|
match result {
|
||||||
|
Ok(DownloadOutcome::Complete) => {
|
||||||
|
let progress = model::artifact_download_progress(artifact, &models_path());
|
||||||
|
self.model_download = ModelDownload::Complete(artifact, operation, progress);
|
||||||
|
self.error = None;
|
||||||
|
}
|
||||||
|
Ok(DownloadOutcome::Stopped) => {
|
||||||
|
self.model_download = ModelDownload::Idle;
|
||||||
|
self.error = None;
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
self.error = Some(error.clone());
|
||||||
|
self.model_download = ModelDownload::Failed(artifact, error, progress);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
476
src/app/preferences.rs
Normal file
476
src/app/preferences.rs
Normal file
@@ -0,0 +1,476 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub(super) struct PreferenceDraft {
|
||||||
|
pub(super) model: ModelChoice,
|
||||||
|
pub(super) dspark_enabled: bool,
|
||||||
|
pub(super) idle_timeout_minutes: String,
|
||||||
|
pub(super) endpoint_port: String,
|
||||||
|
pub(super) context_tokens: String,
|
||||||
|
pub(super) max_generated_tokens: String,
|
||||||
|
pub(super) system_prompt: String,
|
||||||
|
pub(super) temperature: String,
|
||||||
|
pub(super) top_p: String,
|
||||||
|
pub(super) min_p: String,
|
||||||
|
pub(super) seed: String,
|
||||||
|
pub(super) reasoning_mode: ReasoningMode,
|
||||||
|
pub(super) cpu_threads: String,
|
||||||
|
pub(super) power_percent: String,
|
||||||
|
pub(super) prefill_chunk: String,
|
||||||
|
pub(super) quality: bool,
|
||||||
|
pub(super) warm_weights: bool,
|
||||||
|
pub(super) mtp_draft_tokens: String,
|
||||||
|
pub(super) mtp_margin: String,
|
||||||
|
pub(super) glm_mtp: bool,
|
||||||
|
pub(super) glm_mtp_timing: bool,
|
||||||
|
pub(super) dspark_confidence_threshold: String,
|
||||||
|
pub(super) dspark_strict: bool,
|
||||||
|
pub(super) ssd_streaming: bool,
|
||||||
|
pub(super) ssd_streaming_cold: bool,
|
||||||
|
pub(super) ssd_cache: String,
|
||||||
|
pub(super) ssd_full_layers: String,
|
||||||
|
pub(super) ssd_preload_experts: String,
|
||||||
|
pub(super) directional_steering_file: String,
|
||||||
|
pub(super) directional_steering_ffn: String,
|
||||||
|
pub(super) directional_steering_attn: String,
|
||||||
|
pub(super) simulated_used_memory_gib: String,
|
||||||
|
pub(super) expert_profile_path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PreferenceDraft {
|
||||||
|
pub(super) fn from_saved(preferences: &AppPreferences) -> Result<Self, String> {
|
||||||
|
let model = ModelChoice::from_id(&preferences.selected_model)
|
||||||
|
.ok_or_else(|| format!("Unsupported model: {}", preferences.selected_model))?;
|
||||||
|
let generation = preferences.generation()?;
|
||||||
|
let runtime = preferences.runtime()?;
|
||||||
|
runtime.validate(model)?;
|
||||||
|
let execution = &runtime.execution;
|
||||||
|
let speculative = &runtime.speculative;
|
||||||
|
Ok(Self {
|
||||||
|
model,
|
||||||
|
dspark_enabled: speculative.dspark_enabled,
|
||||||
|
idle_timeout_minutes: preferences.idle_timeout_minutes.to_string(),
|
||||||
|
endpoint_port: preferences.endpoint_port.to_string(),
|
||||||
|
context_tokens: generation.context_tokens.to_string(),
|
||||||
|
max_generated_tokens: generation.max_generated_tokens.to_string(),
|
||||||
|
system_prompt: generation.system_prompt,
|
||||||
|
temperature: generation
|
||||||
|
.temperature
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
top_p: generation
|
||||||
|
.top_p
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
min_p: generation
|
||||||
|
.min_p
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
seed: generation
|
||||||
|
.seed
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
reasoning_mode: generation.reasoning_mode,
|
||||||
|
cpu_threads: execution
|
||||||
|
.cpu_threads
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
power_percent: execution
|
||||||
|
.power_percent
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
prefill_chunk: execution
|
||||||
|
.prefill_chunk
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
quality: execution.quality,
|
||||||
|
warm_weights: execution.warm_weights,
|
||||||
|
mtp_draft_tokens: speculative.mtp_draft_tokens.to_string(),
|
||||||
|
mtp_margin: speculative.mtp_margin.to_string(),
|
||||||
|
glm_mtp: speculative.glm_mtp,
|
||||||
|
glm_mtp_timing: speculative.glm_mtp_timing,
|
||||||
|
dspark_confidence_threshold: speculative
|
||||||
|
.dspark_confidence_threshold
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
dspark_strict: speculative.dspark_strict,
|
||||||
|
ssd_streaming: runtime.ssd.enabled,
|
||||||
|
ssd_streaming_cold: runtime.ssd.cold,
|
||||||
|
ssd_cache: runtime
|
||||||
|
.ssd
|
||||||
|
.cache
|
||||||
|
.map_or_else(String::new, |cache| match cache {
|
||||||
|
StreamingCacheBudget::Experts(experts) => experts.to_string(),
|
||||||
|
StreamingCacheBudget::Gib(gib) => format!("{gib}GB"),
|
||||||
|
}),
|
||||||
|
ssd_full_layers: runtime
|
||||||
|
.ssd
|
||||||
|
.full_layers
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
ssd_preload_experts: runtime
|
||||||
|
.ssd
|
||||||
|
.preload_experts
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
directional_steering_file: runtime.steering.file.unwrap_or_default(),
|
||||||
|
directional_steering_ffn: runtime
|
||||||
|
.steering
|
||||||
|
.ffn_scale
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
directional_steering_attn: runtime
|
||||||
|
.steering
|
||||||
|
.attention_scale
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
simulated_used_memory_gib: runtime
|
||||||
|
.diagnostics
|
||||||
|
.simulated_used_memory_gib
|
||||||
|
.map_or_else(String::new, |value| value.to_string()),
|
||||||
|
expert_profile_path: runtime.diagnostics.expert_profile_path.unwrap_or_default(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn generation(&self) -> Result<GenerationPreferences, String> {
|
||||||
|
let preferences = GenerationPreferences {
|
||||||
|
context_tokens: parse_positive_i32("Context tokens", &self.context_tokens)?,
|
||||||
|
max_generated_tokens: parse_positive_i32(
|
||||||
|
"Maximum generated tokens",
|
||||||
|
&self.max_generated_tokens,
|
||||||
|
)?,
|
||||||
|
system_prompt: self.system_prompt.clone(),
|
||||||
|
temperature: parse_optional_f32("Temperature", &self.temperature)?,
|
||||||
|
top_p: parse_optional_f32("Top-p", &self.top_p)?,
|
||||||
|
min_p: parse_optional_f32("Min-p", &self.min_p)?,
|
||||||
|
seed: parse_optional_u64("Seed", &self.seed)?,
|
||||||
|
reasoning_mode: self.reasoning_mode,
|
||||||
|
};
|
||||||
|
preferences.validate()?;
|
||||||
|
Ok(preferences)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn reset(&mut self) {
|
||||||
|
let defaults = GenerationPreferences::default();
|
||||||
|
let execution = ExecutionPreferences::default();
|
||||||
|
let speculative = SpeculativePreferences::default();
|
||||||
|
let ssd = SsdPreferences::default();
|
||||||
|
self.model = ModelChoice::default();
|
||||||
|
self.dspark_enabled = false;
|
||||||
|
self.idle_timeout_minutes = "10".into();
|
||||||
|
self.endpoint_port = "4000".into();
|
||||||
|
self.context_tokens = defaults.context_tokens.to_string();
|
||||||
|
self.max_generated_tokens = defaults.max_generated_tokens.to_string();
|
||||||
|
self.system_prompt = defaults.system_prompt;
|
||||||
|
self.temperature.clear();
|
||||||
|
self.top_p.clear();
|
||||||
|
self.min_p.clear();
|
||||||
|
self.seed.clear();
|
||||||
|
self.reasoning_mode = defaults.reasoning_mode;
|
||||||
|
self.cpu_threads.clear();
|
||||||
|
self.power_percent.clear();
|
||||||
|
self.prefill_chunk.clear();
|
||||||
|
self.quality = execution.quality;
|
||||||
|
self.warm_weights = execution.warm_weights;
|
||||||
|
self.mtp_draft_tokens = speculative.mtp_draft_tokens.to_string();
|
||||||
|
self.mtp_margin = speculative.mtp_margin.to_string();
|
||||||
|
self.glm_mtp = speculative.glm_mtp;
|
||||||
|
self.glm_mtp_timing = speculative.glm_mtp_timing;
|
||||||
|
self.dspark_confidence_threshold.clear();
|
||||||
|
self.dspark_strict = speculative.dspark_strict;
|
||||||
|
self.ssd_streaming = ssd.enabled;
|
||||||
|
self.ssd_streaming_cold = ssd.cold;
|
||||||
|
self.ssd_cache.clear();
|
||||||
|
self.ssd_full_layers.clear();
|
||||||
|
self.ssd_preload_experts.clear();
|
||||||
|
self.directional_steering_file.clear();
|
||||||
|
self.directional_steering_ffn.clear();
|
||||||
|
self.directional_steering_attn.clear();
|
||||||
|
self.simulated_used_memory_gib.clear();
|
||||||
|
self.expert_profile_path.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn execution(&self) -> Result<ExecutionPreferences, String> {
|
||||||
|
Ok(ExecutionPreferences {
|
||||||
|
cpu_threads: parse_optional_u32("CPU helper threads", &self.cpu_threads)?,
|
||||||
|
power_percent: parse_optional_u8("GPU power", &self.power_percent)?,
|
||||||
|
prefill_chunk: parse_optional_u32("Prefill chunk", &self.prefill_chunk)?,
|
||||||
|
quality: self.quality,
|
||||||
|
warm_weights: self.warm_weights,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn speculative(&self) -> Result<SpeculativePreferences, String> {
|
||||||
|
Ok(SpeculativePreferences {
|
||||||
|
mtp_draft_tokens: parse_positive_i32("MTP draft tokens", &self.mtp_draft_tokens)?,
|
||||||
|
mtp_margin: parse_f32("MTP margin", &self.mtp_margin)?,
|
||||||
|
glm_mtp: self.glm_mtp,
|
||||||
|
glm_mtp_timing: self.glm_mtp_timing,
|
||||||
|
dspark_enabled: self.dspark_enabled,
|
||||||
|
dspark_confidence_threshold: parse_optional_f32(
|
||||||
|
"DSpark confidence",
|
||||||
|
&self.dspark_confidence_threshold,
|
||||||
|
)?,
|
||||||
|
dspark_strict: self.dspark_strict,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn runtime(&self) -> Result<RuntimePreferences, String> {
|
||||||
|
Ok(RuntimePreferences {
|
||||||
|
execution: self.execution()?,
|
||||||
|
speculative: self.speculative()?,
|
||||||
|
ssd: SsdPreferences {
|
||||||
|
enabled: self.ssd_streaming,
|
||||||
|
cold: self.ssd_streaming_cold,
|
||||||
|
cache: parse_streaming_cache(&self.ssd_cache)?,
|
||||||
|
full_layers: parse_optional_u32("SSD full-layer count", &self.ssd_full_layers)?,
|
||||||
|
preload_experts: parse_optional_u32(
|
||||||
|
"SSD preload experts",
|
||||||
|
&self.ssd_preload_experts,
|
||||||
|
)?,
|
||||||
|
},
|
||||||
|
steering: SteeringPreferences {
|
||||||
|
file: optional_text(&self.directional_steering_file),
|
||||||
|
ffn_scale: parse_optional_f32(
|
||||||
|
"Directional FFN scale",
|
||||||
|
&self.directional_steering_ffn,
|
||||||
|
)?,
|
||||||
|
attention_scale: parse_optional_f32(
|
||||||
|
"Directional attention scale",
|
||||||
|
&self.directional_steering_attn,
|
||||||
|
)?,
|
||||||
|
},
|
||||||
|
diagnostics: DiagnosticPreferences {
|
||||||
|
simulated_used_memory_gib: parse_optional_gib(
|
||||||
|
"Simulated used memory",
|
||||||
|
&self.simulated_used_memory_gib,
|
||||||
|
)?,
|
||||||
|
expert_profile_path: optional_text(&self.expert_profile_path),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_positive_i32(name: &str, value: &str) -> Result<i32, String> {
|
||||||
|
value
|
||||||
|
.trim()
|
||||||
|
.parse::<i32>()
|
||||||
|
.ok()
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.ok_or_else(|| format!("{name} must be a positive whole number."))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_optional_f32(name: &str, value: &str) -> Result<Option<f32>, String> {
|
||||||
|
let value = value.trim();
|
||||||
|
if value.is_empty() {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
value
|
||||||
|
.parse()
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| format!("{name} must be a number or left blank for the DS4 default."))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_f32(name: &str, value: &str) -> Result<f32, String> {
|
||||||
|
value
|
||||||
|
.trim()
|
||||||
|
.parse()
|
||||||
|
.map_err(|_| format!("{name} must be a number."))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_optional_u64(name: &str, value: &str) -> Result<Option<u64>, String> {
|
||||||
|
let value = value.trim();
|
||||||
|
if value.is_empty() {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
value
|
||||||
|
.parse()
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| format!("{name} must be a positive whole number or left blank."))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_optional_u32(name: &str, value: &str) -> Result<Option<u32>, String> {
|
||||||
|
parse_optional_number(name, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_optional_u8(name: &str, value: &str) -> Result<Option<u8>, String> {
|
||||||
|
parse_optional_number(name, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_optional_number<T: std::str::FromStr>(
|
||||||
|
name: &str,
|
||||||
|
value: &str,
|
||||||
|
) -> Result<Option<T>, String> {
|
||||||
|
let value = value.trim();
|
||||||
|
if value.is_empty() {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
value
|
||||||
|
.parse()
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| format!("{name} must be a positive whole number or left blank."))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn parse_streaming_cache(value: &str) -> Result<Option<StreamingCacheBudget>, String> {
|
||||||
|
let value = value.trim();
|
||||||
|
if value.is_empty() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
if value.len() > 2
|
||||||
|
&& value
|
||||||
|
.get(value.len() - 2..)
|
||||||
|
.is_some_and(|suffix| suffix.eq_ignore_ascii_case("gb"))
|
||||||
|
{
|
||||||
|
return parse_gib("SSD cache budget", value)
|
||||||
|
.map(|gib| Some(StreamingCacheBudget::Gib(gib)));
|
||||||
|
}
|
||||||
|
if !value.chars().all(|character| character.is_ascii_digit()) {
|
||||||
|
return Err(
|
||||||
|
"SSD cache budget must be a positive expert count or whole GiB value such as 64GB."
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
value
|
||||||
|
.parse::<u32>()
|
||||||
|
.ok()
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.map(StreamingCacheBudget::Experts)
|
||||||
|
.map(Some)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
"SSD cache budget must be a positive expert count or whole GiB value such as 64GB."
|
||||||
|
.into()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn parse_optional_gib(name: &str, value: &str) -> Result<Option<u64>, String> {
|
||||||
|
let value = value.trim();
|
||||||
|
if value.is_empty() {
|
||||||
|
Ok(None)
|
||||||
|
} else {
|
||||||
|
parse_gib(name, value).map(Some)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_gib(name: &str, value: &str) -> Result<u64, String> {
|
||||||
|
let value = value
|
||||||
|
.get(value.len().saturating_sub(2)..)
|
||||||
|
.filter(|suffix| suffix.eq_ignore_ascii_case("gb"))
|
||||||
|
.map_or(value, |_| &value[..value.len() - 2]);
|
||||||
|
if !value.chars().all(|character| character.is_ascii_digit()) {
|
||||||
|
return Err(format!("{name} must be a positive whole GiB value."));
|
||||||
|
}
|
||||||
|
value
|
||||||
|
.parse::<u64>()
|
||||||
|
.ok()
|
||||||
|
.filter(|value| *value > 0 && *value <= u64::MAX / GIB)
|
||||||
|
.ok_or_else(|| format!("{name} must be a positive whole GiB value."))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn optional_text(value: &str) -> Option<String> {
|
||||||
|
let value = value.trim();
|
||||||
|
(!value.is_empty()).then(|| value.to_owned())
|
||||||
|
}
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn open_preferences(&mut self) {
|
||||||
|
if self.database.is_none() || self.pending_project_path.is_some() || self.choosing_folder {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
match PreferenceDraft::from_saved(&self.preferences) {
|
||||||
|
Ok(draft) => {
|
||||||
|
self.preference_draft = draft;
|
||||||
|
self.preference_error = None;
|
||||||
|
self.preferences_open = true;
|
||||||
|
}
|
||||||
|
Err(error) => self.error = Some(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn save_preferences(&mut self) {
|
||||||
|
let Ok(idle_timeout_minutes) = self
|
||||||
|
.preference_draft
|
||||||
|
.idle_timeout_minutes
|
||||||
|
.trim()
|
||||||
|
.parse::<i32>()
|
||||||
|
else {
|
||||||
|
self.preference_error = Some("Idle timeout must be a whole number.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if !(1..=1440).contains(&idle_timeout_minutes) {
|
||||||
|
self.preference_error = Some("Idle timeout must be between 1 and 1440 minutes.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Ok(endpoint_port) = self.preference_draft.endpoint_port.trim().parse::<u16>() else {
|
||||||
|
self.preference_error = Some("Endpoint port must be a whole number.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if endpoint_port == 0 {
|
||||||
|
self.preference_error = Some("Endpoint port must be between 1 and 65535.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let generation = match self.preference_draft.generation() {
|
||||||
|
Ok(generation) => generation,
|
||||||
|
Err(error) => {
|
||||||
|
self.preference_error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let model = self.preference_draft.model;
|
||||||
|
let runtime = match self.preference_draft.runtime() {
|
||||||
|
Ok(runtime) => runtime,
|
||||||
|
Err(error) => {
|
||||||
|
self.preference_error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if let Err(error) = runtime.validate(model) {
|
||||||
|
self.preference_error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
let pending_endpoint = if self.preferences.endpoint_port != i32::from(endpoint_port)
|
||||||
|
|| self._endpoint.is_none()
|
||||||
|
{
|
||||||
|
let Some(generation) = &self.generation_service else {
|
||||||
|
self.preference_error = Some("The model runtime is unavailable.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
match crate::server::ServerHandle::spawn(
|
||||||
|
generation.clone(),
|
||||||
|
Arc::clone(&self.runtime_preferences),
|
||||||
|
models_path(),
|
||||||
|
application_support_path().join("kv-cache").join("http"),
|
||||||
|
endpoint_port,
|
||||||
|
Arc::clone(&self.metrics),
|
||||||
|
) {
|
||||||
|
Ok(endpoint) => Some(endpoint),
|
||||||
|
Err(error) => {
|
||||||
|
self.preference_error = Some(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let Some(database) = &mut self.database else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
match database.update_preferences(
|
||||||
|
model.id(),
|
||||||
|
idle_timeout_minutes,
|
||||||
|
i32::from(endpoint_port),
|
||||||
|
&generation,
|
||||||
|
&runtime,
|
||||||
|
) {
|
||||||
|
Ok(preferences) => {
|
||||||
|
self.preferences = preferences;
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
if let Ok(mut runtime_preferences) = self.runtime_preferences.write() {
|
||||||
|
*runtime_preferences = self.preferences.clone();
|
||||||
|
}
|
||||||
|
#[cfg(target_os = "macos")]
|
||||||
|
if let Some(endpoint) = pending_endpoint {
|
||||||
|
self._endpoint = Some(endpoint);
|
||||||
|
}
|
||||||
|
self.preference_draft = PreferenceDraft::from_saved(&self.preferences)
|
||||||
|
.expect("the saved model was selected from the supported catalog");
|
||||||
|
self.preferences_open = false;
|
||||||
|
self.preference_error = None;
|
||||||
|
self.error = None;
|
||||||
|
}
|
||||||
|
Err(error) => self.preference_error = Some(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
102
src/app/projects.rs
Normal file
102
src/app/projects.rs
Normal file
@@ -0,0 +1,102 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn prepare_project(&mut self, path: PathBuf) {
|
||||||
|
let Ok(path) = fs::canonicalize(path) else {
|
||||||
|
self.error = Some("The selected folder is no longer available.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Some(path_text) = path.to_str() else {
|
||||||
|
self.error = Some("The selected folder path is not valid UTF-8.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if self
|
||||||
|
.projects
|
||||||
|
.iter()
|
||||||
|
.any(|item| item.project.path == path_text)
|
||||||
|
{
|
||||||
|
self.error = Some("That project is already in the sidebar.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
self.project_name_input = path
|
||||||
|
.file_name()
|
||||||
|
.and_then(|name| name.to_str())
|
||||||
|
.unwrap_or("Project")
|
||||||
|
.to_owned();
|
||||||
|
self.pending_project_path = Some(path);
|
||||||
|
self.error = None;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn create_project(&mut self) {
|
||||||
|
let name = self.project_name_input.trim();
|
||||||
|
if name.is_empty() {
|
||||||
|
self.error = Some("Project name cannot be empty.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Some(path) = &self.pending_project_path else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Some(path) = path.to_str() else {
|
||||||
|
self.error = Some("The selected folder path is not valid UTF-8.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Some(database) = &mut self.database else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
match database.create_project(name, path) {
|
||||||
|
Ok(project) => {
|
||||||
|
self.selected_project = Some(project.id);
|
||||||
|
self.selected_session = None;
|
||||||
|
self.pending_project_path = None;
|
||||||
|
self.project_name_input.clear();
|
||||||
|
self.error = None;
|
||||||
|
self.reload_projects();
|
||||||
|
}
|
||||||
|
Err(error) => self.error = Some(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn create_session(&mut self) {
|
||||||
|
if self.generating {
|
||||||
|
self.error = Some("Stop the active generation before creating a session.".into());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Some(project_id) = self.selected_project else {
|
||||||
|
self.error = Some("Select a project first.".into());
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let default_number = self
|
||||||
|
.selected_project()
|
||||||
|
.map(|project| project.sessions.len() + 1)
|
||||||
|
.unwrap_or(1);
|
||||||
|
let title = format!("Session {default_number}");
|
||||||
|
let Some(database) = &mut self.database else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
match database.create_session(project_id, &title) {
|
||||||
|
Ok(session) => {
|
||||||
|
self.selected_session = Some(session.id);
|
||||||
|
self.conversation.clear();
|
||||||
|
self.composer.clear();
|
||||||
|
self.context_used = 0;
|
||||||
|
self.context_limit = self.preferences.context_tokens.max(0) as u32;
|
||||||
|
self.tokens_per_second = None;
|
||||||
|
self.error = None;
|
||||||
|
self.reload_projects();
|
||||||
|
}
|
||||||
|
Err(error) => self.error = Some(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn reload_projects(&mut self) {
|
||||||
|
if let Some(database) = &mut self.database {
|
||||||
|
match database.load_projects() {
|
||||||
|
Ok(projects) => self.projects = projects,
|
||||||
|
Err(error) => self.error = Some(error),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
1308
src/app/view.rs
1308
src/app/view.rs
File diff suppressed because it is too large
Load Diff
210
src/app/view/chat.rs
Normal file
210
src/app/view/chat.rs
Normal file
@@ -0,0 +1,210 @@
|
|||||||
|
use super::*;
|
||||||
|
use iced::widget::column;
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn chat_detail(&self) -> Element<'_, Message> {
|
||||||
|
let Some(item) = self.selected_project() else {
|
||||||
|
let open_project_content = row![icon(ICON_FOLDER_PLUS, 17), text("Open project…"),]
|
||||||
|
.spacing(8)
|
||||||
|
.align_y(Alignment::Center);
|
||||||
|
let open_project = if self.database.is_some() && !self.choosing_folder {
|
||||||
|
action_button(open_project_content).on_press(Message::ChooseProjectFolder)
|
||||||
|
} else {
|
||||||
|
action_button(open_project_content)
|
||||||
|
};
|
||||||
|
return container(
|
||||||
|
column![
|
||||||
|
icon(ICON_SPARK, 36),
|
||||||
|
text("Start a local coding session").size(28),
|
||||||
|
text("Choose a project folder to create your first session.").size(14),
|
||||||
|
Space::with_height(10),
|
||||||
|
open_project,
|
||||||
|
]
|
||||||
|
.spacing(10)
|
||||||
|
.align_x(Alignment::Center),
|
||||||
|
)
|
||||||
|
.center_x(Length::Fill)
|
||||||
|
.center_y(Length::Fill)
|
||||||
|
.into();
|
||||||
|
};
|
||||||
|
|
||||||
|
let project = &item.project;
|
||||||
|
let selected_title = self
|
||||||
|
.selected_session(item)
|
||||||
|
.map(|session| session.title.as_str())
|
||||||
|
.unwrap_or(project.name.as_str());
|
||||||
|
let header = row![
|
||||||
|
icon(ICON_FOLDER, 19),
|
||||||
|
text(selected_title).size(18),
|
||||||
|
icon(ICON_MORE, 18),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
]
|
||||||
|
.spacing(10)
|
||||||
|
.align_y(Alignment::Center);
|
||||||
|
|
||||||
|
let body: Element<'_, Message> = if let Some(session) = self.selected_session(item) {
|
||||||
|
let mut messages = column![].spacing(12);
|
||||||
|
if self.conversation.is_empty() {
|
||||||
|
messages = messages.push(
|
||||||
|
column![
|
||||||
|
text(&session.title).size(26),
|
||||||
|
text("Run DeepSeek locally with the Rust Metal engine.").size(14),
|
||||||
|
]
|
||||||
|
.spacing(8),
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
let markdown_style = markdown::Style::from_palette(app_theme().palette());
|
||||||
|
for (index, message) in self.conversation.iter().enumerate() {
|
||||||
|
let label = if message.user { "You" } else { "DS4" };
|
||||||
|
let active = self.generating && index + 1 == self.conversation.len();
|
||||||
|
let mut body = column![text(label).size(11)].spacing(5);
|
||||||
|
if let Some(reasoning) = &message.reasoning {
|
||||||
|
let reasoning_label =
|
||||||
|
match (message.reasoning_open, message.reasoning_complete, active) {
|
||||||
|
(true, false, true) => "▾ Thinking",
|
||||||
|
(false, false, true) => "› Thinking",
|
||||||
|
(true, false, false) => "▾ Reasoning (stopped)",
|
||||||
|
(false, false, false) => "› Reasoning (stopped)",
|
||||||
|
(true, true, _) => "▾ Reasoning",
|
||||||
|
(false, true, _) => "› Reasoning",
|
||||||
|
};
|
||||||
|
body = body.push(
|
||||||
|
button(text(reasoning_label).size(12))
|
||||||
|
.padding(0)
|
||||||
|
.style(button::text)
|
||||||
|
.on_press(Message::ToggleReasoning(index)),
|
||||||
|
);
|
||||||
|
if message.reasoning_open {
|
||||||
|
body = body.push(
|
||||||
|
text(if reasoning.is_empty() && active {
|
||||||
|
"Thinking…"
|
||||||
|
} else {
|
||||||
|
reasoning
|
||||||
|
})
|
||||||
|
.size(13)
|
||||||
|
.color(muted_text()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !message.content.is_empty() {
|
||||||
|
if message.user || message.markdown.is_empty() {
|
||||||
|
let content = if message.reasoning.is_some() {
|
||||||
|
message.content.trim_start()
|
||||||
|
} else {
|
||||||
|
&message.content
|
||||||
|
};
|
||||||
|
body = body.push(text(content).size(14));
|
||||||
|
} else {
|
||||||
|
body = body.push(
|
||||||
|
markdown::view(
|
||||||
|
&message.markdown,
|
||||||
|
markdown::Settings::with_text_size(14),
|
||||||
|
markdown_style,
|
||||||
|
)
|
||||||
|
.map(Message::OpenLink),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
} else if active && message.reasoning.is_none() {
|
||||||
|
body = body.push(text("Loading model…").size(14));
|
||||||
|
}
|
||||||
|
let user = message.user;
|
||||||
|
messages = messages.push(
|
||||||
|
container(body)
|
||||||
|
.padding(14)
|
||||||
|
.width(Length::Fill)
|
||||||
|
.style(move |theme| chat_message_style(theme, user)),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let composer = text_input("Ask DS4Server anything…", &self.composer)
|
||||||
|
.on_input(Message::ComposerChanged)
|
||||||
|
.on_submit(Message::SubmitPrompt)
|
||||||
|
.padding(12)
|
||||||
|
.size(14);
|
||||||
|
let action = if self.generating {
|
||||||
|
action_button(text("Stop").size(12)).on_press(Message::StopGeneration)
|
||||||
|
} else if self.composer.trim().is_empty() {
|
||||||
|
action_button(icon(ICON_SEND, 18)).padding(8)
|
||||||
|
} else {
|
||||||
|
action_button(icon(ICON_SEND, 18))
|
||||||
|
.padding(8)
|
||||||
|
.on_press(Message::SubmitPrompt)
|
||||||
|
};
|
||||||
|
let context_fraction = if self.context_limit == 0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
self.context_used.min(self.context_limit) as f32 / self.context_limit as f32
|
||||||
|
};
|
||||||
|
let conversation = column![
|
||||||
|
scrollable(messages)
|
||||||
|
.id(chat_scroll_id())
|
||||||
|
.height(Length::Fill),
|
||||||
|
container(
|
||||||
|
column![
|
||||||
|
composer,
|
||||||
|
progress_bar(0.0..=1.0, context_fraction).height(3),
|
||||||
|
row![
|
||||||
|
icon(ICON_PAPERCLIP, 19),
|
||||||
|
text(format!(
|
||||||
|
"{} / {} tokens ({:.0}%) • {}",
|
||||||
|
self.context_used,
|
||||||
|
self.context_limit,
|
||||||
|
context_fraction * 100.0,
|
||||||
|
self.tokens_per_second.map_or_else(
|
||||||
|
|| "— tok/s".to_owned(),
|
||||||
|
|speed| format!("{speed:.1} tok/s")
|
||||||
|
)
|
||||||
|
))
|
||||||
|
.size(11)
|
||||||
|
.color(muted_text()),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
icon(ICON_MODEL, 16),
|
||||||
|
text(
|
||||||
|
ModelChoice::from_id(&self.preferences.selected_model)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.to_string(),
|
||||||
|
)
|
||||||
|
.size(12),
|
||||||
|
action,
|
||||||
|
]
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
]
|
||||||
|
.spacing(8),
|
||||||
|
)
|
||||||
|
.padding(16)
|
||||||
|
.width(Length::Fill)
|
||||||
|
.style(overview_style),
|
||||||
|
]
|
||||||
|
.height(Length::Fill)
|
||||||
|
.spacing(8);
|
||||||
|
container(conversation)
|
||||||
|
.max_width(860)
|
||||||
|
.center_x(Length::Fill)
|
||||||
|
.height(Length::Fill)
|
||||||
|
.into()
|
||||||
|
} else {
|
||||||
|
container(
|
||||||
|
column![
|
||||||
|
text(if item.sessions.is_empty() {
|
||||||
|
"No sessions yet"
|
||||||
|
} else {
|
||||||
|
"Choose a session"
|
||||||
|
})
|
||||||
|
.size(24),
|
||||||
|
text("Create a session above or select one from the sidebar.").size(14),
|
||||||
|
]
|
||||||
|
.spacing(8)
|
||||||
|
.align_x(Alignment::Center),
|
||||||
|
)
|
||||||
|
.center_x(Length::Fill)
|
||||||
|
.center_y(Length::Fill)
|
||||||
|
.into()
|
||||||
|
};
|
||||||
|
|
||||||
|
container(column![header, body].spacing(24))
|
||||||
|
.width(Length::Fill)
|
||||||
|
.height(Length::Fill)
|
||||||
|
.padding(24)
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
}
|
||||||
296
src/app/view/model_manager.rs
Normal file
296
src/app/view/model_manager.rs
Normal file
@@ -0,0 +1,296 @@
|
|||||||
|
use super::*;
|
||||||
|
use iced::widget::column;
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn model_manager(&self) -> Element<'_, Message> {
|
||||||
|
let busy = matches!(self.model_download, ModelDownload::Active(_));
|
||||||
|
let mut artifacts = column![];
|
||||||
|
for (index, artifact) in model::managed_artifacts(&models_path()).iter().enumerate() {
|
||||||
|
if index > 0 {
|
||||||
|
artifacts = artifacts.push(horizontal_rule(1));
|
||||||
|
}
|
||||||
|
artifacts = artifacts.push(model_artifact_row(artifact, busy));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut content = column![
|
||||||
|
text("Model Manager").size(26),
|
||||||
|
text("Download, verify, or remove locally stored model files.")
|
||||||
|
.size(14)
|
||||||
|
.color(muted_text()),
|
||||||
|
]
|
||||||
|
.spacing(8);
|
||||||
|
if !matches!(self.model_download, ModelDownload::Idle) {
|
||||||
|
content = content.push(
|
||||||
|
container(model_download_status(&self.model_download))
|
||||||
|
.padding(16)
|
||||||
|
.style(overview_style),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if let Some(error) = &self.error {
|
||||||
|
content = content.push(text(error).style(iced::widget::text::danger));
|
||||||
|
}
|
||||||
|
content = content.push(Space::with_height(8)).push(
|
||||||
|
container(scrollable(artifacts).height(Length::Fill))
|
||||||
|
.height(Length::Fill)
|
||||||
|
.style(overview_style),
|
||||||
|
);
|
||||||
|
|
||||||
|
let base: Element<'_, Message> = container(content)
|
||||||
|
.width(Length::Fill)
|
||||||
|
.height(Length::Fill)
|
||||||
|
.padding(28)
|
||||||
|
.into();
|
||||||
|
let Some(artifact) = self.pending_model_delete else {
|
||||||
|
return base;
|
||||||
|
};
|
||||||
|
let confirmation = container(
|
||||||
|
column![
|
||||||
|
text("Delete model file?").size(22),
|
||||||
|
text(format!(
|
||||||
|
"Delete {artifact}, including any resumable partial download?"
|
||||||
|
))
|
||||||
|
.size(13),
|
||||||
|
row![
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
action_button("Cancel").on_press(Message::CancelDeleteArtifact),
|
||||||
|
danger_button("Delete").on_press(Message::ConfirmDeleteArtifact),
|
||||||
|
]
|
||||||
|
.spacing(8),
|
||||||
|
]
|
||||||
|
.spacing(14),
|
||||||
|
)
|
||||||
|
.padding(22)
|
||||||
|
.width(460)
|
||||||
|
.style(overview_style);
|
||||||
|
stack![
|
||||||
|
base,
|
||||||
|
opaque(
|
||||||
|
container(confirmation)
|
||||||
|
.center_x(Length::Fill)
|
||||||
|
.center_y(Length::Fill)
|
||||||
|
.style(|_| container::Style::default()
|
||||||
|
.background(Color::from_rgba8(0, 0, 0, 0.68)))
|
||||||
|
)
|
||||||
|
]
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model_artifact_row(artifact: &ManagedArtifact, busy: bool) -> Element<'static, Message> {
|
||||||
|
let status = match artifact.state {
|
||||||
|
ManagedArtifactState::Missing => "Not downloaded",
|
||||||
|
ManagedArtifactState::Partial => "Partial download",
|
||||||
|
ManagedArtifactState::NeedsVerification => "Downloaded; verification required",
|
||||||
|
ManagedArtifactState::Ready => "Ready and verified",
|
||||||
|
};
|
||||||
|
let download_label = match artifact.state {
|
||||||
|
ManagedArtifactState::Ready => "Downloaded",
|
||||||
|
_ if artifact.stored > 0 => "Resume",
|
||||||
|
_ => "Download",
|
||||||
|
};
|
||||||
|
let download = if artifact.state == ManagedArtifactState::Ready || busy {
|
||||||
|
action_button(download_label)
|
||||||
|
} else {
|
||||||
|
action_button(download_label).on_press(Message::DownloadArtifact(artifact.id))
|
||||||
|
};
|
||||||
|
let validate = if artifact.can_validate() && !busy {
|
||||||
|
action_button("Validate").on_press(Message::ValidateArtifact(artifact.id))
|
||||||
|
} else {
|
||||||
|
action_button("Validate")
|
||||||
|
};
|
||||||
|
let delete = if artifact.stored > 0 && !busy {
|
||||||
|
danger_button("Delete").on_press(Message::DeleteArtifact(artifact.id))
|
||||||
|
} else {
|
||||||
|
danger_button("Delete")
|
||||||
|
};
|
||||||
|
|
||||||
|
container(
|
||||||
|
row![
|
||||||
|
icon(ICON_MODEL, 28),
|
||||||
|
column![
|
||||||
|
text(artifact.id.to_string()).size(16),
|
||||||
|
text(status).size(13).color(muted_text()),
|
||||||
|
text(format!(
|
||||||
|
"{} on disk • {} expected",
|
||||||
|
format_bytes(artifact.stored),
|
||||||
|
format_bytes(artifact.expected),
|
||||||
|
))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
]
|
||||||
|
.spacing(5),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
row![download, validate, delete]
|
||||||
|
.spacing(8)
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
]
|
||||||
|
.spacing(16)
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
)
|
||||||
|
.width(Length::Fill)
|
||||||
|
.padding([18, 20])
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn download_status_bar(download: &ActiveDownload) -> Element<'_, Message> {
|
||||||
|
let progress = &download.progress;
|
||||||
|
let percent = progress.fraction() * 100.0;
|
||||||
|
let measurement = if progress.verification.is_some() {
|
||||||
|
format!(
|
||||||
|
"{} verified of {} ({percent:.1}%)",
|
||||||
|
format_bytes(progress.completed()),
|
||||||
|
format_bytes(progress.active_total()),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"{} of {} ({percent:.1}%)",
|
||||||
|
format_bytes(progress.completed()),
|
||||||
|
format_bytes(progress.active_total()),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let transfer = if download.bytes_per_second > 0.0 && progress.remaining() > 0 {
|
||||||
|
format!(
|
||||||
|
"{}/s • about {} remaining",
|
||||||
|
format_bytes(download.bytes_per_second as u64),
|
||||||
|
format_duration(progress.remaining() as f64 / download.bytes_per_second),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
"Calculating time remaining…".to_owned()
|
||||||
|
};
|
||||||
|
let stop = if download.stopping {
|
||||||
|
danger_button("Stopping…")
|
||||||
|
} else {
|
||||||
|
danger_button("Stop").on_press(Message::StopModelDownload)
|
||||||
|
};
|
||||||
|
|
||||||
|
container(
|
||||||
|
row![
|
||||||
|
text(phase_text(progress.phase)).size(12),
|
||||||
|
progress_bar(0.0..=1.0, progress.fraction())
|
||||||
|
.width(180)
|
||||||
|
.height(7),
|
||||||
|
text(measurement).size(12),
|
||||||
|
text(transfer).size(12),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
stop,
|
||||||
|
]
|
||||||
|
.spacing(12)
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
)
|
||||||
|
.width(Length::Fill)
|
||||||
|
.padding([7, 14])
|
||||||
|
.style(sidebar_style)
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model_download_status(download: &ModelDownload) -> Element<'_, Message> {
|
||||||
|
let (progress, heading, speed, failed) = match download {
|
||||||
|
ModelDownload::Idle => unreachable!("idle operations are not displayed"),
|
||||||
|
ModelDownload::Active(active) => (
|
||||||
|
&active.progress,
|
||||||
|
if active.stopping {
|
||||||
|
format!("Stopping {}…", active.artifact)
|
||||||
|
} else {
|
||||||
|
phase_text(active.progress.phase)
|
||||||
|
},
|
||||||
|
active.bytes_per_second,
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
ModelDownload::Complete(artifact, operation, progress) => (
|
||||||
|
progress,
|
||||||
|
match operation {
|
||||||
|
ModelOperation::Download => format!("{artifact} is downloaded and verified."),
|
||||||
|
ModelOperation::Validate => format!("{artifact} passed validation."),
|
||||||
|
},
|
||||||
|
0.0,
|
||||||
|
false,
|
||||||
|
),
|
||||||
|
ModelDownload::Failed(artifact, error, progress) => (
|
||||||
|
progress,
|
||||||
|
format!("{artifact} operation failed: {error}"),
|
||||||
|
0.0,
|
||||||
|
true,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
let percent = progress.fraction() * 100.0;
|
||||||
|
let measurement = if progress.verification.is_some() {
|
||||||
|
format!(
|
||||||
|
"{} verified of {} ({percent:.1}%) • {} remaining",
|
||||||
|
format_bytes(progress.completed()),
|
||||||
|
format_bytes(progress.active_total()),
|
||||||
|
format_bytes(progress.remaining()),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"{} of {} ({percent:.1}%) • {} remaining",
|
||||||
|
format_bytes(progress.completed()),
|
||||||
|
format_bytes(progress.active_total()),
|
||||||
|
format_bytes(progress.remaining()),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let heading = if failed {
|
||||||
|
text(heading).size(12).style(iced::widget::text::danger)
|
||||||
|
} else {
|
||||||
|
text(heading).size(12)
|
||||||
|
};
|
||||||
|
let mut heading_row = row![heading, Space::with_width(Length::Fill)]
|
||||||
|
.align_y(Alignment::Center)
|
||||||
|
.spacing(8);
|
||||||
|
if let ModelDownload::Active(active) = download {
|
||||||
|
heading_row = heading_row.push(if active.stopping {
|
||||||
|
danger_button("Stopping…")
|
||||||
|
} else {
|
||||||
|
danger_button("Stop").on_press(Message::StopModelDownload)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let mut status = column![
|
||||||
|
heading_row,
|
||||||
|
progress_bar(0.0..=1.0, progress.fraction()).height(8),
|
||||||
|
text(measurement).size(12),
|
||||||
|
]
|
||||||
|
.spacing(6);
|
||||||
|
if speed > 0.0 && progress.remaining() > 0 {
|
||||||
|
status = status.push(
|
||||||
|
text(format!(
|
||||||
|
"{}/s • about {} remaining",
|
||||||
|
format_bytes(speed as u64),
|
||||||
|
format_duration(progress.remaining() as f64 / speed),
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
status.into()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn phase_text(phase: DownloadPhase) -> String {
|
||||||
|
match phase {
|
||||||
|
DownloadPhase::Pending(artifact) => format!("Ready to download {artifact}."),
|
||||||
|
DownloadPhase::Downloading(artifact) => format!("Downloading {artifact}…"),
|
||||||
|
DownloadPhase::Verifying(artifact) => format!("Verifying {artifact}…"),
|
||||||
|
DownloadPhase::Complete => "Download complete.".to_owned(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn format_bytes(bytes: u64) -> String {
|
||||||
|
const GB: f64 = 1_000_000_000.0;
|
||||||
|
const MB: f64 = 1_000_000.0;
|
||||||
|
if bytes >= 1_000_000_000 {
|
||||||
|
format!("{:.1} GB", bytes as f64 / GB)
|
||||||
|
} else {
|
||||||
|
format!("{:.1} MB", bytes as f64 / MB)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn format_duration(seconds: f64) -> String {
|
||||||
|
let seconds = seconds.max(0.0).round() as u64;
|
||||||
|
let hours = seconds / 3600;
|
||||||
|
let minutes = seconds % 3600 / 60;
|
||||||
|
if hours > 0 {
|
||||||
|
format!("{hours}h {minutes}m")
|
||||||
|
} else if minutes > 0 {
|
||||||
|
format!("{minutes}m {}s", seconds % 60)
|
||||||
|
} else {
|
||||||
|
format!("{seconds}s")
|
||||||
|
}
|
||||||
|
}
|
||||||
440
src/app/view/preferences.rs
Normal file
440
src/app/view/preferences.rs
Normal file
@@ -0,0 +1,440 @@
|
|||||||
|
use super::*;
|
||||||
|
use iced::widget::column;
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn preferences_panel(&self) -> Element<'_, Message> {
|
||||||
|
let dspark_toggle: Option<fn(bool) -> Message> = self
|
||||||
|
.preference_draft
|
||||||
|
.model
|
||||||
|
.supports_dspark()
|
||||||
|
.then_some(Message::PreferenceDsparkChanged);
|
||||||
|
let dspark = checkbox(
|
||||||
|
"Enable DSpark for this model",
|
||||||
|
self.preference_draft.dspark_enabled,
|
||||||
|
)
|
||||||
|
.on_toggle_maybe(dspark_toggle);
|
||||||
|
let glm_mtp_toggle: Option<fn(bool) -> Message> = (self.preference_draft.model
|
||||||
|
== ModelChoice::Glm52)
|
||||||
|
.then_some(Message::PreferenceGlmMtpChanged);
|
||||||
|
let glm_mtp_timing_toggle: Option<fn(bool) -> Message> = (self.preference_draft.model
|
||||||
|
== ModelChoice::Glm52)
|
||||||
|
.then_some(Message::PreferenceGlmMtpTimingChanged);
|
||||||
|
let dspark_strict_toggle: Option<fn(bool) -> Message> = self
|
||||||
|
.preference_draft
|
||||||
|
.model
|
||||||
|
.supports_dspark()
|
||||||
|
.then_some(Message::PreferenceDsparkStrictChanged);
|
||||||
|
let effective = self
|
||||||
|
.preference_draft
|
||||||
|
.generation()
|
||||||
|
.and_then(|generation| {
|
||||||
|
self.preference_draft.runtime().and_then(|runtime| {
|
||||||
|
crate::settings::effective_settings(
|
||||||
|
self.preference_draft.model,
|
||||||
|
&generation,
|
||||||
|
&runtime,
|
||||||
|
&models_path(),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.ok();
|
||||||
|
let engine = effective.as_ref().map(|settings| &settings.engine);
|
||||||
|
let turn = effective.as_ref().map(|settings| &settings.turn);
|
||||||
|
let mut power = text_input("100", &self.preference_draft.power_percent);
|
||||||
|
let mut prefill = text_input("Automatic", &self.preference_draft.prefill_chunk);
|
||||||
|
let mut ssd_full_layers = text_input("Automatic", &self.preference_draft.ssd_full_layers);
|
||||||
|
let mut steering_file = text_input(
|
||||||
|
"Direction-vector file path",
|
||||||
|
&self.preference_draft.directional_steering_file,
|
||||||
|
);
|
||||||
|
let mut steering_ffn =
|
||||||
|
text_input("Automatic", &self.preference_draft.directional_steering_ffn);
|
||||||
|
let mut steering_attn = text_input("0", &self.preference_draft.directional_steering_attn);
|
||||||
|
let mut dspark_confidence = text_input(
|
||||||
|
"0.9 (DS4 default)",
|
||||||
|
&self.preference_draft.dspark_confidence_threshold,
|
||||||
|
);
|
||||||
|
if self.preference_draft.model != ModelChoice::Glm52 {
|
||||||
|
power = power.on_input(Message::PreferencePowerChanged);
|
||||||
|
prefill = prefill.on_input(Message::PreferencePrefillChunkChanged);
|
||||||
|
steering_file = steering_file.on_input(Message::PreferenceSteeringFileChanged);
|
||||||
|
steering_ffn = steering_ffn.on_input(Message::PreferenceSteeringFfnChanged);
|
||||||
|
steering_attn = steering_attn.on_input(Message::PreferenceSteeringAttnChanged);
|
||||||
|
} else {
|
||||||
|
ssd_full_layers = ssd_full_layers.on_input(Message::PreferenceSsdFullLayersChanged);
|
||||||
|
}
|
||||||
|
if self.preference_draft.model.supports_dspark() {
|
||||||
|
dspark_confidence =
|
||||||
|
dspark_confidence.on_input(Message::PreferenceDsparkConfidenceChanged);
|
||||||
|
}
|
||||||
|
|
||||||
|
let model_group = preference_group(
|
||||||
|
"MODEL & LIFECYCLE",
|
||||||
|
column![
|
||||||
|
pick_list(
|
||||||
|
&MODEL_CHOICES[..],
|
||||||
|
Some(self.preference_draft.model),
|
||||||
|
Message::PreferenceModelChanged,
|
||||||
|
)
|
||||||
|
.width(Length::Fill),
|
||||||
|
text(format!(
|
||||||
|
"Main: {}{}",
|
||||||
|
engine.map_or_else(
|
||||||
|
|| "Invalid settings".to_owned(),
|
||||||
|
|engine| engine.artifacts.model.display().to_string(),
|
||||||
|
),
|
||||||
|
engine
|
||||||
|
.and_then(|engine| engine.artifacts.mtp.as_ref())
|
||||||
|
.map_or_else(String::new, |path| format!(
|
||||||
|
" • support: {}",
|
||||||
|
path.display()
|
||||||
|
)),
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
row![
|
||||||
|
text_input("10", &self.preference_draft.idle_timeout_minutes)
|
||||||
|
.on_input(Message::PreferenceTimeoutChanged)
|
||||||
|
.width(90)
|
||||||
|
.padding(9),
|
||||||
|
text("minutes before unloading the model").size(13),
|
||||||
|
]
|
||||||
|
.spacing(10)
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
text("Enter a whole number from 1 to 1440.").size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let endpoint_group = preference_group(
|
||||||
|
"LOCAL ENDPOINT",
|
||||||
|
column![
|
||||||
|
preference_input_row(
|
||||||
|
"Port",
|
||||||
|
text_input("4000", &self.preference_draft.endpoint_port)
|
||||||
|
.on_input(Message::PreferenceEndpointPortChanged),
|
||||||
|
),
|
||||||
|
text("Listens on 127.0.0.1. Saving a changed port restarts the local endpoint.")
|
||||||
|
.size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let generation_group = preference_group(
|
||||||
|
"GENERATION",
|
||||||
|
column![
|
||||||
|
preference_input_row(
|
||||||
|
"Context tokens",
|
||||||
|
text_input("32768", &self.preference_draft.context_tokens)
|
||||||
|
.on_input(Message::PreferenceContextChanged),
|
||||||
|
),
|
||||||
|
preference_input_row(
|
||||||
|
"Maximum generated tokens",
|
||||||
|
text_input("50000", &self.preference_draft.max_generated_tokens)
|
||||||
|
.on_input(Message::PreferenceMaxTokensChanged),
|
||||||
|
),
|
||||||
|
text("System prompt").size(13),
|
||||||
|
text_input(
|
||||||
|
"You are a helpful assistant",
|
||||||
|
&self.preference_draft.system_prompt,
|
||||||
|
)
|
||||||
|
.on_input(Message::PreferenceSystemPromptChanged)
|
||||||
|
.padding(9),
|
||||||
|
Space::with_height(4),
|
||||||
|
text("SAMPLING & REASONING").size(11).color(muted_text()),
|
||||||
|
preference_input_row(
|
||||||
|
"Temperature",
|
||||||
|
text_input("DS4 default", &self.preference_draft.temperature)
|
||||||
|
.on_input(Message::PreferenceTemperatureChanged),
|
||||||
|
),
|
||||||
|
preference_input_row(
|
||||||
|
"Top-p",
|
||||||
|
text_input("DS4 default", &self.preference_draft.top_p)
|
||||||
|
.on_input(Message::PreferenceTopPChanged),
|
||||||
|
),
|
||||||
|
preference_input_row(
|
||||||
|
"Min-p",
|
||||||
|
text_input("DS4 default", &self.preference_draft.min_p)
|
||||||
|
.on_input(Message::PreferenceMinPChanged),
|
||||||
|
),
|
||||||
|
preference_input_row(
|
||||||
|
"Seed",
|
||||||
|
text_input("Random", &self.preference_draft.seed)
|
||||||
|
.on_input(Message::PreferenceSeedChanged),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text("Reasoning").size(13).width(Length::Fill),
|
||||||
|
pick_list(
|
||||||
|
&REASONING_MODES[..],
|
||||||
|
Some(self.preference_draft.reasoning_mode),
|
||||||
|
Message::PreferenceReasoningChanged,
|
||||||
|
)
|
||||||
|
.width(240),
|
||||||
|
]
|
||||||
|
.spacing(12)
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
text("Blank sampling values retain DS4's model-family defaults. Think Max needs at least 393216 context tokens.")
|
||||||
|
.size(12),
|
||||||
|
text(turn.map_or_else(
|
||||||
|
|| "Effective settings will appear after valid values are entered.".to_owned(),
|
||||||
|
|settings| format!(
|
||||||
|
"Effective: {} context • {} max • temp {} • top-p {} • min-p {} • seed {} • {} • system prompt {}",
|
||||||
|
settings.context_tokens,
|
||||||
|
settings.max_generated_tokens,
|
||||||
|
settings.temperature,
|
||||||
|
settings.top_p,
|
||||||
|
settings.min_p,
|
||||||
|
settings.seed.map_or_else(|| "random".to_owned(), |seed| seed.to_string()),
|
||||||
|
settings.reasoning_mode,
|
||||||
|
if settings.system_prompt.is_empty() { "off" } else { "on" },
|
||||||
|
),
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let execution_group = preference_group(
|
||||||
|
"EXECUTION",
|
||||||
|
column![
|
||||||
|
preference_input_row(
|
||||||
|
"CPU helper threads",
|
||||||
|
text_input("Automatic", &self.preference_draft.cpu_threads)
|
||||||
|
.on_input(Message::PreferenceCpuThreadsChanged),
|
||||||
|
),
|
||||||
|
preference_input_row("GPU power percent", power),
|
||||||
|
preference_input_row("Prefill chunk", prefill),
|
||||||
|
checkbox("Prefer exact quality kernels", self.preference_draft.quality)
|
||||||
|
.on_toggle(Message::PreferenceQualityChanged),
|
||||||
|
checkbox("Warm mapped weights at load time", self.preference_draft.warm_weights)
|
||||||
|
.on_toggle(Message::PreferenceWarmWeightsChanged),
|
||||||
|
text(if self.preference_draft.model == ModelChoice::Glm52 {
|
||||||
|
"GLM 5.2 uses full GPU power and selects prefill chunks automatically."
|
||||||
|
} else {
|
||||||
|
"Blank numeric values preserve DS4's automatic engine behavior."
|
||||||
|
})
|
||||||
|
.size(12),
|
||||||
|
text(engine.as_ref().map_or_else(
|
||||||
|
|| "Effective execution settings will appear after valid values are entered."
|
||||||
|
.to_owned(),
|
||||||
|
|engine| {
|
||||||
|
let settings = engine.execution;
|
||||||
|
format!(
|
||||||
|
"Metal engine: threads {} • power {}% • prefill {} • quality {} • warm weights {}",
|
||||||
|
if settings.cpu_threads == 0 { "auto".to_owned() } else { settings.cpu_threads.to_string() },
|
||||||
|
if settings.power_percent == 0 { 100 } else { settings.power_percent },
|
||||||
|
if settings.prefill_chunk == 0 { "auto".to_owned() } else { settings.prefill_chunk.to_string() },
|
||||||
|
if settings.quality { "on" } else { "off" },
|
||||||
|
if settings.warm_weights { "on" } else { "off" },
|
||||||
|
)
|
||||||
|
},
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let acceleration_group = preference_group(
|
||||||
|
"ACCELERATION & MEMORY",
|
||||||
|
column![
|
||||||
|
text("SPECULATIVE DECODING").size(11).color(muted_text()),
|
||||||
|
preference_input_row(
|
||||||
|
"MTP draft tokens",
|
||||||
|
text_input("1", &self.preference_draft.mtp_draft_tokens)
|
||||||
|
.on_input(Message::PreferenceMtpDraftChanged),
|
||||||
|
),
|
||||||
|
preference_input_row(
|
||||||
|
"MTP verifier margin",
|
||||||
|
text_input("3", &self.preference_draft.mtp_margin)
|
||||||
|
.on_input(Message::PreferenceMtpMarginChanged),
|
||||||
|
),
|
||||||
|
checkbox("Enable integrated GLM MTP", self.preference_draft.glm_mtp)
|
||||||
|
.on_toggle_maybe(glm_mtp_toggle),
|
||||||
|
checkbox(
|
||||||
|
"Log GLM MTP timing counters",
|
||||||
|
self.preference_draft.glm_mtp_timing,
|
||||||
|
)
|
||||||
|
.on_toggle_maybe(glm_mtp_timing_toggle),
|
||||||
|
dspark,
|
||||||
|
preference_input_row("DSpark confidence threshold", dspark_confidence),
|
||||||
|
checkbox(
|
||||||
|
"DSpark target-only decode",
|
||||||
|
self.preference_draft.dspark_strict,
|
||||||
|
)
|
||||||
|
.on_toggle_maybe(dspark_strict_toggle),
|
||||||
|
text(if self.preference_draft.model.supports_dspark() {
|
||||||
|
"DSpark uses the managed support artifact; entering a threshold or enabling strict mode also enables DSpark."
|
||||||
|
} else if self.preference_draft.model == ModelChoice::Glm52 {
|
||||||
|
"GLM MTP is integrated; DSpark is unavailable for this model."
|
||||||
|
} else {
|
||||||
|
"No managed MTP support artifact is available for this model."
|
||||||
|
})
|
||||||
|
.size(12),
|
||||||
|
text(engine.as_ref().map_or_else(
|
||||||
|
|| "Effective speculative settings will appear after valid values are entered."
|
||||||
|
.to_owned(),
|
||||||
|
|engine| {
|
||||||
|
let settings = engine.speculative;
|
||||||
|
format!(
|
||||||
|
"Engine: MTP draft {} • margin {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}",
|
||||||
|
settings.mtp_draft_tokens,
|
||||||
|
settings.mtp_margin,
|
||||||
|
if settings.glm_mtp { "on" } else { "off" },
|
||||||
|
if settings.glm_mtp_timing { "on" } else { "off" },
|
||||||
|
if settings.dspark { "on" } else { "off" },
|
||||||
|
settings.dspark_confidence_threshold,
|
||||||
|
if settings.dspark_confidence_threshold_set { " explicit" } else { " default" },
|
||||||
|
if settings.dspark_strict { "on" } else { "off" },
|
||||||
|
)
|
||||||
|
},
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
Space::with_height(6),
|
||||||
|
text("SSD STREAMING").size(11).color(muted_text()),
|
||||||
|
checkbox("Enable SSD-backed model streaming", self.preference_draft.ssd_streaming)
|
||||||
|
.on_toggle(Message::PreferenceSsdChanged),
|
||||||
|
checkbox("Skip automatic expert preload", self.preference_draft.ssd_streaming_cold)
|
||||||
|
.on_toggle(Message::PreferenceSsdColdChanged),
|
||||||
|
preference_input_row(
|
||||||
|
"Expert cache count or GiB",
|
||||||
|
text_input("Automatic, 128, or 64GB", &self.preference_draft.ssd_cache)
|
||||||
|
.on_input(Message::PreferenceSsdCacheChanged),
|
||||||
|
),
|
||||||
|
preference_input_row("Fully resident GLM layers", ssd_full_layers),
|
||||||
|
preference_input_row(
|
||||||
|
"Explicit expert preload count",
|
||||||
|
text_input("Automatic", &self.preference_draft.ssd_preload_experts)
|
||||||
|
.on_input(Message::PreferenceSsdPreloadChanged),
|
||||||
|
),
|
||||||
|
text("A blank full-layer value is automatic; an explicit 0 disables fully resident GLM layers. SSD streaming and DSpark are mutually exclusive.")
|
||||||
|
.size(12),
|
||||||
|
text(engine.as_ref().map_or_else(
|
||||||
|
|| "Effective SSD settings will appear after valid values are entered."
|
||||||
|
.to_owned(),
|
||||||
|
|engine| {
|
||||||
|
let settings = engine.ssd;
|
||||||
|
let cache = if settings.cache_bytes > 0 {
|
||||||
|
format!("{} GiB", settings.cache_bytes / GIB)
|
||||||
|
} else if settings.cache_experts > 0 {
|
||||||
|
format!("{} experts", settings.cache_experts)
|
||||||
|
} else {
|
||||||
|
"auto".to_owned()
|
||||||
|
};
|
||||||
|
format!(
|
||||||
|
"Engine: streaming {} • cold {} • cache {} • full layers {}{} • preload {}",
|
||||||
|
if settings.enabled { "on" } else { "off" },
|
||||||
|
if settings.cold { "on" } else { "off" },
|
||||||
|
cache,
|
||||||
|
settings.full_layers,
|
||||||
|
if settings.full_layers_set { " explicit" } else { " auto" },
|
||||||
|
if settings.preload_experts == 0 { "auto".to_owned() } else { settings.preload_experts.to_string() },
|
||||||
|
)
|
||||||
|
},
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let steering_group = preference_group(
|
||||||
|
"STEERING & DIAGNOSTICS",
|
||||||
|
column![
|
||||||
|
text("DIRECTIONAL STEERING").size(11).color(muted_text()),
|
||||||
|
text("Direction-vector file").size(13),
|
||||||
|
steering_file.padding(9),
|
||||||
|
preference_input_row("FFN scale", steering_ffn),
|
||||||
|
preference_input_row("Attention scale", steering_attn),
|
||||||
|
text(if self.preference_draft.model == ModelChoice::Glm52 {
|
||||||
|
"Directional steering is not supported for GLM 5.2."
|
||||||
|
} else {
|
||||||
|
"With a file and no explicit scale, DS4 defaults the FFN scale to 1. Scales accept -100 through 100."
|
||||||
|
})
|
||||||
|
.size(12),
|
||||||
|
text(engine.as_ref().map_or_else(
|
||||||
|
|| "Effective steering settings will appear after valid values are entered."
|
||||||
|
.to_owned(),
|
||||||
|
|engine| format!(
|
||||||
|
"Engine: file {} • FFN scale {} • attention scale {}",
|
||||||
|
if engine.steering.file.is_some() { "set" } else { "off" },
|
||||||
|
engine.steering.ffn_scale,
|
||||||
|
engine.steering.attention_scale,
|
||||||
|
),
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
Space::with_height(6),
|
||||||
|
text("ADVANCED DIAGNOSTICS").size(11).color(muted_text()),
|
||||||
|
preference_input_row(
|
||||||
|
"Simulated used memory (GiB)",
|
||||||
|
text_input("Disabled", &self.preference_draft.simulated_used_memory_gib)
|
||||||
|
.on_input(Message::PreferenceSimulatedMemoryChanged),
|
||||||
|
),
|
||||||
|
text("Routed expert profile output").size(13),
|
||||||
|
text_input("Output file path", &self.preference_draft.expert_profile_path)
|
||||||
|
.on_input(Message::PreferenceExpertProfileChanged)
|
||||||
|
.padding(9),
|
||||||
|
text(engine.as_ref().map_or_else(
|
||||||
|
|| "Effective diagnostic settings will appear after valid values are entered."
|
||||||
|
.to_owned(),
|
||||||
|
|engine| format!(
|
||||||
|
"{} load: simulated memory {} • expert profile {}",
|
||||||
|
engine.model,
|
||||||
|
if engine.diagnostics.simulated_used_memory_bytes == 0 {
|
||||||
|
"off".to_owned()
|
||||||
|
} else {
|
||||||
|
format!("{} GiB", engine.diagnostics.simulated_used_memory_bytes / GIB)
|
||||||
|
},
|
||||||
|
if engine.diagnostics.expert_profile_path.is_some() { "set" } else { "off" },
|
||||||
|
),
|
||||||
|
))
|
||||||
|
.size(12),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
);
|
||||||
|
let mut fields = column![
|
||||||
|
model_group,
|
||||||
|
endpoint_group,
|
||||||
|
generation_group,
|
||||||
|
execution_group,
|
||||||
|
acceleration_group,
|
||||||
|
steering_group,
|
||||||
|
]
|
||||||
|
.spacing(12);
|
||||||
|
|
||||||
|
if let Some(error) = &self.preference_error {
|
||||||
|
fields = fields.push(text(error).style(iced::widget::text::danger));
|
||||||
|
}
|
||||||
|
let header = row![
|
||||||
|
icon(ICON_SETTINGS, 22),
|
||||||
|
text("Preferences").size(24),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text("⌘,").size(12),
|
||||||
|
]
|
||||||
|
.spacing(10)
|
||||||
|
.align_y(Alignment::Center);
|
||||||
|
let footer = row![
|
||||||
|
action_button("Reset DS4 defaults").on_press(Message::ResetPreferences),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
action_button("Cancel").on_press(Message::DismissPanel),
|
||||||
|
action_button("Save").on_press(Message::SavePreferences),
|
||||||
|
]
|
||||||
|
.spacing(8);
|
||||||
|
|
||||||
|
let panel = container(
|
||||||
|
column![
|
||||||
|
header,
|
||||||
|
scrollable(container(fields).padding(iced::Padding::ZERO.right(18)))
|
||||||
|
.height(Length::Fill),
|
||||||
|
footer
|
||||||
|
]
|
||||||
|
.spacing(16),
|
||||||
|
)
|
||||||
|
.padding(24)
|
||||||
|
.width(700)
|
||||||
|
.height(Length::Fill)
|
||||||
|
.max_height(660)
|
||||||
|
.style(overview_style);
|
||||||
|
opaque(
|
||||||
|
container(panel)
|
||||||
|
.padding(24)
|
||||||
|
.center_x(Length::Fill)
|
||||||
|
.center_y(Length::Fill)
|
||||||
|
.style(|_| {
|
||||||
|
container::Style::default().background(Color::from_rgba8(0, 0, 0, 0.68))
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
371
src/app/view/stats.rs
Normal file
371
src/app/view/stats.rs
Normal file
@@ -0,0 +1,371 @@
|
|||||||
|
use super::*;
|
||||||
|
use iced::widget::column;
|
||||||
|
|
||||||
|
impl App {
|
||||||
|
pub(super) fn stats_dashboard(&self) -> Element<'_, Message> {
|
||||||
|
let stats = &self.metrics_snapshot;
|
||||||
|
let context_fraction = if stats.context_limit == 0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
stats.context_used.min(stats.context_limit) as f32 / stats.context_limit as f32
|
||||||
|
};
|
||||||
|
let cache_fraction = if stats.last_prompt_tokens == 0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
stats.last_cached_tokens as f32 / stats.last_prompt_tokens as f32
|
||||||
|
};
|
||||||
|
let cache_hit_fraction = if stats.kv_lookups == 0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
stats.kv_hits as f32 / stats.kv_lookups as f32
|
||||||
|
};
|
||||||
|
let endpoint = if stats.server_listening {
|
||||||
|
format!("Listening · 127.0.0.1:{}", stats.server_port)
|
||||||
|
} else {
|
||||||
|
"Stopped".to_owned()
|
||||||
|
};
|
||||||
|
let phase_color = match stats.phase {
|
||||||
|
crate::metrics::RuntimePhase::Generating => Color::from_rgb8(84, 170, 255),
|
||||||
|
crate::metrics::RuntimePhase::Prefilling | crate::metrics::RuntimePhase::Loading => {
|
||||||
|
Color::from_rgb8(240, 180, 70)
|
||||||
|
}
|
||||||
|
crate::metrics::RuntimePhase::Ready => Color::from_rgb8(72, 176, 112),
|
||||||
|
crate::metrics::RuntimePhase::Failed => Color::from_rgb8(220, 80, 86),
|
||||||
|
crate::metrics::RuntimePhase::Unloaded => muted_text(),
|
||||||
|
};
|
||||||
|
let heading = container(
|
||||||
|
row![
|
||||||
|
column![
|
||||||
|
text("Runtime observability").size(24),
|
||||||
|
text(format!(
|
||||||
|
"{} · {} · uptime {}",
|
||||||
|
stats.model,
|
||||||
|
stats.source.label(),
|
||||||
|
format_duration(stats.uptime_seconds as f64)
|
||||||
|
))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
]
|
||||||
|
.spacing(5),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
container(text(stats.phase.label()).size(12).color(phase_color))
|
||||||
|
.padding([7, 11])
|
||||||
|
.style(move |_| status_badge_style(phase_color)),
|
||||||
|
]
|
||||||
|
.align_y(Alignment::Center),
|
||||||
|
)
|
||||||
|
.padding(16)
|
||||||
|
.style(overview_style);
|
||||||
|
|
||||||
|
let headline = column![
|
||||||
|
row![
|
||||||
|
metric_card(
|
||||||
|
"DECODE",
|
||||||
|
format!("{:.1} tok/s", stats.decode_tokens_per_second),
|
||||||
|
format!(
|
||||||
|
"{} completion tokens total",
|
||||||
|
format_count(stats.completion_tokens)
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_card(
|
||||||
|
"PREFILL",
|
||||||
|
format!("{:.1} tok/s", stats.prefill_tokens_per_second),
|
||||||
|
format!("{} prompt tokens total", format_count(stats.prompt_tokens)),
|
||||||
|
),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
row![
|
||||||
|
metric_card(
|
||||||
|
"CONTEXT",
|
||||||
|
format!(
|
||||||
|
"{} / {}",
|
||||||
|
format_count(u64::from(stats.context_used)),
|
||||||
|
format_count(u64::from(stats.context_limit))
|
||||||
|
),
|
||||||
|
format!("{:.0}% occupied", context_fraction * 100.0),
|
||||||
|
),
|
||||||
|
metric_card(
|
||||||
|
"WORK",
|
||||||
|
format!(
|
||||||
|
"{} active · {} queued",
|
||||||
|
stats.http_active, stats.queue_depth
|
||||||
|
),
|
||||||
|
format!("{} runtime requests", format_count(stats.runtime_requests)),
|
||||||
|
),
|
||||||
|
]
|
||||||
|
.spacing(10),
|
||||||
|
]
|
||||||
|
.spacing(10);
|
||||||
|
|
||||||
|
let throughput = stats_panel(
|
||||||
|
"MODEL ACTIVITY · LAST 24 SECONDS",
|
||||||
|
column![
|
||||||
|
mini_chart(
|
||||||
|
&self.metrics_history,
|
||||||
|
|point| point.decode_tokens_per_second,
|
||||||
|
Color::from_rgb8(84, 170, 255),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text("Decode")
|
||||||
|
.size(12)
|
||||||
|
.color(Color::from_rgb8(84, 170, 255)),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text(format!("{:.1} tok/s", stats.decode_tokens_per_second))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
],
|
||||||
|
mini_chart(
|
||||||
|
&self.metrics_history,
|
||||||
|
|point| point.prefill_tokens_per_second,
|
||||||
|
Color::from_rgb8(157, 119, 255),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text("Prefill")
|
||||||
|
.size(12)
|
||||||
|
.color(Color::from_rgb8(157, 119, 255)),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text(format!("{:.1} tok/s", stats.prefill_tokens_per_second))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
],
|
||||||
|
]
|
||||||
|
.spacing(7)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
let requests = stats_panel(
|
||||||
|
"SERVER REQUEST RATE · LAST 24 SECONDS",
|
||||||
|
column![
|
||||||
|
mini_chart(
|
||||||
|
&self.metrics_history,
|
||||||
|
|point| point.http_requests_per_second,
|
||||||
|
Color::from_rgb8(72, 176, 112),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text(format!("{} requests", format_count(stats.http_requests))).size(12),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text(format!(
|
||||||
|
"{} errors · {} streaming",
|
||||||
|
stats.http_errors, stats.http_streaming_requests
|
||||||
|
))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
],
|
||||||
|
]
|
||||||
|
.spacing(7)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
let latest = self.metrics_history.back().copied().unwrap_or_default();
|
||||||
|
let kv_io = stats_panel(
|
||||||
|
"KV CHECKPOINT I/O · LAST 24 SECONDS",
|
||||||
|
column![
|
||||||
|
mini_chart(
|
||||||
|
&self.metrics_history,
|
||||||
|
|point| point.kv_read_bytes_per_second,
|
||||||
|
Color::from_rgb8(67, 194, 203),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text(if stats.kv_read_active {
|
||||||
|
"Disk read · active"
|
||||||
|
} else {
|
||||||
|
"Disk read"
|
||||||
|
})
|
||||||
|
.size(12)
|
||||||
|
.color(Color::from_rgb8(67, 194, 203)),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text(format_rate(latest.kv_read_bytes_per_second))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
],
|
||||||
|
mini_chart(
|
||||||
|
&self.metrics_history,
|
||||||
|
|point| point.kv_write_bytes_per_second,
|
||||||
|
Color::from_rgb8(240, 180, 70),
|
||||||
|
),
|
||||||
|
row![
|
||||||
|
text(if stats.kv_write_active {
|
||||||
|
"Disk write · active"
|
||||||
|
} else {
|
||||||
|
"Disk write"
|
||||||
|
})
|
||||||
|
.size(12)
|
||||||
|
.color(Color::from_rgb8(240, 180, 70)),
|
||||||
|
Space::with_width(Length::Fill),
|
||||||
|
text(format_rate(latest.kv_write_bytes_per_second))
|
||||||
|
.size(12)
|
||||||
|
.color(muted_text()),
|
||||||
|
],
|
||||||
|
]
|
||||||
|
.spacing(7)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let model = stats_panel(
|
||||||
|
"MODEL CORE",
|
||||||
|
column![
|
||||||
|
metric_row("State", stats.phase.label()),
|
||||||
|
metric_row("Loaded model", stats.model),
|
||||||
|
metric_row("Mapped weights", format_bytes(stats.model_bytes)),
|
||||||
|
metric_row("Tensors", format_count(stats.tensor_count)),
|
||||||
|
metric_row("Vocabulary", format_count(stats.vocabulary_size)),
|
||||||
|
metric_row("Last load", format_milliseconds(stats.model_load_ms)),
|
||||||
|
metric_row(
|
||||||
|
"Lifecycle",
|
||||||
|
format!(
|
||||||
|
"{} loads · {} unloads",
|
||||||
|
stats.model_loads, stats.model_unloads
|
||||||
|
),
|
||||||
|
),
|
||||||
|
]
|
||||||
|
.spacing(9)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
let runtime = stats_panel(
|
||||||
|
"GENERATION",
|
||||||
|
column![
|
||||||
|
metric_row("Last runtime", format_milliseconds(stats.last_runtime_ms)),
|
||||||
|
metric_row(
|
||||||
|
"Average runtime",
|
||||||
|
format_milliseconds(stats.average_runtime_ms)
|
||||||
|
),
|
||||||
|
metric_row("Last prompt", format_count(stats.last_prompt_tokens)),
|
||||||
|
metric_row("Last reused", format_count(stats.last_cached_tokens)),
|
||||||
|
metric_row(
|
||||||
|
"Last completion",
|
||||||
|
format_count(stats.last_completion_tokens)
|
||||||
|
),
|
||||||
|
metric_row("Cached tokens total", format_count(stats.cached_tokens)),
|
||||||
|
metric_row(
|
||||||
|
"Cache reuse",
|
||||||
|
format!("{:.0}% of last prompt", cache_fraction * 100.0),
|
||||||
|
),
|
||||||
|
metric_row(
|
||||||
|
"Results",
|
||||||
|
format!(
|
||||||
|
"{} completed · {} failed",
|
||||||
|
stats.completed_requests, stats.failed_requests
|
||||||
|
),
|
||||||
|
),
|
||||||
|
]
|
||||||
|
.spacing(9)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
let cache = stats_panel(
|
||||||
|
"KV CACHE",
|
||||||
|
column![
|
||||||
|
metric_row(
|
||||||
|
"Total",
|
||||||
|
format!(
|
||||||
|
"{} · {} files",
|
||||||
|
format_bytes(stats.kv_bytes),
|
||||||
|
stats.kv_files
|
||||||
|
)
|
||||||
|
),
|
||||||
|
metric_row(
|
||||||
|
"Local sessions",
|
||||||
|
format!(
|
||||||
|
"{} · {} files",
|
||||||
|
format_bytes(stats.local_kv_bytes),
|
||||||
|
stats.local_kv_files
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_row(
|
||||||
|
"HTTP transient",
|
||||||
|
format!(
|
||||||
|
"{} · {} files",
|
||||||
|
format_bytes(stats.http_kv_bytes),
|
||||||
|
stats.http_kv_files
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_row("Checkpoint writes", format_count(stats.checkpoint_writes)),
|
||||||
|
metric_row(
|
||||||
|
"Exact hits",
|
||||||
|
format!(
|
||||||
|
"{} · {} memory / {} disk",
|
||||||
|
stats.kv_hits, stats.kv_memory_hits, stats.kv_disk_hits
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_row(
|
||||||
|
"Misses",
|
||||||
|
format!("{} · {} invalid", stats.kv_misses, stats.kv_invalid),
|
||||||
|
),
|
||||||
|
metric_row("Lookups", format_count(stats.kv_lookups)),
|
||||||
|
metric_row(
|
||||||
|
"Exact hit rate",
|
||||||
|
format!("{:.1}%", cache_hit_fraction * 100.0)
|
||||||
|
),
|
||||||
|
metric_row("Prefix hits", format_count(stats.kv_prefix_hits)),
|
||||||
|
metric_row(
|
||||||
|
"Reads",
|
||||||
|
format!(
|
||||||
|
"{} · {} · {} errors · last {}",
|
||||||
|
stats.kv_read_operations,
|
||||||
|
format_bytes(stats.kv_read_bytes),
|
||||||
|
stats.kv_read_errors,
|
||||||
|
format_milliseconds(stats.last_kv_read_ms),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_row(
|
||||||
|
"Writes",
|
||||||
|
format!(
|
||||||
|
"{} · {} · {} errors · last {}",
|
||||||
|
stats.kv_write_operations,
|
||||||
|
format_bytes(stats.kv_write_bytes),
|
||||||
|
stats.kv_write_errors,
|
||||||
|
format_milliseconds(stats.last_kv_write_ms),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
progress_bar(0.0..=1.0, cache_fraction.min(1.0)).height(4),
|
||||||
|
]
|
||||||
|
.spacing(9)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
let server = stats_panel(
|
||||||
|
"LOCAL SERVER",
|
||||||
|
column![
|
||||||
|
metric_row("Endpoint", endpoint),
|
||||||
|
metric_row("Active", format_count(u64::from(stats.http_active))),
|
||||||
|
metric_row("Completed", format_count(stats.http_completed)),
|
||||||
|
metric_row("Chat completions", format_count(stats.http_chat_requests)),
|
||||||
|
metric_row("Model queries", format_count(stats.http_model_requests)),
|
||||||
|
metric_row(
|
||||||
|
"Runtime sources",
|
||||||
|
format!(
|
||||||
|
"{} local · {} HTTP",
|
||||||
|
stats.local_requests, stats.endpoint_generations
|
||||||
|
),
|
||||||
|
),
|
||||||
|
metric_row("Received", format_bytes(stats.http_bytes_received)),
|
||||||
|
metric_row("Last latency", format_milliseconds(stats.last_http_ms)),
|
||||||
|
metric_row(
|
||||||
|
"Average latency",
|
||||||
|
format_milliseconds(stats.average_http_ms)
|
||||||
|
),
|
||||||
|
]
|
||||||
|
.spacing(9)
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
|
||||||
|
scrollable(
|
||||||
|
container(
|
||||||
|
column![
|
||||||
|
heading,
|
||||||
|
headline,
|
||||||
|
throughput,
|
||||||
|
kv_io,
|
||||||
|
requests,
|
||||||
|
row![model, runtime].spacing(10),
|
||||||
|
row![cache, server].spacing(10),
|
||||||
|
text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.")
|
||||||
|
.size(11)
|
||||||
|
.color(muted_text()),
|
||||||
|
]
|
||||||
|
.spacing(12),
|
||||||
|
)
|
||||||
|
.padding(24)
|
||||||
|
.max_width(960)
|
||||||
|
.center_x(Length::Fill),
|
||||||
|
)
|
||||||
|
.height(Length::Fill)
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
}
|
||||||
997
src/engine.rs
997
src/engine.rs
File diff suppressed because it is too large
Load Diff
1286
src/engine/metal.rs
1286
src/engine/metal.rs
File diff suppressed because it is too large
Load Diff
367
src/engine/metal/checkpoint.rs
Normal file
367
src/engine/metal/checkpoint.rs
Normal file
@@ -0,0 +1,367 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
impl Executor {
|
||||||
|
pub(in crate::engine) fn save_checkpoint(
|
||||||
|
&mut self,
|
||||||
|
path: &Path,
|
||||||
|
tag: [u8; 32],
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<(), String> {
|
||||||
|
if let Some(parent) = path.parent() {
|
||||||
|
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
|
||||||
|
}
|
||||||
|
let temporary = path.with_extension("tmp");
|
||||||
|
let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
|
||||||
|
let mut reported = 0_u64;
|
||||||
|
let mut pending = 0_u64;
|
||||||
|
self.write_checkpoint(&mut file, tag, &mut |bytes| {
|
||||||
|
reported += bytes;
|
||||||
|
pending += bytes;
|
||||||
|
if pending >= CHECKPOINT_IO_CHUNK as u64 {
|
||||||
|
progress(pending);
|
||||||
|
pending = 0;
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
progress(pending);
|
||||||
|
let total = file.metadata().map_err(|error| error.to_string())?.len();
|
||||||
|
progress(total.saturating_sub(reported));
|
||||||
|
file.sync_all().map_err(|error| error.to_string())?;
|
||||||
|
fs::rename(&temporary, path).map_err(|error| error.to_string())?;
|
||||||
|
self.checkpoint_tag = tag;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(in crate::engine) fn load_checkpoint(
|
||||||
|
&mut self,
|
||||||
|
path: &Path,
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<bool, String> {
|
||||||
|
let mut file = match File::open(path) {
|
||||||
|
Ok(file) => file,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
||||||
|
Err(error) => return Err(error.to_string()),
|
||||||
|
};
|
||||||
|
let mut reported = 0_u64;
|
||||||
|
let mut pending = 0_u64;
|
||||||
|
self.read_checkpoint(&mut file, &mut |bytes| {
|
||||||
|
reported += bytes;
|
||||||
|
pending += bytes;
|
||||||
|
if pending >= CHECKPOINT_IO_CHUNK as u64 {
|
||||||
|
progress(pending);
|
||||||
|
pending = 0;
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
progress(pending);
|
||||||
|
let total = file.metadata().map_err(|error| error.to_string())?.len();
|
||||||
|
progress(total.saturating_sub(reported));
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_checkpoint(
|
||||||
|
&self,
|
||||||
|
file: &mut File,
|
||||||
|
tag: [u8; 32],
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let shape = self.model.shape;
|
||||||
|
file.write_all(CHECKPOINT_MAGIC)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
for value in [
|
||||||
|
CHECKPOINT_VERSION,
|
||||||
|
self.session.context,
|
||||||
|
self.session.raw_cap,
|
||||||
|
shape.layers,
|
||||||
|
shape.head_dim as u32,
|
||||||
|
shape.indexer_head_dim as u32,
|
||||||
|
shape.vocab as u32,
|
||||||
|
u32::from(self.quality),
|
||||||
|
] {
|
||||||
|
write_u32(file, value)?;
|
||||||
|
}
|
||||||
|
write_u64(file, self.model.main.len())?;
|
||||||
|
write_u64(file, self.model_modified.0)?;
|
||||||
|
write_u32(file, self.model_modified.1)?;
|
||||||
|
for weight in [self.weights.token_embedding, self.weights.output] {
|
||||||
|
write_u64(file, weight.offset)?;
|
||||||
|
write_u64(file, weight.bytes)?;
|
||||||
|
write_u32(file, weight.kind)?;
|
||||||
|
}
|
||||||
|
file.write_all(&tag).map_err(|error| error.to_string())?;
|
||||||
|
let token_count = u32::try_from(self.tokens.len())
|
||||||
|
.map_err(|_| "KV checkpoint has too many tokens".to_owned())?;
|
||||||
|
let raw_live = token_count.min(self.session.raw_cap);
|
||||||
|
write_u32(file, token_count)?;
|
||||||
|
write_u32(file, raw_live)?;
|
||||||
|
for &token in &self.tokens {
|
||||||
|
write_u32(file, token as u32)?;
|
||||||
|
}
|
||||||
|
for &logit in &self.logits {
|
||||||
|
write_u32(file, logit.to_bits())?;
|
||||||
|
}
|
||||||
|
for layer in &self.session.layers {
|
||||||
|
write_u32(
|
||||||
|
file,
|
||||||
|
layer.compression.as_ref().map_or(0, |state| state.rows),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
for layer in &self.session.layers {
|
||||||
|
write_u32(file, layer.indexer.as_ref().map_or(0, |state| state.rows))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
||||||
|
let raw_first = token_count - raw_live;
|
||||||
|
for layer in &self.session.layers {
|
||||||
|
for position in raw_first..token_count {
|
||||||
|
let physical = position % self.session.raw_cap;
|
||||||
|
write_buffer(
|
||||||
|
file,
|
||||||
|
&layer.raw_cache,
|
||||||
|
u64::from(physical) * shape.head_dim * 4,
|
||||||
|
shape.head_dim * 4,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if let Some(state) = &layer.compression {
|
||||||
|
write_buffer(
|
||||||
|
file,
|
||||||
|
&state.cache,
|
||||||
|
0,
|
||||||
|
u64::from(state.rows) * shape.head_dim * 2,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
|
||||||
|
write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
||||||
|
write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
||||||
|
}
|
||||||
|
if let Some(state) = &layer.indexer {
|
||||||
|
write_buffer(
|
||||||
|
file,
|
||||||
|
&state.cache,
|
||||||
|
0,
|
||||||
|
u64::from(state.rows) * shape.indexer_head_dim * 4,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
|
||||||
|
write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
||||||
|
write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_checkpoint(
|
||||||
|
&mut self,
|
||||||
|
file: &mut File,
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let mut magic = [0; 8];
|
||||||
|
file.read_exact(&mut magic)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
if &magic != CHECKPOINT_MAGIC {
|
||||||
|
return Err("KV checkpoint has an invalid signature".into());
|
||||||
|
}
|
||||||
|
let shape = self.model.shape;
|
||||||
|
let header = [
|
||||||
|
CHECKPOINT_VERSION,
|
||||||
|
self.session.context,
|
||||||
|
self.session.raw_cap,
|
||||||
|
shape.layers,
|
||||||
|
shape.head_dim as u32,
|
||||||
|
shape.indexer_head_dim as u32,
|
||||||
|
shape.vocab as u32,
|
||||||
|
u32::from(self.quality),
|
||||||
|
];
|
||||||
|
for expected in header {
|
||||||
|
if read_u32(file)? != expected {
|
||||||
|
return Err("KV checkpoint does not match the current executor".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if read_u64(file)? != self.model.main.len() {
|
||||||
|
return Err("KV checkpoint was written for a different model".into());
|
||||||
|
}
|
||||||
|
if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 {
|
||||||
|
return Err("KV checkpoint model file has changed".into());
|
||||||
|
}
|
||||||
|
for weight in [self.weights.token_embedding, self.weights.output] {
|
||||||
|
if read_u64(file)? != weight.offset
|
||||||
|
|| read_u64(file)? != weight.bytes
|
||||||
|
|| read_u32(file)? != weight.kind
|
||||||
|
{
|
||||||
|
return Err("KV checkpoint was written for a different model layout".into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut checkpoint_tag = [0; 32];
|
||||||
|
file.read_exact(&mut checkpoint_tag)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
let token_count = read_u32(file)?;
|
||||||
|
let raw_live = read_u32(file)?;
|
||||||
|
if token_count > self.session.context || raw_live != token_count.min(self.session.raw_cap) {
|
||||||
|
return Err("KV checkpoint token count is invalid".into());
|
||||||
|
}
|
||||||
|
let mut tokens = Vec::with_capacity(token_count as usize);
|
||||||
|
for _ in 0..token_count {
|
||||||
|
let token = read_u32(file)?;
|
||||||
|
if u64::from(token) >= shape.vocab {
|
||||||
|
return Err("KV checkpoint contains an invalid token".into());
|
||||||
|
}
|
||||||
|
tokens.push(token as i32);
|
||||||
|
}
|
||||||
|
let mut logits = Vec::with_capacity(shape.vocab as usize);
|
||||||
|
for _ in 0..shape.vocab {
|
||||||
|
logits.push(f32::from_bits(read_u32(file)?));
|
||||||
|
}
|
||||||
|
let mut compressed_rows = Vec::with_capacity(shape.layers as usize);
|
||||||
|
let mut indexer_rows = Vec::with_capacity(shape.layers as usize);
|
||||||
|
for layer in 0..shape.layers {
|
||||||
|
let rows = read_u32(file)?;
|
||||||
|
let ratio = compression_ratio(layer);
|
||||||
|
if rows != token_count.checked_div(ratio).unwrap_or(0) {
|
||||||
|
return Err("KV checkpoint compressed row count is invalid".into());
|
||||||
|
}
|
||||||
|
compressed_rows.push(rows);
|
||||||
|
}
|
||||||
|
for layer in 0..shape.layers {
|
||||||
|
let rows = read_u32(file)?;
|
||||||
|
let expected = if compression_ratio(layer) == 4 {
|
||||||
|
token_count / 4
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
if rows != expected {
|
||||||
|
return Err("KV checkpoint indexer row count is invalid".into());
|
||||||
|
}
|
||||||
|
indexer_rows.push(rows);
|
||||||
|
}
|
||||||
|
|
||||||
|
self.reset()?;
|
||||||
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
||||||
|
let raw_first = token_count - raw_live;
|
||||||
|
for (index, layer) in self.session.layers.iter_mut().enumerate() {
|
||||||
|
for position in raw_first..token_count {
|
||||||
|
let physical = position % self.session.raw_cap;
|
||||||
|
read_buffer(
|
||||||
|
file,
|
||||||
|
&layer.raw_cache,
|
||||||
|
u64::from(physical) * shape.head_dim * 4,
|
||||||
|
shape.head_dim * 4,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if let Some(state) = &mut layer.compression {
|
||||||
|
state.rows = compressed_rows[index];
|
||||||
|
read_buffer(
|
||||||
|
file,
|
||||||
|
&state.cache,
|
||||||
|
0,
|
||||||
|
u64::from(state.rows) * shape.head_dim * 2,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
|
||||||
|
read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
||||||
|
read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
||||||
|
}
|
||||||
|
if let Some(state) = &mut layer.indexer {
|
||||||
|
state.rows = indexer_rows[index];
|
||||||
|
read_buffer(
|
||||||
|
file,
|
||||||
|
&state.cache,
|
||||||
|
0,
|
||||||
|
u64::from(state.rows) * shape.indexer_head_dim * 4,
|
||||||
|
&mut chunk,
|
||||||
|
progress,
|
||||||
|
)?;
|
||||||
|
let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
|
||||||
|
read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
||||||
|
read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut trailing = [0];
|
||||||
|
if file
|
||||||
|
.read(&mut trailing)
|
||||||
|
.map_err(|error| error.to_string())?
|
||||||
|
!= 0
|
||||||
|
{
|
||||||
|
self.reset()?;
|
||||||
|
return Err("KV checkpoint has trailing data".into());
|
||||||
|
}
|
||||||
|
self.session.position = token_count;
|
||||||
|
self.tokens = tokens;
|
||||||
|
self.logits = logits;
|
||||||
|
self.checkpoint_tag = checkpoint_tag;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 {
|
||||||
|
let coefficient = if ratio == 4 { 2 } else { 1 };
|
||||||
|
coefficient * head_dim * coefficient * u64::from(ratio) * 4
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_buffer(
|
||||||
|
file: &mut File,
|
||||||
|
buffer: &Buffer,
|
||||||
|
mut offset: u64,
|
||||||
|
mut bytes: u64,
|
||||||
|
chunk: &mut [u8],
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<(), String> {
|
||||||
|
while bytes != 0 {
|
||||||
|
let length = bytes.min(chunk.len() as u64) as usize;
|
||||||
|
buffer.read(offset, &mut chunk[..length])?;
|
||||||
|
file.write_all(&chunk[..length])
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
progress(length as u64);
|
||||||
|
offset += length as u64;
|
||||||
|
bytes -= length as u64;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_buffer(
|
||||||
|
file: &mut File,
|
||||||
|
buffer: &Buffer,
|
||||||
|
mut offset: u64,
|
||||||
|
mut bytes: u64,
|
||||||
|
chunk: &mut [u8],
|
||||||
|
progress: &mut impl FnMut(u64),
|
||||||
|
) -> Result<(), String> {
|
||||||
|
while bytes != 0 {
|
||||||
|
let length = bytes.min(chunk.len() as u64) as usize;
|
||||||
|
file.read_exact(&mut chunk[..length])
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
buffer.write(offset, &chunk[..length])?;
|
||||||
|
progress(length as u64);
|
||||||
|
offset += length as u64;
|
||||||
|
bytes -= length as u64;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
|
||||||
|
file.write_all(&value.to_le_bytes())
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
|
||||||
|
file.write_all(&value.to_le_bytes())
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_u32(file: &mut File) -> Result<u32, String> {
|
||||||
|
let mut bytes = [0; 4];
|
||||||
|
file.read_exact(&mut bytes)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
Ok(u32::from_le_bytes(bytes))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_u64(file: &mut File) -> Result<u64, String> {
|
||||||
|
let mut bytes = [0; 8];
|
||||||
|
file.read_exact(&mut bytes)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
Ok(u64::from_le_bytes(bytes))
|
||||||
|
}
|
||||||
926
src/engine/metal/gpu.rs
Normal file
926
src/engine/metal/gpu.rs
Normal file
@@ -0,0 +1,926 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[repr(C)]
|
||||||
|
pub(super) struct GpuTensor {
|
||||||
|
_private: [u8; 0],
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe extern "C" {
|
||||||
|
pub(super) fn ds4_gpu_init() -> i32;
|
||||||
|
pub(super) fn ds4_gpu_cleanup();
|
||||||
|
pub(super) fn ds4_gpu_set_model_map_range(
|
||||||
|
model_map: *const c_void,
|
||||||
|
model_size: u64,
|
||||||
|
map_offset: u64,
|
||||||
|
map_size: u64,
|
||||||
|
max_tensor_bytes: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_set_quality(quality: bool);
|
||||||
|
pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
|
||||||
|
pub(super) fn ds4_gpu_tensor_view(
|
||||||
|
base: *const GpuTensor,
|
||||||
|
offset: u64,
|
||||||
|
bytes: u64,
|
||||||
|
) -> *mut GpuTensor;
|
||||||
|
pub(super) fn ds4_gpu_tensor_free(tensor: *mut GpuTensor);
|
||||||
|
pub(super) fn ds4_gpu_tensor_fill_f32(tensor: *mut GpuTensor, value: f32, count: u64) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_tensor_read(
|
||||||
|
tensor: *const GpuTensor,
|
||||||
|
offset: u64,
|
||||||
|
data: *mut c_void,
|
||||||
|
bytes: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_tensor_write(
|
||||||
|
tensor: *mut GpuTensor,
|
||||||
|
offset: u64,
|
||||||
|
data: *const c_void,
|
||||||
|
bytes: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_tensor_copy(
|
||||||
|
dst: *mut GpuTensor,
|
||||||
|
dst_offset: u64,
|
||||||
|
src: *const GpuTensor,
|
||||||
|
src_offset: u64,
|
||||||
|
bytes: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_tensor_copy_f32_to_f16(
|
||||||
|
dst: *mut GpuTensor,
|
||||||
|
dst_offset: u64,
|
||||||
|
src: *const GpuTensor,
|
||||||
|
src_offset: u64,
|
||||||
|
count: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_begin_commands() -> i32;
|
||||||
|
pub(super) fn ds4_gpu_end_commands() -> i32;
|
||||||
|
|
||||||
|
pub(super) fn ds4_gpu_embed_tokens_hc_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
tokens: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
vocab: u32,
|
||||||
|
rows: u32,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
|
||||||
|
pub(super) fn ds4_gpu_embed_token_hc_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
vocab: u32,
|
||||||
|
token: u32,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_rms_norm_plain_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
n: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_rms_norm_weight_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
n: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
scale: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u32,
|
||||||
|
output: u32,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_f16_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_q8_0_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_q8_0_pair_tensor(
|
||||||
|
out_a: *mut GpuTensor,
|
||||||
|
out_b: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight_a: u64,
|
||||||
|
weight_b: u64,
|
||||||
|
input: u64,
|
||||||
|
output_a: u64,
|
||||||
|
output_b: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_f16_pair_tensor(
|
||||||
|
out_a: *mut GpuTensor,
|
||||||
|
out_b: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight_a: u64,
|
||||||
|
weight_b: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u64,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_f16_pair_compressor_store_tensor(
|
||||||
|
out_kv: *mut GpuTensor,
|
||||||
|
out_score: *mut GpuTensor,
|
||||||
|
state_kv: *mut GpuTensor,
|
||||||
|
state_score: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight_kv: u64,
|
||||||
|
weight_score: u64,
|
||||||
|
ape: u64,
|
||||||
|
ape_type: u32,
|
||||||
|
input: u64,
|
||||||
|
width: u32,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
ratio: u32,
|
||||||
|
pos: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_split_weighted_sum_norm_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
norm: *mut GpuTensor,
|
||||||
|
split: *mut GpuTensor,
|
||||||
|
mix: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
scale: u64,
|
||||||
|
base: u64,
|
||||||
|
norm_weight: u64,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
iterations: u32,
|
||||||
|
eps: f32,
|
||||||
|
norm_eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
|
||||||
|
q_out: *mut GpuTensor,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
q_weight: u64,
|
||||||
|
q_n: u32,
|
||||||
|
kv_out: *mut GpuTensor,
|
||||||
|
kv: *const GpuTensor,
|
||||||
|
kv_weight: u64,
|
||||||
|
kv_n: u32,
|
||||||
|
rows: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
half: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
rot: u32,
|
||||||
|
pos: u32,
|
||||||
|
original_context: u32,
|
||||||
|
inverse: bool,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_head_rms_norm_tensor(
|
||||||
|
x: *mut GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_head_rms_norm_rope_tail_tensor(
|
||||||
|
x: *mut GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
rot: u32,
|
||||||
|
pos: u32,
|
||||||
|
original_context: u32,
|
||||||
|
inverse: bool,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_rope_tail_tensor(
|
||||||
|
x: *mut GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
rot: u32,
|
||||||
|
pos: u32,
|
||||||
|
original_context: u32,
|
||||||
|
inverse: bool,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_decode_heads_tensor(
|
||||||
|
heads_out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw_kv: *const GpuTensor,
|
||||||
|
n_raw: u32,
|
||||||
|
raw_cap: u32,
|
||||||
|
raw_start: u32,
|
||||||
|
comp_kv: *const GpuTensor,
|
||||||
|
comp_f16: u32,
|
||||||
|
n_comp: u32,
|
||||||
|
comp_mask: *const GpuTensor,
|
||||||
|
use_mask: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_decode_raw_batch_heads_tensor(
|
||||||
|
heads: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw_kv: *const GpuTensor,
|
||||||
|
tokens: u32,
|
||||||
|
pos: u32,
|
||||||
|
n_raw: u32,
|
||||||
|
raw_cap: u32,
|
||||||
|
raw_start: u32,
|
||||||
|
window: u32,
|
||||||
|
heads_count: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_decode_mixed_batch_heads_tensor(
|
||||||
|
heads: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw_kv: *const GpuTensor,
|
||||||
|
compressed: *const GpuTensor,
|
||||||
|
compressed_f16: u32,
|
||||||
|
compressed_mask: *const GpuTensor,
|
||||||
|
use_mask: u32,
|
||||||
|
tokens: u32,
|
||||||
|
pos: u32,
|
||||||
|
n_raw: u32,
|
||||||
|
raw_cap: u32,
|
||||||
|
raw_start: u32,
|
||||||
|
n_comp: u32,
|
||||||
|
window: u32,
|
||||||
|
ratio: u32,
|
||||||
|
heads_count: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
|
||||||
|
heads: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw_kv: *const GpuTensor,
|
||||||
|
comp_kv: *const GpuTensor,
|
||||||
|
comp_f16: u32,
|
||||||
|
topk: *const GpuTensor,
|
||||||
|
tokens: u32,
|
||||||
|
pos: u32,
|
||||||
|
n_raw: u32,
|
||||||
|
raw_cap: u32,
|
||||||
|
raw_start: u32,
|
||||||
|
n_comp: u32,
|
||||||
|
top_k: u32,
|
||||||
|
window: u32,
|
||||||
|
ratio: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_indexer_score_one_tensor(
|
||||||
|
scores: *mut GpuTensor,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
index_comp: *const GpuTensor,
|
||||||
|
n_comp: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
scale: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_indexer_scores_decode_batch_tensor(
|
||||||
|
scores: *mut GpuTensor,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
index_comp: *const GpuTensor,
|
||||||
|
n_comp: u32,
|
||||||
|
tokens: u32,
|
||||||
|
pos: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
ratio: u32,
|
||||||
|
scale: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_indexer_topk_tensor(
|
||||||
|
selected: *mut GpuTensor,
|
||||||
|
scores: *const GpuTensor,
|
||||||
|
n_comp: u32,
|
||||||
|
tokens: u32,
|
||||||
|
top_k: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_dsv4_indexer_qat_tensor(
|
||||||
|
x: *mut GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_compressor_update_tensor(
|
||||||
|
kv: *const GpuTensor,
|
||||||
|
score: *const GpuTensor,
|
||||||
|
state_kv: *mut GpuTensor,
|
||||||
|
state_score: *mut GpuTensor,
|
||||||
|
cache: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
ape: u64,
|
||||||
|
ape_type: u32,
|
||||||
|
norm: u64,
|
||||||
|
norm_type: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
ratio: u32,
|
||||||
|
pos: u32,
|
||||||
|
row: u32,
|
||||||
|
rot: u32,
|
||||||
|
original_context: u32,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
rms_eps: f32,
|
||||||
|
state_already_stored: bool,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_compressor_prefill_state_ratio4_tensor(
|
||||||
|
state_kv: *mut GpuTensor,
|
||||||
|
state_score: *mut GpuTensor,
|
||||||
|
kv_tail: *const GpuTensor,
|
||||||
|
score_tail: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
ape: u64,
|
||||||
|
ape_type: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
pos: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_dsv4_fp8_kv_quantize_tensor(
|
||||||
|
x: *mut GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
rot: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_kv_fp8_store_raw_tensor(
|
||||||
|
kv: *mut GpuTensor,
|
||||||
|
raw_cache: *mut GpuTensor,
|
||||||
|
raw_cap: u32,
|
||||||
|
row: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
rot: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_store_raw_kv_batch_tensor(
|
||||||
|
raw_cache: *mut GpuTensor,
|
||||||
|
kv: *const GpuTensor,
|
||||||
|
raw_cap: u32,
|
||||||
|
pos: u32,
|
||||||
|
rows: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_compressor_prefill_tensor(
|
||||||
|
cache: *mut GpuTensor,
|
||||||
|
state_kv: *mut GpuTensor,
|
||||||
|
state_score: *mut GpuTensor,
|
||||||
|
kv: *const GpuTensor,
|
||||||
|
score: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
ape: u64,
|
||||||
|
ape_type: u32,
|
||||||
|
norm: u64,
|
||||||
|
norm_type: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
ratio: u32,
|
||||||
|
pos: u32,
|
||||||
|
rows: u32,
|
||||||
|
rot: u32,
|
||||||
|
original_context: u32,
|
||||||
|
quantize_fp8: bool,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
rms_eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_compressor_prefill_ratio4_replay_tensor(
|
||||||
|
cache: *mut GpuTensor,
|
||||||
|
state_kv: *mut GpuTensor,
|
||||||
|
state_score: *mut GpuTensor,
|
||||||
|
kv: *const GpuTensor,
|
||||||
|
score: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
ape: u64,
|
||||||
|
ape_type: u32,
|
||||||
|
norm: u64,
|
||||||
|
norm_type: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
pos: u32,
|
||||||
|
rows: u32,
|
||||||
|
rot: u32,
|
||||||
|
original_context: u32,
|
||||||
|
quantize_fp8: bool,
|
||||||
|
freq_base: f32,
|
||||||
|
freq_scale: f32,
|
||||||
|
ext_factor: f32,
|
||||||
|
attn_factor: f32,
|
||||||
|
beta_fast: f32,
|
||||||
|
beta_slow: f32,
|
||||||
|
rms_eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_prefill_raw_heads_tensor(
|
||||||
|
heads: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw: *const GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
window: u32,
|
||||||
|
heads_count: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_prefill_static_mixed_heads_tensor(
|
||||||
|
heads: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
sinks: u64,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
raw: *const GpuTensor,
|
||||||
|
compressed: *const GpuTensor,
|
||||||
|
compressed_f16: u32,
|
||||||
|
rows: u32,
|
||||||
|
compressed_rows: u32,
|
||||||
|
window: u32,
|
||||||
|
ratio: u32,
|
||||||
|
heads_count: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_indexer_scores_prefill_tensor(
|
||||||
|
scores: *mut GpuTensor,
|
||||||
|
q: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
compressed: *const GpuTensor,
|
||||||
|
compressed_rows: u32,
|
||||||
|
rows: u32,
|
||||||
|
heads: u32,
|
||||||
|
head_dim: u32,
|
||||||
|
ratio: u32,
|
||||||
|
scale: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_output_q8_batch_f16_tensor(
|
||||||
|
out_half: *mut GpuTensor,
|
||||||
|
low: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight_a: u64,
|
||||||
|
weight_b: u64,
|
||||||
|
group_dim: u64,
|
||||||
|
rank: u64,
|
||||||
|
groups: u32,
|
||||||
|
output: u64,
|
||||||
|
heads: *const GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_output_q8_batch_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
low: *mut GpuTensor,
|
||||||
|
group_scratch: *mut GpuTensor,
|
||||||
|
low_scratch: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight_a: u64,
|
||||||
|
weight_b: u64,
|
||||||
|
group_dim: u64,
|
||||||
|
rank: u64,
|
||||||
|
groups: u32,
|
||||||
|
output: u64,
|
||||||
|
heads: *const GpuTensor,
|
||||||
|
rows: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_router_select_batch_tensor(
|
||||||
|
selected: *mut GpuTensor,
|
||||||
|
weights: *mut GpuTensor,
|
||||||
|
probs: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
bias: u64,
|
||||||
|
hash: u64,
|
||||||
|
hash_rows: u32,
|
||||||
|
expert_groups: u32,
|
||||||
|
groups_used: u32,
|
||||||
|
has_bias: bool,
|
||||||
|
hash_mode: bool,
|
||||||
|
logits: *const GpuTensor,
|
||||||
|
tokens: *const GpuTensor,
|
||||||
|
experts: u32,
|
||||||
|
experts_used: u32,
|
||||||
|
scale: f32,
|
||||||
|
rows: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_routed_moe_batch_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
gate: *mut GpuTensor,
|
||||||
|
up: *mut GpuTensor,
|
||||||
|
mid: *mut GpuTensor,
|
||||||
|
experts_out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
gate_weight: u64,
|
||||||
|
up_weight: u64,
|
||||||
|
down_weight: u64,
|
||||||
|
gate_type: u32,
|
||||||
|
down_type: u32,
|
||||||
|
gate_expert_bytes: u64,
|
||||||
|
gate_row_bytes: u64,
|
||||||
|
down_expert_bytes: u64,
|
||||||
|
down_row_bytes: u64,
|
||||||
|
input: u32,
|
||||||
|
middle: u32,
|
||||||
|
output: u32,
|
||||||
|
selected: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
total_experts: u32,
|
||||||
|
used_experts: u32,
|
||||||
|
clamp: f32,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
layer: u32,
|
||||||
|
rows: u32,
|
||||||
|
mid_f16: *mut bool,
|
||||||
|
force_resident: bool,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_swiglu_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
gate: *const GpuTensor,
|
||||||
|
up: *const GpuTensor,
|
||||||
|
count: u32,
|
||||||
|
clamp: f32,
|
||||||
|
scale: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_expand_split_half_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
block_half: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
split: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_expand_split_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
block: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
split: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_expand_add_split_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
block: *const GpuTensor,
|
||||||
|
add: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
split: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_attention_output_low_q8_tensor(
|
||||||
|
low: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
group_dim: u64,
|
||||||
|
rank: u64,
|
||||||
|
groups: u32,
|
||||||
|
heads: *const GpuTensor,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_matmul_q8_0_hc_expand_tensor(
|
||||||
|
out_hc: *mut GpuTensor,
|
||||||
|
block_out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
split: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_router_select_tensor(
|
||||||
|
selected: *mut GpuTensor,
|
||||||
|
weights: *mut GpuTensor,
|
||||||
|
probs: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
bias: u64,
|
||||||
|
hash: u64,
|
||||||
|
hash_rows: u32,
|
||||||
|
token: u32,
|
||||||
|
experts: u32,
|
||||||
|
used: u32,
|
||||||
|
scale: f32,
|
||||||
|
expert_groups: u32,
|
||||||
|
groups_used: u32,
|
||||||
|
has_bias: bool,
|
||||||
|
hash_mode: bool,
|
||||||
|
logits: *const GpuTensor,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_routed_moe_one_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
gate: *mut GpuTensor,
|
||||||
|
up: *mut GpuTensor,
|
||||||
|
mid: *mut GpuTensor,
|
||||||
|
experts_out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
gate_weight: u64,
|
||||||
|
up_weight: u64,
|
||||||
|
down_weight: u64,
|
||||||
|
gate_type: u32,
|
||||||
|
down_type: u32,
|
||||||
|
gate_expert_bytes: u64,
|
||||||
|
gate_row_bytes: u64,
|
||||||
|
down_expert_bytes: u64,
|
||||||
|
down_row_bytes: u64,
|
||||||
|
input: u32,
|
||||||
|
middle: u32,
|
||||||
|
output: u32,
|
||||||
|
selected: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
total_experts: u32,
|
||||||
|
used_experts: u32,
|
||||||
|
clamp: f32,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
add: *const GpuTensor,
|
||||||
|
layer: u32,
|
||||||
|
force_resident: bool,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
|
||||||
|
gate: *mut GpuTensor,
|
||||||
|
up: *mut GpuTensor,
|
||||||
|
mid: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
gate_weight: u64,
|
||||||
|
up_weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
x: *const GpuTensor,
|
||||||
|
clamp: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_shared_down_hc_expand_q8_0_tensor(
|
||||||
|
out_hc: *mut GpuTensor,
|
||||||
|
shared_out: *mut GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
weight: u64,
|
||||||
|
input: u64,
|
||||||
|
output: u64,
|
||||||
|
middle: *const GpuTensor,
|
||||||
|
routed: *const GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
split: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_output_hc_weights_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
pre: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
scale: u64,
|
||||||
|
base: u64,
|
||||||
|
hc: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_weighted_sum_norm_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
norm: *mut GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
map: *const c_void,
|
||||||
|
size: u64,
|
||||||
|
norm_weight: u64,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
eps: f32,
|
||||||
|
) -> i32;
|
||||||
|
pub(super) fn ds4_gpu_hc_weighted_sum_tensor(
|
||||||
|
out: *mut GpuTensor,
|
||||||
|
residual: *const GpuTensor,
|
||||||
|
weights: *const GpuTensor,
|
||||||
|
embd: u32,
|
||||||
|
hc: u32,
|
||||||
|
) -> i32;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) struct Context;
|
||||||
|
|
||||||
|
impl Context {
|
||||||
|
pub(super) fn open(model: &Model, quality: bool) -> Result<Self, String> {
|
||||||
|
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
|
||||||
|
let data_offset = model.main.data_offset();
|
||||||
|
if let Err(error) = check(
|
||||||
|
unsafe {
|
||||||
|
ds4_gpu_set_model_map_range(
|
||||||
|
model.main.map_ptr().cast(),
|
||||||
|
model.main.len(),
|
||||||
|
data_offset,
|
||||||
|
model.main.len() - data_offset,
|
||||||
|
model.main.max_tensor_bytes(),
|
||||||
|
)
|
||||||
|
},
|
||||||
|
"model mapping",
|
||||||
|
) {
|
||||||
|
unsafe { ds4_gpu_cleanup() };
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
unsafe { ds4_gpu_set_quality(quality) };
|
||||||
|
Ok(Self)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for Context {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
unsafe { ds4_gpu_cleanup() };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) struct Commands(bool);
|
||||||
|
|
||||||
|
impl Commands {
|
||||||
|
pub(super) fn begin() -> Result<Self, String> {
|
||||||
|
check(
|
||||||
|
unsafe { ds4_gpu_begin_commands() },
|
||||||
|
"beginning Metal commands",
|
||||||
|
)?;
|
||||||
|
Ok(Self(true))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn finish(mut self) -> Result<(), String> {
|
||||||
|
self.0 = false;
|
||||||
|
check(
|
||||||
|
unsafe { ds4_gpu_end_commands() },
|
||||||
|
"executing Metal commands",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for Commands {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if self.0 {
|
||||||
|
unsafe { ds4_gpu_end_commands() };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) struct Buffer(NonNull<GpuTensor>);
|
||||||
|
|
||||||
|
impl Buffer {
|
||||||
|
pub(super) fn floats(count: u64) -> Result<Self, String> {
|
||||||
|
Self::bytes(count * 4)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn bytes(bytes: u64) -> Result<Self, String> {
|
||||||
|
NonNull::new(unsafe { ds4_gpu_tensor_alloc(bytes) })
|
||||||
|
.map(Self)
|
||||||
|
.ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
|
||||||
|
NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) })
|
||||||
|
.map(Self)
|
||||||
|
.ok_or_else(|| "Metal could not create a tensor view".to_owned())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn read_f32(&self, values: &mut [f32]) -> Result<(), String> {
|
||||||
|
check(
|
||||||
|
unsafe {
|
||||||
|
ds4_gpu_tensor_read(
|
||||||
|
self.raw(),
|
||||||
|
0,
|
||||||
|
values.as_mut_ptr().cast(),
|
||||||
|
std::mem::size_of_val(values) as u64,
|
||||||
|
)
|
||||||
|
},
|
||||||
|
"reading Metal output",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
|
||||||
|
check(
|
||||||
|
unsafe {
|
||||||
|
ds4_gpu_tensor_read(
|
||||||
|
self.raw(),
|
||||||
|
offset,
|
||||||
|
values.as_mut_ptr().cast(),
|
||||||
|
values.len() as u64,
|
||||||
|
)
|
||||||
|
},
|
||||||
|
"reading a Metal buffer",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn write(&self, offset: u64, values: &[u8]) -> Result<(), String> {
|
||||||
|
check(
|
||||||
|
unsafe {
|
||||||
|
ds4_gpu_tensor_write(
|
||||||
|
self.raw(),
|
||||||
|
offset,
|
||||||
|
values.as_ptr().cast(),
|
||||||
|
values.len() as u64,
|
||||||
|
)
|
||||||
|
},
|
||||||
|
"restoring a Metal buffer",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn write_i32(&self, values: &[i32]) -> Result<(), String> {
|
||||||
|
check(
|
||||||
|
unsafe {
|
||||||
|
ds4_gpu_tensor_write(
|
||||||
|
self.raw(),
|
||||||
|
0,
|
||||||
|
values.as_ptr().cast(),
|
||||||
|
std::mem::size_of_val(values) as u64,
|
||||||
|
)
|
||||||
|
},
|
||||||
|
"uploading tokens",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> {
|
||||||
|
call(
|
||||||
|
unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },
|
||||||
|
"initializing a Metal buffer",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn raw(&self) -> *mut GpuTensor {
|
||||||
|
self.0.as_ptr()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for Buffer {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
unsafe { ds4_gpu_tensor_free(self.raw()) };
|
||||||
|
}
|
||||||
|
}
|
||||||
994
src/engine/validation.rs
Normal file
994
src/engine/validation.rs
Normal file
@@ -0,0 +1,994 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
pub(crate) fn validate_model_artifact(
|
||||||
|
path: &Path,
|
||||||
|
expected: ModelChoice,
|
||||||
|
support: bool,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
if support {
|
||||||
|
let model = Gguf::open(path)?;
|
||||||
|
validate_dspark(&model, &FLASH)
|
||||||
|
} else {
|
||||||
|
let model = Model::open_main(path, expected)?;
|
||||||
|
let summary = model.summary();
|
||||||
|
if summary.tensor_count == 0
|
||||||
|
|| model
|
||||||
|
.render_prompt("", "", ReasoningMode::Direct)
|
||||||
|
.is_empty()
|
||||||
|
{
|
||||||
|
return Err("model intake produced an empty tensor directory or prompt".into());
|
||||||
|
}
|
||||||
|
let _ = model.tokenize("");
|
||||||
|
let eos = model.eos_token();
|
||||||
|
let _ = model.token_bytes(eos);
|
||||||
|
if !model.is_stop_token(eos) {
|
||||||
|
return Err("tokenizer EOS marker is not a stop token".into());
|
||||||
|
}
|
||||||
|
model.tensor_data("token_embd.weight")?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> {
|
||||||
|
let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") {
|
||||||
|
ModelFamily::Glm
|
||||||
|
} else {
|
||||||
|
ModelFamily::DeepSeek
|
||||||
|
};
|
||||||
|
let shape = match family {
|
||||||
|
ModelFamily::Glm => GLM,
|
||||||
|
ModelFamily::DeepSeek => match model.u32("deepseek4.block_count")? {
|
||||||
|
43 => FLASH,
|
||||||
|
61 => PRO,
|
||||||
|
layers => return Err(format!("unsupported DeepSeek layer count: {layers}")),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
if shape.model != expected {
|
||||||
|
return Err(format!(
|
||||||
|
"{} contains {}, but the selected model is {expected}",
|
||||||
|
model.path().display(),
|
||||||
|
shape.model
|
||||||
|
));
|
||||||
|
}
|
||||||
|
validate_metadata(model, &shape)?;
|
||||||
|
validate_tensors(model, &shape)?;
|
||||||
|
Ok(shape)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||||||
|
let prefix = if shape.family == ModelFamily::Glm {
|
||||||
|
"glm-dsa"
|
||||||
|
} else {
|
||||||
|
"deepseek4"
|
||||||
|
};
|
||||||
|
for (key, expected) in [
|
||||||
|
("block_count", u64::from(shape.layers)),
|
||||||
|
("embedding_length", shape.embd),
|
||||||
|
("vocab_size", shape.vocab),
|
||||||
|
("attention.head_count", shape.heads),
|
||||||
|
("attention.head_count_kv", shape.head_kv),
|
||||||
|
("attention.key_length", shape.head_dim),
|
||||||
|
("attention.value_length", shape.value_dim),
|
||||||
|
("rope.dimension_count", shape.rot),
|
||||||
|
("attention.q_lora_rank", shape.lora_q),
|
||||||
|
("expert_count", shape.experts),
|
||||||
|
("expert_used_count", shape.experts_used),
|
||||||
|
("expert_feed_forward_length", shape.ff_expert),
|
||||||
|
("expert_shared_count", shape.expert_shared),
|
||||||
|
("attention.indexer.head_count", shape.indexer_heads),
|
||||||
|
("attention.indexer.key_length", shape.indexer_head_dim),
|
||||||
|
("attention.indexer.top_k", shape.indexer_top_k),
|
||||||
|
] {
|
||||||
|
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||||||
|
}
|
||||||
|
expect_float(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.attention.layer_norm_rms_epsilon"),
|
||||||
|
shape.rms_epsilon,
|
||||||
|
)?;
|
||||||
|
expect_float(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.expert_weights_scale"),
|
||||||
|
shape.expert_weight_scale,
|
||||||
|
)?;
|
||||||
|
if !model.boolean(&format!("{prefix}.expert_weights_norm"))? {
|
||||||
|
return Err(format!("{prefix}.expert_weights_norm must be true"));
|
||||||
|
}
|
||||||
|
expect_float(model, &format!("{prefix}.rope.freq_base"), shape.rope_base)?;
|
||||||
|
|
||||||
|
if shape.family == ModelFamily::Glm {
|
||||||
|
for (key, expected) in [
|
||||||
|
("context_length", shape.original_context),
|
||||||
|
("feed_forward_length", shape.ff_dense),
|
||||||
|
("attention.kv_lora_rank", shape.kv_lora),
|
||||||
|
("attention.key_length_mla", shape.key_mla),
|
||||||
|
("attention.value_length_mla", shape.value_mla),
|
||||||
|
("expert_group_count", 1),
|
||||||
|
("expert_group_used_count", 1),
|
||||||
|
("expert_gating_func", 2),
|
||||||
|
("leading_dense_block_count", u64::from(shape.leading_dense)),
|
||||||
|
("nextn_predict_layers", u64::from(shape.nextn)),
|
||||||
|
] {
|
||||||
|
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||||||
|
}
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
for (key, expected) in [
|
||||||
|
("attention.output_group_count", shape.out_groups),
|
||||||
|
("attention.output_lora_rank", shape.lora_o),
|
||||||
|
("hash_layer_count", u64::from(shape.hash_layers)),
|
||||||
|
("attention.sliding_window", shape.sliding_window),
|
||||||
|
("hyper_connection.count", shape.hc),
|
||||||
|
("hyper_connection.sinkhorn_iterations", shape.hc_sinkhorn),
|
||||||
|
] {
|
||||||
|
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||||||
|
}
|
||||||
|
for key in ["expert_group_count", "expert_group_used_count"] {
|
||||||
|
if let Some(Value::U32(value)) = model.metadata.get(&format!("{prefix}.{key}"))
|
||||||
|
&& *value != 0
|
||||||
|
{
|
||||||
|
return Err(format!("{prefix}.{key} must be zero"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (key, expected) in [
|
||||||
|
("hyper_connection.epsilon", shape.hc_epsilon),
|
||||||
|
(
|
||||||
|
"attention.compress_rope_freq_base",
|
||||||
|
shape.compress_rope_base,
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
expect_float(model, &format!("{prefix}.{key}"), expected)?;
|
||||||
|
}
|
||||||
|
for (key, expected) in [
|
||||||
|
("rope.scaling.factor", shape.rope_scale),
|
||||||
|
("rope.scaling.yarn_beta_fast", shape.rope_beta_fast),
|
||||||
|
("rope.scaling.yarn_beta_slow", shape.rope_beta_slow),
|
||||||
|
] {
|
||||||
|
if model.metadata.contains_key(&format!("{prefix}.{key}")) {
|
||||||
|
expect_float(model, &format!("{prefix}.{key}"), expected)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if model
|
||||||
|
.metadata
|
||||||
|
.contains_key("deepseek4.rope.scaling.original_context_length")
|
||||||
|
{
|
||||||
|
expect_u64(
|
||||||
|
model,
|
||||||
|
"deepseek4.rope.scaling.original_context_length",
|
||||||
|
shape.original_context,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let ratios = model.u32s("deepseek4.attention.compress_ratios")?;
|
||||||
|
let clamps = model.f32s("deepseek4.swiglu_clamp_exp")?;
|
||||||
|
if ratios.len() < shape.layers as usize || clamps.len() < shape.layers as usize {
|
||||||
|
return Err("DeepSeek per-layer metadata is shorter than the layer count".into());
|
||||||
|
}
|
||||||
|
for layer in 0..shape.layers as usize {
|
||||||
|
let expected = compression_ratio(shape, layer as u32);
|
||||||
|
if ratios[layer] != expected {
|
||||||
|
return Err(format!(
|
||||||
|
"layer {layer} compression ratio is {}, expected {expected}",
|
||||||
|
ratios[layer]
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if !float_eq(clamps[layer], shape.swiglu_clamp) {
|
||||||
|
return Err(format!("layer {layer} has an invalid SwiGLU clamp"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||||||
|
match shape.family {
|
||||||
|
ModelFamily::DeepSeek => validate_deepseek_tensors(model, shape),
|
||||||
|
ModelFamily::Glm => validate_glm_tensors(model, shape),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_deepseek_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||||||
|
let hc_dim = shape.embd * shape.hc;
|
||||||
|
let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
|
||||||
|
let q_dim = shape.heads * shape.head_dim;
|
||||||
|
let output_low = shape.out_groups * shape.lora_o;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
"token_embd.weight",
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, shape.vocab],
|
||||||
|
)?;
|
||||||
|
expect(model, "output_hc_base.weight", &[F32], &[shape.hc])?;
|
||||||
|
expect(model, "output_hc_fn.weight", &[F16], &[hc_dim, shape.hc])?;
|
||||||
|
expect(model, "output_hc_scale.weight", &[F32], &[1])?;
|
||||||
|
expect(model, "output_norm.weight", &[F32], &[shape.embd])?;
|
||||||
|
expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?;
|
||||||
|
|
||||||
|
for layer in 0..shape.layers {
|
||||||
|
let name = |suffix: &str| format!("blk.{layer}.{suffix}");
|
||||||
|
expect(model, &name("hc_attn_fn.weight"), &[F16], &[hc_dim, hc_mix])?;
|
||||||
|
expect(model, &name("hc_attn_scale.weight"), &[F32], &[3])?;
|
||||||
|
expect(model, &name("hc_attn_base.weight"), &[F32], &[hc_mix])?;
|
||||||
|
expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_a.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.lora_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_a_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.lora_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.lora_q, q_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_kv.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.head_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_kv_a_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.head_dim],
|
||||||
|
)?;
|
||||||
|
expect(model, &name("attn_sinks.weight"), &[F32], &[shape.heads])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_output_a.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[
|
||||||
|
shape.head_dim * (shape.heads / shape.out_groups),
|
||||||
|
output_low,
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_output_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[output_low, shape.embd],
|
||||||
|
)?;
|
||||||
|
|
||||||
|
let ratio = compression_ratio(shape, layer);
|
||||||
|
if ratio != 0 {
|
||||||
|
let compression_width = if ratio == 4 { 2 } else { 1 } * shape.head_dim;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_compressor_ape.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[compression_width, u64::from(ratio)],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_compressor_kv.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, compression_width],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_compressor_gate.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, compression_width],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_compressor_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.head_dim],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if ratio == 4 {
|
||||||
|
let index_q = shape.indexer_heads * shape.indexer_head_dim;
|
||||||
|
let index_width = 2 * shape.indexer_head_dim;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.attn_q_b.weight"),
|
||||||
|
&[F16, Q8_0],
|
||||||
|
&[shape.lora_q, index_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.proj.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, shape.indexer_heads],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer_compressor_ape.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[index_width, 4],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer_compressor_kv.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, index_width],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer_compressor_gate.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, index_width],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer_compressor_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.indexer_head_dim],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
expect(model, &name("hc_ffn_fn.weight"), &[F16], &[hc_dim, hc_mix])?;
|
||||||
|
expect(model, &name("hc_ffn_scale.weight"), &[F32], &[3])?;
|
||||||
|
expect(model, &name("hc_ffn_base.weight"), &[F32], &[hc_mix])?;
|
||||||
|
expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_inp.weight"),
|
||||||
|
&[F16],
|
||||||
|
&[shape.embd, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect_optional(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_up_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_down_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.ff_expert, shape.embd, shape.experts],
|
||||||
|
)?;
|
||||||
|
same_type(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_exps.weight"),
|
||||||
|
&name("ffn_up_exps.weight"),
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_expert],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_up_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_expert],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_down_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.ff_expert, shape.embd],
|
||||||
|
)?;
|
||||||
|
if layer < shape.hash_layers {
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_tid2eid.weight"),
|
||||||
|
&[I32],
|
||||||
|
&[shape.experts_used, shape.vocab],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||||||
|
let q_dim = shape.heads * shape.key_mla;
|
||||||
|
let q_nope = shape.key_mla - shape.rot;
|
||||||
|
let index_q = shape.indexer_heads * shape.indexer_head_dim;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
"token_embd.weight",
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.vocab],
|
||||||
|
)?;
|
||||||
|
expect(model, "output_norm.weight", &[F32], &[shape.embd])?;
|
||||||
|
expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?;
|
||||||
|
for layer in 0..shape.layers {
|
||||||
|
let name = |suffix: &str| format!("blk.{layer}.{suffix}");
|
||||||
|
expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_a.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.lora_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_a_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.lora_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_q_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.lora_q, q_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_kv_a_mqa.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.head_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_kv_a_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.kv_lora],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_k_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[q_nope, shape.kv_lora, shape.heads],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_v_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.kv_lora, shape.value_mla, shape.heads],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("attn_output.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.heads * shape.value_mla, shape.embd],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.attn_k.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.indexer_head_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.attn_q_b.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.lora_q, index_q],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.k_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.indexer_head_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.k_norm.bias"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.indexer_head_dim],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("indexer.proj.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.embd, shape.indexer_heads],
|
||||||
|
)?;
|
||||||
|
expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
if layer < shape.leading_dense {
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_dense],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_up.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_dense],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_down.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.ff_dense, shape.embd],
|
||||||
|
)?;
|
||||||
|
} else {
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_inp.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.embd, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_up_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_down_exps.weight"),
|
||||||
|
ROUTED,
|
||||||
|
&[shape.ff_expert, shape.embd, shape.experts],
|
||||||
|
)?;
|
||||||
|
same_type(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_exps.weight"),
|
||||||
|
&name("ffn_up_exps.weight"),
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_expert],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_up_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.embd, shape.ff_expert],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("ffn_down_shexp.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[shape.ff_expert, shape.embd],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if layer + shape.nextn >= shape.layers {
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("nextn.eh_proj.weight"),
|
||||||
|
DENSE,
|
||||||
|
&[2 * shape.embd, shape.embd],
|
||||||
|
)?;
|
||||||
|
expect(model, &name("nextn.enorm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
expect(model, &name("nextn.hnorm.weight"), &[F32], &[shape.embd])?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&name("nextn.shared_head_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.embd],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||||||
|
if shape.model != ModelChoice::DeepSeekV4Flash {
|
||||||
|
return Err("DSpark support is available only for DeepSeek V4 Flash".into());
|
||||||
|
}
|
||||||
|
let block_size = first_u32(
|
||||||
|
model,
|
||||||
|
&[
|
||||||
|
"deepseek4.dspark.block_size",
|
||||||
|
"deepseek4.dspark_block_size",
|
||||||
|
"dspark.block_size",
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
let markov_rank = first_u32(
|
||||||
|
model,
|
||||||
|
&[
|
||||||
|
"deepseek4.dspark.markov_rank",
|
||||||
|
"deepseek4.dspark_markov_rank",
|
||||||
|
"dspark.markov_rank",
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
let noise_token = first_u32(
|
||||||
|
model,
|
||||||
|
&[
|
||||||
|
"deepseek4.dspark.noise_token_id",
|
||||||
|
"deepseek4.dspark_noise_token_id",
|
||||||
|
"dspark.noise_token_id",
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
let targets = first_u32s(
|
||||||
|
model,
|
||||||
|
&[
|
||||||
|
"deepseek4.dspark.target_layer_ids",
|
||||||
|
"deepseek4.dspark_target_layer_ids",
|
||||||
|
"dspark.target_layer_ids",
|
||||||
|
],
|
||||||
|
)?;
|
||||||
|
if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 {
|
||||||
|
return Err("invalid DSpark block, Markov, or noise-token metadata".into());
|
||||||
|
}
|
||||||
|
if targets.is_empty()
|
||||||
|
|| targets.len() > 8
|
||||||
|
|| targets.windows(2).any(|pair| pair[0] >= pair[1])
|
||||||
|
|| targets.iter().any(|layer| *layer >= shape.layers)
|
||||||
|
{
|
||||||
|
return Err("invalid DSpark target-layer metadata".into());
|
||||||
|
}
|
||||||
|
let stages = model
|
||||||
|
.tensors
|
||||||
|
.keys()
|
||||||
|
.filter_map(|name| {
|
||||||
|
name.strip_prefix("mtp.")?
|
||||||
|
.split('.')
|
||||||
|
.next()?
|
||||||
|
.parse::<u32>()
|
||||||
|
.ok()
|
||||||
|
})
|
||||||
|
.max()
|
||||||
|
.map_or(0, |stage| stage + 1);
|
||||||
|
if !(1..=8).contains(&stages) {
|
||||||
|
return Err(format!("invalid DSpark stage count: {stages}"));
|
||||||
|
}
|
||||||
|
for stage in 0..stages {
|
||||||
|
validate_dspark_block(model, shape, stage)?;
|
||||||
|
if stage == 0 {
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("mtp.{stage}.main_proj.weight"),
|
||||||
|
DSPARK_DENSE,
|
||||||
|
&[targets.len() as u64 * shape.embd, shape.embd],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("mtp.{stage}.main_norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.embd],
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let prefix = format!("mtp.{}", stages - 1);
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.norm.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.embd],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.hc_head_base.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[shape.hc],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.hc_head_fn.weight"),
|
||||||
|
PLAIN,
|
||||||
|
&[shape.embd * shape.hc, shape.hc],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.hc_head_scale.weight"),
|
||||||
|
&[F32],
|
||||||
|
&[1],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.markov_head.markov_w1.weight"),
|
||||||
|
DSPARK_DENSE,
|
||||||
|
&[u64::from(markov_rank), shape.vocab],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.markov_head.markov_w2.weight"),
|
||||||
|
DSPARK_DENSE,
|
||||||
|
&[u64::from(markov_rank), shape.vocab],
|
||||||
|
)?;
|
||||||
|
expect(
|
||||||
|
model,
|
||||||
|
&format!("{prefix}.confidence_head.proj.weight"),
|
||||||
|
DSPARK_DENSE,
|
||||||
|
&[shape.embd + u64::from(markov_rank), 1],
|
||||||
|
)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_dspark_block(model: &Gguf, shape: &Shape, stage: u32) -> Result<(), String> {
|
||||||
|
let hc_dim = shape.embd * shape.hc;
|
||||||
|
let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
|
||||||
|
let q_dim = shape.heads * shape.head_dim;
|
||||||
|
let output_low = shape.out_groups * shape.lora_o;
|
||||||
|
let name = |suffix: &str| format!("mtp.{stage}.{suffix}");
|
||||||
|
for (suffix, types, dims) in [
|
||||||
|
("hc_attn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
|
||||||
|
("hc_attn_scale.weight", &[F32][..], vec![3]),
|
||||||
|
("hc_attn_base.weight", &[F32][..], vec![hc_mix]),
|
||||||
|
("attn_norm.weight", &[F32][..], vec![shape.embd]),
|
||||||
|
(
|
||||||
|
"attn_q_a.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.embd, shape.lora_q],
|
||||||
|
),
|
||||||
|
("attn_q_a_norm.weight", &[F32][..], vec![shape.lora_q]),
|
||||||
|
("attn_q_b.weight", DSPARK_DENSE, vec![shape.lora_q, q_dim]),
|
||||||
|
(
|
||||||
|
"attn_kv.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.embd, shape.head_dim],
|
||||||
|
),
|
||||||
|
("attn_kv_a_norm.weight", &[F32][..], vec![shape.head_dim]),
|
||||||
|
("attn_sinks.weight", &[F32][..], vec![shape.heads]),
|
||||||
|
(
|
||||||
|
"attn_output_a.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![
|
||||||
|
shape.head_dim * (shape.heads / shape.out_groups),
|
||||||
|
output_low,
|
||||||
|
],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"attn_output_b.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![output_low, shape.embd],
|
||||||
|
),
|
||||||
|
("hc_ffn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
|
||||||
|
("hc_ffn_scale.weight", &[F32][..], vec![3]),
|
||||||
|
("hc_ffn_base.weight", &[F32][..], vec![hc_mix]),
|
||||||
|
("ffn_norm.weight", &[F32][..], vec![shape.embd]),
|
||||||
|
(
|
||||||
|
"ffn_gate_inp.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.embd, shape.experts],
|
||||||
|
),
|
||||||
|
("exp_probs_b.bias", &[F32][..], vec![shape.experts]),
|
||||||
|
(
|
||||||
|
"ffn_gate_exps.weight",
|
||||||
|
ROUTED,
|
||||||
|
vec![shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"ffn_up_exps.weight",
|
||||||
|
ROUTED,
|
||||||
|
vec![shape.embd, shape.ff_expert, shape.experts],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"ffn_down_exps.weight",
|
||||||
|
ROUTED,
|
||||||
|
vec![shape.ff_expert, shape.embd, shape.experts],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"ffn_gate_shexp.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.embd, shape.ff_expert],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"ffn_up_shexp.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.embd, shape.ff_expert],
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"ffn_down_shexp.weight",
|
||||||
|
DSPARK_DENSE,
|
||||||
|
vec![shape.ff_expert, shape.embd],
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
expect(model, &name(suffix), types, &dims)?;
|
||||||
|
}
|
||||||
|
same_type(
|
||||||
|
model,
|
||||||
|
&name("ffn_gate_exps.weight"),
|
||||||
|
&name("ffn_up_exps.weight"),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn expect(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||||||
|
validate_tensor(name, model.tensor(name)?, types, dims)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn expect_optional(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||||||
|
match model.tensors.get(name) {
|
||||||
|
Some(tensor) => validate_tensor(name, tensor, types, dims),
|
||||||
|
None => Ok(()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_tensor(name: &str, tensor: &Tensor, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||||||
|
if !types.contains(&tensor.kind) {
|
||||||
|
return Err(format!(
|
||||||
|
"tensor {name} has unsupported type {}",
|
||||||
|
tensor.kind
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if tensor.dims != dims {
|
||||||
|
return Err(format!(
|
||||||
|
"tensor {name} has dimensions {:?}, expected {dims:?}",
|
||||||
|
tensor.dims
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn same_type(model: &Gguf, first: &str, second: &str) -> Result<(), String> {
|
||||||
|
if model.tensor(first)?.kind != model.tensor(second)?.kind {
|
||||||
|
Err(format!(
|
||||||
|
"tensors {first} and {second} use different quantizations"
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn expect_u64(model: &Gguf, key: &str, expected: u64) -> Result<(), String> {
|
||||||
|
let actual = model.u64(key)?;
|
||||||
|
if actual == expected {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(format!("{key} is {actual}, expected {expected}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn expect_float(model: &Gguf, key: &str, expected: f32) -> Result<(), String> {
|
||||||
|
let actual = model.f32(key)?;
|
||||||
|
if float_eq(actual, expected) {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(format!("{key} is {actual}, expected {expected}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn float_eq(actual: f32, expected: f32) -> bool {
|
||||||
|
actual.is_finite() && (actual - expected).abs() <= expected.abs().max(1.0) * 1.0e-6
|
||||||
|
}
|
||||||
|
|
||||||
|
fn compression_ratio(shape: &Shape, layer: u32) -> u32 {
|
||||||
|
match shape.model {
|
||||||
|
ModelChoice::DeepSeekV4Flash if layer < 2 => 0,
|
||||||
|
ModelChoice::DeepSeekV4Pro if layer < 2 => 128,
|
||||||
|
ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro if layer.is_multiple_of(2) => 4,
|
||||||
|
ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro => 128,
|
||||||
|
ModelChoice::Glm52 => 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn first_u32(model: &Gguf, keys: &[&str]) -> Result<u32, String> {
|
||||||
|
keys.iter()
|
||||||
|
.find_map(|key| model.u32(key).ok())
|
||||||
|
.ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0]))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn first_u32s<'a>(model: &'a Gguf, keys: &[&str]) -> Result<&'a [u32], String> {
|
||||||
|
keys.iter()
|
||||||
|
.find_map(|key| model.u32s(key).ok())
|
||||||
|
.ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0]))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn installed_ds4_fixture_opens_and_renders_a_prompt() {
|
||||||
|
let path = Path::new(
|
||||||
|
"../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
|
||||||
|
);
|
||||||
|
if !path.exists() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let model = Model::open_main(path, ModelChoice::DeepSeekV4Flash).unwrap();
|
||||||
|
let summary = model.summary();
|
||||||
|
assert_eq!(summary.model, ModelChoice::DeepSeekV4Flash);
|
||||||
|
assert_eq!(summary.vocabulary_size, 129_280);
|
||||||
|
assert_eq!(
|
||||||
|
model.tokenize("Hello, world! 1234\nint café = 7;\n中文テスト"),
|
||||||
|
[
|
||||||
|
19_923, 14, 2_058, 3, 223, 6_895, 22, 201, 650, 57_664, 438, 223, 25, 510, 21_134,
|
||||||
|
109_288,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model.render_prompt("", "Hello", ReasoningMode::Direct),
|
||||||
|
[0, 128_803, 19_923, 128_804, 128_822]
|
||||||
|
);
|
||||||
|
assert_eq!(model.tokenizer.tokenize_rendered("|DSML|").len(), 1);
|
||||||
|
assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct));
|
||||||
|
assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High));
|
||||||
|
let think_start = *model
|
||||||
|
.render_prompt("", "Hello", ReasoningMode::High)
|
||||||
|
.last()
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
model.render_conversation(
|
||||||
|
"",
|
||||||
|
&[
|
||||||
|
ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
ChatTurn {
|
||||||
|
user: false,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
ReasoningMode::Direct,
|
||||||
|
),
|
||||||
|
[
|
||||||
|
0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model.render_continuation("Hello", ReasoningMode::Direct, false),
|
||||||
|
[1, 128_803, 19_923, 128_804, 128_822]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model.render_continuation("Hello", ReasoningMode::Direct, true),
|
||||||
|
[128_803, 19_923, 128_804, 128_822]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model.render_conversation(
|
||||||
|
"",
|
||||||
|
&[
|
||||||
|
ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
ChatTurn {
|
||||||
|
user: false,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: Some("Hello".into()),
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: "Hello".into(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
ReasoningMode::Direct,
|
||||||
|
),
|
||||||
|
[
|
||||||
|
0,
|
||||||
|
128_803,
|
||||||
|
19_923,
|
||||||
|
128_804,
|
||||||
|
think_start,
|
||||||
|
19_923,
|
||||||
|
128_822,
|
||||||
|
19_923,
|
||||||
|
1,
|
||||||
|
128_803,
|
||||||
|
19_923,
|
||||||
|
128_804,
|
||||||
|
128_822,
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn installed_dspark_fixture_passes_the_target_layout() {
|
||||||
|
let path = Path::new("../ds4/gguf/DeepSeek-V4-Flash-DSpark-support.gguf");
|
||||||
|
if path.exists() {
|
||||||
|
validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
516
src/model.rs
516
src/model.rs
@@ -1,10 +1,12 @@
|
|||||||
use std::fmt;
|
mod transfer;
|
||||||
use std::fs::{self, File, OpenOptions};
|
|
||||||
use std::io::{Read, Write};
|
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
||||||
|
|
||||||
use sha2::{Digest, Sha256};
|
pub(crate) use transfer::{
|
||||||
|
delete_managed_artifact, download_managed_artifact, validate_managed_artifact,
|
||||||
|
};
|
||||||
|
|
||||||
|
use std::fmt;
|
||||||
|
use std::fs;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [
|
pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [
|
||||||
ModelChoice::DeepSeekV4Flash,
|
ModelChoice::DeepSeekV4Flash,
|
||||||
@@ -379,505 +381,3 @@ pub(crate) fn artifact_verification_progress(
|
|||||||
}),
|
}),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn download_managed_artifact(
|
|
||||||
id: ManagedArtifactId,
|
|
||||||
models_path: &Path,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
verified_bytes: &AtomicU64,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
download_artifact_with_cancel(
|
|
||||||
id.model(),
|
|
||||||
id.artifact(),
|
|
||||||
models_path,
|
|
||||||
cancel,
|
|
||||||
verified_bytes,
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) fn validate_managed_artifact(
|
|
||||||
id: ManagedArtifactId,
|
|
||||||
models_path: &Path,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
verified_bytes: &AtomicU64,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
let model = id.model();
|
|
||||||
let artifact = id.artifact();
|
|
||||||
let destination = artifact.path(model, models_path);
|
|
||||||
let partial = artifact.partial_path(model, models_path);
|
|
||||||
let (path, promote) = if destination.exists() {
|
|
||||||
(destination.clone(), false)
|
|
||||||
} else if partial.exists() {
|
|
||||||
(partial.clone(), true)
|
|
||||||
} else {
|
|
||||||
return Err(format!("{} is not downloaded", artifact.label));
|
|
||||||
};
|
|
||||||
|
|
||||||
match verify(&path, artifact, model, cancel, verified_bytes) {
|
|
||||||
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
|
||||||
Ok(DownloadOutcome::Complete) => {}
|
|
||||||
Err(error) => {
|
|
||||||
let marker = artifact.verification_path(model, models_path);
|
|
||||||
if let Err(remove_error) = fs::remove_file(marker)
|
|
||||||
&& remove_error.kind() != std::io::ErrorKind::NotFound
|
|
||||||
{
|
|
||||||
return Err(format!(
|
|
||||||
"{error}; could not remove checksum marker: {remove_error}"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
return Err(error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if promote {
|
|
||||||
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
|
||||||
}
|
|
||||||
mark_verified(model, artifact, models_path)?;
|
|
||||||
Ok(DownloadOutcome::Complete)
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) fn delete_managed_artifact(
|
|
||||||
id: ManagedArtifactId,
|
|
||||||
models_path: &Path,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
let model = id.model();
|
|
||||||
let artifact = id.artifact();
|
|
||||||
for path in [
|
|
||||||
artifact.path(model, models_path),
|
|
||||||
artifact.partial_path(model, models_path),
|
|
||||||
artifact.verification_path(model, models_path),
|
|
||||||
] {
|
|
||||||
match fs::remove_file(path) {
|
|
||||||
Ok(()) => {}
|
|
||||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
|
|
||||||
Err(error) => return Err(error.to_string()),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
fn download_artifact(
|
|
||||||
model: ModelChoice,
|
|
||||||
artifact: &Artifact,
|
|
||||||
models_path: &Path,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
download_artifact_with_cancel(
|
|
||||||
model,
|
|
||||||
artifact,
|
|
||||||
models_path,
|
|
||||||
&AtomicBool::new(false),
|
|
||||||
&AtomicU64::new(0),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn download_artifact_with_cancel(
|
|
||||||
model: ModelChoice,
|
|
||||||
artifact: &Artifact,
|
|
||||||
models_path: &Path,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
verified_bytes: &AtomicU64,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
if cancel.load(Ordering::Relaxed) {
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
let destination = artifact.path(model, models_path);
|
|
||||||
if artifact.is_installed(model, models_path) {
|
|
||||||
return Ok(DownloadOutcome::Complete);
|
|
||||||
}
|
|
||||||
if destination.exists() {
|
|
||||||
if verify(&destination, artifact, model, cancel, verified_bytes)?
|
|
||||||
== DownloadOutcome::Stopped
|
|
||||||
{
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
mark_verified(model, artifact, models_path)?;
|
|
||||||
return Ok(DownloadOutcome::Complete);
|
|
||||||
}
|
|
||||||
|
|
||||||
let directory = destination
|
|
||||||
.parent()
|
|
||||||
.ok_or_else(|| "model artifact path has no parent directory".to_owned())?;
|
|
||||||
fs::create_dir_all(directory).map_err(|error| error.to_string())?;
|
|
||||||
let partial = artifact.partial_path(model, models_path);
|
|
||||||
let partial_size = partial.metadata().map_or(0, |metadata| metadata.len());
|
|
||||||
if partial_size > artifact.size {
|
|
||||||
File::create(&partial).map_err(|error| error.to_string())?;
|
|
||||||
}
|
|
||||||
if partial.metadata().map_or(0, |metadata| metadata.len()) != artifact.size
|
|
||||||
&& download_to_partial(artifact, &partial, cancel)? == DownloadOutcome::Stopped
|
|
||||||
{
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
|
|
||||||
if cancel.load(Ordering::Relaxed) {
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
match verify(&partial, artifact, model, cancel, verified_bytes) {
|
|
||||||
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
|
||||||
Ok(DownloadOutcome::Complete) => {}
|
|
||||||
Err(error) => {
|
|
||||||
fs::remove_file(&partial).map_err(|remove_error| {
|
|
||||||
format!("{error}; could not remove partial file: {remove_error}")
|
|
||||||
})?;
|
|
||||||
return Err(error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
|
||||||
mark_verified(model, artifact, models_path)?;
|
|
||||||
Ok(DownloadOutcome::Complete)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn mark_verified(
|
|
||||||
model: ModelChoice,
|
|
||||||
artifact: &Artifact,
|
|
||||||
models_path: &Path,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
fs::write(
|
|
||||||
artifact.verification_path(model, models_path),
|
|
||||||
artifact.sha256,
|
|
||||||
)
|
|
||||||
.map_err(|error| error.to_string())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn verify(
|
|
||||||
path: &Path,
|
|
||||||
artifact: &Artifact,
|
|
||||||
model: ModelChoice,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
verified_bytes: &AtomicU64,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
verified_bytes.store(0, Ordering::Relaxed);
|
|
||||||
let size = path.metadata().map_err(|error| error.to_string())?.len();
|
|
||||||
if size != artifact.size {
|
|
||||||
return Err(format!(
|
|
||||||
"{} has size {size}, expected {}",
|
|
||||||
path.display(),
|
|
||||||
artifact.size
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut file = File::open(path).map_err(|error| error.to_string())?;
|
|
||||||
let mut hasher = Sha256::new();
|
|
||||||
let mut buffer = vec![0; 1024 * 1024];
|
|
||||||
loop {
|
|
||||||
if cancel.load(Ordering::Relaxed) {
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
let count = file.read(&mut buffer).map_err(|error| error.to_string())?;
|
|
||||||
if count == 0 {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
hasher.update(&buffer[..count]);
|
|
||||||
verified_bytes.fetch_add(count as u64, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
let actual = hex(&hasher.finalize());
|
|
||||||
if actual != artifact.sha256 {
|
|
||||||
return Err(format!(
|
|
||||||
"Checksum verification failed for {}",
|
|
||||||
path.display()
|
|
||||||
));
|
|
||||||
}
|
|
||||||
if let Some(support) = artifact.support {
|
|
||||||
crate::engine::validate_model_artifact(path, model, support)?;
|
|
||||||
}
|
|
||||||
Ok(DownloadOutcome::Complete)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn hex(bytes: &[u8]) -> String {
|
|
||||||
const DIGITS: &[u8; 16] = b"0123456789abcdef";
|
|
||||||
let mut result = String::with_capacity(bytes.len() * 2);
|
|
||||||
for byte in bytes {
|
|
||||||
result.push(DIGITS[(byte >> 4) as usize] as char);
|
|
||||||
result.push(DIGITS[(byte & 0x0f) as usize] as char);
|
|
||||||
}
|
|
||||||
result
|
|
||||||
}
|
|
||||||
|
|
||||||
fn download_to_partial(
|
|
||||||
artifact: &Artifact,
|
|
||||||
partial: &Path,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
download_url_to_partial(&artifact.url(), partial, cancel)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn download_url_to_partial(
|
|
||||||
url: &str,
|
|
||||||
partial: &Path,
|
|
||||||
cancel: &AtomicBool,
|
|
||||||
) -> Result<DownloadOutcome, String> {
|
|
||||||
let offset = partial.metadata().map_or(0, |metadata| metadata.len());
|
|
||||||
let agent: ureq::Agent = ureq::Agent::config_builder()
|
|
||||||
.https_only(url.starts_with("https://"))
|
|
||||||
.build()
|
|
||||||
.into();
|
|
||||||
let mut request = agent.get(url);
|
|
||||||
if offset > 0 {
|
|
||||||
request = request.header("Range", format!("bytes={offset}-"));
|
|
||||||
}
|
|
||||||
let mut response = request
|
|
||||||
.call()
|
|
||||||
.map_err(|error| format!("Model download failed: {error}"))?;
|
|
||||||
let status = response.status().as_u16();
|
|
||||||
let append = offset > 0 && status == 206;
|
|
||||||
if offset > 0 && status != 200 && status != 206 {
|
|
||||||
return Err(format!(
|
|
||||||
"Model server returned HTTP {status} while resuming at byte {offset}"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
if append {
|
|
||||||
validate_content_range(&response, offset)?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut output = OpenOptions::new()
|
|
||||||
.create(true)
|
|
||||||
.write(true)
|
|
||||||
.append(append)
|
|
||||||
.truncate(!append)
|
|
||||||
.open(partial)
|
|
||||||
.map_err(|error| error.to_string())?;
|
|
||||||
let mut body = response.body_mut().as_reader();
|
|
||||||
let mut buffer = vec![0; 1024 * 1024];
|
|
||||||
loop {
|
|
||||||
if cancel.load(Ordering::Relaxed) {
|
|
||||||
return Ok(DownloadOutcome::Stopped);
|
|
||||||
}
|
|
||||||
let count = body
|
|
||||||
.read(&mut buffer)
|
|
||||||
.map_err(|error| format!("Model download failed: {error}"))?;
|
|
||||||
if count == 0 {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
output
|
|
||||||
.write_all(&buffer[..count])
|
|
||||||
.map_err(|error| error.to_string())?;
|
|
||||||
}
|
|
||||||
Ok(DownloadOutcome::Complete)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn validate_content_range(
|
|
||||||
response: &ureq::http::Response<ureq::Body>,
|
|
||||||
offset: u64,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
let expected = format!("bytes {offset}-");
|
|
||||||
let content_range = response
|
|
||||||
.headers()
|
|
||||||
.get("content-range")
|
|
||||||
.and_then(|value| value.to_str().ok())
|
|
||||||
.unwrap_or_default();
|
|
||||||
if content_range.starts_with(&expected) {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err(format!(
|
|
||||||
"Model server returned an invalid Content-Range while resuming at byte {offset}"
|
|
||||||
))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use std::net::TcpListener;
|
|
||||||
use std::thread;
|
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn catalog_and_checksum_verification_are_explicit() {
|
|
||||||
assert_eq!(ModelChoice::from_id("glm-5.2"), Some(ModelChoice::Glm52));
|
|
||||||
assert!(ModelChoice::from_id("unknown").is_none());
|
|
||||||
assert_eq!(
|
|
||||||
ModelChoice::DeepSeekV4Flash.main_artifact().size,
|
|
||||||
86_720_111_488
|
|
||||||
);
|
|
||||||
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
|
|
||||||
assert_eq!(ModelChoice::DeepSeekV4Flash.artifacts(true).count(), 2);
|
|
||||||
assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1);
|
|
||||||
|
|
||||||
let id = SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos();
|
|
||||||
let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}"));
|
|
||||||
let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash, true, &models_path);
|
|
||||||
assert_eq!(
|
|
||||||
engine.model.file_name(),
|
|
||||||
Some(std::ffi::OsStr::new(FLASH.file_name))
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
engine.mtp.as_deref().and_then(Path::file_name),
|
|
||||||
Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name))
|
|
||||||
);
|
|
||||||
let empty = Artifact {
|
|
||||||
label: "empty",
|
|
||||||
file_name: "empty",
|
|
||||||
repository: "",
|
|
||||||
size: 0,
|
|
||||||
sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
|
|
||||||
support: None,
|
|
||||||
};
|
|
||||||
let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash, &models_path);
|
|
||||||
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
||||||
fs::write(&partial, []).unwrap();
|
|
||||||
download_artifact(ModelChoice::DeepSeekV4Flash, &empty, &models_path).unwrap();
|
|
||||||
assert!(empty.is_installed(ModelChoice::DeepSeekV4Flash, &models_path));
|
|
||||||
assert!(!partial.exists());
|
|
||||||
assert_eq!(
|
|
||||||
fs::read_to_string(empty.verification_path(ModelChoice::DeepSeekV4Flash, &models_path))
|
|
||||||
.unwrap(),
|
|
||||||
empty.sha256
|
|
||||||
);
|
|
||||||
fs::remove_dir_all(models_path).unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn verification_reports_bytes_read() {
|
|
||||||
let path = std::env::temp_dir().join(format!(
|
|
||||||
"ds4-server-verify-{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
));
|
|
||||||
fs::write(&path, b"abc").unwrap();
|
|
||||||
let artifact = Artifact {
|
|
||||||
label: "test model",
|
|
||||||
file_name: "unused",
|
|
||||||
repository: "unused",
|
|
||||||
size: 3,
|
|
||||||
sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
|
|
||||||
support: None,
|
|
||||||
};
|
|
||||||
let verified_bytes = AtomicU64::new(999);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
verify(
|
|
||||||
&path,
|
|
||||||
&artifact,
|
|
||||||
ModelChoice::DeepSeekV4Flash,
|
|
||||||
&AtomicBool::new(false),
|
|
||||||
&verified_bytes,
|
|
||||||
)
|
|
||||||
.unwrap(),
|
|
||||||
DownloadOutcome::Complete
|
|
||||||
);
|
|
||||||
assert_eq!(verified_bytes.load(Ordering::Relaxed), 3);
|
|
||||||
fs::remove_file(path).unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn managed_artifact_inventory_and_delete_include_partial_files() {
|
|
||||||
let models_path = std::env::temp_dir().join(format!(
|
|
||||||
"ds4-server-manager-{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
));
|
|
||||||
let id = ManagedArtifactId::DeepSeekV4Flash;
|
|
||||||
let partial = id.artifact().partial_path(id.model(), &models_path);
|
|
||||||
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
||||||
fs::write(&partial, b"part").unwrap();
|
|
||||||
|
|
||||||
let managed = managed_artifacts(&models_path)
|
|
||||||
.into_iter()
|
|
||||||
.find(|artifact| artifact.id == id)
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(managed.stored, 4);
|
|
||||||
assert_eq!(managed.state, ManagedArtifactState::Partial);
|
|
||||||
|
|
||||||
delete_managed_artifact(id, &models_path).unwrap();
|
|
||||||
assert!(!partial.exists());
|
|
||||||
fs::remove_dir_all(models_path).unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn restart_resumes_at_the_existing_partial_byte() {
|
|
||||||
let content = b"restart-resume works";
|
|
||||||
let offset = 8;
|
|
||||||
let directory = std::env::temp_dir().join(format!(
|
|
||||||
"ds4-server-resume-{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
));
|
|
||||||
fs::create_dir_all(&directory).unwrap();
|
|
||||||
let partial = directory.join("model.gguf.part");
|
|
||||||
fs::write(&partial, &content[..offset]).unwrap();
|
|
||||||
|
|
||||||
let listener = match TcpListener::bind("127.0.0.1:0") {
|
|
||||||
Ok(listener) => listener,
|
|
||||||
Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => {
|
|
||||||
fs::remove_dir_all(directory).unwrap();
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
Err(error) => panic!("could not start test server: {error}"),
|
|
||||||
};
|
|
||||||
let address = listener.local_addr().unwrap();
|
|
||||||
let server = thread::spawn(move || {
|
|
||||||
let (mut connection, _) = listener.accept().unwrap();
|
|
||||||
let mut request = [0; 2048];
|
|
||||||
let count = connection.read(&mut request).unwrap();
|
|
||||||
let request = String::from_utf8_lossy(&request[..count]).to_ascii_lowercase();
|
|
||||||
assert!(request.contains("range: bytes=8-"));
|
|
||||||
let remaining = &content[offset..];
|
|
||||||
write!(
|
|
||||||
connection,
|
|
||||||
"HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n",
|
|
||||||
remaining.len(),
|
|
||||||
content.len() - 1,
|
|
||||||
content.len(),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
connection.write_all(remaining).unwrap();
|
|
||||||
});
|
|
||||||
|
|
||||||
let outcome = download_url_to_partial(
|
|
||||||
&format!("http://{address}/model.gguf"),
|
|
||||||
&partial,
|
|
||||||
&AtomicBool::new(false),
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
server.join().unwrap();
|
|
||||||
assert_eq!(outcome, DownloadOutcome::Complete);
|
|
||||||
assert_eq!(fs::read(&partial).unwrap(), content);
|
|
||||||
fs::remove_dir_all(directory).unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn cancellation_keeps_the_partial_file_for_the_next_run() {
|
|
||||||
let directory = std::env::temp_dir().join(format!(
|
|
||||||
"ds4-server-cancel-{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap()
|
|
||||||
.as_nanos()
|
|
||||||
));
|
|
||||||
let artifact = Artifact {
|
|
||||||
label: "test model",
|
|
||||||
file_name: "model.gguf",
|
|
||||||
repository: "unused",
|
|
||||||
size: 10,
|
|
||||||
sha256: "unused",
|
|
||||||
support: None,
|
|
||||||
};
|
|
||||||
let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash, &directory);
|
|
||||||
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
|
||||||
fs::write(&partial, b"part").unwrap();
|
|
||||||
let cancel = AtomicBool::new(true);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
download_artifact_with_cancel(
|
|
||||||
ModelChoice::DeepSeekV4Flash,
|
|
||||||
&artifact,
|
|
||||||
&directory,
|
|
||||||
&cancel,
|
|
||||||
&AtomicU64::new(0),
|
|
||||||
)
|
|
||||||
.unwrap(),
|
|
||||||
DownloadOutcome::Stopped
|
|
||||||
);
|
|
||||||
assert_eq!(fs::read(&partial).unwrap(), b"part");
|
|
||||||
fs::remove_dir_all(directory).unwrap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
507
src/model/transfer.rs
Normal file
507
src/model/transfer.rs
Normal file
@@ -0,0 +1,507 @@
|
|||||||
|
use super::*;
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
use std::fs::{self, File, OpenOptions};
|
||||||
|
use std::io::{Read, Write};
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||||
|
|
||||||
|
pub(crate) fn download_managed_artifact(
|
||||||
|
id: ManagedArtifactId,
|
||||||
|
models_path: &Path,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
verified_bytes: &AtomicU64,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
download_artifact_with_cancel(
|
||||||
|
id.model(),
|
||||||
|
id.artifact(),
|
||||||
|
models_path,
|
||||||
|
cancel,
|
||||||
|
verified_bytes,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn validate_managed_artifact(
|
||||||
|
id: ManagedArtifactId,
|
||||||
|
models_path: &Path,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
verified_bytes: &AtomicU64,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
let model = id.model();
|
||||||
|
let artifact = id.artifact();
|
||||||
|
let destination = artifact.path(model, models_path);
|
||||||
|
let partial = artifact.partial_path(model, models_path);
|
||||||
|
let (path, promote) = if destination.exists() {
|
||||||
|
(destination.clone(), false)
|
||||||
|
} else if partial.exists() {
|
||||||
|
(partial.clone(), true)
|
||||||
|
} else {
|
||||||
|
return Err(format!("{} is not downloaded", artifact.label));
|
||||||
|
};
|
||||||
|
|
||||||
|
match verify(&path, artifact, model, cancel, verified_bytes) {
|
||||||
|
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
||||||
|
Ok(DownloadOutcome::Complete) => {}
|
||||||
|
Err(error) => {
|
||||||
|
let marker = artifact.verification_path(model, models_path);
|
||||||
|
if let Err(remove_error) = fs::remove_file(marker)
|
||||||
|
&& remove_error.kind() != std::io::ErrorKind::NotFound
|
||||||
|
{
|
||||||
|
return Err(format!(
|
||||||
|
"{error}; could not remove checksum marker: {remove_error}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if promote {
|
||||||
|
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
||||||
|
}
|
||||||
|
mark_verified(model, artifact, models_path)?;
|
||||||
|
Ok(DownloadOutcome::Complete)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn delete_managed_artifact(
|
||||||
|
id: ManagedArtifactId,
|
||||||
|
models_path: &Path,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let model = id.model();
|
||||||
|
let artifact = id.artifact();
|
||||||
|
for path in [
|
||||||
|
artifact.path(model, models_path),
|
||||||
|
artifact.partial_path(model, models_path),
|
||||||
|
artifact.verification_path(model, models_path),
|
||||||
|
] {
|
||||||
|
match fs::remove_file(path) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
|
||||||
|
Err(error) => return Err(error.to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
fn download_artifact(
|
||||||
|
model: ModelChoice,
|
||||||
|
artifact: &Artifact,
|
||||||
|
models_path: &Path,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
download_artifact_with_cancel(
|
||||||
|
model,
|
||||||
|
artifact,
|
||||||
|
models_path,
|
||||||
|
&AtomicBool::new(false),
|
||||||
|
&AtomicU64::new(0),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn download_artifact_with_cancel(
|
||||||
|
model: ModelChoice,
|
||||||
|
artifact: &Artifact,
|
||||||
|
models_path: &Path,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
verified_bytes: &AtomicU64,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
if cancel.load(Ordering::Relaxed) {
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
let destination = artifact.path(model, models_path);
|
||||||
|
if artifact.is_installed(model, models_path) {
|
||||||
|
return Ok(DownloadOutcome::Complete);
|
||||||
|
}
|
||||||
|
if destination.exists() {
|
||||||
|
if verify(&destination, artifact, model, cancel, verified_bytes)?
|
||||||
|
== DownloadOutcome::Stopped
|
||||||
|
{
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
mark_verified(model, artifact, models_path)?;
|
||||||
|
return Ok(DownloadOutcome::Complete);
|
||||||
|
}
|
||||||
|
|
||||||
|
let directory = destination
|
||||||
|
.parent()
|
||||||
|
.ok_or_else(|| "model artifact path has no parent directory".to_owned())?;
|
||||||
|
fs::create_dir_all(directory).map_err(|error| error.to_string())?;
|
||||||
|
let partial = artifact.partial_path(model, models_path);
|
||||||
|
let partial_size = partial.metadata().map_or(0, |metadata| metadata.len());
|
||||||
|
if partial_size > artifact.size {
|
||||||
|
File::create(&partial).map_err(|error| error.to_string())?;
|
||||||
|
}
|
||||||
|
if partial.metadata().map_or(0, |metadata| metadata.len()) != artifact.size
|
||||||
|
&& download_to_partial(artifact, &partial, cancel)? == DownloadOutcome::Stopped
|
||||||
|
{
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
|
||||||
|
if cancel.load(Ordering::Relaxed) {
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
match verify(&partial, artifact, model, cancel, verified_bytes) {
|
||||||
|
Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped),
|
||||||
|
Ok(DownloadOutcome::Complete) => {}
|
||||||
|
Err(error) => {
|
||||||
|
fs::remove_file(&partial).map_err(|remove_error| {
|
||||||
|
format!("{error}; could not remove partial file: {remove_error}")
|
||||||
|
})?;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
fs::rename(partial, destination).map_err(|error| error.to_string())?;
|
||||||
|
mark_verified(model, artifact, models_path)?;
|
||||||
|
Ok(DownloadOutcome::Complete)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mark_verified(
|
||||||
|
model: ModelChoice,
|
||||||
|
artifact: &Artifact,
|
||||||
|
models_path: &Path,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
fs::write(
|
||||||
|
artifact.verification_path(model, models_path),
|
||||||
|
artifact.sha256,
|
||||||
|
)
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn verify(
|
||||||
|
path: &Path,
|
||||||
|
artifact: &Artifact,
|
||||||
|
model: ModelChoice,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
verified_bytes: &AtomicU64,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
verified_bytes.store(0, Ordering::Relaxed);
|
||||||
|
let size = path.metadata().map_err(|error| error.to_string())?.len();
|
||||||
|
if size != artifact.size {
|
||||||
|
return Err(format!(
|
||||||
|
"{} has size {size}, expected {}",
|
||||||
|
path.display(),
|
||||||
|
artifact.size
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut file = File::open(path).map_err(|error| error.to_string())?;
|
||||||
|
let mut hasher = Sha256::new();
|
||||||
|
let mut buffer = vec![0; 1024 * 1024];
|
||||||
|
loop {
|
||||||
|
if cancel.load(Ordering::Relaxed) {
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
let count = file.read(&mut buffer).map_err(|error| error.to_string())?;
|
||||||
|
if count == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
hasher.update(&buffer[..count]);
|
||||||
|
verified_bytes.fetch_add(count as u64, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
let actual = hex(&hasher.finalize());
|
||||||
|
if actual != artifact.sha256 {
|
||||||
|
return Err(format!(
|
||||||
|
"Checksum verification failed for {}",
|
||||||
|
path.display()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if let Some(support) = artifact.support {
|
||||||
|
crate::engine::validate_model_artifact(path, model, support)?;
|
||||||
|
}
|
||||||
|
Ok(DownloadOutcome::Complete)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn hex(bytes: &[u8]) -> String {
|
||||||
|
const DIGITS: &[u8; 16] = b"0123456789abcdef";
|
||||||
|
let mut result = String::with_capacity(bytes.len() * 2);
|
||||||
|
for byte in bytes {
|
||||||
|
result.push(DIGITS[(byte >> 4) as usize] as char);
|
||||||
|
result.push(DIGITS[(byte & 0x0f) as usize] as char);
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn download_to_partial(
|
||||||
|
artifact: &Artifact,
|
||||||
|
partial: &Path,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
download_url_to_partial(&artifact.url(), partial, cancel)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn download_url_to_partial(
|
||||||
|
url: &str,
|
||||||
|
partial: &Path,
|
||||||
|
cancel: &AtomicBool,
|
||||||
|
) -> Result<DownloadOutcome, String> {
|
||||||
|
let offset = partial.metadata().map_or(0, |metadata| metadata.len());
|
||||||
|
let agent: ureq::Agent = ureq::Agent::config_builder()
|
||||||
|
.https_only(url.starts_with("https://"))
|
||||||
|
.build()
|
||||||
|
.into();
|
||||||
|
let mut request = agent.get(url);
|
||||||
|
if offset > 0 {
|
||||||
|
request = request.header("Range", format!("bytes={offset}-"));
|
||||||
|
}
|
||||||
|
let mut response = request
|
||||||
|
.call()
|
||||||
|
.map_err(|error| format!("Model download failed: {error}"))?;
|
||||||
|
let status = response.status().as_u16();
|
||||||
|
let append = offset > 0 && status == 206;
|
||||||
|
if offset > 0 && status != 200 && status != 206 {
|
||||||
|
return Err(format!(
|
||||||
|
"Model server returned HTTP {status} while resuming at byte {offset}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if append {
|
||||||
|
validate_content_range(&response, offset)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut output = OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.write(true)
|
||||||
|
.append(append)
|
||||||
|
.truncate(!append)
|
||||||
|
.open(partial)
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
let mut body = response.body_mut().as_reader();
|
||||||
|
let mut buffer = vec![0; 1024 * 1024];
|
||||||
|
loop {
|
||||||
|
if cancel.load(Ordering::Relaxed) {
|
||||||
|
return Ok(DownloadOutcome::Stopped);
|
||||||
|
}
|
||||||
|
let count = body
|
||||||
|
.read(&mut buffer)
|
||||||
|
.map_err(|error| format!("Model download failed: {error}"))?;
|
||||||
|
if count == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
output
|
||||||
|
.write_all(&buffer[..count])
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
}
|
||||||
|
Ok(DownloadOutcome::Complete)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_content_range(
|
||||||
|
response: &ureq::http::Response<ureq::Body>,
|
||||||
|
offset: u64,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let expected = format!("bytes {offset}-");
|
||||||
|
let content_range = response
|
||||||
|
.headers()
|
||||||
|
.get("content-range")
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or_default();
|
||||||
|
if content_range.starts_with(&expected) {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(format!(
|
||||||
|
"Model server returned an invalid Content-Range while resuming at byte {offset}"
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::net::TcpListener;
|
||||||
|
use std::thread;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn catalog_and_checksum_verification_are_explicit() {
|
||||||
|
assert_eq!(ModelChoice::from_id("glm-5.2"), Some(ModelChoice::Glm52));
|
||||||
|
assert!(ModelChoice::from_id("unknown").is_none());
|
||||||
|
assert_eq!(
|
||||||
|
ModelChoice::DeepSeekV4Flash.main_artifact().size,
|
||||||
|
86_720_111_488
|
||||||
|
);
|
||||||
|
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
|
||||||
|
assert_eq!(ModelChoice::DeepSeekV4Flash.artifacts(true).count(), 2);
|
||||||
|
assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1);
|
||||||
|
|
||||||
|
let id = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos();
|
||||||
|
let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}"));
|
||||||
|
let engine = engine_artifacts(ModelChoice::DeepSeekV4Flash, true, &models_path);
|
||||||
|
assert_eq!(
|
||||||
|
engine.model.file_name(),
|
||||||
|
Some(std::ffi::OsStr::new(FLASH.file_name))
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
engine.mtp.as_deref().and_then(Path::file_name),
|
||||||
|
Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name))
|
||||||
|
);
|
||||||
|
let empty = Artifact {
|
||||||
|
label: "empty",
|
||||||
|
file_name: "empty",
|
||||||
|
repository: "",
|
||||||
|
size: 0,
|
||||||
|
sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
|
||||||
|
support: None,
|
||||||
|
};
|
||||||
|
let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash, &models_path);
|
||||||
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
||||||
|
fs::write(&partial, []).unwrap();
|
||||||
|
download_artifact(ModelChoice::DeepSeekV4Flash, &empty, &models_path).unwrap();
|
||||||
|
assert!(empty.is_installed(ModelChoice::DeepSeekV4Flash, &models_path));
|
||||||
|
assert!(!partial.exists());
|
||||||
|
assert_eq!(
|
||||||
|
fs::read_to_string(empty.verification_path(ModelChoice::DeepSeekV4Flash, &models_path))
|
||||||
|
.unwrap(),
|
||||||
|
empty.sha256
|
||||||
|
);
|
||||||
|
fs::remove_dir_all(models_path).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn verification_reports_bytes_read() {
|
||||||
|
let path = std::env::temp_dir().join(format!(
|
||||||
|
"ds4-server-verify-{}",
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos()
|
||||||
|
));
|
||||||
|
fs::write(&path, b"abc").unwrap();
|
||||||
|
let artifact = Artifact {
|
||||||
|
label: "test model",
|
||||||
|
file_name: "unused",
|
||||||
|
repository: "unused",
|
||||||
|
size: 3,
|
||||||
|
sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
|
||||||
|
support: None,
|
||||||
|
};
|
||||||
|
let verified_bytes = AtomicU64::new(999);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
verify(
|
||||||
|
&path,
|
||||||
|
&artifact,
|
||||||
|
ModelChoice::DeepSeekV4Flash,
|
||||||
|
&AtomicBool::new(false),
|
||||||
|
&verified_bytes,
|
||||||
|
)
|
||||||
|
.unwrap(),
|
||||||
|
DownloadOutcome::Complete
|
||||||
|
);
|
||||||
|
assert_eq!(verified_bytes.load(Ordering::Relaxed), 3);
|
||||||
|
fs::remove_file(path).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn managed_artifact_inventory_and_delete_include_partial_files() {
|
||||||
|
let models_path = std::env::temp_dir().join(format!(
|
||||||
|
"ds4-server-manager-{}",
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos()
|
||||||
|
));
|
||||||
|
let id = ManagedArtifactId::DeepSeekV4Flash;
|
||||||
|
let partial = id.artifact().partial_path(id.model(), &models_path);
|
||||||
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
||||||
|
fs::write(&partial, b"part").unwrap();
|
||||||
|
|
||||||
|
let managed = managed_artifacts(&models_path)
|
||||||
|
.into_iter()
|
||||||
|
.find(|artifact| artifact.id == id)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(managed.stored, 4);
|
||||||
|
assert_eq!(managed.state, ManagedArtifactState::Partial);
|
||||||
|
|
||||||
|
delete_managed_artifact(id, &models_path).unwrap();
|
||||||
|
assert!(!partial.exists());
|
||||||
|
fs::remove_dir_all(models_path).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn restart_resumes_at_the_existing_partial_byte() {
|
||||||
|
let content = b"restart-resume works";
|
||||||
|
let offset = 8;
|
||||||
|
let directory = std::env::temp_dir().join(format!(
|
||||||
|
"ds4-server-resume-{}",
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos()
|
||||||
|
));
|
||||||
|
fs::create_dir_all(&directory).unwrap();
|
||||||
|
let partial = directory.join("model.gguf.part");
|
||||||
|
fs::write(&partial, &content[..offset]).unwrap();
|
||||||
|
|
||||||
|
let listener = match TcpListener::bind("127.0.0.1:0") {
|
||||||
|
Ok(listener) => listener,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => {
|
||||||
|
fs::remove_dir_all(directory).unwrap();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
Err(error) => panic!("could not start test server: {error}"),
|
||||||
|
};
|
||||||
|
let address = listener.local_addr().unwrap();
|
||||||
|
let server = thread::spawn(move || {
|
||||||
|
let (mut connection, _) = listener.accept().unwrap();
|
||||||
|
let mut request = [0; 2048];
|
||||||
|
let count = connection.read(&mut request).unwrap();
|
||||||
|
let request = String::from_utf8_lossy(&request[..count]).to_ascii_lowercase();
|
||||||
|
assert!(request.contains("range: bytes=8-"));
|
||||||
|
let remaining = &content[offset..];
|
||||||
|
write!(
|
||||||
|
connection,
|
||||||
|
"HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n",
|
||||||
|
remaining.len(),
|
||||||
|
content.len() - 1,
|
||||||
|
content.len(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
connection.write_all(remaining).unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let outcome = download_url_to_partial(
|
||||||
|
&format!("http://{address}/model.gguf"),
|
||||||
|
&partial,
|
||||||
|
&AtomicBool::new(false),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
server.join().unwrap();
|
||||||
|
assert_eq!(outcome, DownloadOutcome::Complete);
|
||||||
|
assert_eq!(fs::read(&partial).unwrap(), content);
|
||||||
|
fs::remove_dir_all(directory).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cancellation_keeps_the_partial_file_for_the_next_run() {
|
||||||
|
let directory = std::env::temp_dir().join(format!(
|
||||||
|
"ds4-server-cancel-{}",
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_nanos()
|
||||||
|
));
|
||||||
|
let artifact = Artifact {
|
||||||
|
label: "test model",
|
||||||
|
file_name: "model.gguf",
|
||||||
|
repository: "unused",
|
||||||
|
size: 10,
|
||||||
|
sha256: "unused",
|
||||||
|
support: None,
|
||||||
|
};
|
||||||
|
let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash, &directory);
|
||||||
|
fs::create_dir_all(partial.parent().unwrap()).unwrap();
|
||||||
|
fs::write(&partial, b"part").unwrap();
|
||||||
|
let cancel = AtomicBool::new(true);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
download_artifact_with_cancel(
|
||||||
|
ModelChoice::DeepSeekV4Flash,
|
||||||
|
&artifact,
|
||||||
|
&directory,
|
||||||
|
&cancel,
|
||||||
|
&AtomicU64::new(0),
|
||||||
|
)
|
||||||
|
.unwrap(),
|
||||||
|
DownloadOutcome::Stopped
|
||||||
|
);
|
||||||
|
assert_eq!(fs::read(&partial).unwrap(), b"part");
|
||||||
|
fs::remove_dir_all(directory).unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
2490
src/server.rs
2490
src/server.rs
File diff suppressed because it is too large
Load Diff
178
src/server/http.rs
Normal file
178
src/server/http.rs
Normal file
@@ -0,0 +1,178 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
pub(super) fn send_sse_headers(stream: &mut impl Write) -> Result<(), String> {
|
||||||
|
stream
|
||||||
|
.write_all(
|
||||||
|
b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n",
|
||||||
|
)
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_sse(stream: &mut impl Write, value: &Value) -> Result<(), String> {
|
||||||
|
let body = serde_json::to_vec(value).map_err(|error| error.to_string())?;
|
||||||
|
stream
|
||||||
|
.write_all(b"data: ")
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
stream.write_all(&body).map_err(|error| error.to_string())?;
|
||||||
|
stream.write_all(b"\n\n").map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_sse_error(stream: &mut impl Write, message: &str) -> Result<(), String> {
|
||||||
|
let body = serde_json::to_vec(&json!({"error": {"message": message, "type": "server_error"}}))
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
stream
|
||||||
|
.write_all(b"event: error\ndata: ")
|
||||||
|
.map_err(|error| error.to_string())?;
|
||||||
|
stream.write_all(&body).map_err(|error| error.to_string())?;
|
||||||
|
stream.write_all(b"\n\n").map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_json(stream: &mut TcpStream, code: u16, value: &Value) -> Result<(), String> {
|
||||||
|
let mut body = serde_json::to_vec(value).map_err(|error| error.to_string())?;
|
||||||
|
body.push(b'\n');
|
||||||
|
send_response(stream, code, Some("application/json"), &body)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_error(stream: &mut TcpStream, code: u16, message: &str) -> Result<(), String> {
|
||||||
|
send_json(
|
||||||
|
stream,
|
||||||
|
code,
|
||||||
|
&json!({"error": {"message": message, "type": "invalid_request_error"}}),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
code: u16,
|
||||||
|
content_type: Option<&str>,
|
||||||
|
body: &[u8],
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let reason = match code {
|
||||||
|
200 => "OK",
|
||||||
|
204 => "No Content",
|
||||||
|
400 => "Bad Request",
|
||||||
|
404 => "Not Found",
|
||||||
|
409 => "Conflict",
|
||||||
|
500 => "Internal Server Error",
|
||||||
|
_ => "Error",
|
||||||
|
};
|
||||||
|
let mut header = format!(
|
||||||
|
"HTTP/1.1 {code} {reason}\r\nContent-Length: {}\r\n",
|
||||||
|
body.len()
|
||||||
|
);
|
||||||
|
if let Some(content_type) = content_type {
|
||||||
|
header.push_str("Content-Type: ");
|
||||||
|
header.push_str(content_type);
|
||||||
|
header.push_str("\r\n");
|
||||||
|
}
|
||||||
|
header.push_str("Connection: close\r\n\r\n");
|
||||||
|
stream
|
||||||
|
.write_all(header.as_bytes())
|
||||||
|
.and_then(|()| stream.write_all(body))
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, String> {
|
||||||
|
let mut bytes = Vec::new();
|
||||||
|
let header_end = loop {
|
||||||
|
if bytes.len() >= MAX_HEADER_BYTES {
|
||||||
|
return Err("HTTP headers are too large".into());
|
||||||
|
}
|
||||||
|
let mut chunk = [0_u8; 4096];
|
||||||
|
let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
|
||||||
|
if read == 0 {
|
||||||
|
return Err("bad HTTP request".into());
|
||||||
|
}
|
||||||
|
bytes.extend_from_slice(&chunk[..read]);
|
||||||
|
if let Some(end) = find_header_end(&bytes) {
|
||||||
|
break end;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let header = std::str::from_utf8(&bytes[..header_end])
|
||||||
|
.map_err(|_| "HTTP headers are not UTF-8".to_owned())?;
|
||||||
|
let mut lines = header.lines();
|
||||||
|
let request_line = lines.next().ok_or_else(|| "bad HTTP request".to_owned())?;
|
||||||
|
let mut parts = request_line.split_whitespace();
|
||||||
|
let method = parts
|
||||||
|
.next()
|
||||||
|
.ok_or_else(|| "bad HTTP request".to_owned())?
|
||||||
|
.to_owned();
|
||||||
|
let path = parts
|
||||||
|
.next()
|
||||||
|
.ok_or_else(|| "bad HTTP request".to_owned())?
|
||||||
|
.to_owned();
|
||||||
|
let length = lines
|
||||||
|
.find_map(|line| {
|
||||||
|
line.split_once(':').and_then(|(name, value)| {
|
||||||
|
name.eq_ignore_ascii_case("content-length")
|
||||||
|
.then(|| value.trim().parse::<usize>().ok())
|
||||||
|
.flatten()
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.unwrap_or(0);
|
||||||
|
if length > MAX_BODY_BYTES {
|
||||||
|
return Err("HTTP body is too large".into());
|
||||||
|
}
|
||||||
|
while bytes.len() < header_end + length {
|
||||||
|
let mut chunk = [0_u8; 8192];
|
||||||
|
let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
|
||||||
|
if read == 0 {
|
||||||
|
return Err("incomplete HTTP body".into());
|
||||||
|
}
|
||||||
|
bytes.extend_from_slice(&chunk[..read]);
|
||||||
|
}
|
||||||
|
Ok(HttpRequest {
|
||||||
|
method,
|
||||||
|
path: path.split('?').next().unwrap_or(&path).to_owned(),
|
||||||
|
body: bytes[header_end..header_end + length].to_vec(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn find_header_end(bytes: &[u8]) -> Option<usize> {
|
||||||
|
bytes
|
||||||
|
.windows(4)
|
||||||
|
.position(|window| window == b"\r\n\r\n")
|
||||||
|
.map(|position| position + 4)
|
||||||
|
.or_else(|| {
|
||||||
|
bytes
|
||||||
|
.windows(2)
|
||||||
|
.position(|window| window == b"\n\n")
|
||||||
|
.map(|position| position + 2)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn unix_time() -> u64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn random_id(prefix: &str) -> String {
|
||||||
|
random_hex_id::<12>(prefix)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn random_tool_id(prefix: &str) -> String {
|
||||||
|
random_hex_id::<16>(prefix)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn random_hex_id<const N: usize>(prefix: &str) -> String {
|
||||||
|
let mut bytes = [0_u8; N];
|
||||||
|
if File::open("/dev/urandom")
|
||||||
|
.and_then(|mut file| file.read_exact(&mut bytes))
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
let fallback = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_nanos();
|
||||||
|
bytes.copy_from_slice(&fallback.to_le_bytes()[..N]);
|
||||||
|
}
|
||||||
|
let mut id = String::with_capacity(prefix.len() + N * 2);
|
||||||
|
id.push_str(prefix);
|
||||||
|
for byte in bytes {
|
||||||
|
use std::fmt::Write as _;
|
||||||
|
let _ = write!(id, "{byte:02x}");
|
||||||
|
}
|
||||||
|
id
|
||||||
|
}
|
||||||
595
src/server/request.rs
Normal file
595
src/server/request.rs
Normal file
@@ -0,0 +1,595 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
pub(super) fn completion_request(mut value: Value) -> Result<ChatRequest, (u16, String)> {
|
||||||
|
let object = value
|
||||||
|
.as_object_mut()
|
||||||
|
.ok_or_else(|| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
let prompt = object
|
||||||
|
.remove("prompt")
|
||||||
|
.ok_or_else(|| (400, "missing prompt".to_owned()))?;
|
||||||
|
let prompt = match prompt {
|
||||||
|
Value::String(text) => text,
|
||||||
|
Value::Array(values) => values
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
.and_then(|value| value.as_str().map(str::to_owned))
|
||||||
|
.unwrap_or_default(),
|
||||||
|
_ => String::new(),
|
||||||
|
};
|
||||||
|
object.insert(
|
||||||
|
"messages".into(),
|
||||||
|
json!([
|
||||||
|
{"role": "system", "content": "You are a helpful assistant"},
|
||||||
|
{"role": "user", "content": prompt}
|
||||||
|
]),
|
||||||
|
);
|
||||||
|
decode_chat_request(value)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn anthropic_request(value: Value) -> Result<ChatRequest, (u16, String)> {
|
||||||
|
let object = value
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
let source = object
|
||||||
|
.get("messages")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.ok_or_else(|| (400, "missing messages".to_owned()))?;
|
||||||
|
let mut messages = Vec::new();
|
||||||
|
let system = object.get("system").map(content_text).unwrap_or_default();
|
||||||
|
if !system.is_empty() {
|
||||||
|
messages.push(json!({"role": "system", "content": system}));
|
||||||
|
}
|
||||||
|
for message in source {
|
||||||
|
let role = message
|
||||||
|
.get("role")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("user");
|
||||||
|
let Some(blocks) = message.get("content").and_then(Value::as_array) else {
|
||||||
|
messages.push(json!({"role": role, "content": message.get("content").cloned().unwrap_or(Value::Null)}));
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let mut text = String::new();
|
||||||
|
let mut reasoning = String::new();
|
||||||
|
let mut calls = Vec::new();
|
||||||
|
for block in blocks {
|
||||||
|
match block.get("type").and_then(Value::as_str).unwrap_or("text") {
|
||||||
|
"text" => text.push_str(block.get("text").and_then(Value::as_str).unwrap_or("")),
|
||||||
|
"thinking" | "redacted_thinking" => reasoning.push_str(
|
||||||
|
block
|
||||||
|
.get("thinking")
|
||||||
|
.or_else(|| block.get("data"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or(""),
|
||||||
|
),
|
||||||
|
"tool_use" => calls.push(json!({
|
||||||
|
"id": block.get("id").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": block.get("name").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"arguments": block.get("input").cloned().unwrap_or_else(|| json!({})).to_string()
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
"tool_result" => {
|
||||||
|
if !text.is_empty() {
|
||||||
|
messages.push(json!({"role": role, "content": std::mem::take(&mut text)}));
|
||||||
|
}
|
||||||
|
messages.push(json!({
|
||||||
|
"role": "tool",
|
||||||
|
"content": block.get("content").map(content_text).unwrap_or_default(),
|
||||||
|
"tool_call_id": block.get("tool_use_id").and_then(Value::as_str).unwrap_or("")
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !text.is_empty() || !reasoning.is_empty() || !calls.is_empty() {
|
||||||
|
messages.push(json!({
|
||||||
|
"role": role,
|
||||||
|
"content": text,
|
||||||
|
"reasoning_content": reasoning,
|
||||||
|
"tool_calls": calls
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut request = Map::new();
|
||||||
|
copy_request_fields(object, &mut request);
|
||||||
|
request.insert("messages".into(), Value::Array(messages));
|
||||||
|
if let Some(stops) = object.get("stop_sequences") {
|
||||||
|
request.insert("stop".into(), stops.clone());
|
||||||
|
}
|
||||||
|
if let Some(effort) = object
|
||||||
|
.get("output_config")
|
||||||
|
.and_then(|value| value.get("effort"))
|
||||||
|
{
|
||||||
|
request.insert("reasoning_effort".into(), effort.clone());
|
||||||
|
}
|
||||||
|
request.insert(
|
||||||
|
"tools".into(),
|
||||||
|
normalize_anthropic_tools(object.get("tools")),
|
||||||
|
);
|
||||||
|
decode_chat_request(Value::Object(request))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn responses_request(value: Value) -> Result<ChatRequest, (u16, String)> {
|
||||||
|
let object = value
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
for key in ["previous_response_id", "conversation"] {
|
||||||
|
if object.get(key).is_some_and(|value| !value.is_null()) {
|
||||||
|
return Err((
|
||||||
|
400,
|
||||||
|
format!("{key} is not supported; replay full input instead"),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(choice) = object.get("tool_choice") {
|
||||||
|
match choice {
|
||||||
|
Value::String(choice) if choice == "none" || choice == "auto" => {}
|
||||||
|
Value::String(choice) => {
|
||||||
|
return Err((400, format!("tool_choice={choice} not supported")));
|
||||||
|
}
|
||||||
|
Value::Object(_) => return Err((400, "forced tool_choice not supported".into())),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let input = object
|
||||||
|
.get("input")
|
||||||
|
.ok_or_else(|| (400, "missing input".to_owned()))?;
|
||||||
|
let mut messages = Vec::new();
|
||||||
|
match input {
|
||||||
|
Value::String(text) => messages.push(json!({"role": "user", "content": text})),
|
||||||
|
Value::Array(items) => {
|
||||||
|
let mut pending_reasoning = String::new();
|
||||||
|
for item in items {
|
||||||
|
let item = item
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
if item
|
||||||
|
.get("status")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.is_some_and(|status| status != "completed")
|
||||||
|
{
|
||||||
|
return Err((400, "invalid JSON request".into()));
|
||||||
|
}
|
||||||
|
let kind = item
|
||||||
|
.get("type")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("message");
|
||||||
|
let role = item.get("role").and_then(Value::as_str).unwrap_or("user");
|
||||||
|
let consumes_reasoning = (kind == "message" && role == "assistant")
|
||||||
|
|| matches!(
|
||||||
|
kind,
|
||||||
|
"function_call"
|
||||||
|
| "custom_tool_call"
|
||||||
|
| "local_shell_call"
|
||||||
|
| "web_search_call"
|
||||||
|
| "tool_search_call"
|
||||||
|
| "image_generation_call"
|
||||||
|
);
|
||||||
|
let bookkeeping = matches!(kind, "compaction" | "context_compaction");
|
||||||
|
if !consumes_reasoning && !bookkeeping && !pending_reasoning.is_empty() {
|
||||||
|
messages.push(reasoning_message(std::mem::take(&mut pending_reasoning)));
|
||||||
|
}
|
||||||
|
match kind {
|
||||||
|
"message" => {
|
||||||
|
let mut message = json!({
|
||||||
|
"role": role,
|
||||||
|
"content": responses_content_text(item.get("content").unwrap_or(&Value::Null))?
|
||||||
|
});
|
||||||
|
if role == "assistant" && !pending_reasoning.is_empty() {
|
||||||
|
message["reasoning_content"] =
|
||||||
|
Value::String(std::mem::take(&mut pending_reasoning));
|
||||||
|
}
|
||||||
|
messages.push(message);
|
||||||
|
}
|
||||||
|
"function_call" | "custom_tool_call" => {
|
||||||
|
let name = if kind == "function_call" {
|
||||||
|
format!(
|
||||||
|
"{}{}",
|
||||||
|
item.get("namespace").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
item.get("name").and_then(Value::as_str).unwrap_or("")
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
item.get("name")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_owned()
|
||||||
|
};
|
||||||
|
push_responses_call(
|
||||||
|
&mut messages,
|
||||||
|
&mut pending_reasoning,
|
||||||
|
item,
|
||||||
|
&name,
|
||||||
|
item.get("arguments").or_else(|| item.get("input")),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
"function_call_output" | "custom_tool_call_output" => messages.push(json!({
|
||||||
|
"role": "tool",
|
||||||
|
"content": responses_output_text(item.get("output"))?,
|
||||||
|
"tool_call_id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or("")
|
||||||
|
})),
|
||||||
|
"reasoning" => {
|
||||||
|
for value in [item.get("summary"), item.get("content")]
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
{
|
||||||
|
let text = responses_content_text(value)?;
|
||||||
|
if !pending_reasoning.is_empty() && !text.is_empty() {
|
||||||
|
pending_reasoning.push('\n');
|
||||||
|
}
|
||||||
|
pending_reasoning.push_str(&text);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"local_shell_call"
|
||||||
|
| "web_search_call"
|
||||||
|
| "tool_search_call"
|
||||||
|
| "image_generation_call" => {
|
||||||
|
let name = match kind {
|
||||||
|
"local_shell_call" => "local_shell",
|
||||||
|
"tool_search_call" => "tool_search",
|
||||||
|
other => other,
|
||||||
|
};
|
||||||
|
push_responses_call(
|
||||||
|
&mut messages,
|
||||||
|
&mut pending_reasoning,
|
||||||
|
item,
|
||||||
|
name,
|
||||||
|
item.get("action")
|
||||||
|
.or_else(|| item.get("arguments"))
|
||||||
|
.or_else(|| item.get("input")),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
"local_shell_call_output"
|
||||||
|
| "web_search_call_output"
|
||||||
|
| "tool_search_output"
|
||||||
|
| "tool_search_call_output"
|
||||||
|
| "image_generation_call_output" => {
|
||||||
|
let content = if let Some(value) =
|
||||||
|
item.get("output").or_else(|| item.get("result"))
|
||||||
|
{
|
||||||
|
responses_output_text(Some(value))?
|
||||||
|
} else {
|
||||||
|
item.get("tools").map(Value::to_string).unwrap_or_default()
|
||||||
|
};
|
||||||
|
messages.push(json!({
|
||||||
|
"role": "tool",
|
||||||
|
"content": content,
|
||||||
|
"tool_call_id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or("")
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
"compaction" | "context_compaction" => {}
|
||||||
|
_ => return Err((400, "invalid JSON request".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !pending_reasoning.is_empty() {
|
||||||
|
messages.push(reasoning_message(pending_reasoning));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => return Err((400, "invalid JSON request".into())),
|
||||||
|
}
|
||||||
|
if let Some(instructions) = object.get("instructions") {
|
||||||
|
let instructions = match instructions {
|
||||||
|
Value::Null => String::new(),
|
||||||
|
Value::String(text) => text.clone(),
|
||||||
|
_ => return Err((400, "invalid JSON request".into())),
|
||||||
|
};
|
||||||
|
if !instructions.is_empty() {
|
||||||
|
messages.insert(0, json!({"role": "system", "content": instructions}));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut request = Map::new();
|
||||||
|
copy_request_fields(object, &mut request);
|
||||||
|
request.insert("messages".into(), Value::Array(messages));
|
||||||
|
if let Some(tokens) = object.get("max_output_tokens") {
|
||||||
|
request.insert("max_tokens".into(), tokens.clone());
|
||||||
|
}
|
||||||
|
if let Some(effort) = object
|
||||||
|
.get("reasoning")
|
||||||
|
.and_then(|value| value.get("effort"))
|
||||||
|
{
|
||||||
|
request.insert("reasoning_effort".into(), effort.clone());
|
||||||
|
}
|
||||||
|
request.insert(
|
||||||
|
"tools".into(),
|
||||||
|
normalize_responses_tools(object.get("tools")),
|
||||||
|
);
|
||||||
|
decode_chat_request(Value::Object(request))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reasoning_message(reasoning: String) -> Value {
|
||||||
|
json!({"role": "assistant", "content": "", "reasoning_content": reasoning})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn push_responses_call(
|
||||||
|
messages: &mut Vec<Value>,
|
||||||
|
pending_reasoning: &mut String,
|
||||||
|
item: &Map<String, Value>,
|
||||||
|
name: &str,
|
||||||
|
arguments: Option<&Value>,
|
||||||
|
) {
|
||||||
|
let call = json!({
|
||||||
|
"id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"type": "function",
|
||||||
|
"function": {"name": name, "arguments": raw_json_text(arguments)}
|
||||||
|
});
|
||||||
|
if let Some(last) = messages.last_mut()
|
||||||
|
&& last.get("role").and_then(Value::as_str) == Some("assistant")
|
||||||
|
{
|
||||||
|
if !pending_reasoning.is_empty()
|
||||||
|
&& last
|
||||||
|
.get("reasoning_content")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.is_none_or(str::is_empty)
|
||||||
|
{
|
||||||
|
last["reasoning_content"] = Value::String(std::mem::take(pending_reasoning));
|
||||||
|
}
|
||||||
|
if !last.get("tool_calls").is_some_and(Value::is_array) {
|
||||||
|
last["tool_calls"] = json!([]);
|
||||||
|
}
|
||||||
|
last["tool_calls"].as_array_mut().unwrap().push(call);
|
||||||
|
} else {
|
||||||
|
messages.push(json!({
|
||||||
|
"role": "assistant",
|
||||||
|
"content": "",
|
||||||
|
"reasoning_content": std::mem::take(pending_reasoning),
|
||||||
|
"tool_calls": [call]
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn raw_json_text(value: Option<&Value>) -> String {
|
||||||
|
match value {
|
||||||
|
Some(Value::String(text)) => text.clone(),
|
||||||
|
Some(value) => value.to_string(),
|
||||||
|
None => "{}".into(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn responses_output_text(value: Option<&Value>) -> Result<String, (u16, String)> {
|
||||||
|
match value {
|
||||||
|
Some(Value::Array(_)) => responses_content_text(value.unwrap()),
|
||||||
|
Some(Value::String(text)) => Ok(text.clone()),
|
||||||
|
Some(value) => Ok(value.to_string()),
|
||||||
|
None => Ok(String::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn responses_content_text(value: &Value) -> Result<String, (u16, String)> {
|
||||||
|
match value {
|
||||||
|
Value::String(text) => Ok(text.clone()),
|
||||||
|
Value::Null => Ok(String::new()),
|
||||||
|
Value::Array(parts) => {
|
||||||
|
let mut text = String::new();
|
||||||
|
for part in parts {
|
||||||
|
match part {
|
||||||
|
Value::String(part) => text.push_str(part),
|
||||||
|
Value::Object(part)
|
||||||
|
if matches!(
|
||||||
|
part.get("type").and_then(Value::as_str),
|
||||||
|
Some(
|
||||||
|
"input_text"
|
||||||
|
| "output_text"
|
||||||
|
| "text"
|
||||||
|
| "summary_text"
|
||||||
|
| "reasoning_text"
|
||||||
|
)
|
||||||
|
) && matches!(
|
||||||
|
part.get("text"),
|
||||||
|
Some(Value::String(_) | Value::Null)
|
||||||
|
) =>
|
||||||
|
{
|
||||||
|
if let Some(value) = part.get("text").and_then(Value::as_str) {
|
||||||
|
text.push_str(value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => return Err((400, "invalid JSON request".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
|
_ => Err((400, "invalid JSON request".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn copy_request_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
|
||||||
|
for key in [
|
||||||
|
"model",
|
||||||
|
"max_tokens",
|
||||||
|
"max_completion_tokens",
|
||||||
|
"temperature",
|
||||||
|
"top_p",
|
||||||
|
"min_p",
|
||||||
|
"top_k",
|
||||||
|
"seed",
|
||||||
|
"stream",
|
||||||
|
"stream_options",
|
||||||
|
"thinking",
|
||||||
|
"think",
|
||||||
|
"reasoning_effort",
|
||||||
|
"tool_choice",
|
||||||
|
"stop",
|
||||||
|
] {
|
||||||
|
if let Some(value) = source.get(key) {
|
||||||
|
target.insert(key.into(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_anthropic_tools(tools: Option<&Value>) -> Value {
|
||||||
|
Value::Array(
|
||||||
|
tools
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
.map(|tool| {
|
||||||
|
json!({"type": "function", "function": {
|
||||||
|
"name": tool.get("name").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"description": tool.get("description").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"parameters": tool.get("input_schema").cloned().unwrap_or_else(|| json!({"type": "object"}))
|
||||||
|
}})
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_responses_tools(tools: Option<&Value>) -> Value {
|
||||||
|
Value::Array(
|
||||||
|
tools
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
.filter(|tool| tool.get("type").and_then(Value::as_str) == Some("function"))
|
||||||
|
.map(|tool| {
|
||||||
|
json!({"type": "function", "function": {
|
||||||
|
"name": tool.get("name").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"description": tool.get("description").and_then(Value::as_str).unwrap_or(""),
|
||||||
|
"parameters": tool.get("parameters").cloned().unwrap_or_else(|| json!({"type": "object"}))
|
||||||
|
}})
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn decode_chat_request(value: Value) -> Result<ChatRequest, (u16, String)> {
|
||||||
|
serde_json::from_value(value).map_err(|_| (400, "invalid JSON request".to_owned()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn raw_tool_schemas(
|
||||||
|
body: &[u8],
|
||||||
|
protocol: Protocol,
|
||||||
|
) -> Result<Vec<String>, (u16, String)> {
|
||||||
|
let request: RawToolsRequest<'_> =
|
||||||
|
serde_json::from_slice(body).map_err(|_| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
request
|
||||||
|
.tools
|
||||||
|
.into_iter()
|
||||||
|
.map(|tool| {
|
||||||
|
if protocol == Protocol::Chat {
|
||||||
|
let wrapper: RawOpenAiTool<'_> = serde_json::from_str(tool.get())
|
||||||
|
.map_err(|_| (400, "invalid JSON request".to_owned()))?;
|
||||||
|
Ok(wrapper.function.unwrap_or(tool).get().to_owned())
|
||||||
|
} else {
|
||||||
|
Ok(tool.get().to_owned())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn parse_chat_request(
|
||||||
|
state: &State,
|
||||||
|
request: ChatRequest,
|
||||||
|
protocol: Protocol,
|
||||||
|
) -> Result<ParsedRequest, (u16, String)> {
|
||||||
|
if request.messages.is_empty() {
|
||||||
|
return Err((400, "missing messages".into()));
|
||||||
|
}
|
||||||
|
let requested_id = request.model.as_deref().unwrap_or_default();
|
||||||
|
let preferences = state
|
||||||
|
.preferences
|
||||||
|
.read()
|
||||||
|
.map_err(|_| (500, "preferences are unavailable".to_owned()))?
|
||||||
|
.clone();
|
||||||
|
let model = if requested_id.is_empty() {
|
||||||
|
ModelChoice::from_id(&preferences.selected_model)
|
||||||
|
} else {
|
||||||
|
model_alias(requested_id)
|
||||||
|
}
|
||||||
|
.ok_or_else(|| (400, format!("unknown model: {requested_id}")))?;
|
||||||
|
if !installed_endpoint_models(&state.models_path).contains(&model) {
|
||||||
|
return Err((400, format!("model is not installed and verified: {model}")));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut generation = preferences.generation().map_err(|error| (500, error))?;
|
||||||
|
generation.system_prompt.clear();
|
||||||
|
generation.max_generated_tokens = request
|
||||||
|
.max_completion_tokens
|
||||||
|
.or(request.max_tokens)
|
||||||
|
.unwrap_or(generation.max_generated_tokens);
|
||||||
|
if generation.max_generated_tokens <= 0 {
|
||||||
|
return Err((400, "max_tokens must be positive".into()));
|
||||||
|
}
|
||||||
|
generation.temperature = request.temperature.or(Some(1.0));
|
||||||
|
generation.top_p = request.top_p.or(Some(1.0));
|
||||||
|
generation.min_p = request.min_p.or(Some(0.05));
|
||||||
|
generation.seed = request.seed.filter(|seed| *seed > 0);
|
||||||
|
generation.reasoning_mode = request_reasoning(&request, requested_id)?;
|
||||||
|
|
||||||
|
let tools_enabled = (!request.tools.is_empty() || !request.tool_schemas.is_empty())
|
||||||
|
&& request.tool_choice.as_ref().and_then(Value::as_str) != Some("none");
|
||||||
|
let (system, messages) = render_messages(
|
||||||
|
state,
|
||||||
|
&request.messages,
|
||||||
|
&request.tools,
|
||||||
|
&request.tool_schemas,
|
||||||
|
tools_enabled,
|
||||||
|
protocol,
|
||||||
|
)?;
|
||||||
|
generation.system_prompt = system;
|
||||||
|
let runtime = preferences.runtime().map_err(|error| (500, error))?;
|
||||||
|
let mut effective = effective_settings(model, &generation, &runtime, &state.models_path)
|
||||||
|
.map_err(|error| (400, error))?;
|
||||||
|
effective.turn.top_k = request.top_k.unwrap_or(0);
|
||||||
|
if effective.turn.top_k < 0 {
|
||||||
|
return Err((400, "top_k must not be negative".into()));
|
||||||
|
}
|
||||||
|
effective.turn.stops = match request.stop {
|
||||||
|
Some(OneOrMany::One(stop)) => vec![stop],
|
||||||
|
Some(OneOrMany::Many(stops)) => stops,
|
||||||
|
None => Vec::new(),
|
||||||
|
};
|
||||||
|
effective.turn.stops.retain(|stop| !stop.is_empty());
|
||||||
|
|
||||||
|
Ok(ParsedRequest {
|
||||||
|
protocol,
|
||||||
|
model_id: if requested_id.is_empty() {
|
||||||
|
model.id().to_owned()
|
||||||
|
} else {
|
||||||
|
requested_id.to_owned()
|
||||||
|
},
|
||||||
|
messages,
|
||||||
|
turn: effective.turn,
|
||||||
|
engine: effective.engine,
|
||||||
|
idle_timeout: Duration::from_secs(preferences.idle_timeout_minutes.max(1) as u64 * 60),
|
||||||
|
stream: request.stream,
|
||||||
|
include_usage: request
|
||||||
|
.stream_options
|
||||||
|
.is_some_and(|options| options.include_usage),
|
||||||
|
has_tools: tools_enabled,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn request_reasoning(
|
||||||
|
request: &ChatRequest,
|
||||||
|
model_id: &str,
|
||||||
|
) -> Result<ReasoningMode, (u16, String)> {
|
||||||
|
let explicit_thinking = request
|
||||||
|
.think
|
||||||
|
.or_else(|| request.thinking.as_ref().and_then(thinking_enabled));
|
||||||
|
let mut reasoning = match request.reasoning_effort.as_deref() {
|
||||||
|
Some("max") => ReasoningMode::Max,
|
||||||
|
Some("none") => ReasoningMode::Direct,
|
||||||
|
Some("xhigh" | "high" | "medium" | "low" | "minimal") | None => ReasoningMode::High,
|
||||||
|
Some(value) => return Err((400, format!("unsupported reasoning_effort: {value}"))),
|
||||||
|
};
|
||||||
|
if explicit_thinking == Some(false)
|
||||||
|
|| (explicit_thinking.is_none()
|
||||||
|
&& matches!(
|
||||||
|
model_id,
|
||||||
|
"deepseek-chat" | "glm-5.2-chat" | "glm-5.2-no-think" | "glm-5.2-nothink"
|
||||||
|
))
|
||||||
|
{
|
||||||
|
reasoning = ReasoningMode::Direct;
|
||||||
|
}
|
||||||
|
Ok(reasoning)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn thinking_enabled(value: &Value) -> Option<bool> {
|
||||||
|
value.as_bool().or_else(|| {
|
||||||
|
value.as_str().map(|value| value != "disabled").or_else(|| {
|
||||||
|
value
|
||||||
|
.get("type")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(|value| value != "disabled")
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
974
src/server/response.rs
Normal file
974
src/server/response.rs
Normal file
@@ -0,0 +1,974 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
pub(super) fn final_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
let output = wait_for_output(active)?;
|
||||||
|
let (content, calls) = parse_generated_tools(state, &output.message.content, request.protocol);
|
||||||
|
let finish = if calls.is_empty() {
|
||||||
|
output.finish_reason
|
||||||
|
} else {
|
||||||
|
"tool_calls"
|
||||||
|
};
|
||||||
|
let reasoning = output.message.reasoning.filter(|value| !value.is_empty());
|
||||||
|
let usage = usage_json(
|
||||||
|
output.prompt_tokens,
|
||||||
|
output.cached_tokens,
|
||||||
|
output.completion_tokens,
|
||||||
|
);
|
||||||
|
let body = match request.protocol {
|
||||||
|
Protocol::Chat => {
|
||||||
|
let mut message = json!({"role": "assistant", "content": content});
|
||||||
|
if let Some(reasoning) = reasoning {
|
||||||
|
message["reasoning_content"] = Value::String(reasoning);
|
||||||
|
}
|
||||||
|
if !calls.is_empty() {
|
||||||
|
message["tool_calls"] = tool_calls_json(&calls);
|
||||||
|
}
|
||||||
|
json!({
|
||||||
|
"id": id, "object": "chat.completion", "created": unix_time(),
|
||||||
|
"model": request.model_id,
|
||||||
|
"choices": [{"index": 0, "message": message, "finish_reason": finish}],
|
||||||
|
"usage": usage,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
Protocol::Completion => json!({
|
||||||
|
"id": id, "object": "text_completion", "created": unix_time(),
|
||||||
|
"model": request.model_id,
|
||||||
|
"choices": [{"text": content, "index": 0, "finish_reason": finish}],
|
||||||
|
"usage": usage,
|
||||||
|
}),
|
||||||
|
Protocol::Anthropic => {
|
||||||
|
let mut blocks = Vec::new();
|
||||||
|
if let Some(reasoning) = reasoning {
|
||||||
|
blocks.push(json!({"type": "thinking", "thinking": reasoning, "signature": id}));
|
||||||
|
}
|
||||||
|
if !content.is_empty() {
|
||||||
|
blocks.push(json!({"type": "text", "text": content}));
|
||||||
|
}
|
||||||
|
for call in &calls {
|
||||||
|
blocks.push(json!({
|
||||||
|
"type": "tool_use", "id": call.id, "name": call.function.name,
|
||||||
|
"input": serde_json::from_str::<Value>(&call.function.arguments).unwrap_or_else(|_| json!({}))
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
if blocks.is_empty() || (blocks.iter().all(|block| block["type"] == "thinking")) {
|
||||||
|
blocks.push(json!({"type": "text", "text": ""}));
|
||||||
|
}
|
||||||
|
let cached = output.cached_tokens.min(output.prompt_tokens);
|
||||||
|
let written = output.prompt_tokens - cached;
|
||||||
|
json!({
|
||||||
|
"id": id, "type": "message", "role": "assistant", "model": request.model_id,
|
||||||
|
"content": blocks,
|
||||||
|
"stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" },
|
||||||
|
"stop_sequence": Value::Null,
|
||||||
|
"usage": {
|
||||||
|
"input_tokens": output.prompt_tokens - cached - written,
|
||||||
|
"output_tokens": output.completion_tokens,
|
||||||
|
"cache_read_input_tokens": cached,
|
||||||
|
"cache_creation_input_tokens": written
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
Protocol::Responses => {
|
||||||
|
let status = if finish == "length" {
|
||||||
|
"incomplete"
|
||||||
|
} else if finish == "error" {
|
||||||
|
"failed"
|
||||||
|
} else {
|
||||||
|
"completed"
|
||||||
|
};
|
||||||
|
let mut items = Vec::new();
|
||||||
|
if let Some(reasoning) = reasoning {
|
||||||
|
items.push(json!({
|
||||||
|
"id": random_id("rs_"), "type": "reasoning", "status": status,
|
||||||
|
"summary": [{"type": "summary_text", "text": reasoning}]
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
if !content.is_empty() {
|
||||||
|
items.push(json!({
|
||||||
|
"id": random_id("msg_"), "type": "message", "status": status,
|
||||||
|
"role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}]
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
for call in &calls {
|
||||||
|
items.push(json!({
|
||||||
|
"id": random_id("fc_"), "type": "function_call", "status": status,
|
||||||
|
"name": call.function.name, "call_id": call.id, "arguments": call.function.arguments
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
json!({
|
||||||
|
"id": id, "object": "response", "created_at": unix_time(), "status": status,
|
||||||
|
"model": request.model_id, "output": items,
|
||||||
|
"usage": {
|
||||||
|
"input_tokens": output.prompt_tokens,
|
||||||
|
"input_tokens_details": {"cached_tokens": output.cached_tokens.min(output.prompt_tokens), "cache_write_tokens": output.prompt_tokens - output.cached_tokens.min(output.prompt_tokens)},
|
||||||
|
"output_tokens": output.completion_tokens,
|
||||||
|
"output_tokens_details": {"reasoning_tokens": 0},
|
||||||
|
"total_tokens": output.prompt_tokens + output.completion_tokens
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
};
|
||||||
|
send_json(stream, 200, &body).map_err(|error| (500, error))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn stream_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
if request.protocol != Protocol::Chat {
|
||||||
|
return structured_stream_response(stream, state, request, active, id);
|
||||||
|
}
|
||||||
|
stream_response_with_keepalive(
|
||||||
|
stream,
|
||||||
|
state,
|
||||||
|
request,
|
||||||
|
active,
|
||||||
|
id,
|
||||||
|
PREFILL_KEEPALIVE_INTERVAL,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn structured_stream_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
match request.protocol {
|
||||||
|
Protocol::Completion => completion_stream_response(stream, request, active, id),
|
||||||
|
Protocol::Anthropic => anthropic_stream_response(stream, state, request, active, id),
|
||||||
|
Protocol::Responses => responses_stream_response(stream, state, request, active, id),
|
||||||
|
Protocol::Chat => unreachable!(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn completion_stream_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
send_sse_headers(stream).map_err(|error| (500, error))?;
|
||||||
|
while let Ok(event) = active.events.recv() {
|
||||||
|
match event {
|
||||||
|
GenerationEvent::Chunk {
|
||||||
|
reasoning: false,
|
||||||
|
content,
|
||||||
|
} if !content.is_empty() => send_sse(
|
||||||
|
stream,
|
||||||
|
&json!({
|
||||||
|
"id": id, "object": "text_completion", "created": unix_time(),
|
||||||
|
"model": request.model_id,
|
||||||
|
"choices": [{"text": content, "index": 0, "finish_reason": Value::Null}]
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?,
|
||||||
|
GenerationEvent::Finished(Ok(output)) => {
|
||||||
|
send_sse(
|
||||||
|
stream,
|
||||||
|
&json!({
|
||||||
|
"id": id, "object": "text_completion", "created": unix_time(),
|
||||||
|
"model": request.model_id,
|
||||||
|
"choices": [{"text": "", "index": 0, "finish_reason": output.finish_reason}]
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?;
|
||||||
|
if request.include_usage {
|
||||||
|
send_sse(
|
||||||
|
stream,
|
||||||
|
&json!({
|
||||||
|
"id": id, "object": "text_completion", "created": unix_time(),
|
||||||
|
"model": request.model_id, "choices": [],
|
||||||
|
"usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens)
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?;
|
||||||
|
}
|
||||||
|
return stream
|
||||||
|
.write_all(b"data: [DONE]\n\n")
|
||||||
|
.map_err(|error| (500, error.to_string()));
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(Err(error)) => {
|
||||||
|
let _ = send_sse_error(stream, &error);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn anthropic_stream_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
send_sse_headers(stream).map_err(|error| (500, error))?;
|
||||||
|
let mut prompt_tokens = 0;
|
||||||
|
let mut started = false;
|
||||||
|
let mut block = None::<(usize, bool)>;
|
||||||
|
let mut next_index = 0;
|
||||||
|
let mut projector = ToolProjector::new();
|
||||||
|
let mut tool_indices = Vec::new();
|
||||||
|
while let Ok(event) = active.events.recv() {
|
||||||
|
match event {
|
||||||
|
GenerationEvent::Context {
|
||||||
|
used,
|
||||||
|
tokens_per_second,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
if tokens_per_second.is_none() {
|
||||||
|
prompt_tokens = used;
|
||||||
|
} else if !started {
|
||||||
|
anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?;
|
||||||
|
started = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GenerationEvent::Chunk { reasoning, content } => {
|
||||||
|
if !started {
|
||||||
|
anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?;
|
||||||
|
started = true;
|
||||||
|
}
|
||||||
|
if !reasoning && request.has_tools {
|
||||||
|
let events = projector.push(&content, false, "toolu_");
|
||||||
|
send_anthropic_projection_events(
|
||||||
|
stream,
|
||||||
|
events,
|
||||||
|
&mut block,
|
||||||
|
&mut next_index,
|
||||||
|
&mut tool_indices,
|
||||||
|
)?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if block.is_some_and(|(_, current_reasoning)| current_reasoning != reasoning) {
|
||||||
|
let (index, _) = block.take().unwrap();
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": index}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let index = if let Some((index, _)) = block {
|
||||||
|
index
|
||||||
|
} else {
|
||||||
|
let index = next_index;
|
||||||
|
next_index += 1;
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_start",
|
||||||
|
&json!({
|
||||||
|
"type": "content_block_start", "index": index,
|
||||||
|
"content_block": if reasoning { json!({"type": "thinking", "thinking": "", "signature": ""}) } else { json!({"type": "text", "text": ""}) }
|
||||||
|
}),
|
||||||
|
)?;
|
||||||
|
block = Some((index, reasoning));
|
||||||
|
index
|
||||||
|
};
|
||||||
|
let delta = if reasoning {
|
||||||
|
json!({"type": "thinking_delta", "thinking": content})
|
||||||
|
} else {
|
||||||
|
json!({"type": "text_delta", "text": content})
|
||||||
|
};
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_delta",
|
||||||
|
&json!({"type": "content_block_delta", "index": index, "delta": delta}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(Ok(output)) => {
|
||||||
|
if !started {
|
||||||
|
anthropic_stream_start(
|
||||||
|
stream,
|
||||||
|
&request,
|
||||||
|
id,
|
||||||
|
output.prompt_tokens,
|
||||||
|
output.cached_tokens,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if request.has_tools {
|
||||||
|
let events = projector.push("", true, "toolu_");
|
||||||
|
send_anthropic_projection_events(
|
||||||
|
stream,
|
||||||
|
events,
|
||||||
|
&mut block,
|
||||||
|
&mut next_index,
|
||||||
|
&mut tool_indices,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
if let Some((index, _)) = block.take() {
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": index}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let (_, calls) = parse_generated_tools_with_ids(
|
||||||
|
state,
|
||||||
|
&output.message.content,
|
||||||
|
Protocol::Anthropic,
|
||||||
|
&projector.ids,
|
||||||
|
);
|
||||||
|
if projector.ids.is_empty() {
|
||||||
|
for call in &calls {
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_start",
|
||||||
|
&json!({"type": "content_block_start", "index": next_index, "content_block": {"type": "tool_use", "id": call.id, "name": call.function.name, "input": {}}}),
|
||||||
|
)?;
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_delta",
|
||||||
|
&json!({"type": "content_block_delta", "index": next_index, "delta": {"type": "input_json_delta", "partial_json": call.function.arguments}}),
|
||||||
|
)?;
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": next_index}),
|
||||||
|
)?;
|
||||||
|
next_index += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let finish = if calls.is_empty() {
|
||||||
|
output.finish_reason
|
||||||
|
} else {
|
||||||
|
"tool_calls"
|
||||||
|
};
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"message_delta",
|
||||||
|
&json!({"type": "message_delta", "delta": {"stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" }, "stop_sequence": Value::Null}, "usage": {"output_tokens": output.completion_tokens}}),
|
||||||
|
)?;
|
||||||
|
return send_named_sse(stream, "message_stop", &json!({"type": "message_stop"}));
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(Err(error)) => {
|
||||||
|
return send_named_sse(
|
||||||
|
stream,
|
||||||
|
"error",
|
||||||
|
&json!({"type": "error", "error": {"type": "api_error", "message": error}}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn send_anthropic_projection_events(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
events: Vec<ToolProjectionEvent>,
|
||||||
|
block: &mut Option<(usize, bool)>,
|
||||||
|
next_index: &mut usize,
|
||||||
|
tool_indices: &mut Vec<usize>,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
for event in events {
|
||||||
|
match event {
|
||||||
|
ToolProjectionEvent::Text(text) if !text.is_empty() => {
|
||||||
|
if block.is_some_and(|(_, reasoning)| reasoning) {
|
||||||
|
let (index, _) = block.take().unwrap();
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": index}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let index = if let Some((index, _)) = *block {
|
||||||
|
index
|
||||||
|
} else {
|
||||||
|
let index = *next_index;
|
||||||
|
*next_index += 1;
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_start",
|
||||||
|
&json!({"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}),
|
||||||
|
)?;
|
||||||
|
*block = Some((index, false));
|
||||||
|
index
|
||||||
|
};
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_delta",
|
||||||
|
&json!({"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
ToolProjectionEvent::Start { index, id, name } => {
|
||||||
|
if let Some((open_index, _)) = block.take() {
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": open_index}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
let content_index = *next_index;
|
||||||
|
*next_index += 1;
|
||||||
|
if tool_indices.len() == index {
|
||||||
|
tool_indices.push(content_index);
|
||||||
|
}
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_start",
|
||||||
|
&json!({"type": "content_block_start", "index": content_index, "content_block": {"type": "tool_use", "id": id, "name": name, "input": {}}}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
ToolProjectionEvent::Arguments { index, fragment } => {
|
||||||
|
if let Some(content_index) = tool_indices.get(index) {
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_delta",
|
||||||
|
&json!({"type": "content_block_delta", "index": content_index, "delta": {"type": "input_json_delta", "partial_json": fragment}}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ToolProjectionEvent::End { index } => {
|
||||||
|
if let Some(content_index) = tool_indices.get(index) {
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"content_block_stop",
|
||||||
|
&json!({"type": "content_block_stop", "index": content_index}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ToolProjectionEvent::Text(_) => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn anthropic_stream_start(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
request: &ResponseOptions,
|
||||||
|
id: &str,
|
||||||
|
prompt_tokens: u32,
|
||||||
|
cached_tokens: u32,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
let cached = cached_tokens.min(prompt_tokens);
|
||||||
|
let written = prompt_tokens - cached;
|
||||||
|
send_named_sse(
|
||||||
|
stream,
|
||||||
|
"message_start",
|
||||||
|
&json!({"type": "message_start", "message": {"id": id, "type": "message", "role": "assistant", "model": request.model_id, "content": [], "stop_reason": Value::Null, "stop_sequence": Value::Null, "usage": {"input_tokens": prompt_tokens - cached - written, "output_tokens": 0, "cache_read_input_tokens": cached, "cache_creation_input_tokens": written}}}),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn responses_stream_response(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
send_sse_headers(stream).map_err(|error| (500, error))?;
|
||||||
|
let created = unix_time();
|
||||||
|
let message_id = random_id("msg_");
|
||||||
|
let reasoning_id = random_id("rs_");
|
||||||
|
let mut sequence = 0;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.created", "response": {"id": id, "object": "response", "created_at": created, "status": "in_progress", "model": request.model_id, "output": []}}),
|
||||||
|
)?;
|
||||||
|
let mut reasoning_open = false;
|
||||||
|
let mut message_open = false;
|
||||||
|
let mut reasoning = String::new();
|
||||||
|
let mut content = String::new();
|
||||||
|
while let Ok(event) = active.events.recv() {
|
||||||
|
match event {
|
||||||
|
GenerationEvent::Chunk {
|
||||||
|
reasoning: true,
|
||||||
|
content: chunk,
|
||||||
|
} => {
|
||||||
|
if !request.reasoning_summary {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !reasoning_open {
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.added", "output_index": 0, "item": {"id": reasoning_id, "type": "reasoning", "status": "in_progress", "summary": []}}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.reasoning_summary_part.added", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "part": {"type": "summary_text", "text": ""}}),
|
||||||
|
)?;
|
||||||
|
reasoning_open = true;
|
||||||
|
}
|
||||||
|
reasoning.push_str(&chunk);
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.reasoning_summary_text.delta", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "delta": chunk}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
GenerationEvent::Chunk {
|
||||||
|
reasoning: false,
|
||||||
|
content: chunk,
|
||||||
|
} => {
|
||||||
|
if request.has_tools {
|
||||||
|
content.push_str(&chunk);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !message_open {
|
||||||
|
let output_index = usize::from(reasoning_open);
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}),
|
||||||
|
)?;
|
||||||
|
message_open = true;
|
||||||
|
}
|
||||||
|
content.push_str(&chunk);
|
||||||
|
let output_index = usize::from(reasoning_open);
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": chunk}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(Ok(output)) => {
|
||||||
|
let (parsed_content, calls) =
|
||||||
|
parse_generated_tools(state, &output.message.content, Protocol::Responses);
|
||||||
|
if request.has_tools {
|
||||||
|
content = parsed_content;
|
||||||
|
}
|
||||||
|
let finish = if calls.is_empty() {
|
||||||
|
output.finish_reason
|
||||||
|
} else {
|
||||||
|
"tool_calls"
|
||||||
|
};
|
||||||
|
let status = if finish == "length" {
|
||||||
|
"incomplete"
|
||||||
|
} else if finish == "error" {
|
||||||
|
"failed"
|
||||||
|
} else {
|
||||||
|
"completed"
|
||||||
|
};
|
||||||
|
let mut terminal_items = Vec::new();
|
||||||
|
let mut output_index = 0;
|
||||||
|
if reasoning_open {
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.reasoning_summary_text.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "text": reasoning}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.reasoning_summary_part.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "part": {"type": "summary_text", "text": reasoning}}),
|
||||||
|
)?;
|
||||||
|
let item = json!({"id": reasoning_id, "type": "reasoning", "status": status, "summary": [{"type": "summary_text", "text": reasoning}]});
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
||||||
|
)?;
|
||||||
|
terminal_items.push(item);
|
||||||
|
output_index += 1;
|
||||||
|
}
|
||||||
|
if !content.is_empty() {
|
||||||
|
if !message_open {
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": content}),
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_text.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "text": content}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.content_part.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": content, "annotations": []}}),
|
||||||
|
)?;
|
||||||
|
let item = json!({"id": message_id, "type": "message", "status": status, "role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}]});
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
||||||
|
)?;
|
||||||
|
terminal_items.push(item);
|
||||||
|
output_index += 1;
|
||||||
|
}
|
||||||
|
for call in &calls {
|
||||||
|
let item_id = random_id("fc_");
|
||||||
|
let mut item = json!({"id": item_id, "type": "function_call", "status": status, "name": call.function.name, "call_id": call.id, "arguments": call.function.arguments});
|
||||||
|
let mut added = item.clone();
|
||||||
|
added["status"] = Value::String("in_progress".into());
|
||||||
|
added["arguments"] = Value::String(String::new());
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.added", "output_index": output_index, "item": added}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.function_call_arguments.delta", "item_id": item_id, "output_index": output_index, "delta": call.function.arguments}),
|
||||||
|
)?;
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.function_call_arguments.done", "item_id": item_id, "output_index": output_index, "name": call.function.name, "arguments": call.function.arguments}),
|
||||||
|
)?;
|
||||||
|
item["id"] = Value::String(item_id);
|
||||||
|
send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": "response.output_item.done", "output_index": output_index, "item": item}),
|
||||||
|
)?;
|
||||||
|
terminal_items.push(item);
|
||||||
|
output_index += 1;
|
||||||
|
}
|
||||||
|
let event_type = if finish == "length" {
|
||||||
|
"response.incomplete"
|
||||||
|
} else if finish == "error" {
|
||||||
|
"response.failed"
|
||||||
|
} else {
|
||||||
|
"response.completed"
|
||||||
|
};
|
||||||
|
let cached = output.cached_tokens.min(output.prompt_tokens);
|
||||||
|
return send_responses_sse(
|
||||||
|
stream,
|
||||||
|
&mut sequence,
|
||||||
|
json!({"type": event_type, "response": {"id": id, "object": "response", "created_at": created, "status": status, "model": request.model_id, "output": terminal_items, "usage": {"input_tokens": output.prompt_tokens, "input_tokens_details": {"cached_tokens": cached, "cache_write_tokens": output.prompt_tokens - cached}, "output_tokens": output.completion_tokens, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": output.prompt_tokens + output.completion_tokens}}}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(Err(error)) => {
|
||||||
|
let _ = send_sse_error(stream, &error);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
fn send_named_sse(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
event: &str,
|
||||||
|
value: &Value,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
let body = serde_json::to_vec(value).map_err(|error| (500, error.to_string()))?;
|
||||||
|
stream
|
||||||
|
.write_all(b"event: ")
|
||||||
|
.and_then(|()| stream.write_all(event.as_bytes()))
|
||||||
|
.and_then(|()| stream.write_all(b"\ndata: "))
|
||||||
|
.and_then(|()| stream.write_all(&body))
|
||||||
|
.and_then(|()| stream.write_all(b"\n\n"))
|
||||||
|
.map_err(|error| (500, error.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn send_responses_sse(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
sequence: &mut u32,
|
||||||
|
value: Value,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
let mut object = value
|
||||||
|
.as_object()
|
||||||
|
.cloned()
|
||||||
|
.ok_or_else(|| (500, "Responses event is not an object".to_owned()))?;
|
||||||
|
let event_type = object.shift_remove("type").unwrap_or(Value::Null);
|
||||||
|
let mut ordered = Map::new();
|
||||||
|
ordered.insert("type".into(), event_type);
|
||||||
|
ordered.insert("sequence_number".into(), Value::from(*sequence));
|
||||||
|
ordered.extend(object);
|
||||||
|
*sequence += 1;
|
||||||
|
send_sse(stream, &Value::Object(ordered)).map_err(|error| (500, error))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stream_response_with_keepalive(
|
||||||
|
stream: &mut TcpStream,
|
||||||
|
state: &State,
|
||||||
|
request: ResponseOptions,
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
id: &str,
|
||||||
|
keepalive_interval: Duration,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
let mut projector = ToolProjector::new();
|
||||||
|
let mut output = None;
|
||||||
|
let mut prefilling = true;
|
||||||
|
let mut headers_sent = false;
|
||||||
|
let mut role_sent = false;
|
||||||
|
let mut last_keepalive = Instant::now();
|
||||||
|
loop {
|
||||||
|
let event = match receive_stream_event(
|
||||||
|
stream,
|
||||||
|
&active,
|
||||||
|
prefilling && headers_sent,
|
||||||
|
&mut last_keepalive,
|
||||||
|
keepalive_interval,
|
||||||
|
) {
|
||||||
|
Ok(Some(event)) => event,
|
||||||
|
Ok(None) => break,
|
||||||
|
Err(_) => return Ok(()),
|
||||||
|
};
|
||||||
|
match event {
|
||||||
|
GenerationEvent::Loading => {}
|
||||||
|
GenerationEvent::Context {
|
||||||
|
tokens_per_second, ..
|
||||||
|
} => {
|
||||||
|
prefilling = tokens_per_second.is_none();
|
||||||
|
if prefilling && !headers_sent {
|
||||||
|
send_sse_headers(stream).map_err(|error| (500, error))?;
|
||||||
|
headers_sent = true;
|
||||||
|
last_keepalive = Instant::now();
|
||||||
|
} else if !prefilling {
|
||||||
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GenerationEvent::Chunk { reasoning, content } => {
|
||||||
|
prefilling = false;
|
||||||
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
||||||
|
if request.has_tools && !reasoning {
|
||||||
|
let events = projector.push(&content, false, "call_");
|
||||||
|
if send_chat_projection_events(stream, &request, id, events).is_err() {
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
if TOOL_SYNTAXES
|
||||||
|
.iter()
|
||||||
|
.any(|syntax| projector.raw.contains(syntax.tool_end))
|
||||||
|
{
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
} else if !content.is_empty() {
|
||||||
|
let field = if reasoning {
|
||||||
|
"reasoning_content"
|
||||||
|
} else {
|
||||||
|
"content"
|
||||||
|
};
|
||||||
|
let chunk = chunk_json(id, &request.model_id, json!({field: content}), None);
|
||||||
|
if send_sse(stream, &chunk).is_err() {
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(result) => {
|
||||||
|
match result {
|
||||||
|
Ok(result) => {
|
||||||
|
send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?;
|
||||||
|
output = Some(result);
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
if headers_sent {
|
||||||
|
let _ = send_sse_error(stream, &error);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
return Err((500, error));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let output = match output {
|
||||||
|
Some(output) => output,
|
||||||
|
None if headers_sent => {
|
||||||
|
let _ = send_sse_error(stream, "The model runtime stopped unexpectedly.");
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
None => return Err((500, "The model runtime stopped unexpectedly.".into())),
|
||||||
|
};
|
||||||
|
if request.has_tools {
|
||||||
|
let events = projector.push("", true, "call_");
|
||||||
|
send_chat_projection_events(stream, &request, id, events)?;
|
||||||
|
}
|
||||||
|
let (_, calls) = parse_generated_tools_with_ids(
|
||||||
|
state,
|
||||||
|
&output.message.content,
|
||||||
|
Protocol::Chat,
|
||||||
|
&projector.ids,
|
||||||
|
);
|
||||||
|
if request.has_tools && projector.ids.is_empty() && !calls.is_empty() {
|
||||||
|
send_sse(
|
||||||
|
stream,
|
||||||
|
&chunk_json(
|
||||||
|
id,
|
||||||
|
&request.model_id,
|
||||||
|
json!({"tool_calls": tool_calls_json(&calls)}),
|
||||||
|
None,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?;
|
||||||
|
}
|
||||||
|
let finish = if calls.is_empty() {
|
||||||
|
output.finish_reason
|
||||||
|
} else {
|
||||||
|
"tool_calls"
|
||||||
|
};
|
||||||
|
send_sse(
|
||||||
|
stream,
|
||||||
|
&chunk_json(id, &request.model_id, json!({}), Some(finish)),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?;
|
||||||
|
if request.include_usage {
|
||||||
|
let usage = json!({
|
||||||
|
"id": id,
|
||||||
|
"object": "chat.completion.chunk",
|
||||||
|
"created": unix_time(),
|
||||||
|
"model": request.model_id,
|
||||||
|
"choices": [],
|
||||||
|
"usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens),
|
||||||
|
});
|
||||||
|
send_sse(stream, &usage).map_err(|error| (500, error))?;
|
||||||
|
}
|
||||||
|
stream
|
||||||
|
.write_all(b"data: [DONE]\n\n")
|
||||||
|
.map_err(|error| (500, error.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn send_chat_projection_events(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
request: &ResponseOptions,
|
||||||
|
response_id: &str,
|
||||||
|
events: Vec<ToolProjectionEvent>,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
for event in events {
|
||||||
|
let delta = match event {
|
||||||
|
ToolProjectionEvent::Text(content) if !content.is_empty() => {
|
||||||
|
json!({"content": content})
|
||||||
|
}
|
||||||
|
ToolProjectionEvent::Start { index, id, name } => json!({"tool_calls": [{
|
||||||
|
"index": index, "id": id, "type": "function",
|
||||||
|
"function": {"name": name, "arguments": ""}
|
||||||
|
}]}),
|
||||||
|
ToolProjectionEvent::Arguments { index, fragment } => json!({
|
||||||
|
"tool_calls": [{"index": index, "function": {"arguments": fragment}}]
|
||||||
|
}),
|
||||||
|
ToolProjectionEvent::End { .. } | ToolProjectionEvent::Text(_) => continue,
|
||||||
|
};
|
||||||
|
send_sse(
|
||||||
|
stream,
|
||||||
|
&chunk_json(response_id, &request.model_id, delta, None),
|
||||||
|
)
|
||||||
|
.map_err(|error| (500, error))?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn send_stream_start(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
request: &ResponseOptions,
|
||||||
|
id: &str,
|
||||||
|
headers_sent: &mut bool,
|
||||||
|
role_sent: &mut bool,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
if !*headers_sent {
|
||||||
|
send_sse_headers(stream).map_err(|error| (500, error))?;
|
||||||
|
*headers_sent = true;
|
||||||
|
}
|
||||||
|
if !*role_sent {
|
||||||
|
let role = chunk_json(id, &request.model_id, json!({"role": "assistant"}), None);
|
||||||
|
send_sse(stream, &role).map_err(|error| (500, error))?;
|
||||||
|
*role_sent = true;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn receive_stream_event(
|
||||||
|
stream: &mut impl Write,
|
||||||
|
active: &crate::runtime::ActiveGeneration,
|
||||||
|
prefilling: bool,
|
||||||
|
last_keepalive: &mut Instant,
|
||||||
|
keepalive_interval: Duration,
|
||||||
|
) -> Result<Option<GenerationEvent>, String> {
|
||||||
|
if !prefilling {
|
||||||
|
return Ok(active.events.recv().ok());
|
||||||
|
}
|
||||||
|
loop {
|
||||||
|
if last_keepalive.elapsed() >= keepalive_interval {
|
||||||
|
if let Err(error) = stream.write_all(b": prefill\n\n") {
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
return Err(error.to_string());
|
||||||
|
}
|
||||||
|
*last_keepalive = Instant::now();
|
||||||
|
}
|
||||||
|
let remaining = keepalive_interval.saturating_sub(last_keepalive.elapsed());
|
||||||
|
match active.events.recv_timeout(remaining) {
|
||||||
|
Ok(event) => return Ok(Some(event)),
|
||||||
|
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
|
||||||
|
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn wait_for_output(
|
||||||
|
active: crate::runtime::ActiveGeneration,
|
||||||
|
) -> Result<crate::engine::GenerationOutput, (u16, String)> {
|
||||||
|
let mut content = String::new();
|
||||||
|
while let Ok(event) = active.events.recv() {
|
||||||
|
match event {
|
||||||
|
GenerationEvent::Chunk {
|
||||||
|
reasoning: false,
|
||||||
|
content: chunk,
|
||||||
|
} => {
|
||||||
|
content.push_str(&chunk);
|
||||||
|
if TOOL_SYNTAXES
|
||||||
|
.iter()
|
||||||
|
.any(|syntax| content.contains(syntax.tool_end))
|
||||||
|
{
|
||||||
|
active.cancel.store(true, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GenerationEvent::Finished(result) => {
|
||||||
|
return result.map_err(|error| (500, error));
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err((500, "The model runtime stopped unexpectedly.".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn usage_json(prompt: u32, cached: u32, completion: u32) -> Value {
|
||||||
|
let cached = cached.min(prompt);
|
||||||
|
json!({
|
||||||
|
"prompt_tokens": prompt,
|
||||||
|
"completion_tokens": completion,
|
||||||
|
"total_tokens": prompt + completion,
|
||||||
|
"prompt_tokens_details": {
|
||||||
|
"cached_tokens": cached,
|
||||||
|
"cache_write_tokens": prompt - cached
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn chunk_json(id: &str, model: &str, delta: Value, finish: Option<&str>) -> Value {
|
||||||
|
json!({
|
||||||
|
"id": id,
|
||||||
|
"object": "chat.completion.chunk",
|
||||||
|
"created": unix_time(),
|
||||||
|
"model": model,
|
||||||
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}]
|
||||||
|
})
|
||||||
|
}
|
||||||
733
src/server/tools.rs
Normal file
733
src/server/tools.rs
Normal file
@@ -0,0 +1,733 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
enum ToolProjectionState {
|
||||||
|
Seeking,
|
||||||
|
Invokes,
|
||||||
|
Parameters,
|
||||||
|
Value,
|
||||||
|
Done,
|
||||||
|
Failed,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) enum ToolProjectionEvent {
|
||||||
|
Text(String),
|
||||||
|
Start {
|
||||||
|
index: usize,
|
||||||
|
id: String,
|
||||||
|
name: String,
|
||||||
|
},
|
||||||
|
Arguments {
|
||||||
|
index: usize,
|
||||||
|
fragment: String,
|
||||||
|
},
|
||||||
|
End {
|
||||||
|
index: usize,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) struct ToolProjector {
|
||||||
|
pub(super) raw: String,
|
||||||
|
position: usize,
|
||||||
|
text_emitted: usize,
|
||||||
|
state: ToolProjectionState,
|
||||||
|
index: usize,
|
||||||
|
pub(super) ids: Vec<String>,
|
||||||
|
first_parameter: bool,
|
||||||
|
string_parameter: bool,
|
||||||
|
syntax: Option<ToolSyntax>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolProjector {
|
||||||
|
pub(super) fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
raw: String::new(),
|
||||||
|
position: 0,
|
||||||
|
text_emitted: 0,
|
||||||
|
state: ToolProjectionState::Seeking,
|
||||||
|
index: 0,
|
||||||
|
ids: Vec::new(),
|
||||||
|
first_parameter: true,
|
||||||
|
string_parameter: false,
|
||||||
|
syntax: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn push(
|
||||||
|
&mut self,
|
||||||
|
chunk: &str,
|
||||||
|
final_chunk: bool,
|
||||||
|
prefix: &str,
|
||||||
|
) -> Vec<ToolProjectionEvent> {
|
||||||
|
self.raw.push_str(chunk);
|
||||||
|
let mut events = Vec::new();
|
||||||
|
loop {
|
||||||
|
match self.state {
|
||||||
|
ToolProjectionState::Seeking => {
|
||||||
|
if let Some((start, syntax)) = TOOL_SYNTAXES
|
||||||
|
.iter()
|
||||||
|
.filter_map(|syntax| {
|
||||||
|
self.raw
|
||||||
|
.find(syntax.tool_start)
|
||||||
|
.map(|start| (start, *syntax))
|
||||||
|
})
|
||||||
|
.min_by_key(|(start, _)| *start)
|
||||||
|
{
|
||||||
|
if start > self.text_emitted {
|
||||||
|
let text = self.raw[self.text_emitted..start].trim_end();
|
||||||
|
if !text.is_empty() {
|
||||||
|
events.push(ToolProjectionEvent::Text(text.to_owned()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.position = start + syntax.tool_start.len();
|
||||||
|
self.text_emitted = start;
|
||||||
|
self.syntax = Some(syntax);
|
||||||
|
self.state = ToolProjectionState::Invokes;
|
||||||
|
} else {
|
||||||
|
let limit = if final_chunk {
|
||||||
|
self.raw.len()
|
||||||
|
} else {
|
||||||
|
TOOL_SYNTAXES
|
||||||
|
.iter()
|
||||||
|
.map(|syntax| {
|
||||||
|
safe_before_partial_marker(&self.raw, syntax.tool_start)
|
||||||
|
})
|
||||||
|
.min()
|
||||||
|
.unwrap_or(self.raw.len())
|
||||||
|
};
|
||||||
|
if limit > self.text_emitted {
|
||||||
|
let text = &self.raw[self.text_emitted..limit];
|
||||||
|
if !text.trim().is_empty() {
|
||||||
|
events.push(ToolProjectionEvent::Text(text.to_owned()));
|
||||||
|
self.text_emitted = limit;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ToolProjectionState::Invokes => {
|
||||||
|
let syntax = self.syntax.unwrap();
|
||||||
|
self.skip_whitespace();
|
||||||
|
if self.full_at(syntax.tool_end) {
|
||||||
|
self.position += syntax.tool_end.len();
|
||||||
|
self.state = ToolProjectionState::Done;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if self.partial_at(syntax.tool_end) || self.partial_at(syntax.invoke_start) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if !self.full_at(syntax.invoke_start) {
|
||||||
|
self.state = ToolProjectionState::Failed;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let Some(tag_end) = self.raw[self.position..].find('>') else {
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
let tag_end = self.position + tag_end + 1;
|
||||||
|
let Some(name) = dsml_attribute(&self.raw[self.position..tag_end], "name")
|
||||||
|
else {
|
||||||
|
self.state = ToolProjectionState::Failed;
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
let id = random_tool_id(prefix);
|
||||||
|
self.ids.push(id.clone());
|
||||||
|
events.push(ToolProjectionEvent::Start {
|
||||||
|
index: self.index,
|
||||||
|
id,
|
||||||
|
name,
|
||||||
|
});
|
||||||
|
events.push(ToolProjectionEvent::Arguments {
|
||||||
|
index: self.index,
|
||||||
|
fragment: "{".into(),
|
||||||
|
});
|
||||||
|
self.position = tag_end;
|
||||||
|
self.first_parameter = true;
|
||||||
|
self.state = ToolProjectionState::Parameters;
|
||||||
|
}
|
||||||
|
ToolProjectionState::Parameters => {
|
||||||
|
let syntax = self.syntax.unwrap();
|
||||||
|
self.skip_whitespace();
|
||||||
|
if self.full_at(syntax.invoke_end) {
|
||||||
|
events.push(ToolProjectionEvent::Arguments {
|
||||||
|
index: self.index,
|
||||||
|
fragment: "}".into(),
|
||||||
|
});
|
||||||
|
events.push(ToolProjectionEvent::End { index: self.index });
|
||||||
|
self.position += syntax.invoke_end.len();
|
||||||
|
self.index += 1;
|
||||||
|
self.state = ToolProjectionState::Invokes;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if self.partial_at(syntax.invoke_end) || self.partial_at(syntax.parameter_start)
|
||||||
|
{
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if !self.full_at(syntax.parameter_start) {
|
||||||
|
self.state = ToolProjectionState::Failed;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let Some(tag_end) = self.raw[self.position..].find('>') else {
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
let tag_end = self.position + tag_end + 1;
|
||||||
|
let tag = &self.raw[self.position..tag_end];
|
||||||
|
let Some(name) = dsml_attribute(tag, "name") else {
|
||||||
|
self.state = ToolProjectionState::Failed;
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
self.string_parameter =
|
||||||
|
dsml_attribute(tag, "string").as_deref() != Some("false");
|
||||||
|
let mut fragment = if self.first_parameter {
|
||||||
|
String::new()
|
||||||
|
} else {
|
||||||
|
",".into()
|
||||||
|
};
|
||||||
|
self.first_parameter = false;
|
||||||
|
fragment
|
||||||
|
.push_str(&serde_json::to_string(&name).unwrap_or_else(|_| "\"\"".into()));
|
||||||
|
fragment.push(':');
|
||||||
|
if self.string_parameter {
|
||||||
|
fragment.push('"');
|
||||||
|
}
|
||||||
|
events.push(ToolProjectionEvent::Arguments {
|
||||||
|
index: self.index,
|
||||||
|
fragment,
|
||||||
|
});
|
||||||
|
self.position = tag_end;
|
||||||
|
self.state = ToolProjectionState::Value;
|
||||||
|
}
|
||||||
|
ToolProjectionState::Value => {
|
||||||
|
let syntax = self.syntax.unwrap();
|
||||||
|
if let Some(relative_end) = self.raw[self.position..].find(syntax.parameter_end)
|
||||||
|
{
|
||||||
|
let end = self.position + relative_end;
|
||||||
|
self.emit_value(end, &mut events);
|
||||||
|
if self.string_parameter {
|
||||||
|
events.push(ToolProjectionEvent::Arguments {
|
||||||
|
index: self.index,
|
||||||
|
fragment: "\"".into(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
self.position = end + syntax.parameter_end.len();
|
||||||
|
self.state = ToolProjectionState::Parameters;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let limit = safe_parameter_value_limit(
|
||||||
|
&self.raw,
|
||||||
|
self.position,
|
||||||
|
syntax.parameter_end,
|
||||||
|
self.string_parameter,
|
||||||
|
);
|
||||||
|
self.emit_value(limit, &mut events);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
ToolProjectionState::Done | ToolProjectionState::Failed => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
events
|
||||||
|
}
|
||||||
|
|
||||||
|
fn emit_value(&mut self, end: usize, events: &mut Vec<ToolProjectionEvent>) {
|
||||||
|
if end <= self.position {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let raw = &self.raw[self.position..end];
|
||||||
|
let fragment = if self.string_parameter {
|
||||||
|
let value = unescape_dsml(raw);
|
||||||
|
let encoded = serde_json::to_string(&value).unwrap_or_else(|_| "\"\"".into());
|
||||||
|
encoded[1..encoded.len() - 1].to_owned()
|
||||||
|
} else {
|
||||||
|
raw.to_owned()
|
||||||
|
};
|
||||||
|
events.push(ToolProjectionEvent::Arguments {
|
||||||
|
index: self.index,
|
||||||
|
fragment,
|
||||||
|
});
|
||||||
|
self.position = end;
|
||||||
|
}
|
||||||
|
|
||||||
|
fn skip_whitespace(&mut self) {
|
||||||
|
while self.raw[self.position..]
|
||||||
|
.chars()
|
||||||
|
.next()
|
||||||
|
.is_some_and(char::is_whitespace)
|
||||||
|
{
|
||||||
|
self.position += self.raw[self.position..].chars().next().unwrap().len_utf8();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn full_at(&self, marker: &str) -> bool {
|
||||||
|
self.raw.as_bytes()[self.position..].starts_with(marker.as_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn partial_at(&self, marker: &str) -> bool {
|
||||||
|
let tail = &self.raw.as_bytes()[self.position..];
|
||||||
|
tail.len() < marker.len() && marker.as_bytes().starts_with(tail)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dsml_attribute(tag: &str, name: &str) -> Option<String> {
|
||||||
|
let start = tag.find(&format!("{name}=\""))? + name.len() + 2;
|
||||||
|
let end = start + tag[start..].find('"')?;
|
||||||
|
Some(unescape_dsml(&tag[start..end]))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn safe_before_partial_marker(text: &str, marker: &str) -> usize {
|
||||||
|
let mut limit = text.len().saturating_sub(marker.len().saturating_sub(1));
|
||||||
|
while !text.is_char_boundary(limit) {
|
||||||
|
limit -= 1;
|
||||||
|
}
|
||||||
|
limit
|
||||||
|
}
|
||||||
|
|
||||||
|
fn safe_parameter_value_limit(text: &str, start: usize, end_marker: &str, string: bool) -> usize {
|
||||||
|
let bytes = text.as_bytes();
|
||||||
|
let marker = end_marker.as_bytes();
|
||||||
|
let mut limit = bytes.len();
|
||||||
|
for length in (1..marker.len().min(bytes.len().saturating_sub(start) + 1)).rev() {
|
||||||
|
if bytes[start..].ends_with(&marker[..length]) {
|
||||||
|
limit -= length;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if string {
|
||||||
|
for entity in ["&", "<", ">", """, "'"] {
|
||||||
|
let entity = entity.as_bytes();
|
||||||
|
for length in 1..entity.len() {
|
||||||
|
if bytes[start..limit].ends_with(&entity[..length]) {
|
||||||
|
limit -= length;
|
||||||
|
return limit;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
limit
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn render_messages(
|
||||||
|
state: &State,
|
||||||
|
messages: &[ApiMessage],
|
||||||
|
tools: &[Value],
|
||||||
|
tool_schemas: &[String],
|
||||||
|
tools_enabled: bool,
|
||||||
|
protocol: Protocol,
|
||||||
|
) -> Result<(String, Vec<ChatTurn>), (u16, String)> {
|
||||||
|
validate_tool_results(state, messages, protocol)?;
|
||||||
|
let preserve_reasoning = tools_enabled
|
||||||
|
|| messages.iter().any(|message| {
|
||||||
|
matches!(message.role.as_str(), "tool" | "function") || !message.tool_calls.is_empty()
|
||||||
|
});
|
||||||
|
let mut system = String::new();
|
||||||
|
if tools_enabled {
|
||||||
|
system.push_str(TOOLS_PROMPT);
|
||||||
|
if tool_schemas.is_empty() {
|
||||||
|
for tool in tools {
|
||||||
|
let schema = tool.get("function").unwrap_or(tool);
|
||||||
|
if !system.ends_with("\n\n") {
|
||||||
|
system.push('\n');
|
||||||
|
}
|
||||||
|
system.push_str(
|
||||||
|
&serde_json::to_string(schema)
|
||||||
|
.map_err(|error| (400, format!("invalid tool schema: {error}")))?,
|
||||||
|
);
|
||||||
|
system.push('\n');
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for schema in tool_schemas {
|
||||||
|
if !system.ends_with("\n\n") {
|
||||||
|
system.push('\n');
|
||||||
|
}
|
||||||
|
system.push_str(schema);
|
||||||
|
system.push('\n');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
system.push_str(
|
||||||
|
"\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. Use the exact parameter names from the schemas.",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut turns = Vec::<ChatTurn>::new();
|
||||||
|
for message in messages {
|
||||||
|
let content = content_text(&message.content);
|
||||||
|
match message.role.as_str() {
|
||||||
|
"system" | "developer" => {
|
||||||
|
if !system.is_empty() {
|
||||||
|
system.push_str("\n\n");
|
||||||
|
}
|
||||||
|
system.push_str(&content);
|
||||||
|
}
|
||||||
|
"user" => turns.push(ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content,
|
||||||
|
}),
|
||||||
|
"tool" | "function" => {
|
||||||
|
let wrapped = format!(
|
||||||
|
"<tool_result>{}</tool_result>",
|
||||||
|
escape_tool_result(&content)
|
||||||
|
);
|
||||||
|
if let Some(previous) = turns.last_mut()
|
||||||
|
&& previous.user
|
||||||
|
&& previous.content.starts_with("<tool_result>")
|
||||||
|
{
|
||||||
|
previous.content.push_str(&wrapped);
|
||||||
|
} else {
|
||||||
|
turns.push(ChatTurn {
|
||||||
|
user: true,
|
||||||
|
skip_previous_eos: protocol == Protocol::Responses,
|
||||||
|
reasoning: None,
|
||||||
|
reasoning_complete: true,
|
||||||
|
content: wrapped,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"assistant" => {
|
||||||
|
let mut content = content;
|
||||||
|
if !message.tool_calls.is_empty() {
|
||||||
|
content.push_str(&replayed_or_canonical_tools(state, &message.tool_calls));
|
||||||
|
}
|
||||||
|
let reasoning = content_text(&message.reasoning_content);
|
||||||
|
turns.push(ChatTurn {
|
||||||
|
user: false,
|
||||||
|
skip_previous_eos: false,
|
||||||
|
reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning),
|
||||||
|
reasoning_complete: true,
|
||||||
|
content,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
role => return Err((400, format!("unsupported message role: {role}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok((system, turns))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn validate_tool_results(
|
||||||
|
state: &State,
|
||||||
|
messages: &[ApiMessage],
|
||||||
|
protocol: Protocol,
|
||||||
|
) -> Result<(), (u16, String)> {
|
||||||
|
if !matches!(protocol, Protocol::Anthropic | Protocol::Responses) {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let memory = state.tool_memory.lock().ok();
|
||||||
|
for (index, message) in messages.iter().enumerate() {
|
||||||
|
if !matches!(message.role.as_str(), "tool" | "function") || message.tool_call_id.is_empty()
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let id = &message.tool_call_id;
|
||||||
|
let live = memory
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|memory| memory.contains_key(id));
|
||||||
|
let replayed = messages[..index].iter().any(|message| {
|
||||||
|
message.role == "assistant" && message.tool_calls.iter().any(|call| call.id == *id)
|
||||||
|
});
|
||||||
|
if live || replayed {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let message = match protocol {
|
||||||
|
Protocol::Anthropic => format!(
|
||||||
|
"Anthropic continuation state is not available for tool_use_id {id}; retry by replaying the full messages history"
|
||||||
|
),
|
||||||
|
Protocol::Responses => format!(
|
||||||
|
"Responses continuation state is not available for call_id {id}; retry by replaying the full input history"
|
||||||
|
),
|
||||||
|
_ => unreachable!(),
|
||||||
|
};
|
||||||
|
return Err((400, message));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn replayed_or_canonical_tools(state: &State, calls: &[ApiToolCall]) -> String {
|
||||||
|
if let Ok(memory) = state.tool_memory.lock()
|
||||||
|
&& let Some(raw) = calls.iter().find_map(|call| {
|
||||||
|
(!call.id.is_empty())
|
||||||
|
.then(|| memory.get(&call.id))
|
||||||
|
.flatten()
|
||||||
|
})
|
||||||
|
{
|
||||||
|
return raw.clone();
|
||||||
|
}
|
||||||
|
canonical_tools(calls)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn canonical_tools(calls: &[ApiToolCall]) -> String {
|
||||||
|
let mut output = String::from("\n\n<|DSML|tool_calls>\n");
|
||||||
|
for call in calls {
|
||||||
|
output.push_str("<|DSML|invoke name=\"");
|
||||||
|
output.push_str(&escape_attribute(&call.function.name));
|
||||||
|
output.push_str("\">\n");
|
||||||
|
match serde_json::from_str::<Value>(&call.function.arguments) {
|
||||||
|
Ok(Value::Object(arguments)) => {
|
||||||
|
for (name, value) in arguments {
|
||||||
|
output.push_str("<|DSML|parameter name=\"");
|
||||||
|
output.push_str(&escape_attribute(&name));
|
||||||
|
let string = value.as_str();
|
||||||
|
output.push_str(if string.is_some() {
|
||||||
|
"\" string=\"true\">"
|
||||||
|
} else {
|
||||||
|
"\" string=\"false\">"
|
||||||
|
});
|
||||||
|
if let Some(value) = string {
|
||||||
|
output.push_str(&escape_parameter(value));
|
||||||
|
} else {
|
||||||
|
output.push_str(&escape_json_parameter(&value.to_string()));
|
||||||
|
}
|
||||||
|
output.push_str("</|DSML|parameter>\n");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
output.push_str("<|DSML|parameter name=\"arguments\" string=\"true\">");
|
||||||
|
output.push_str(&escape_parameter(&call.function.arguments));
|
||||||
|
output.push_str("</|DSML|parameter>\n");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
output.push_str("</|DSML|invoke>\n");
|
||||||
|
}
|
||||||
|
output.push_str("</|DSML|tool_calls>");
|
||||||
|
output
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn parse_generated_tools(
|
||||||
|
state: &State,
|
||||||
|
text: &str,
|
||||||
|
protocol: Protocol,
|
||||||
|
) -> (String, Vec<ApiToolCall>) {
|
||||||
|
parse_generated_tools_with_ids(state, text, protocol, &[])
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn parse_generated_tools_with_ids(
|
||||||
|
state: &State,
|
||||||
|
text: &str,
|
||||||
|
protocol: Protocol,
|
||||||
|
streamed_ids: &[String],
|
||||||
|
) -> (String, Vec<ApiToolCall>) {
|
||||||
|
let Some((start, syntax)) = TOOL_SYNTAXES
|
||||||
|
.iter()
|
||||||
|
.filter_map(|syntax| text.find(syntax.tool_start).map(|start| (start, *syntax)))
|
||||||
|
.min_by_key(|(start, _)| *start)
|
||||||
|
else {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
};
|
||||||
|
let Some(relative_end) = text[start..].find(syntax.tool_end) else {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
};
|
||||||
|
let end = start + relative_end + syntax.tool_end.len();
|
||||||
|
let content = text[..start].trim_end();
|
||||||
|
let raw = &text[content.len()..end];
|
||||||
|
let mut calls = Vec::new();
|
||||||
|
let mut cursor = raw.find(syntax.tool_start).unwrap() + syntax.tool_start.len();
|
||||||
|
loop {
|
||||||
|
skip_text_whitespace(raw, &mut cursor);
|
||||||
|
if raw[cursor..].starts_with(syntax.tool_end) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if !raw[cursor..].starts_with(syntax.invoke_start) {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
}
|
||||||
|
let Some(tag_end) = raw[cursor..].find('>').map(|end| cursor + end + 1) else {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
};
|
||||||
|
let Some(name) = dsml_attribute(&raw[cursor..tag_end], "name") else {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
};
|
||||||
|
cursor = tag_end;
|
||||||
|
let mut arguments = Map::new();
|
||||||
|
loop {
|
||||||
|
skip_text_whitespace(raw, &mut cursor);
|
||||||
|
if raw[cursor..].starts_with(syntax.invoke_end) {
|
||||||
|
cursor += syntax.invoke_end.len();
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let Some((name, value)) = parse_tool_parameter(raw, &mut cursor, syntax) else {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
};
|
||||||
|
arguments.insert(name, value);
|
||||||
|
}
|
||||||
|
calls.push(ApiToolCall {
|
||||||
|
id: String::new(),
|
||||||
|
function: ApiFunction {
|
||||||
|
name,
|
||||||
|
arguments: Value::Object(arguments).to_string(),
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if calls.is_empty() {
|
||||||
|
return (text.to_owned(), Vec::new());
|
||||||
|
}
|
||||||
|
let prefix = if protocol == Protocol::Anthropic {
|
||||||
|
"toolu_"
|
||||||
|
} else {
|
||||||
|
"call_"
|
||||||
|
};
|
||||||
|
for (index, call) in calls.iter_mut().enumerate() {
|
||||||
|
call.id = streamed_ids
|
||||||
|
.get(index)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| random_tool_id(prefix));
|
||||||
|
}
|
||||||
|
if let Ok(mut memory) = state.tool_memory.lock() {
|
||||||
|
// ponytail: one process-local replay table; add LRU eviction if 100k live tool ids is measured insufficient.
|
||||||
|
if memory.len() >= 100_000 {
|
||||||
|
memory.clear();
|
||||||
|
}
|
||||||
|
for call in &calls {
|
||||||
|
memory.insert(call.id.clone(), raw.to_owned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(content.to_owned(), calls)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
pub(super) struct ToolSyntax {
|
||||||
|
pub(super) tool_start: &'static str,
|
||||||
|
pub(super) tool_end: &'static str,
|
||||||
|
pub(super) invoke_start: &'static str,
|
||||||
|
pub(super) invoke_end: &'static str,
|
||||||
|
pub(super) parameter_start: &'static str,
|
||||||
|
pub(super) parameter_end: &'static str,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) const TOOL_SYNTAXES: [ToolSyntax; 3] = [
|
||||||
|
ToolSyntax {
|
||||||
|
tool_start: "<|DSML|tool_calls>",
|
||||||
|
tool_end: "</|DSML|tool_calls>",
|
||||||
|
invoke_start: "<|DSML|invoke",
|
||||||
|
invoke_end: "</|DSML|invoke>",
|
||||||
|
parameter_start: "<|DSML|parameter",
|
||||||
|
parameter_end: "</|DSML|parameter>",
|
||||||
|
},
|
||||||
|
ToolSyntax {
|
||||||
|
tool_start: "<DSML|tool_calls>",
|
||||||
|
tool_end: "</DSML|tool_calls>",
|
||||||
|
invoke_start: "<DSML|invoke",
|
||||||
|
invoke_end: "</DSML|invoke>",
|
||||||
|
parameter_start: "<DSML|parameter",
|
||||||
|
parameter_end: "</DSML|parameter>",
|
||||||
|
},
|
||||||
|
ToolSyntax {
|
||||||
|
tool_start: "<tool_calls>",
|
||||||
|
tool_end: "</tool_calls>",
|
||||||
|
invoke_start: "<invoke",
|
||||||
|
invoke_end: "</invoke>",
|
||||||
|
parameter_start: "<parameter",
|
||||||
|
parameter_end: "</parameter>",
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
fn parse_tool_parameter(
|
||||||
|
text: &str,
|
||||||
|
cursor: &mut usize,
|
||||||
|
syntax: ToolSyntax,
|
||||||
|
) -> Option<(String, Value)> {
|
||||||
|
if !text[*cursor..].starts_with(syntax.parameter_start) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let tag_end = text[*cursor..].find('>').map(|end| *cursor + end + 1)?;
|
||||||
|
let tag = &text[*cursor..tag_end];
|
||||||
|
let name = dsml_attribute(tag, "name")?;
|
||||||
|
let is_string = dsml_attribute(tag, "string");
|
||||||
|
*cursor = tag_end;
|
||||||
|
let mut nested_start = *cursor;
|
||||||
|
skip_text_whitespace(text, &mut nested_start);
|
||||||
|
if is_string.is_none() && text[nested_start..].starts_with(syntax.parameter_start) {
|
||||||
|
*cursor = nested_start;
|
||||||
|
let mut nested = Map::new();
|
||||||
|
loop {
|
||||||
|
skip_text_whitespace(text, cursor);
|
||||||
|
if !text[*cursor..].starts_with(syntax.parameter_start) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let (name, value) = parse_tool_parameter(text, cursor, syntax)?;
|
||||||
|
nested.insert(name, value);
|
||||||
|
}
|
||||||
|
skip_text_whitespace(text, cursor);
|
||||||
|
if !text[*cursor..].starts_with(syntax.parameter_end) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
*cursor += syntax.parameter_end.len();
|
||||||
|
return Some((name, Value::Object(nested)));
|
||||||
|
}
|
||||||
|
let value_end = text[*cursor..]
|
||||||
|
.find(syntax.parameter_end)
|
||||||
|
.map(|end| *cursor + end)?;
|
||||||
|
let raw = &text[*cursor..value_end];
|
||||||
|
*cursor = value_end + syntax.parameter_end.len();
|
||||||
|
let value = if is_string.as_deref().unwrap_or("true") == "true" {
|
||||||
|
Value::String(unescape_dsml(raw))
|
||||||
|
} else {
|
||||||
|
serde_json::from_str(raw).unwrap_or(Value::Null)
|
||||||
|
};
|
||||||
|
Some((name, value))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn skip_text_whitespace(text: &str, cursor: &mut usize) {
|
||||||
|
while text[*cursor..]
|
||||||
|
.chars()
|
||||||
|
.next()
|
||||||
|
.is_some_and(char::is_whitespace)
|
||||||
|
{
|
||||||
|
*cursor += text[*cursor..].chars().next().unwrap().len_utf8();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn tool_calls_json(calls: &[ApiToolCall]) -> Value {
|
||||||
|
Value::Array(
|
||||||
|
calls
|
||||||
|
.iter()
|
||||||
|
.map(|call| {
|
||||||
|
json!({
|
||||||
|
"id": call.id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": call.function.name,
|
||||||
|
"arguments": call.function.arguments,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn content_text(value: &Value) -> String {
|
||||||
|
match value {
|
||||||
|
Value::String(text) => text.clone(),
|
||||||
|
Value::Array(parts) => parts
|
||||||
|
.iter()
|
||||||
|
.filter_map(|part| match part {
|
||||||
|
Value::String(text) => Some(text.as_str()),
|
||||||
|
Value::Object(object) => object.get("text").and_then(Value::as_str),
|
||||||
|
_ => None,
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
_ => String::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn escape_attribute(text: &str) -> String {
|
||||||
|
text.replace('&', "&")
|
||||||
|
.replace('<', "<")
|
||||||
|
.replace('>', ">")
|
||||||
|
.replace('"', """)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn escape_parameter(text: &str) -> String {
|
||||||
|
text.replace("</|DSML|parameter>", "</|DSML|parameter>")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn escape_json_parameter(text: &str) -> String {
|
||||||
|
text.replace("</|DSML|parameter>", "\\u003c/|DSML|parameter>")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn escape_tool_result(text: &str) -> String {
|
||||||
|
text.replace("</tool_result>", "</tool_result>")
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn unescape_dsml(text: &str) -> String {
|
||||||
|
text.replace(""", "\"")
|
||||||
|
.replace(">", ">")
|
||||||
|
.replace("<", "<")
|
||||||
|
.replace("&", "&")
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user