valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve gpt-oss compatibility (#2461)

Browse Source 
The gpt-oss models require reasoning with subsequent Chat Completions
requests because otherwise the model forgets why the tools were called.
This change fixes that and also adds some additional missing
documentation around how to handle context windows in Ollama and how to
show the CoT if you desire to.
This commit is contained in:
Dominik Kundel
2025-09-02 19:49:03 -07:00
committed by GitHub
parent a93a907c7e
commit b127a3643f
5 changed files with 886 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
codex-rs/core/src/chat_completions.rs
View File

@@ -43,7 +43,107 @@ pub(crate) async fn stream_chat_completions(
let input = prompt.get_formatted_input();
// Pre-scan: map Reasoning blocks to the adjacent assistant anchor after the last user.
// - If the last emitted message is a user message, drop all reasoning.
// - Otherwise, for each Reasoning item after the last user message, attach it
// to the immediate previous assistant message (stop turns) or the immediate
// next assistant anchor (tool-call turns: function/local shell call, or assistant message).
let mut reasoning_by_anchor_index: std::collections::HashMap<usize, String> =
std::collections::HashMap::new();
// Determine the last role that would be emitted to Chat Completions.
let mut last_emitted_role: Option<&str> = None;
for item in &input {
match item {
ResponseItem::Message { role, .. } => last_emitted_role = Some(role.as_str()),
ResponseItem::FunctionCall { .. } | ResponseItem::LocalShellCall { .. } => {
last_emitted_role = Some("assistant")
}
ResponseItem::FunctionCallOutput { .. } => last_emitted_role = Some("tool"),
ResponseItem::Reasoning { .. } | ResponseItem::Other => {}
ResponseItem::CustomToolCall { .. } => {}
ResponseItem::CustomToolCallOutput { .. } => {}
ResponseItem::WebSearchCall { .. } => {}
}
}
// Find the last user message index in the input.
let mut last_user_index: Option<usize> = None;
for (idx, item) in input.iter().enumerate() {
if let ResponseItem::Message { role, .. } = item
&& role == "user"
{
last_user_index = Some(idx);
}
}
// Attach reasoning only if the conversation does not end with a user message.
if !matches!(last_emitted_role, Some("user")) {
for (idx, item) in input.iter().enumerate() {
// Only consider reasoning that appears after the last user message.
if let Some(u_idx) = last_user_index
&& idx <= u_idx
{
continue;
}
if let ResponseItem::Reasoning {
content: Some(items),
..
} = item
{
let mut text = String::new();
for c in items {
match c {
ReasoningItemContent::ReasoningText { text: t }
| ReasoningItemContent::Text { text: t } => text.push_str(t),
}
}
if text.trim().is_empty() {
continue;
}
// Prefer immediate previous assistant message (stop turns)
let mut attached = false;
if idx > 0
&& let ResponseItem::Message { role, .. } = &input[idx - 1]
&& role == "assistant"
{
reasoning_by_anchor_index
.entry(idx - 1)
.and_modify(|v| v.push_str(&text))
.or_insert(text.clone());
attached = true;
}
// Otherwise, attach to immediate next assistant anchor (tool-calls or assistant message)
if !attached && idx + 1 < input.len() {
match &input[idx + 1] {
ResponseItem::FunctionCall { .. } | ResponseItem::LocalShellCall { .. } => {
reasoning_by_anchor_index
.entry(idx + 1)
.and_modify(|v| v.push_str(&text))
.or_insert(text.clone());
}
ResponseItem::Message { role, .. } if role == "assistant" => {
reasoning_by_anchor_index
.entry(idx + 1)
.and_modify(|v| v.push_str(&text))
.or_insert(text.clone());
}
_ => {}
}
}
}
}
}
// Track last assistant text we emitted to avoid duplicate assistant messages
// in the outbound Chat Completions payload (can happen if a final
// aggregated assistant message was recorded alongside an earlier partial).
let mut last_assistant_text: Option<String> = None;
for (idx, item) in input.iter().enumerate() {
match item {
ResponseItem::Message { role, content, .. } => {
let mut text = String::new();
@@ -56,7 +156,24 @@ pub(crate) async fn stream_chat_completions(
_ => {}
}
}
messages.push(json!({"role": role, "content": text}));
// Skip exact-duplicate assistant messages.
if role == "assistant" {
if let Some(prev) = &last_assistant_text
&& prev == &text
{
continue;
}
last_assistant_text = Some(text.clone());
}
let mut msg = json!({"role": role, "content": text});
if role == "assistant"
&& let Some(reasoning) = reasoning_by_anchor_index.get(&idx)
&& let Some(obj) = msg.as_object_mut()
{
obj.insert("reasoning".to_string(), json!(reasoning));
}
messages.push(msg);
}
ResponseItem::FunctionCall {
name,
@@ -64,7 +181,7 @@ pub(crate) async fn stream_chat_completions(
call_id,
..
} => {
messages.push(json!({
let mut msg = json!({
"role": "assistant",
"content": null,
"tool_calls": [{
@@ -75,7 +192,13 @@ pub(crate) async fn stream_chat_completions(
"arguments": arguments,
}
}]
}));
});
if let Some(reasoning) = reasoning_by_anchor_index.get(&idx)
&& let Some(obj) = msg.as_object_mut()
{
obj.insert("reasoning".to_string(), json!(reasoning));
}
messages.push(msg);
}
ResponseItem::LocalShellCall {
id,
@@ -84,7 +207,7 @@ pub(crate) async fn stream_chat_completions(
action,
} => {
// Confirm with API team.
messages.push(json!({
let mut msg = json!({
"role": "assistant",
"content": null,
"tool_calls": [{
@@ -93,7 +216,13 @@ pub(crate) async fn stream_chat_completions(
"status": status,
"action": action,
}]
}));
});
if let Some(reasoning) = reasoning_by_anchor_index.get(&idx)
&& let Some(obj) = msg.as_object_mut()
{
obj.insert("reasoning".to_string(), json!(reasoning));
}
messages.push(msg);
}
ResponseItem::FunctionCallOutput { call_id, output } => {
messages.push(json!({
@@ -331,7 +460,10 @@ async fn process_chat_sse<S>(
// Some providers stream `reasoning` as a plain string while others
// nest the text under an object (e.g. `{ "reasoning": { "text": "…" } }`).
if let Some(reasoning_val) = choice.get("delta").and_then(|d| d.get("reasoning")) {
let mut maybe_text = reasoning_val.as_str().map(|s| s.to_string());
let mut maybe_text = reasoning_val
.as_str()
.map(|s| s.to_string())
.filter(|s| !s.is_empty());
if maybe_text.is_none() && reasoning_val.is_object() {
if let Some(s) = reasoning_val
@@ -350,12 +482,39 @@ async fn process_chat_sse<S>(
}
if let Some(reasoning) = maybe_text {
// Accumulate so we can emit a terminal Reasoning item at the end.
reasoning_text.push_str(&reasoning);
let _ = tx_event
.send(Ok(ResponseEvent::ReasoningContentDelta(reasoning)))
.await;
}
}
// Some providers only include reasoning on the final message object.
if let Some(message_reasoning) = choice.get("message").and_then(|m| m.get("reasoning"))
{
// Accept either a plain string or an object with { text | content }
if let Some(s) = message_reasoning.as_str() {
if !s.is_empty() {
reasoning_text.push_str(s);
let _ = tx_event
.send(Ok(ResponseEvent::ReasoningContentDelta(s.to_string())))
.await;
}
} else if let Some(obj) = message_reasoning.as_object()
&& let Some(s) = obj
.get("text")
.and_then(|v| v.as_str())
.or_else(|| obj.get("content").and_then(|v| v.as_str()))
&& !s.is_empty()
{
reasoning_text.push_str(s);
let _ = tx_event
.send(Ok(ResponseEvent::ReasoningContentDelta(s.to_string())))
.await;
}
}
// Handle streaming function / tool calls.
if let Some(tool_calls) = choice
.get("delta")
@@ -511,27 +670,47 @@ where
// do NOT emit yet. Forward any other item (e.g. FunctionCall) right
// away so downstream consumers see it.
let is_assistant_delta = matches!(&item, codex_protocol::models::ResponseItem::Message { role, .. } if role == "assistant");
let is_assistant_message = matches!(
&item,
codex_protocol::models::ResponseItem::Message { role, .. } if role == "assistant"
);
if is_assistant_delta {
// Only use the final assistant message if we have not
// seen any deltas; otherwise, deltas already built the
// cumulative text and this would duplicate it.
if this.cumulative.is_empty()
&& let codex_protocol::models::ResponseItem::Message { content, .. } =
&item
&& let Some(text) = content.iter().find_map(|c| match c {
codex_protocol::models::ContentItem::OutputText { text } => {
Some(text)
if is_assistant_message {
match this.mode {
AggregateMode::AggregatedOnly => {
// Only use the final assistant message if we have not
// seen any deltas; otherwise, deltas already built the
// cumulative text and this would duplicate it.
if this.cumulative.is_empty()
&& let codex_protocol::models::ResponseItem::Message {
content,
..
} = &item
&& let Some(text) = content.iter().find_map(|c| match c {
codex_protocol::models::ContentItem::OutputText {
text,
} => Some(text),
_ => None,
})
{
this.cumulative.push_str(text);
}
_ => None,
})
{
this.cumulative.push_str(text);
// Swallow assistant message here; emit on Completed.
continue;
}
AggregateMode::Streaming => {
// In streaming mode, if we have not seen any deltas, forward
// the final assistant message directly. If deltas were seen,
// suppress the final message to avoid duplication.
if this.cumulative.is_empty() {
return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(
item,
))));
} else {
continue;
}
}
}
// Swallow assistant message here; emit on Completed.
continue;
}
// Not an assistant message – forward immediately.
@@ -563,6 +742,11 @@ where
emitted_any = true;
}
// Always emit the final aggregated assistant message when any
// content deltas have been observed. In AggregatedOnly mode this
// is the sole assistant output; in Streaming mode this finalizes
// the streamed deltas into a terminal OutputItemDone so callers
// can persist/render the message once per turn.
if !this.cumulative.is_empty() {
let aggregated_message = codex_protocol::models::ResponseItem::Message {
id: None,

4
codex-rs/core/src/client_common.rs
View File

@@ -35,7 +35,7 @@ pub struct Prompt {
/// Tools available to the model, including additional tools sourced from
/// external MCP servers.
pub tools: Vec<OpenAiTool>,
pub(crate) tools: Vec<OpenAiTool>,
/// Optional override for the built-in BASE_INSTRUCTIONS.
pub base_instructions_override: Option<String>,
@@ -174,7 +174,7 @@ pub(crate) fn create_text_param_for_request(
})
}
pub(crate) struct ResponseStream {
pub struct ResponseStream {
pub(crate) rx_event: mpsc::Receiver<Result<ResponseEvent>>,
}

11
codex-rs/core/src/lib.rs
View File

@@ -69,3 +69,14 @@ pub use codex_protocol::protocol;
// Re-export protocol config enums to ensure call sites can use the same types
// as those in the protocol crate when constructing protocol messages.
pub use codex_protocol::config_types as protocol_config_types;
pub use client::ModelClient;
pub use client_common::Prompt;
pub use client_common::ResponseEvent;
pub use client_common::ResponseStream;
pub use codex_protocol::models::ContentItem;
pub use codex_protocol::models::LocalShellAction;
pub use codex_protocol::models::LocalShellExecAction;
pub use codex_protocol::models::LocalShellStatus;
pub use codex_protocol::models::ReasoningItemContent;
pub use codex_protocol::models::ResponseItem;