ds4-winhip / crates /ds4-cli /src /token_printer.rs

ds4-winhip: Windows/HIP 7.2 build tree, launcher, supervisor, bench + notes

fda8632 verified 17 days ago

7.36 kB

//! CLI output: preserve C thinking formatting and project Inkling channel IDs.

use std::io::{self, IsTerminal, Write};

use ds4_core::ModelFamily;

const THINK_OPEN: &[u8] = b"<think>";
const THINK_CLOSE: &[u8] = b"</think>";
const INKLING_MODEL: i32 = 200001;
const INKLING_TEXT: i32 = 200004;
const INKLING_THINK: i32 = 200008;
const INKLING_END: i32 = 200010;

/// C `token_printer_set_grey`: SGR 90 bright-black.
const GREY: &[u8] = b"\x1b[90m";
/// C `token_printer_reset_color`.
const RESET: &[u8] = b"\x1b[0m";

pub(crate) struct TokenPrinter {
    family: ModelFamily,
    format_thinking: bool,
    in_think: bool,
    color_open: bool,
    use_color: bool,
    pending: Vec<u8>,
    last_output_newline: bool,
}

impl TokenPrinter {
    pub(crate) fn new(family: ModelFamily, format_thinking: bool) -> Self {
        let mut printer = Self::with_color(format_thinking, io::stdout().is_terminal());
        printer.family = family;
        if family == ModelFamily::Inkling {
            printer.in_think = false;
        }
        printer
    }

    pub(crate) fn with_color(format_thinking: bool, use_color: bool) -> Self {
        Self {
            family: ModelFamily::DeepSeek4,
            format_thinking,
            in_think: format_thinking,
            color_open: false,
            use_color,
            pending: Vec::new(),
            last_output_newline: true,
        }
    }

    pub(crate) fn write_token<W: Write>(
        &mut self,
        out: &mut W,
        token: i32,
        text: &[u8],
    ) -> io::Result<()> {
        if self.family != ModelFamily::Inkling {
            return self.write_text(out, text);
        }
        // Use special IDs so ordinary text that spells a marker stays literal.
        // Inkling emits these boundaries even when thinking effort is zero.
        match token {
            INKLING_MODEL => return Ok(()),
            INKLING_THINK => {
                self.in_think = self.format_thinking;
                return Ok(());
            }
            INKLING_TEXT | INKLING_END => {
                self.in_think = false;
                self.reset_color(out)?;
                if token == INKLING_END && !self.last_output_newline {
                    out.write_all(b"\n")?;
                    self.last_output_newline = true;
                }
                return Ok(());
            }
            _ => {}
        }
        for &byte in text {
            self.write_char(out, byte)?;
        }
        Ok(())
    }

    pub(crate) fn write_text<W: Write>(&mut self, out: &mut W, text: &[u8]) -> io::Result<()> {
        if !self.format_thinking {
            out.write_all(text)?;
            if let Some(last) = text.last() {
                self.last_output_newline = *last == b'\n';
            }
            return Ok(());
        }
        self.process(out, text, false)
    }

