Merge branch 'openai:main' into main

This commit is contained in:
Yehor Smoliakov 2025-08-05 20:23:23 +03:00 committed by GitHub
commit f68e3ccebd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 225 additions and 449 deletions

2
Cargo.lock generated
View file

@ -1317,7 +1317,7 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
[[package]]
name = "openai-harmony"
version = "0.0.1"
version = "0.0.2"
dependencies = [
"anyhow",
"base64",

View file

@ -1,6 +1,6 @@
[package]
name = "openai-harmony"
version = "0.0.1"
version = "0.0.2"
edition = "2021"
license = "Apache-2.0"
repository = "https://github.com/openai/harmony"

View file

@ -86,7 +86,7 @@ convo = Conversation.from_messages([
Message.from_role_and_content(
Role.DEVELOPER,
DeveloperContent.new().with_instructions("Talk like a pirate!")
)
),
Message.from_role_and_content(Role.USER, "Arrr, how be you?"),
])
tokens = enc.render_conversation_for_completion(convo, Role.ASSISTANT)

View file

@ -12,6 +12,8 @@ classifiers = [
]
dynamic = ["version"]
dependencies = ["pydantic>=2.11.7"]
description = "OpenAI's response format for its open-weight model series gpt-oss"
readme = "README.md"
[project.optional-dependencies]
demo = ["uvicorn", "fastapi"]
@ -20,3 +22,7 @@ demo = ["uvicorn", "fastapi"]
features = ["pyo3/extension-module"]
module-name = "openai_harmony"
python-source = "python"
[tool.pytest.ini_options]
# Only collect tests from the top-level tests directory
testpaths = ["tests"]

View file

@ -29,7 +29,7 @@ from typing import (
)
import re
from pydantic import BaseModel, Field, root_validator, validator
from pydantic import BaseModel, Field
# Re-export the low-level Rust bindings under a private name so that we can
# keep the *public* namespace clean and purely Pythonic.
@ -612,6 +612,10 @@ class StreamableParser:
self._inner.process(token)
return self
def process_eos(self) -> "StreamableParser":
self._inner.process_eos()
return self
@property
def current_content(self) -> str:
return self._inner.current_content

View file

