This commit is contained in:
jif-oai
2025-11-11 15:16:07 +00:00
parent 0c609d441b
commit b7dcc8ef5c
7 changed files with 133 additions and 62 deletions

View File

@@ -9,6 +9,7 @@ use crate::error::Error;
use crate::error::Result;
use crate::stream::ResponseEvent;
use crate::stream::ResponseStream;
use crate::stream::WireResponseStream;
use codex_provider_config::ModelProviderInfo;
#[derive(Clone)]
@@ -25,6 +26,7 @@ pub struct ChatCompletionsApiClientConfig {
pub model: String,
pub otel_event_manager: OtelEventManager,
pub session_source: SessionSource,
pub extra_headers: Vec<(String, String)>,
}
#[derive(Clone)]
@@ -56,12 +58,18 @@ impl PayloadClient for ChatCompletionsApiClient {
}
let auth = crate::client::http::resolve_auth(&None).await;
let extra_headers: Vec<(&str, String)> = self
.config
.extra_headers
.iter()
.map(|(k, v)| (k.as_str(), v.clone()))
.collect();
let mut req_builder = crate::client::http::build_request(
&self.config.http_client,
&self.config.provider,
&auth,
session_source,
&[],
&extra_headers,
)
.await?;
@@ -95,3 +103,25 @@ impl PayloadClient for ChatCompletionsApiClient {
Ok(crate::stream::EventStream::from_receiver(rx_event))
}
}
impl ChatCompletionsApiClient {
pub async fn stream_payload_wire(
&self,
payload_json: &serde_json::Value,
session_source: Option<&codex_protocol::protocol::SessionSource>,
) -> Result<WireResponseStream> {
use futures::StreamExt;
let legacy = self.stream_payload(payload_json, session_source).await?;
let (tx, rx) = tokio::sync::mpsc::channel(1600);
tokio::spawn(async move {
futures::pin_mut!(legacy);
while let Some(item) = legacy.next().await {
let converted = item.and_then(|ev| crate::wire::map_response_event_to_wire(ev));
if tx.send(converted).await.is_err() {
break;
}
}
});
Ok(crate::stream::EventStream::from_receiver(rx))
}
}

View File

@@ -1,4 +1,3 @@
pub mod aggregate;
pub mod api;
pub mod auth;
pub mod chat;
@@ -10,8 +9,8 @@ pub mod error;
pub mod responses;
pub mod routed_client;
pub mod stream;
mod wire;
pub use crate::aggregate::AggregateStreamExt;
pub use crate::auth::AuthContext;
pub use crate::auth::AuthProvider;
pub use crate::chat::ChatCompletionsApiClient;
@@ -32,9 +31,5 @@ pub use crate::stream::TextFormat;
pub use crate::stream::TextFormatType;
pub use crate::stream::WireEvent;
pub use crate::stream::WireResponseStream;
pub use codex_provider_config::BUILT_IN_OSS_MODEL_PROVIDER_ID;
pub use codex_provider_config::ModelProviderInfo;
pub use codex_provider_config::WireApi;
pub use codex_provider_config::built_in_model_providers;
pub use codex_provider_config::create_oss_provider;
pub use codex_provider_config::create_oss_provider_with_base_url;

View File

