use crate::{
chat::{Author, Content, Message, ReasoningEffort, Role, SystemContent, TextContent},
tiktoken::{CoreBPE, Rank},
};
use anyhow::Context as _;
use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ParsedHeader {
author: Author,
recipient: Option<String>,
channel: Option<String>,
content_type: Option<String>,
}
#[derive(thiserror::Error, Debug)]
pub(crate) enum RenderFormattingTokenError {
#[error("tried to render unmapped formatting token {0}")]
UnmappedToken(FormattingToken),
#[error(
"Expected encoding of formatting token {token} to be a single token, but got {encoding:?}"
)]
InvalidEncoding {
token: FormattingToken,
encoding: Vec<Rank>,
},
}
#[allow(dead_code)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum FormattingToken {
Start,
Message,
EndMessage,
EndMessageDoneSampling,
EndMessageAssistantToTool,
Refusal,
ConstrainedFormat,
Channel,
BeginUntrusted,
EndUntrusted,
MetaSep,
MetaEnd,
}
impl FormattingToken {
fn as_str(&self) -> &str {
match self {
FormattingToken::Start => "<|start|>",
FormattingToken::Message => "<|message|>",
FormattingToken::EndMessage => "<|end|>",
FormattingToken::EndMessageDoneSampling => "<|return|>",
FormattingToken::EndMessageAssistantToTool => "<|call|>",
FormattingToken::Refusal => "<|refusal|>",
FormattingToken::ConstrainedFormat => "<|constrain|>",
FormattingToken::Channel => "<|channel|>",
FormattingToken::BeginUntrusted => "<|untrusted|>",
FormattingToken::EndUntrusted => "<|end_untrusted|>",
FormattingToken::MetaSep => "<|channel|>",
FormattingToken::MetaEnd => "<|meta_end|>",
}
}
}
impl std::fmt::Display for FormattingToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[allow(dead_code)]
#[derive(Clone)]
pub struct HarmonyEncoding {
pub(crate) name: String,
pub(crate) n_ctx: usize,
pub(crate) max_message_tokens: usize,
pub(crate) max_action_length: usize,
pub(crate) tokenizer_name: String,
pub(crate) tokenizer: Arc<CoreBPE>,
pub(crate) format_token_mapping: HashMap<FormattingToken, String>,
pub(crate) stop_formatting_tokens: HashSet<FormattingToken>,
pub(crate) stop_formatting_tokens_for_assistant_actions: HashSet<FormattingToken>,
}
impl std::fmt::Debug for HarmonyEncoding {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HarmonyEncoding")
.field("name", &self.name)
.field("tokenizer_name", &self.tokenizer_name)
.field("n_ctx", &self.n_ctx)
.field("max_message_tokens", &self.max_message_tokens)
.field("max_action_length", &self.max_action_length)
.finish()
}
}
impl std::fmt::Display for HarmonyEncoding {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Renderer({})", self.name)
}
}
impl HarmonyEncoding {
pub fn name(&self) -> &str {
&self.name
}
pub fn tokenizer_name(&self) -> &str {
&self.tokenizer_name
}
pub fn max_message_tokens(&self) -> usize {
self.max_message_tokens
}
pub fn tokenizer(&self) -> &CoreBPE {
&self.tokenizer
}
pub fn stop_tokens(&self) -> anyhow::Result<HashSet<Rank>> {
self.stop_formatting_tokens
.iter()
.copied()
.map(|t| match self.render_formatting_token(t) {
Ok(t) => Ok(t),
Err(RenderFormattingTokenError::UnmappedToken(_)) => Err(anyhow::anyhow!(
"token {t} was specified as a stop token, but is not mapped"
)),
Err(e) => Err(anyhow::anyhow!(e).context("could not render stop token")),
})
.collect()
}
pub fn stop_tokens_for_assistant_actions(&self) -> anyhow::Result<HashSet<Rank>> {
self.stop_formatting_tokens_for_assistant_actions
.iter()
.copied()
.map(|t| match self.render_formatting_token(t) {
Ok(t) => Ok(t),
Err(RenderFormattingTokenError::UnmappedToken(_)) => Err(anyhow::anyhow!(
"token {t} was specified as a stop token, but is not mapped"
)),
Err(e) => Err(anyhow::anyhow!(e).context("could not render stop token")),
})
.collect()
}
pub fn render_conversation_into<'a, I, B>(
&self,
conversation: I,
into: &mut B,
config: Option<&RenderConversationConfig>,
) -> anyhow::Result<()>
where
I: IntoIterator<Item = &'a Message>,
B: Extend<Rank>,
{
let messages: Vec<_> = conversation.into_iter().collect();
let has_function_tools = messages.iter().any(|msg| {
msg.content.iter().any(|c| {
if let Content::DeveloperContent(dev) = c {
if let Some(tools) = &dev.tools {
if let Some(ns) = tools.get("functions") {
!ns.tools.is_empty()
} else {
false
}
} else {
false
}
} else {
false
}
})
});
let render_options = RenderOptions {
conversation_has_function_tools: has_function_tools,
};
let last_assistant_is_final = messages
.iter()
.rev()
.find_map(|msg| {
(msg.author.role == Role::Assistant)
.then(|| msg.channel.as_deref() == Some("final"))
})
.unwrap_or(false);
let should_drop_analysis =
config.is_some_and(|c| c.auto_drop_analysis && last_assistant_is_final);
let first_final_idx = messages
.iter()
.position(|msg| msg.channel.as_deref() == Some("final"));
let result = messages
.iter()
.enumerate()
.filter(|(idx, msg)| {
!(should_drop_analysis
&& first_final_idx.is_some_and(|first| *idx < first)
&& msg.channel.as_deref() == Some("analysis"))
})
.try_for_each(|(_, msg)| self.render_into(msg, into, Some(&render_options)));
result?;
Ok(())
}
pub fn render_conversation_for_completion_into<'a, I, B>(
&self,
conversation: I,
next_turn_role: Role,
into: &mut B,
config: Option<&RenderConversationConfig>,
) -> anyhow::Result<()>
where
I: IntoIterator<Item = &'a Message>,
B: Extend<Rank>,
{
let _config = config.unwrap_or(&RenderConversationConfig::default());
self.render_conversation_into(conversation, into, config)?;
self.render_formatting_token_into(FormattingToken::Start, into)?;
self.render_text_into(next_turn_role.as_str(), into)?;
Ok(())
}
pub fn render_conversation_for_completion<'a, I>(
&self,
conversation: I,
next_turn_role: Role,
config: Option<&RenderConversationConfig>,
) -> anyhow::Result<Vec<Rank>>
where
I: IntoIterator<Item = &'a Message>,
{
let mut into = vec![];
self.render_conversation_for_completion_into(
conversation,
next_turn_role,
&mut into,
config,
)?;
Ok(into)
}
pub fn render_conversation_for_training<'a, I>(
&self,
conversation: I,
config: Option<&RenderConversationConfig>,
) -> anyhow::Result<Vec<Rank>>
where
I: IntoIterator<Item = &'a Message>,
{
let messages: Vec<&Message> = conversation.into_iter().collect();
let mut out = vec![];
self.render_conversation_into(messages.iter().copied(), &mut out, config)?;
if let Some(last) = messages.last() {
if last.author.role == Role::Assistant && last.channel.as_deref() == Some("final") {
if let Some(last_token) = out.last_mut() {
*last_token =
self.render_formatting_token(FormattingToken::EndMessageDoneSampling)?;
}
}
}
Ok(out)
}
pub fn render_conversation<'a, I>(
&self,
conversation: I,
config: Option<&RenderConversationConfig>,
) -> anyhow::Result<Vec<Rank>>
where
I: IntoIterator<Item = &'a Message>,
{
let mut out = vec![];
self.render_conversation_into(conversation, &mut out, config)?;
Ok(out)
}
pub fn render(
&self,
message: &Message,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<Vec<Rank>> {
let mut out = vec![];
Render::<Message>::render(self, message, &mut out, render_options)?;
Ok(out)
}
pub fn render_into<B>(
&self,
message: &Message,
into: &mut B,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
Render::<Message>::render(self, message, into, render_options)
}
fn mapped_format_token(&self, t: FormattingToken) -> Option<&str> {
self.format_token_mapping.get(&t).map(|s| s.as_str())
}
fn render_formatting_token(
&self,
t: FormattingToken,
) -> Result<Rank, RenderFormattingTokenError> {
let mapped = self
.mapped_format_token(t)
.ok_or(RenderFormattingTokenError::UnmappedToken(t))?;
let encoded = self.tokenizer.encode_with_special_tokens(mapped);
if encoded.len() != 1 {
return Err(RenderFormattingTokenError::InvalidEncoding {
token: t,
encoding: encoded,
});
}
Ok(encoded[0])
}
fn render_formatting_token_into<B>(
&self,
t: FormattingToken,
into: &mut B,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
let r = self.render_formatting_token(t)?;
into.extend(std::iter::once(r));
Ok(())
}
fn render_text_into<T, B>(&self, text: T, into: &mut B) -> anyhow::Result<()>
where
T: AsRef<str>,
B: Extend<Rank>,
{
into.extend(self.tokenizer.encode_ordinary(text.as_ref()));
Ok(())
}
pub fn parse_messages_from_completion_tokens<I>(
&self,
tokens: I,
role: Option<Role>,
) -> anyhow::Result<Vec<Message>>
where
I: IntoIterator<Item = Rank>,
{
let mut parser = StreamableParser::new(self.clone(), role)?;
for token in tokens {
parser.process(token)?;
}
parser.process_eos()?;
Ok(parser.into_messages())
}
fn template_tools_section(
tools: &std::collections::BTreeMap<String, crate::chat::ToolNamespaceConfig>,
) -> String {
let mut tool_sections = Vec::<String>::new();
tool_sections.push("# Tools".to_string());
for ns_config in tools.values() {
let mut tool_section_content = Vec::<String>::new();
tool_section_content.push(format!("## {}\n", ns_config.name));
if let Some(desc) = &ns_config.description {
for line in desc.lines() {
if !ns_config.tools.is_empty() {
tool_section_content.push(format!("// {line}"));
} else {
tool_section_content.push(line.to_string());
}
}
}
if !ns_config.tools.is_empty() {
tool_section_content.push(format!("namespace {} {{\n", ns_config.name));
for tool in &ns_config.tools {
for line in tool.description.lines() {
tool_section_content.push(format!("// {line}"));
}
if let Some(params) = &tool.parameters {
let param_type = Self::json_schema_to_typescript(params, "");
tool_section_content.push(format!(
"type {} = (_: {}) => any;\n",
tool.name, param_type
));
} else {
tool_section_content.push(format!("type {} = () => any;\n", tool.name));
}
}
tool_section_content.push(format!("}} // namespace {}", ns_config.name));
}
tool_sections.push(tool_section_content.join("\n"));
}
tool_sections.join("\n\n")
}
fn json_schema_to_typescript(schema: &serde_json::Value, indent: &str) -> String {
match schema.get("type").and_then(|v| v.as_str()) {
Some("object") => {
let mut out = String::new();
out.push_str("{\n");
if let Some(props) = schema.get("properties") {
if let Some(props_map) = props.as_object() {
let mut required = std::collections::HashSet::new();
if let Some(req) = schema.get("required") {
if let Some(req_arr) = req.as_array() {
for r in req_arr {
if let Some(s) = r.as_str() {
required.insert(s);
}
}
}
}
for (key, val) in props_map {
out.push_str(&format!(
"{}{}{}: ",
indent,
key,
if required.contains(key.as_str()) {
""
} else {
"?"
}
));
let type_str = Self::json_schema_to_typescript(val, &format!("{indent} "));
out.push_str(&type_str);
out.push_str(",\n");
}
}
}
out.push_str(&format!("{indent}}}"));
out
}
Some("string") => "string".to_string(),
Some("number") | Some("integer") => "number".to_string(),
Some("boolean") => "boolean".to_string(),
Some("array") => {
if let Some(items) = schema.get("items") {
format!("{}[]", Self::json_schema_to_typescript(items, indent))
} else {
"Array<any>".to_string()
}
}
_ => "any".to_string(),
}
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct RenderOptions {
pub conversation_has_function_tools: bool,
}
trait Render<T: ?Sized> {
fn render<B>(
&self,
item: &T,
into: &mut B,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>;
}
impl Render<Message> for HarmonyEncoding {
fn render<B>(
&self,
message: &Message,
into: &mut B,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
self.render_formatting_token_into(FormattingToken::Start, into)?;
if matches!(message.author.role, Role::Tool) {
if let Some(name) = &message.author.name {
self.render_text_into(name, into)?;
} else {
anyhow::bail!("Tools should have a name!");
}
} else {
self.render_text_into(message.author.role.as_str(), into)?;
if let Some(name) = &message.author.name {
self.render_text_into(format!(":{name}"), into)?;
}
};
if let Some(recipient) = &message.recipient {
if recipient != "all" {
self.render_text_into(format!(" to={recipient}"), into)?;
}
}
if let Some(channel) = &message.channel {
self.render_formatting_token_into(FormattingToken::Channel, into)?;
self.render_text_into(channel, into)?;
}
if let Some(content_type) = &message.content_type {
if let Some(constrain_marker) =
self.mapped_format_token(FormattingToken::ConstrainedFormat)
{
if let Some(rest) = content_type.strip_prefix(constrain_marker) {
self.render_text_into(" ", into)?;
self.render_formatting_token_into(FormattingToken::ConstrainedFormat, into)?;
if !rest.is_empty() {
self.render_text_into(rest, into)?;
}
} else {
self.render_text_into(format!(" {content_type}"), into)?;
}
} else {
self.render_text_into(format!(" {content_type}"), into)?;
}
}
self.render_formatting_token_into(FormattingToken::Message, into)?;
for content in message.content.iter() {
if let crate::chat::Content::SystemContent(_) = content {
anyhow::ensure!(
message.author.role == crate::chat::Role::System,
"SystemContent may only appear in system messages, found in {:?}",
message.author.role
);
}
if let crate::chat::Content::DeveloperContent(_) = content {
anyhow::ensure!(
message.author.role == crate::chat::Role::Developer,
"DeveloperContent may only appear in developer messages, found in {:?}",
message.author.role
);
}
Render::<Content>::render(self, content, into, render_options)?;
}
if message.author.role == crate::chat::Role::Assistant && message.recipient.is_some() {
self.render_formatting_token_into(FormattingToken::EndMessageAssistantToTool, into)?;
} else {
self.render_formatting_token_into(FormattingToken::EndMessage, into)?;
}
Ok(())
}
}
impl Render<Content> for HarmonyEncoding {
fn render<B>(
&self,
content: &Content,
into: &mut B,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
match content {
Content::Text(text) => Render::<TextContent>::render(self, text, into, render_options),
Content::SystemContent(sys) => {
Render::<SystemContent>::render(self, sys, into, render_options)
}
Content::DeveloperContent(dev) => {
Render::<crate::chat::DeveloperContent>::render(self, dev, into, render_options)
}
}
}
}
impl Render<TextContent> for HarmonyEncoding {
fn render<B>(
&self,
text: &TextContent,
into: &mut B,
_render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
self.render_text_into(&text.text, into)
}
}
impl Render<SystemContent> for HarmonyEncoding {
fn render<B>(
&self,
sys: &SystemContent,
into: &mut B,
render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
let mut sections = Vec::<String>::new();
let mut top_section = Vec::<String>::new();
if let Some(model_id) = &sys.model_identity {
top_section.push(model_id.clone());
}
if let Some(knowledge_cutoff) = &sys.knowledge_cutoff {
top_section.push(format!("Knowledge cutoff: {knowledge_cutoff}"));
}
if let Some(conversation_start_date) = &sys.conversation_start_date {
top_section.push(format!("Current date: {conversation_start_date}"));
}
if !top_section.is_empty() {
sections.push(top_section.join("\n"));
}
let mut instructions_and_reasoning = Vec::<String>::new();
if let Some(effort) = sys.reasoning_effort {
let effort_str = match effort {
ReasoningEffort::Low => "low",
ReasoningEffort::Medium => "medium",
ReasoningEffort::High => "high",
};
instructions_and_reasoning.push(format!("Reasoning: {effort_str}"));
}
if !instructions_and_reasoning.is_empty() {
sections.push(instructions_and_reasoning.join("\n"));
}
if let Some(tools) = &sys.tools {
if !tools.is_empty() {
sections.push(Self::template_tools_section(tools));
}
}
if let Some(channel_config) = &sys.channel_config {
if !channel_config.valid_channels.is_empty() {
let channels_str = channel_config.valid_channels.join(", ");
let mut channels_header = format!("# Valid channels: {channels_str}.");
if channel_config.channel_required {
channels_header.push_str(" Channel must be included for every message.");
}
if render_options.is_some_and(|o| o.conversation_has_function_tools) {
channels_header.push('\n');
channels_header.push_str(
"Calls to these tools must go to the commentary channel: 'functions'.",
);
}
sections.push(channels_header);
}
}
let formatted = sections.join("\n\n");
self.render_text_into(&formatted, into)?;
Ok(())
}
}
impl Render<crate::chat::DeveloperContent> for HarmonyEncoding {
fn render<B>(
&self,
dev: &crate::chat::DeveloperContent,
into: &mut B,
_render_options: Option<&RenderOptions>,
) -> anyhow::Result<()>
where
B: Extend<Rank>,
{
let mut sections = Vec::<String>::new();
if let Some(instr) = &dev.instructions {
sections.push("# Instructions".to_string());
sections.push(instr.clone());
}
if let Some(tools) = &dev.tools {
if !tools.is_empty() {
sections.push(Self::template_tools_section(tools));
}
}
let formatted = sections.join("\n\n");
self.render_text_into(&formatted, into)?;
Ok(())
}
}
pub struct StreamableParser {
encoding: HarmonyEncoding,
next_role: Option<Role>,
tokens: Vec<Rank>,
messages: Vec<Message>,
state: StreamState,
stop_tokens: HashSet<Rank>,
last_content_delta: Option<String>,
undecoded_tokens: Vec<Rank>,
}
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub enum StreamState {
ExpectStart,
Header {
header_tokens: Vec<Rank>,
},
Content {
header: ParsedHeader,
content_tokens: Vec<Rank>,
},
}
impl StreamableParser {
pub fn new(encoding: HarmonyEncoding, role: Option<Role>) -> anyhow::Result<Self> {
let stop_tokens = encoding.stop_tokens()?;
let (state, next_role) = match role {
Some(role) => (
StreamState::Header {
header_tokens: Vec::new(),
},
Some(role),
),
None => (StreamState::ExpectStart, None),
};
Ok(Self {
encoding,
next_role,
tokens: Vec::new(),
messages: Vec::new(),
state,
stop_tokens,
last_content_delta: None,
undecoded_tokens: Vec::new(),
})
}
fn process_next(&mut self, token: Option<Rank>) -> anyhow::Result<&mut Self> {
if let Some(token) = token {
self.tokens.push(token);
}
let next_role_clone = self.next_role.clone();
match &mut self.state {
StreamState::ExpectStart => {
let start = self
.encoding
.render_formatting_token(FormattingToken::Start)?;
match token {
Some(token) if token == start => {
self.state = StreamState::Header {
header_tokens: Vec::new(),
};
}
Some(token) => {
anyhow::bail!(
"Unexpected token {} while expecting start token {}",
token,
start
);
}
None => {
}
}
}
StreamState::Header { header_tokens } => {
let msg_tok = self
.encoding
.render_formatting_token(FormattingToken::Message)?;
match token {
Some(token) if token == msg_tok => {
let header_tokens_cloned = header_tokens.clone();
let next_role_cloned = next_role_clone;
self.state = StreamState::ExpectStart;
let header =
self.parse_header_from_tokens(&header_tokens_cloned, next_role_cloned)?;
self.next_role = None;
self.state = StreamState::Content {
header,
content_tokens: Vec::new(),
};
}
Some(token) => {
header_tokens.push(token);
}
None => {
anyhow::bail!(
"Unexpected EOS while waiting for message header to complete"
);
}
}
}
StreamState::Content {
header,
content_tokens,
} => {
let is_eos = if let Some(token) = token {
if self.stop_tokens.contains(&token) {
true
} else {
self.undecoded_tokens.push(token);
match self
.encoding
.tokenizer()
.decode_utf8(&self.undecoded_tokens)
{
Ok(decoded) => {
content_tokens.extend(self.undecoded_tokens.iter().copied());
self.last_content_delta = Some(decoded);
self.undecoded_tokens.clear();
}
Err(_) => {
self.last_content_delta = None;
}
}
false
}
} else {
true
};
if is_eos {
let text = self.encoding.tokenizer().decode_utf8(content_tokens)?;
let message = Message {
author: header.author.clone(),
recipient: header.recipient.clone(),
channel: header.channel.clone(),
content_type: header.content_type.clone(),
content: vec![Content::Text(TextContent { text })],
};
self.messages.push(message);
self.state = StreamState::ExpectStart;
self.last_content_delta = None;
self.undecoded_tokens.clear();
}
}
}
Ok(self)
}
pub fn process(&mut self, token: Rank) -> anyhow::Result<&mut Self> {
self.process_next(Some(token))
}
pub fn process_eos(&mut self) -> anyhow::Result<&mut Self> {
self.process_next(None)?;
Ok(self)
}
fn parse_header_from_tokens(
&self,
header_tokens: &[Rank],
role: Option<Role>,
) -> anyhow::Result<ParsedHeader> {
let mut header_string = self
.encoding
.tokenizer()
.decode_utf8(header_tokens)
.context("could not decode header")?;
let mut channel: Option<String> = None;
if let Some(channel_marker) = self.encoding.mapped_format_token(FormattingToken::Channel) {
if let Some(idx) = header_string.find(channel_marker) {
let after_marker = &header_string[idx + channel_marker.len()..];
let channel_end = after_marker
.find(|c: char| c.is_whitespace() || c == '<')
.unwrap_or(after_marker.len());
let channel_value = &after_marker[..channel_end];
if channel_value.is_empty() {
anyhow::bail!("channel marker present but no channel value found in header");
}
channel = Some(channel_value.to_string());
let mut new_header = String::new();
new_header.push_str(&header_string[..idx]);
new_header.push_str(&after_marker[channel_end..]);
header_string = new_header;
}
}
header_string = header_string.trim().to_string();
if let Some(constrain_marker) = self
.encoding
.mapped_format_token(FormattingToken::ConstrainedFormat)
{
if header_string.contains(constrain_marker) {
header_string = header_string
.replace(constrain_marker, &format!(" {constrain_marker}"))
.trim()
.to_string();
}
}
let mut parts: Vec<&str> = header_string.split_ascii_whitespace().collect();
let mut role_str_opt: Option<String> = None;
let role = match role {
Some(r) => r,
None => {
let role_str = parts
.first()
.context("message header did not contain a role")?;
role_str_opt = Some((*role_str).to_string());
let parsed_role = Role::try_from(*role_str);
match parsed_role {
Ok(r) => r,
Err(_) => {
if parts.len() > 1 || (parts.len() == 1 && parts[0].starts_with("to=")) {
parts.remove(0);
Role::Tool
} else {
return Err(anyhow::anyhow!("Unknown role: {}", role_str));
}
}
}
}
};
if let Some(&first) = parts.first() {
if first == role.as_str() {
parts.remove(0);
}
}
let mut recipient: Option<String> = None;
let mut content_type: Option<String> = None;
if !parts.is_empty() {
let num_parts = parts.len();
let last_part = parts.pop().unwrap();
if let Some(stripped) = last_part.strip_prefix("to=") {
recipient = Some(stripped.to_string());
} else if num_parts == 1 {
recipient = Some(last_part.to_string());
} else {
content_type = Some(last_part.to_string());
if let Some(raw_recipient) = parts.pop() {
recipient = if let Some(stripped) = raw_recipient.strip_prefix("to=") {
Some(stripped.to_string())
} else {
Some(raw_recipient.to_string())
};
}
}
}
anyhow::ensure!(
parts.is_empty(),
"unexpected tokens remaining in message header: {:?}",
parts
);
let author = if role == Role::Tool {
let name = role_str_opt;
Author { role, name }
} else {
Author { role, name: None }
};
Ok(ParsedHeader {
author,
recipient,
channel,
content_type,
})
}
pub fn current_content(&self) -> anyhow::Result<String> {
match &self.state {
StreamState::Content { content_tokens, .. } => self
.encoding
.tokenizer()
.decode_utf8(content_tokens)
.map_err(|e| anyhow::anyhow!(e)),
_ => Ok(String::new()),
}
}
pub fn current_role(&self) -> Option<Role> {
match &self.state {
StreamState::Content { header, .. } => Some(header.author.role.clone()),
_ => self.next_role.clone(),
}
}
pub fn current_content_type(&self) -> Option<String> {
match &self.state {
StreamState::Content { header, .. } => header.content_type.clone(),
_ => None,
}
}
pub fn last_content_delta(&self) -> anyhow::Result<Option<String>> {
Ok(self.last_content_delta.clone())
}
pub fn into_messages(self) -> Vec<Message> {
self.messages
}
pub fn messages(&self) -> &[Message] {
&self.messages
}
pub fn tokens(&self) -> &[Rank] {
&self.tokens
}
pub fn current_recipient(&self) -> Option<String> {
match &self.state {
StreamState::Content { header, .. } => header.recipient.clone(),
_ => None,
}
}
pub fn current_channel(&self) -> Option<String> {
match &self.state {
StreamState::Content { header, .. } => header.channel.clone(),
_ => None,
}
}
}
#[derive(Clone, Debug)]
pub struct RenderConversationConfig {
pub auto_drop_analysis: bool,
}
impl Default for RenderConversationConfig {
fn default() -> Self {
Self {
auto_drop_analysis: true,
}
}
}