hugo/DS4Server
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

feat: first cut at openai compatible server

Browse Source
This commit is contained in:
Georg Bauer
2026-07-24 21:12:48 +02:00
parent bbbe65a75f
commit d554b77b9d
18 changed files with 2612 additions and 535 deletions
Download Patch File Download Diff File Expand all files Collapse all files

314
src/engine.rs
View File

@@ -305,6 +305,14 @@ pub(crate) struct ChatTurn {
pub(crate) content: String,
}
pub(crate) struct GenerationOutput {
pub(crate) message: ChatTurn,
pub(crate) prompt_tokens: u32,
pub(crate) cached_tokens: u32,
pub(crate) completion_tokens: u32,
pub(crate) finish_reason: &'static str,
}
#[cfg(target_os = "macos")]
impl Generator {
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
@@ -337,12 +345,12 @@ impl Generator {
cancelled: &AtomicBool,
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32, Option<f32>),
) -> Result<(), String> {
) -> Result<GenerationOutput, String> {
self.select_checkpoint(checkpoint)?;
let (generated, prompt_complete) =
let (output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let mut completed = messages.to_vec();
completed.push(generated);
completed.push(output.message.clone());
self.executor.save_checkpoint(
checkpoint,
if prompt_complete {
@@ -350,7 +358,57 @@ impl Generator {
} else {
[0; 32]
},
)
)?;
Ok(output)
}
pub(crate) fn generate_transient(
&mut self,
directory: &Path,
messages: &[ChatTurn],
settings: &TurnSettings,
cancelled: &AtomicBool,
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32, Option<f32>),
) -> Result<GenerationOutput, String> {
let history = messages
.split_last()
.map_or(messages, |(_, history)| history);
let history_tag =
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history);
let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag)));
if self.executor.checkpoint_tag() != history_tag {
self.select_checkpoint(&checkpoint)?;
if self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
self.checkpoint = None;
let _ = std::fs::remove_file(&checkpoint);
}
}
let (output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
let completed_tag =
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed);
std::fs::create_dir_all(directory).map_err(|error| {
format!(
"Could not create transient KV cache directory {}: {error}",
directory.display()
)
})?;
let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag)));
self.executor.save_checkpoint(
&completed_checkpoint,
if prompt_complete {
completed_tag
} else {
[0; 32]
},
)?;
self.checkpoint = Some(completed_checkpoint);
Ok(output)
}
fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> {
@@ -373,7 +431,7 @@ impl Generator {
cancelled: &AtomicBool,
emit: &mut impl FnMut(bool, String),
progress: &mut impl FnMut(u32, u32, Option<f32>),
) -> Result<(ChatTurn, bool), String> {
) -> Result<(GenerationOutput, bool), String> {
let tokens = match messages.split_last() {
Some((latest, history))
if latest.user
@@ -418,10 +476,21 @@ impl Generator {
reasoning_complete: !reasoning,
content: String::new(),
};
let mut emitted_reasoning = 0;
let mut emitted_content = 0;
let prompt_tokens = tokens.len();
for (index, token) in tokens.into_iter().enumerate().skip(reused) {
if cancelled.load(Ordering::Relaxed) {
return Ok((generated, false));
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: 0,
finish_reason: "stop",
},
false,
));
}
self.executor.eval(token)?;
if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
@@ -436,13 +505,30 @@ impl Generator {
.min((max_context - self.executor.position() as usize) as i32)
{
if cancelled.load(Ordering::Relaxed) {
return Ok((generated, true));
flush_generated(
&mut generated,
&mut emitted_reasoning,
&mut emitted_content,
&settings.stops,
emit,
);
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
},
true,
));
}
let token = sample(
self.executor.logits(),
settings.temperature,
settings.top_p,
settings.min_p,
settings.top_k,
&mut rng,
);
if self
@@ -450,27 +536,89 @@ impl Generator {
.model()
.is_stop_token_for_reasoning(token, settings.reasoning_mode)
{
return Ok((generated, true));
flush_generated(
&mut generated,
&mut emitted_reasoning,
&mut emitted_content,
&settings.stops,
emit,
);
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
},
true,
));
}
if self.executor.model().is_think_start_token(token) {
reasoning = true;
generated.reasoning.get_or_insert_default();
} else if self.executor.model().is_think_end_token(token) {
if let Some(reasoning_text) = &mut generated.reasoning
&& emit_safe_text(
reasoning_text,
&mut emitted_reasoning,
&settings.stops,
true,
true,
emit,
)
{
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
},
false,
));
}
reasoning = false;
generated.reasoning_complete = true;
emit(false, String::new());
} else if let Some(bytes) = self.executor.model().token_bytes(token) {
let content = String::from_utf8_lossy(&bytes).into_owned();
if reasoning {
generated
.reasoning
.get_or_insert_default()
.push_str(&content);
let stopped = if reasoning {
let text = generated.reasoning.get_or_insert_default();
text.push_str(&content);
emit_safe_text(
text,
&mut emitted_reasoning,
&settings.stops,
false,
true,
emit,
)
} else {
generated.reasoning_complete = true;
generated.content.push_str(&content);
emit_safe_text(
&mut generated.content,
&mut emitted_content,
&settings.stops,
false,
false,
emit,
)
};
if stopped {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
},
false,
));
}
emit(reasoning, content);
}
self.executor.eval(token)?;
generated_tokens += 1;
@@ -483,10 +631,103 @@ impl Generator {
),
);
}
Ok((generated, true))
flush_generated(
&mut generated,
&mut emitted_reasoning,
&mut emitted_content,
&settings.stops,
emit,
);
Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "length",
},
true,
))
}
}
#[cfg(target_os = "macos")]
fn flush_generated(
generated: &mut ChatTurn,
emitted_reasoning: &mut usize,
emitted_content: &mut usize,
stops: &[String],
emit: &mut impl FnMut(bool, String),
) {
if let Some(reasoning) = &mut generated.reasoning {
let _ = emit_safe_text(reasoning, emitted_reasoning, stops, true, true, emit);
}
let _ = emit_safe_text(
&mut generated.content,
emitted_content,
stops,
true,
false,
emit,
);
}
#[cfg(target_os = "macos")]
fn emit_safe_text(
text: &mut String,
emitted: &mut usize,
stops: &[String],
final_flush: bool,
reasoning: bool,
emit: &mut impl FnMut(bool, String),
) -> bool {
let stop = stops
.iter()
.filter_map(|stop| {
text[*emitted..]
.find(stop)
.map(|position| *emitted + position)
})
.min();
if let Some(stop) = stop {
if stop > *emitted {
emit(reasoning, text[*emitted..stop].to_owned());
}
text.truncate(stop);
*emitted = stop;
return true;
}
let hold = if final_flush {
0
} else {
stops
.iter()
.map(|stop| stop.len().saturating_sub(1))
.max()
.unwrap_or(0)
};
let mut safe = text.len().saturating_sub(hold);
while safe > *emitted && !text.is_char_boundary(safe) {
safe -= 1;
}
if safe > *emitted {
emit(reasoning, text[*emitted..safe].to_owned());
*emitted = safe;
}
false
}
#[cfg(target_os = "macos")]
fn hex_tag(tag: [u8; 32]) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut output = String::with_capacity(64);
for byte in tag {
output.push(HEX[(byte >> 4) as usize] as char);
output.push(HEX[(byte & 0x0f) as usize] as char);
}
output
}
#[cfg(target_os = "macos")]
fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
fn text(hasher: &mut Sha256, value: &str) {
@@ -518,7 +759,14 @@ fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn
}
#[cfg(target_os = "macos")]
fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rng) -> i32 {
fn sample(
logits: &[f32],
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
) -> i32 {
if temperature <= 0.0 {
return logits
.iter()
@@ -554,12 +802,16 @@ fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rn
.max_by(|a, b| a.1.total_cmp(b.1))
.map_or(0, |(index, _)| index as i32);
}
if top_p < 1.0 {
if top_p < 1.0 || top_k > 0 {
probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
let total: f32 = logits
if top_k > 0 {
probabilities.truncate(probabilities.len().min(top_k as usize));
}
}
if top_p < 1.0 {
let total: f32 = probabilities
.iter()
.filter(|logit| logit.is_finite())
.map(|logit| ((*logit - maximum) / temperature).exp())
.map(|(_, probability)| probability)
.sum();
let mut kept = 0.0;
let count = probabilities
@@ -612,7 +864,27 @@ mod sampling_tests {
#[test]
fn zero_temperature_is_greedy() {
let mut rng = Rng::new(1);
assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, &mut rng), 1);
assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, 0, &mut rng), 1);
}
#[test]
fn top_k_and_stream_stops_are_applied_before_output() {
let mut rng = Rng::new(1);
assert_eq!(sample(&[1.0, 4.0, 2.0], 1.0, 1.0, 0.0, 1, &mut rng), 1);
let mut text = "hello STOP hidden".to_owned();
let mut emitted = 0;
let mut chunks = Vec::new();
assert!(emit_safe_text(
&mut text,
&mut emitted,
&["STOP".into()],
false,
false,
&mut |_, chunk| chunks.push(chunk),
));
assert_eq!(text, "hello ");
assert_eq!(chunks, ["hello "]);
}
#[test]