    pub(crate) fn finish<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
        if self.format_thinking {
            self.process(out, &[], true)?;
            self.reset_color(out)?;
        }
        if !self.last_output_newline {
            out.write_all(b"\n")?;
            self.last_output_newline = true;
        }
        out.flush()
    }

    fn process<W: Write>(&mut self, out: &mut W, text: &[u8], finish: bool) -> io::Result<()> {
        let mut bytes = std::mem::take(&mut self.pending);
        bytes.extend_from_slice(text);
        let mut i = 0;
        while i < bytes.len() {
            let rem = &bytes[i..];
            if rem.starts_with(THINK_OPEN) {
                self.in_think = true;
                i += THINK_OPEN.len();
                continue;
            }
            if rem.starts_with(THINK_CLOSE) {
                self.in_think = false;
                self.reset_color(out)?;
                if !self.last_output_newline {
                    out.write_all(b"\n")?;
                    self.last_output_newline = true;
                }
                i += THINK_CLOSE.len();
                continue;
            }
            if !finish
                && rem[0] == b'<'
                && ((rem.len() < THINK_OPEN.len() && THINK_OPEN.starts_with(rem))
                    || (rem.len() < THINK_CLOSE.len() && THINK_CLOSE.starts_with(rem)))
            {
                self.pending.extend_from_slice(rem);
                break;
            }
            self.write_char(out, rem[0])?;
            i += 1;
        }
        Ok(())
    }

    fn write_char<W: Write>(&mut self, out: &mut W, c: u8) -> io::Result<()> {
        if self.in_think {
            self.set_grey(out)?;
        }
        out.write_all(&[c])?;
        self.last_output_newline = c == b'\n';
        Ok(())
    }

    fn set_grey<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
        if self.use_color && !self.color_open {
            out.write_all(GREY)?;
            self.color_open = true;
        }
        Ok(())
    }

    fn reset_color<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
        if self.use_color && self.color_open {
            out.write_all(RESET)?;
            self.color_open = false;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    enum Style {
        Plain,
        Thinking,
        Color,
    }

    fn inkling_output(style: Style, pieces: &[(i32, &[u8])]) -> Vec<u8> {
        let mut printer = TokenPrinter::new(ModelFamily::Inkling, !matches!(style, Style::Plain));
        printer.use_color = matches!(style, Style::Color);
        let mut output = Vec::new();
        for &(id, text) in pieces {
            printer.write_token(&mut output, id, text).unwrap();
        }
        printer.finish(&mut output).unwrap();
        output
    }

    #[test]
    fn inkling_text_boundaries() {
        let pieces: &[(i32, &[u8])] = &[
            (200004, b"<|content_text|>"),
            (19, b"4"),
            (200010, b"<|end_message|>"),
        ];
        for style in [Style::Plain, Style::Thinking] {
            assert_eq!(inkling_output(style, pieces), b"4\n");
        }
    }

    #[test]
    fn inkling_thinking_transition() {
        let pieces: &[(i32, &[u8])] = &[
            (200008, b"<|content_thinking|>"),
            (100, b"Let me check."),
            (200010, b"<|end_message|>"),
            (200001, b"<|message_model|>"),
            (200004, b"<|content_text|>"),
            (19, b"4"),
            (200010, b"<|end_message|>"),
        ];
        assert_eq!(
            inkling_output(Style::Thinking, pieces),
            b"Let me check.\n4\n"
        );
        assert_eq!(
            inkling_output(Style::Color, pieces),
            b"\x1b[90mLet me check.\x1b[0m\n4\n"
        );
    }

    #[test]
    fn inkling_literal_markers() {
        let text = b"Use <think> and <|content_text|> literally.";
        assert_eq!(
            inkling_output(Style::Thinking, &[(100, text)]),
            [text.as_slice(), b"\n"].concat()
        );
    }

    #[test]
    fn inkling_empty_thinking() {
        assert_eq!(
            inkling_output(
                Style::Color,
                &[
                    (200008, b"<|content_thinking|>"),
                    (200010, b"<|end_message|>"),
                    (200004, b"<|content_text|>"),
                    (19, b"4"),
                ]
            ),
            b"4\n"
        );
    }
}

	//! CLI output: preserve C thinking formatting and project Inkling channel IDs.

	use std::io::{self, IsTerminal, Write};

	use ds4_core::ModelFamily;

	const THINK_OPEN: &[u8] = b"<think>";
	const THINK_CLOSE: &[u8] = b"</think>";
	const INKLING_MODEL: i32 = 200001;
	const INKLING_TEXT: i32 = 200004;
	const INKLING_THINK: i32 = 200008;
	const INKLING_END: i32 = 200010;

	/// C `token_printer_set_grey`: SGR 90 bright-black.
	const GREY: &[u8] = b"\x1b[90m";
	/// C `token_printer_reset_color`.
	const RESET: &[u8] = b"\x1b[0m";

	pub(crate) struct TokenPrinter {
	family: ModelFamily,
	format_thinking: bool,
	in_think: bool,
	color_open: bool,
	use_color: bool,
	pending: Vec<u8>,
	last_output_newline: bool,
	}

	impl TokenPrinter {
	pub(crate) fn new(family: ModelFamily, format_thinking: bool) -> Self {
	let mut printer = Self::with_color(format_thinking, io::stdout().is_terminal());
	printer.family = family;
	if family == ModelFamily::Inkling {
	printer.in_think = false;
	}
	printer
	}

	pub(crate) fn with_color(format_thinking: bool, use_color: bool) -> Self {
	Self {
	family: ModelFamily::DeepSeek4,
	format_thinking,
	in_think: format_thinking,
	color_open: false,
	use_color,
	pending: Vec::new(),
	last_output_newline: true,
	}
	}

	pub(crate) fn write_token<W: Write>(
	&mut self,
	out: &mut W,
	token: i32,
	text: &[u8],
	) -> io::Result<()> {
	if self.family != ModelFamily::Inkling {
	return self.write_text(out, text);
	}
	// Use special IDs so ordinary text that spells a marker stays literal.
	// Inkling emits these boundaries even when thinking effort is zero.
	match token {
	INKLING_MODEL => return Ok(()),
	INKLING_THINK => {
	self.in_think = self.format_thinking;
	return Ok(());
	}
	INKLING_TEXT \| INKLING_END => {
	self.in_think = false;
	self.reset_color(out)?;
	if token == INKLING_END && !self.last_output_newline {
	out.write_all(b"\n")?;
	self.last_output_newline = true;
	}
	return Ok(());
	}
	_ => {}
	}
	for &byte in text {
	self.write_char(out, byte)?;
	}
	Ok(())
	}

	pub(crate) fn write_text<W: Write>(&mut self, out: &mut W, text: &[u8]) -> io::Result<()> {
	if !self.format_thinking {
	out.write_all(text)?;
	if let Some(last) = text.last() {
	self.last_output_newline = *last == b'\n';
	}
	return Ok(());
	}
	self.process(out, text, false)
	}

	pub(crate) fn finish<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
	if self.format_thinking {
	self.process(out, &[], true)?;
	self.reset_color(out)?;
	}
	if !self.last_output_newline {
	out.write_all(b"\n")?;
	self.last_output_newline = true;
	}
	out.flush()
	}

	fn process<W: Write>(&mut self, out: &mut W, text: &[u8], finish: bool) -> io::Result<()> {
	let mut bytes = std::mem::take(&mut self.pending);
	bytes.extend_from_slice(text);
	let mut i = 0;
	while i < bytes.len() {
	let rem = &bytes[i..];
	if rem.starts_with(THINK_OPEN) {
	self.in_think = true;
	i += THINK_OPEN.len();
	continue;
	}
	if rem.starts_with(THINK_CLOSE) {
	self.in_think = false;
	self.reset_color(out)?;
	if !self.last_output_newline {
	out.write_all(b"\n")?;
	self.last_output_newline = true;
	}
	i += THINK_CLOSE.len();
	continue;
	}
	if !finish
	&& rem[0] == b'<'
	&& ((rem.len() < THINK_OPEN.len() && THINK_OPEN.starts_with(rem))
	\|\| (rem.len() < THINK_CLOSE.len() && THINK_CLOSE.starts_with(rem)))
	{
	self.pending.extend_from_slice(rem);
	break;
	}
	self.write_char(out, rem[0])?;
	i += 1;
	}
	Ok(())
	}

	fn write_char<W: Write>(&mut self, out: &mut W, c: u8) -> io::Result<()> {
	if self.in_think {
	self.set_grey(out)?;
	}
	out.write_all(&[c])?;
	self.last_output_newline = c == b'\n';
	Ok(())
	}

	fn set_grey<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
	if self.use_color && !self.color_open {
	out.write_all(GREY)?;
	self.color_open = true;
	}
	Ok(())
	}

	fn reset_color<W: Write>(&mut self, out: &mut W) -> io::Result<()> {
	if self.use_color && self.color_open {
	out.write_all(RESET)?;
	self.color_open = false;
	}
	Ok(())
	}
	}

	#[cfg(test)]
	mod tests {
	use super::*;
	enum Style {
	Plain,
	Thinking,
	Color,
	}

	fn inkling_output(style: Style, pieces: &[(i32, &[u8])]) -> Vec<u8> {
	let mut printer = TokenPrinter::new(ModelFamily::Inkling, !matches!(style, Style::Plain));
	printer.use_color = matches!(style, Style::Color);
	let mut output = Vec::new();
	for &(id, text) in pieces {
	printer.write_token(&mut output, id, text).unwrap();
	}
	printer.finish(&mut output).unwrap();
	output
	}

	#[test]
	fn inkling_text_boundaries() {
	let pieces: &[(i32, &[u8])] = &[
	(200004, b"<\|content_text\|>"),
	(19, b"4"),
	(200010, b"<\|end_message\|>"),
	];
	for style in [Style::Plain, Style::Thinking] {
	assert_eq!(inkling_output(style, pieces), b"4\n");
	}
	}

	#[test]
	fn inkling_thinking_transition() {
	let pieces: &[(i32, &[u8])] = &[
	(200008, b"<\|content_thinking\|>"),
	(100, b"Let me check."),
	(200010, b"<\|end_message\|>"),
	(200001, b"<\|message_model\|>"),
	(200004, b"<\|content_text\|>"),
	(19, b"4"),
	(200010, b"<\|end_message\|>"),
	];
	assert_eq!(
	inkling_output(Style::Thinking, pieces),
	b"Let me check.\n4\n"
	);
	assert_eq!(
	inkling_output(Style::Color, pieces),
	b"\x1b[90mLet me check.\x1b[0m\n4\n"
	);
	}

	#[test]
	fn inkling_literal_markers() {
	let text = b"Use <think> and <\|content_text\|> literally.";
	assert_eq!(
	inkling_output(Style::Thinking, &[(100, text)]),
	[text.as_slice(), b"\n"].concat()
	);
	}

	#[test]
	fn inkling_empty_thinking() {
	assert_eq!(
	inkling_output(
	Style::Color,
	&[
	(200008, b"<\|content_thinking\|>"),
	(200010, b"<\|end_message\|>"),
	(200004, b"<\|content_text\|>"),
	(19, b"4"),
	]
	),
	b"4\n"
	);
	}
	}