@ -217,7 +217,7 @@ impl HarmonyEncoding {
&& first_final_idx.is_some_and(|first| *idx < first)
&& msg.channel.as_deref() == Some("analysis"))
})
.try_for_each(|(_, msg)| self.render_into(*msg, into));
.try_for_each(|(_, msg)| self.render_into(msg, into));
self.conversation_has_function_tools
.store(false, Ordering::Relaxed);
result?;
@ -368,36 +368,27 @@ impl HarmonyEncoding {
pub fn parse_messages_from_completion_tokens<I>(
&self,
tokens: I,
mut role: Option<Role>,
role: Option<Role>,
) -> anyhow::Result<Vec<Message>>
where
I: IntoIterator<Item = Rank>,
{
let mut messages = Vec::<Message>::new();
let mut parser = Parser {
encoding: self,
tokens: tokens.into_iter().peekable(),
};
loop {
let (message, did_reach_end_of_stream) = parser.parse_message(role)?;
messages.push(message);
role = None;
if did_reach_end_of_stream {
break;
}
let mut parser = StreamableParser::new(self.clone(), role)?;
for token in tokens {
parser.process(token)?;
}
anyhow::ensure!(parser.tokens.next().is_none(), "Expected end of stream");
Ok(messages)
parser.process_eos()?;
Ok(parser.into_messages())
}
/// Helper to convert a JSON schema (OpenAPI style) to a TypeScript type definition.
fn json_schema_to_typescript(schema: &serde_json::Value, indent: &str) -> String {
// Helper to check if this schema is an enum
fn is_enum(schema: &serde_json::Value) -> bool {
return schema
schema
.get("enum")
.and_then(|e| e.as_array())
.map_or(false, |arr| !arr.is_empty());
.is_some_and(|arr| !arr.is_empty())
}
// Handle oneOf at the top level
@ -407,30 +398,29 @@ impl HarmonyEncoding {
let mut first = true;
for variant in arr {
if !first {
out.push_str("\n");
out.push_str(&format!("{} | ", indent));
out.push('\n');
out.push_str(&format!("{indent} | "));
} else {
out.push_str(&format!("\n{} | ", indent));
out.push_str(&format!("\n{indent} | "));
first = false;
}
let type_str =
Self::json_schema_to_typescript(variant, &format!("{} ", indent));
Self::json_schema_to_typescript(variant, &format!("{indent} "));
let mut type_str = type_str;
if variant
.get("nullable")
.and_then(|n| n.as_bool())
.unwrap_or(false)
&& !type_str.contains("null")
{
if !type_str.contains("null") {
type_str = format!("{} | null", type_str);
}
type_str = format!("{type_str} | null");
}
out.push_str(&type_str);
// Add trailing comments (description, default)
let mut trailing_comments = Vec::new();
if let Some(desc) = variant.get("description") {
if let Some(desc_str) = desc.as_str() {
trailing_comments.push(format!("{}", desc_str));
trailing_comments.push(desc_str.to_string());
}
}
if let Some(default) = variant.get("default") {
@ -438,7 +428,7 @@ impl HarmonyEncoding {
trailing_comments
.push(format!("default: \"{}\"", default.as_str().unwrap()));
} else {
trailing_comments.push(format!("default: {}", default));
trailing_comments.push(format!("default: {default}"));
}
}
if !trailing_comments.is_empty() {
@ -472,7 +462,7 @@ impl HarmonyEncoding {
// Render object-level description as comment
if let Some(desc) = schema.get("description") {
if let Some(desc_str) = desc.as_str() {
out.push_str(&format!("{}// {}\n", indent, desc_str));
out.push_str(&format!("{indent}// {desc_str}\n"));
}
}
out.push_str("{\n");
@ -495,8 +485,7 @@ impl HarmonyEncoding {
if let Some(title) = val.get("title") {
if let Some(title_str) = title.as_str() {
out.push_str(&format!(
"{0}// {1}\n{0}//\n",
indent, title_str
"{indent}// {title_str}\n{indent}//\n"
));
}
}
@ -504,19 +493,18 @@ impl HarmonyEncoding {
if val.get("oneOf").is_none() {
if let Some(desc) = val.get("description") {
if let Some(desc_str) = desc.as_str() {
out.push_str(&format!("{}// {}\n", indent, desc_str));
out.push_str(&format!("{indent}// {desc_str}\n"));
}
}
}
if let Some(examples) = val.get("examples") {
if let Some(arr) = examples.as_array() {
if !arr.is_empty() {
out.push_str(&format!("{}// Examples:\n", indent));
out.push_str(&format!("{indent}// Examples:\n"));
for ex in arr {
if let Some(ex_str) = ex.as_str() {
out.push_str(&format!(
"{}// - \"{}\"\n",
indent, ex_str
"{indent}// - \"{ex_str}\"\n"
));
}
}
@ -535,7 +523,7 @@ impl HarmonyEncoding {
}
let mut skip_property_desc = false;
if let Some(desc_str) = property_desc {
if let Some(first_variant) = arr.get(0) {
if let Some(first_variant) = arr.first() {
if let Some(variant_desc) =
first_variant.get("description")
{
@ -553,10 +541,7 @@ impl HarmonyEncoding {
let mut rendered_property_desc_above = false;
if !skip_property_desc {
if let Some(desc_str) = property_desc {
out.push_str(&format!(
"{}// {}\n",
indent, desc_str
));
out.push_str(&format!("{indent}// {desc_str}\n"));
rendered_property_desc_above = true;
}
}
@ -575,8 +560,7 @@ impl HarmonyEncoding {
));
} else {
out.push_str(&format!(
"{}// default: {}\n",
indent, default
"{indent}// default: {default}\n"
));
}
}
@ -593,10 +577,10 @@ impl HarmonyEncoding {
));
// Render each variant
for (i, variant) in arr.iter().enumerate() {
out.push_str(&format!("{} | ", indent));
out.push_str(&format!("{indent} | "));
let type_str = Self::json_schema_to_typescript(
variant,
&format!("{} ", indent),
&format!("{indent} "),
);
// Handle nullable in variant
let mut type_str = type_str;
@ -604,10 +588,9 @@ impl HarmonyEncoding {
.get("nullable")
.and_then(|n| n.as_bool())
.unwrap_or(false)
&& !type_str.contains("null")
{
if !type_str.contains("null") {
type_str = format!("{} | null", type_str);
}
type_str = format!("{type_str} | null");
}
out.push_str(&type_str);
// Add variant-level comments after the type
@ -619,7 +602,7 @@ impl HarmonyEncoding {
// Only render if not equal to property-level description
if Some(desc_str) != property_desc {
trailing_comments
.push(format!("{}", desc_str));
.push(desc_str.to_string());
}
}
}
@ -636,7 +619,7 @@ impl HarmonyEncoding {
));
} else {
trailing_comments
.push(format!("default: {}", default));
.push(format!("default: {default}"));
}
}
if !trailing_comments.is_empty() {
@ -645,9 +628,9 @@ impl HarmonyEncoding {
trailing_comments.join(" ")
));
}
out.push_str("\n");
out.push('\n');
}
out.push_str(&format!("{},\n", indent));
out.push_str(&format!("{indent},\n"));
continue;
}
}
@ -663,21 +646,18 @@ impl HarmonyEncoding {
}
));
// Handle nullable
let mut type_str = Self::json_schema_to_typescript(
val,
&format!("{} ", indent),
);
let mut type_str =
Self::json_schema_to_typescript(val, &format!("{indent} "));
if val
.get("nullable")
.and_then(|n| n.as_bool())
.unwrap_or(false)
&& !type_str.contains("null")
{
if !type_str.contains("null") {
type_str = format!("{} | null", type_str);
}
type_str = format!("{type_str} | null");
}
out.push_str(&type_str);
out.push_str(",");
out.push(',');
// Add default as comment if present (and not already handled)
if val.get("oneOf").is_none() {
if let Some(default) = val.get("default") {
@ -692,15 +672,15 @@ impl HarmonyEncoding {
default.as_str().unwrap()
));
} else {
out.push_str(&format!(" // default: {}", default));
out.push_str(&format!(" // default: {default}"));
}
}
}
out.push_str("\n");
out.push('\n');
}
}
}
out.push_str(&format!("{}}}", indent));
out.push_str(&format!("{indent}}}"));
out
}
"string" => {
@ -708,7 +688,7 @@ impl HarmonyEncoding {
if let Some(arr) = enum_vals.as_array() {
let enums: Vec<String> = arr
.iter()
.filter_map(|v| v.as_str().map(|s| format!("\"{}\"", s)))
.filter_map(|v| v.as_str().map(|s| format!("\"{s}\"")))
.collect();
if !enums.is_empty() {
return enums.join(" | ");
@ -756,13 +736,13 @@ impl HarmonyEncoding {
) -> String {
let mut tool_sections = Vec::<String>::new();
tool_sections.push("# Tools".to_string());
for (_namespace, ns_config) in tools {
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));
tool_section_content.push(format!("// {line}"));
} else {
tool_section_content.push(line.to_string());
}
@ -772,7 +752,7 @@ impl HarmonyEncoding {
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));
tool_section_content.push(format!("// {line}"));
}
if let Some(params) = &tool.parameters {
let param_type = Self::json_schema_to_typescript(params, "");
@ -817,14 +797,14 @@ impl Render<Message> for HarmonyEncoding {
// For users and assistants we put both the role, and optionally the user name.
self.render_text_into(message.author.role.as_str(), into)?;
if let Some(name) = &message.author.name {
self.render_text_into(format!(":{}", name), into)?;
self.render_text_into(format!(":{name}"), into)?;
}
};
// next render the header recipient, if there is one
if let Some(recipient) = &message.recipient {
if recipient != "all" {
self.render_text_into(format!(" to={}", recipient), into)?;
self.render_text_into(format!(" to={recipient}"), into)?;
}
}
@ -836,7 +816,7 @@ impl Render<Message> for HarmonyEncoding {
// finally content type
if let Some(content_type) = &message.content_type {
self.render_text_into(format!(" {}", content_type), into)?;
self.render_text_into(format!(" {content_type}"), into)?;
}
self.render_formatting_token_into(FormattingToken::Message, into)?;
@ -944,7 +924,7 @@ impl Render<SystemContent> for HarmonyEncoding {
channels_header.push_str(" Channel must be included for every message.");
}
if self.conversation_has_function_tools.load(Ordering::Relaxed) {
channels_header.push_str("\n");
channels_header.push('\n');
channels_header.push_str(
"Calls to these tools must go to the commentary channel: 'functions'.",
);
@ -982,305 +962,6 @@ impl Render<crate::chat::DeveloperContent> for HarmonyEncoding {
}
}
enum TakeUntilStatus {
Found,
EndOfStream,
}
impl TakeUntilStatus {
fn was_found(&self) -> bool {
matches!(self, TakeUntilStatus::Found)
}
}
struct Parser<'a, I>
where
I: Iterator<Item = Rank>,
{
tokens: std::iter::Peekable<I>,
encoding: &'a HarmonyEncoding,
}
impl<I> Parser<'_, I>
where
I: Iterator<Item = Rank>,
{
fn expect_special(&mut self, token: FormattingToken) -> anyhow::Result<Rank> {
let next = self.tokens.next().context(format!(
"Expected special token ({}), but out of tokens",
token
))?;
let expected = self.encoding.render_formatting_token(token)?;
if next != expected {
anyhow::bail!(
"Expected special token ({}) {} but got {}",
token,
expected,
next,
);
}
Ok(next)
}
fn take_until_any(&mut self, ends: &HashSet<Rank>) -> (Vec<Rank>, TakeUntilStatus) {
let mut out = vec![];
for t in &mut self.tokens {
if ends.contains(&t) {
return (out, TakeUntilStatus::Found);
}
out.push(t);
}
(out, TakeUntilStatus::EndOfStream)
}
fn take_until(&mut self, end: Rank) -> (Vec<Rank>, TakeUntilStatus) {
self.take_until_any(&HashSet::from([end]))
}
fn parse_header(&mut self, role: Option<Role>) -> anyhow::Result<ParsedHeader> {
// FormattingToken::Message marks the end of the header.
// Everything before that belongs to the header.
let message_start_token = self
.encoding
.render_formatting_token(FormattingToken::Message)?;
let (header_tokens, status) = self.take_until(message_start_token);
if !status.was_found() {
anyhow::bail!("Expected message start token but ran out of tokens");
}
// Decode the header into a UTF-8 string so we can reason about its structure.
let mut header_string = self
.encoding
.tokenizer
.decode_utf8(header_tokens)
.context("could not decode header")?;
// --------------------------------------------------------------------
// 1. Extract the channel (if any)
// --------------------------------------------------------------------
// A channel, when present, is encoded as:
// <|channel|>CHANNEL_VALUE
// where <|channel|> is the literal rendering of FormattingToken::Channel
// and CHANNEL_VALUE is a contiguous string (no whitespace) naming the
// channel. The <|channel|> marker can appear before or after the
// recipient part, but *always* before the optional content-type (which
// must be last).
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) {
// Slice parts around the marker
let after_marker = &header_string[idx + channel_marker.len()..];
// The channel value continues until the next ASCII whitespace,
// the start of another special token ("<"), or end-of-string.
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());
// Remove the marker *and* the channel value from the header so
// the remaining pieces can be parsed independently.
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;
}
}
// Trim extraneous whitespace that may have been introduced when we
// removed the channel section.
header_string = header_string.trim().to_string();
// If the constrained format marker is present but not preceded by
// whitespace (e.g. "to=foo<|constrain|>json"), insert a space before
// the marker so that splitting on whitespace treats the content type
// as a separate token.
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();
}
}
// --------------------------------------------------------------------
// 2. Split the remaining header into whitespace-separated tokens.
// --------------------------------------------------------------------
// Debug output for development (only active when the `debug_header_parsing` cfg flag is
// enabled).
// For debugging purposes one might want to inspect the header string
// at this point. To avoid unwanted stdout noise in production use
// the following (commented) line and recompile as needed.
// println!("[DEBUG header] '{}'", header_string);
let mut parts: Vec<&str> = header_string.split_ascii_whitespace().collect();
// --------------------------------------------------------------------
// 3. Determine the role (if not already provided).
// --------------------------------------------------------------------
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);
let out = match parsed_role {
Ok(r) => r,
Err(_) => {
// If recipient is present, treat as tool call
if parts.len() > 1 || (parts.len() == 1 && parts[0].starts_with("to=")) {
parts.remove(0); // Remove the unknown role string
Role::Tool
} else {
return Err(anyhow::anyhow!("Unknown role: {}", role_str));
}
}
};
out
}
};
// If the role was supplied externally but also redundantly present in the
// header itself, strip it off so that it does not interfere with the
// parsing of the remaining fields.
if let Some(first) = parts.first() {
if *first == role.as_str() {
parts.remove(0);
}
}
// --------------------------------------------------------------------
// 4. Identify recipient and content-type.
// --------------------------------------------------------------------
let mut recipient: Option<String> = None;
let mut content_type: Option<String> = None;
if !parts.is_empty() {
// Determine whether the last token is a content-type or part of the
// recipient specification.
let num_tokens_before_pop = parts.len();
let last_token_owned = parts.pop().unwrap().to_string();
if last_token_owned.starts_with("to=") {
// The header contains a recipient but *no* content-type.
recipient = Some(last_token_owned.trim_start_matches("to=").to_string());
} else if num_tokens_before_pop == 1 {
// Only one token total (after potential role removal) and it doesn't start
// with "to=" => interpret it as a standalone recipient.
recipient = Some(last_token_owned);
} else {
// More than one token and the last one is not a recipient -> treat as content-type.
content_type = Some(last_token_owned);
// After removing the content-type there may be exactly one token describing the recipient.
if !parts.is_empty() {
if parts.len() != 1 {
anyhow::bail!("Could not parse header: too many tokens remaining after extracting content-type and recipient");
}
let raw_recipient = parts.pop().unwrap();
recipient = if let Some(stripped) = raw_recipient.strip_prefix("to=") {
Some(stripped.to_string())
} else {
Some(raw_recipient.to_string())
};
}
}
}
// After processing, no unparsed tokens should remain.
anyhow::ensure!(
parts.is_empty(),
"unexpected tokens remaining in message header: {:?}",
parts
);
// We have successfully parsed the header.
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,
})
}
/// Parses a message from the remaining tokens.
///
/// Returns the message and a boolean indicating whether end of stream was reached.
fn parse_message(&mut self, role: Option<Role>) -> anyhow::Result<(Message, bool)> {
let start_token = self
.encoding
.render_formatting_token(FormattingToken::Start)?;
match role {
Some(_) => {
if let Some(&next) = self.tokens.peek() {
if next == start_token {
self.tokens.next();
}
} else {
anyhow::bail!("Expected at least one token while parsing message");
}
}
None => {
self.expect_special(FormattingToken::Start)?;
}
}
let header = self.parse_header(role)?;
let ParsedHeader {
author,
recipient,
channel,
content_type,
} = header;
// TODO other content types
// since we bail on anything other than just the role in the header for now, we can assume
// that the content type is text
let end_tokens = self.encoding.stop_tokens()?;
let (remaining_tokens, status) = self.take_until_any(&end_tokens);
let remaining_text = self
.encoding
.tokenizer
.decode_utf8(remaining_tokens)
.context("could not decode message content")?;
let did_reach_end_of_stream = match status {
TakeUntilStatus::Found => self.tokens.peek().is_none(),
TakeUntilStatus::EndOfStream => true,
};
Ok((
Message {
author,
content: vec![Content::Text(TextContent {
text: remaining_text,
})],
channel,
recipient,
content_type,
},
did_reach_end_of_stream,
))
}
}
// ---------------------------------------------------------------------------
// Streamable parsing ---------------------------------------------------------
// ---------------------------------------------------------------------------
/// Incremental parser that can consume tokens one by one.
///
/// It keeps track of all tokens seen so far, exposes all fully parsed messages
@ -1334,8 +1015,11 @@ impl StreamableParser {
}
/// Consume a single token and update the internal state.
pub fn process(&mut self, token: Rank) -> anyhow::Result<&mut Self> {
self.tokens.push(token);
/// Consume a single token and update the internal state.
fn process_next(&mut self, token: Option<Rank>) -> anyhow::Result<&mut Self> {
if let Some(token) = token {
self.tokens.push(token);
}
// Clone next_role up front to avoid borrow checker issues
let next_role_clone = self.next_role.clone();
match &mut self.state {
@ -1343,44 +1027,89 @@ impl StreamableParser {
let start = self
.encoding
.render_formatting_token(FormattingToken::Start)?;
if token == start {
self.state = StreamState::Header {
header_tokens: Vec::new(),
};
} else {
anyhow::bail!(
"Unexpected token {} while expecting start token {}",
token,
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 => {
// receiving EOS while waiting for start token is actually fine
// as we may have just parsed a stop token. in this case we can
// simple keep state as is
}
}
}
StreamState::Header { header_tokens } => {
let msg_tok = self
.encoding
.render_formatting_token(FormattingToken::Message)?;
if token == msg_tok {
// Clone the tokens and next_role, then clear the state before parsing
let header_tokens_cloned = header_tokens.clone();
let next_role_cloned = next_role_clone;
// Set state to dummy to drop mutable borrow
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(),
};
} else {
header_tokens.push(token);
match token {
Some(token) if token == msg_tok => {
// Clone the tokens and next_role, then clear the state before parsing
let header_tokens_cloned = header_tokens.clone();
let next_role_cloned = next_role_clone;
// Set state to dummy to drop mutable borrow
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,
} => {
if self.stop_tokens.contains(&token) {
let is_eos = if let Some(token) = token {
if self.stop_tokens.contains(&token) {
// this is a stop token, dont parse and mark EOS
true
} else {
self.undecoded_tokens.push(token);
// some tokens might not appropriately decode on their own. If they don't
// we will collect them until they eventually decode
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;
}
}
// this was not an EOS
false
}
} else {
// token = None signals EOS to this function
true
};
if is_eos {
let text = self.encoding.tokenizer().decode_utf8(content_tokens)?;
let message = Message {
author: header.author.clone(),
@ -1393,30 +1122,21 @@ impl StreamableParser {
self.state = StreamState::ExpectStart;
self.last_content_delta = None;
self.undecoded_tokens.clear();
} else {
self.undecoded_tokens.push(token);
// some tokens might not appropriately decode on their own. If they don't
// we will collect them until they eventually decode
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;
}
}
}
}
}
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],
@ -1495,8 +1215,8 @@ impl StreamableParser {
}
};
if let Some(first) = parts.first() {
if *first == role.as_str() {
if let Some(&first) = parts.first() {
if first == role.as_str() {
parts.remove(0);
}
}
@ -1505,23 +1225,25 @@ impl StreamableParser {
let mut content_type: Option<String> = None;
if !parts.is_empty() {
let num_tokens_before_pop = parts.len();
let last_token_owned = parts.pop().unwrap().to_string();
// Determine whether the last token is a content-type or part of the
// recipient specification.
let num_parts = parts.len();
// SAFETY: we know that there is at least one part remaining, because of is_empty check above
let last_part = parts.pop().unwrap();
if last_token_owned.starts_with("to=") {
recipient = Some(last_token_owned.trim_start_matches("to=").to_string());
} else if num_tokens_before_pop == 1 {
recipient = Some(last_token_owned);
if let Some(stripped) = last_part.strip_prefix("to=") {
// The header contains a recipient but *no* content-type.
recipient = Some(stripped.to_string());
} else if num_parts == 1 {
// Only one part total (after potential role removal) and it doesn't start
// with "to=" => interpret it as a standalone recipient.
recipient = Some(last_part.to_string());
} else {
content_type = Some(last_token_owned);
// More than one token and the last one is not a recipient -> treat as content-type.
content_type = Some(last_part.to_string());
if !parts.is_empty() {
if parts.len() != 1 {
anyhow::bail!(
"Could not parse header: too many tokens remaining after extracting content-type and recipient"
);
}
let raw_recipient = parts.pop().unwrap();
// After removing the content-type there may be exactly one token describing the recipient.
if let Some(raw_recipient) = parts.pop() {
recipient = if let Some(stripped) = raw_recipient.strip_prefix("to=") {
Some(stripped.to_string())
} else {
@ -1530,12 +1252,12 @@ impl StreamableParser {
}
}
}
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 }
@ -1583,6 +1305,11 @@ impl StreamableParser {
Ok(self.last_content_delta.clone())
}
/// Consume the parser and return all parsed messages.
pub fn into_messages(self) -> Vec<Message> {
self.messages
}
/// All fully parsed messages so far.
pub fn messages(&self) -> &[Message] {
&self.messages

View file

@ -306,6 +306,13 @@ impl PyStreamableParser {
.map_err(|e| PyErr::new::<HarmonyError, _>(e.to_string()))
}
fn process_eos(&mut self) -> PyResult<()> {
self.inner
.process_eos()
.map(|_| ())
.map_err(|e| PyErr::new::<HarmonyError, _>(e.to_string()))
}
#[getter]
fn current_content(&self) -> PyResult<String> {
self.inner

View file

@ -501,11 +501,11 @@ fn test_preserve_cot() {
fn test_reserved_token_decoding() {
let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap();
assert_eq!(
encoding.tokenizer.decode_utf8(&[200014]).unwrap(),
encoding.tokenizer.decode_utf8([200014]).unwrap(),
"<|reserved_200014|>"
);
assert_eq!(
encoding.tokenizer.decode_utf8(&[201088]).unwrap(),
encoding.tokenizer.decode_utf8([201088]).unwrap(),
"<|reserved_201088|>"
);
}
@ -527,7 +527,7 @@ fn test_render_and_render_conversation_roundtrip() {
#[test]
fn test_decode_utf8_invalid_token() {
let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap();
let result = encoding.tokenizer.decode_utf8(&[99999999]);
let result = encoding.tokenizer.decode_utf8([99999999]);
assert!(result.is_err(), "Expected error for invalid token");
}
@ -663,3 +663,39 @@ fn test_streamable_parser_tool_call_with_constrain_adjacent() {
parser.messages()[0]
);
}
#[test]
fn test_tool_call_with_constrain_marker_adjacent() {
let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap();
let text = "<|start|>assistant to=functions.get_weather<|channel|>commentary<|constrain|>json<|message|>{\"location\": \"Tokyo\"}<|end|>";
let tokens = encoding.tokenizer().encode_with_special_tokens(text);
let parsed = encoding
.parse_messages_from_completion_tokens(tokens, None)
.expect("expected to parse");
let expected =
vec![
Message::from_role_and_content(Role::Assistant, "{\"location\": \"Tokyo\"}")
.with_channel("commentary")
.with_recipient("functions.get_weather")
.with_content_type("<|constrain|>json"),
];
assert_eq!(parsed, expected);
}
#[test]
fn test_tool_call_with_channel_before_recipient_and_constrain_adjacent() {
let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap();
let text = "<|start|>assistant<|channel|>commentary to=functions.get_weather<|constrain|>json<|message|>{\"latitude\":48.8566,\"longitude\":2.3522}<|call|>";
let tokens = encoding.tokenizer().encode_with_special_tokens(text);
let parsed = encoding
.parse_messages_from_completion_tokens(tokens, None)
.expect("expected to parse");
let expected = vec![Message::from_role_and_content(
Role::Assistant,
"{\"latitude\":48.8566,\"longitude\":2.3522}",
)
.with_channel("commentary")
.with_recipient("functions.get_weather")
.with_content_type("<|constrain|>json")];
assert_eq!(parsed, expected);
}

5
test_python.sh Executable file
View file

@ -0,0 +1,5 @@
#!/usr/bin/env bash
set -e
source .venv/bin/activate
maturin develop -F python-binding --release
pytest "$@"

View file

@ -34,7 +34,6 @@ from openai_harmony import ( # noqa: E402
StreamableParser,
SystemContent,
ToolDescription,
ToolNamespaceConfig,
load_harmony_encoding,
)
from pydantic import ValidationError
@ -245,23 +244,18 @@ def test_tool_call_with_constrain_marker_adjacent(encoding_name):
correctly and instead handle it as a separate content type.
"""
encoding = load_harmony_encoding(encoding_name)
text = (
"<|start|>assistant to=functions.get_weather<|channel|>commentary"
'<|constrain|>json<|message|>{"location": "Tokyo"}<|end|>'
)
tokens = encoding.encode(text, allowed_special="all")
parsed = encoding.parse_messages_from_completion_tokens(tokens, role=Role.ASSISTANT)
parsed = encoding.parse_messages_from_completion_tokens(tokens, role=None)
expected = [
Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}')
.with_channel("commentary")
.with_recipient("functions.get_weather")
.with_content_type("<|constrain|>json"),
]
assert parsed == expected
@ -280,11 +274,8 @@ def test_tool_call_with_channel_before_recipient_and_constrain_adjacent(
"<|start|>assistant<|channel|>commentary to=functions.get_weather"
'<|constrain|>json<|message|>{"latitude":48.8566,"longitude":2.3522}<|call|>'
)
tokens = encoding.encode(text, allowed_special="all")
parsed = encoding.parse_messages_from_completion_tokens(tokens, role=Role.ASSISTANT)
parsed = encoding.parse_messages_from_completion_tokens(tokens, role=None)
expected = [
Message.from_role_and_content(
Role.ASSISTANT, '{"latitude":48.8566,"longitude":2.3522}'