@@ -16,6 +16,7 @@ use crate::error::Error;
use crate::error::Result;
use crate::stream::ResponseEvent;
use crate::stream::ResponseStream;
use crate::stream::WireResponseStream;
use codex_provider_config::ModelProviderInfo;
#[derive(Clone)]
@@ -34,6 +35,7 @@ pub struct ResponsesApiClientConfig {
pub conversation_id: ConversationId,
pub auth_provider: Option<Arc<dyn AuthProvider>>,
pub otel_event_manager: OtelEventManager,
pub extra_headers: Vec<(String, String)>,
}
#[derive(Clone)]
@@ -77,10 +79,21 @@ impl PayloadClient for ResponsesApiClient {
.unwrap_or_else(|_| "<unable to serialize payload>".to_string())
);
let extra_headers = vec![
("conversation_id", self.config.conversation_id.to_string()),
("session_id", self.config.conversation_id.to_string()),
let mut owned_headers: Vec<(String, String)> = vec![
(
"conversation_id".to_string(),
self.config.conversation_id.to_string(),
),
(
"session_id".to_string(),
self.config.conversation_id.to_string(),
),
];
owned_headers.extend(self.config.extra_headers.iter().cloned());
let extra_headers: Vec<(&str, String)> = owned_headers
.iter()
.map(|(k, v)| (k.as_str(), v.clone()))
.collect();
let mut req_builder = crate::client::http::build_request(
&self.config.http_client,
&self.config.provider,
@@ -140,3 +153,25 @@ impl PayloadClient for ResponsesApiClient {
Ok(crate::stream::EventStream::from_receiver(rx_event))
}
}
impl ResponsesApiClient {
pub async fn stream_payload_wire(
&self,
payload_json: &Value,
session_source: Option<&codex_protocol::protocol::SessionSource>,
) -> Result<WireResponseStream> {
use futures::StreamExt;
let legacy = self.stream_payload(payload_json, session_source).await?;
let (tx, rx) = tokio::sync::mpsc::channel(1600);
tokio::spawn(async move {
futures::pin_mut!(legacy);
while let Some(item) = legacy.next().await {
let converted = item.and_then(|ev| crate::wire::map_response_event_to_wire(ev));
if tx.send(converted).await.is_err() {
break;
}
}
});
Ok(crate::stream::EventStream::from_receiver(rx))
}
}

View File

@@ -17,7 +17,6 @@ use crate::WireResponseStream;
use crate::api::PayloadClient;
use crate::auth::AuthProvider;
use crate::client::fixtures::stream_from_fixture;
use crate::stream::WireTokenUsage;
use codex_provider_config::ModelProviderInfo;
/// Dispatches to the appropriate API client implementation based on the provider wire API.
@@ -62,6 +61,7 @@ impl RoutedApiClient {
conversation_id: self.config.conversation_id,
auth_provider: self.config.auth_provider.clone(),
otel_event_manager: self.config.otel_event_manager.clone(),
extra_headers: vec![],
};
let client = <ResponsesApiClient as crate::api::PayloadClient>::new(cfg)?;
client
@@ -75,6 +75,7 @@ impl RoutedApiClient {
model: self.config.model.clone(),
otel_event_manager: self.config.otel_event_manager.clone(),
session_source: self.config.session_source.clone(),
extra_headers: vec![],
};
let client = <ChatCompletionsApiClient as crate::api::PayloadClient>::new(cfg)?;
client
@@ -122,52 +123,5 @@ impl PayloadClient for RoutedApiClient {
}
fn map_response_event_to_wire(ev: crate::stream::ResponseEvent) -> Result<WireEvent> {
Ok(match ev {
crate::stream::ResponseEvent::Created => WireEvent::Created,
crate::stream::ResponseEvent::OutputItemDone(item) => {
WireEvent::OutputItemDone(serde_json::to_value(item).unwrap_or(serde_json::Value::Null))
}
crate::stream::ResponseEvent::OutputItemAdded(item) => WireEvent::OutputItemAdded(
serde_json::to_value(item).unwrap_or(serde_json::Value::Null),
),
crate::stream::ResponseEvent::Completed {
response_id,
token_usage,
} => {
let mapped = token_usage.map(|u| WireTokenUsage {
input_tokens: u.input_tokens,
cached_input_tokens: u.cached_input_tokens,
output_tokens: u.output_tokens,
reasoning_output_tokens: u.reasoning_output_tokens,
total_tokens: u.total_tokens,
});
WireEvent::Completed {
response_id,
token_usage: mapped,
}
}
crate::stream::ResponseEvent::OutputTextDelta(s) => WireEvent::OutputTextDelta(s),
crate::stream::ResponseEvent::ReasoningSummaryDelta(s) => {
WireEvent::ReasoningSummaryDelta(s)
}
crate::stream::ResponseEvent::ReasoningContentDelta(s) => {
WireEvent::ReasoningContentDelta(s)
}
crate::stream::ResponseEvent::ReasoningSummaryPartAdded => {
WireEvent::ReasoningSummaryPartAdded
}
crate::stream::ResponseEvent::RateLimits(s) => {
let to_win = |w: Option<codex_protocol::protocol::RateLimitWindow>| -> Option<crate::stream::WireRateLimitWindow> {
w.map(|w| crate::stream::WireRateLimitWindow {
used_percent: Some(w.used_percent),
window_minutes: w.window_minutes,
resets_at: w.resets_at,
})
};
WireEvent::RateLimits(crate::stream::WireRateLimitSnapshot {
primary: to_win(s.primary),
secondary: to_win(s.secondary),
})
}
})
crate::wire::map_response_event_to_wire(ev)
}

View File

@@ -0,0 +1,55 @@
use crate::error::Result;
use crate::stream::WireEvent;
use crate::stream::WireRateLimitSnapshot;
use crate::stream::WireRateLimitWindow;
pub fn map_response_event_to_wire(ev: crate::stream::ResponseEvent) -> Result<WireEvent> {
Ok(match ev {
crate::stream::ResponseEvent::Created => WireEvent::Created,
crate::stream::ResponseEvent::OutputItemDone(item) => {
WireEvent::OutputItemDone(serde_json::to_value(item).unwrap_or(serde_json::Value::Null))
}
crate::stream::ResponseEvent::OutputItemAdded(item) => WireEvent::OutputItemAdded(
serde_json::to_value(item).unwrap_or(serde_json::Value::Null),
),
crate::stream::ResponseEvent::Completed {
response_id,
token_usage,
} => {
let mapped = token_usage.map(|u| crate::stream::WireTokenUsage {
input_tokens: u.input_tokens,
cached_input_tokens: u.cached_input_tokens,
output_tokens: u.output_tokens,
reasoning_output_tokens: u.reasoning_output_tokens,
total_tokens: u.total_tokens,
});
WireEvent::Completed {
response_id,
token_usage: mapped,
}
}
crate::stream::ResponseEvent::OutputTextDelta(s) => WireEvent::OutputTextDelta(s),
crate::stream::ResponseEvent::ReasoningSummaryDelta(s) => {
WireEvent::ReasoningSummaryDelta(s)
}
crate::stream::ResponseEvent::ReasoningContentDelta(s) => {
WireEvent::ReasoningContentDelta(s)
}
crate::stream::ResponseEvent::ReasoningSummaryPartAdded => {
WireEvent::ReasoningSummaryPartAdded
}
crate::stream::ResponseEvent::RateLimits(s) => {
let to_win = |w: Option<codex_protocol::protocol::RateLimitWindow>| -> Option<WireRateLimitWindow> {
w.map(|w| WireRateLimitWindow {
used_percent: Some(w.used_percent),
window_minutes: w.window_minutes,
resets_at: w.resets_at,
})
};
WireEvent::RateLimits(WireRateLimitSnapshot {
primary: to_win(s.primary),
secondary: to_win(s.secondary),
})
}
})
}

View File

@@ -3,12 +3,12 @@ use std::pin::Pin;
use std::task::Context;
use std::task::Poll;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use crate::ContentItem;
use crate::ResponseItem;
use futures::Stream;
use crate::ResponseEvent;
use crate::error::Result;
use crate::stream::ResponseEvent;
pub trait AggregateStreamExt: Stream<Item = Result<ResponseEvent>> + Sized {
fn aggregate(self) -> AggregatedChatStream<Self>

View File

@@ -44,6 +44,7 @@ pub use codex_provider_config::ModelProviderInfo;
pub use codex_provider_config::WireApi;
pub use codex_provider_config::built_in_model_providers;
pub use codex_provider_config::create_oss_provider_with_base_url;
mod aggregate;
mod conversation_manager;
mod event_mapping;
pub mod review_format;
@@ -94,6 +95,7 @@ pub use codex_protocol::protocol;
// as those in the protocol crate when constructing protocol messages.
pub use codex_protocol::config_types as protocol_config_types;
pub use aggregate::AggregateStreamExt;
pub use client::ModelClient;
pub use client_common::Prompt;
pub use client_common::REVIEW_PROMPT;