Moving token_info to ConversationHistory (#5581)
I want to centralize input processing and management to `ConversationHistory`. This would need `ConversationHistory` to have access to `token_info` (i.e. preventing adding a big input to the history). Besides, it makes more sense to have it on `ConversationHistory` than `state`.
This commit is contained in:
@@ -1,5 +1,7 @@
|
|||||||
use codex_protocol::models::FunctionCallOutputPayload;
|
use codex_protocol::models::FunctionCallOutputPayload;
|
||||||
use codex_protocol::models::ResponseItem;
|
use codex_protocol::models::ResponseItem;
|
||||||
|
use codex_protocol::protocol::TokenUsage;
|
||||||
|
use codex_protocol::protocol::TokenUsageInfo;
|
||||||
use tracing::error;
|
use tracing::error;
|
||||||
|
|
||||||
/// Transcript of conversation history
|
/// Transcript of conversation history
|
||||||
@@ -7,11 +9,28 @@ use tracing::error;
|
|||||||
pub(crate) struct ConversationHistory {
|
pub(crate) struct ConversationHistory {
|
||||||
/// The oldest items are at the beginning of the vector.
|
/// The oldest items are at the beginning of the vector.
|
||||||
items: Vec<ResponseItem>,
|
items: Vec<ResponseItem>,
|
||||||
|
token_info: Option<TokenUsageInfo>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ConversationHistory {
|
impl ConversationHistory {
|
||||||
pub(crate) fn new() -> Self {
|
pub(crate) fn new() -> Self {
|
||||||
Self { items: Vec::new() }
|
Self {
|
||||||
|
items: Vec::new(),
|
||||||
|
token_info: TokenUsageInfo::new_or_append(&None, &None, None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn token_info(&self) -> Option<TokenUsageInfo> {
|
||||||
|
self.token_info.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
|
||||||
|
match &mut self.token_info {
|
||||||
|
Some(info) => info.fill_to_context_window(context_window),
|
||||||
|
None => {
|
||||||
|
self.token_info = Some(TokenUsageInfo::full_context_window(context_window));
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `items` is ordered from oldest to newest.
|
/// `items` is ordered from oldest to newest.
|
||||||
@@ -301,6 +320,18 @@ impl ConversationHistory {
|
|||||||
self.items.remove(pos);
|
self.items.remove(pos);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn update_token_info(
|
||||||
|
&mut self,
|
||||||
|
usage: &TokenUsage,
|
||||||
|
model_context_window: Option<i64>,
|
||||||
|
) {
|
||||||
|
self.token_info = TokenUsageInfo::new_or_append(
|
||||||
|
&self.token_info,
|
||||||
|
&Some(usage.clone()),
|
||||||
|
model_context_window,
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
|
|||||||
@@ -12,7 +12,6 @@ use crate::protocol::TokenUsageInfo;
|
|||||||
pub(crate) struct SessionState {
|
pub(crate) struct SessionState {
|
||||||
pub(crate) session_configuration: SessionConfiguration,
|
pub(crate) session_configuration: SessionConfiguration,
|
||||||
pub(crate) history: ConversationHistory,
|
pub(crate) history: ConversationHistory,
|
||||||
pub(crate) token_info: Option<TokenUsageInfo>,
|
|
||||||
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
|
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -22,7 +21,6 @@ impl SessionState {
|
|||||||
Self {
|
Self {
|
||||||
session_configuration,
|
session_configuration,
|
||||||
history: ConversationHistory::new(),
|
history: ConversationHistory::new(),
|
||||||
token_info: None,
|
|
||||||
latest_rate_limits: None,
|
latest_rate_limits: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -54,11 +52,11 @@ impl SessionState {
|
|||||||
usage: &TokenUsage,
|
usage: &TokenUsage,
|
||||||
model_context_window: Option<i64>,
|
model_context_window: Option<i64>,
|
||||||
) {
|
) {
|
||||||
self.token_info = TokenUsageInfo::new_or_append(
|
self.history.update_token_info(usage, model_context_window);
|
||||||
&self.token_info,
|
}
|
||||||
&Some(usage.clone()),
|
|
||||||
model_context_window,
|
pub(crate) fn token_info(&self) -> Option<TokenUsageInfo> {
|
||||||
);
|
self.history.token_info()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) {
|
pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) {
|
||||||
@@ -68,17 +66,10 @@ impl SessionState {
|
|||||||
pub(crate) fn token_info_and_rate_limits(
|
pub(crate) fn token_info_and_rate_limits(
|
||||||
&self,
|
&self,
|
||||||
) -> (Option<TokenUsageInfo>, Option<RateLimitSnapshot>) {
|
) -> (Option<TokenUsageInfo>, Option<RateLimitSnapshot>) {
|
||||||
(self.token_info.clone(), self.latest_rate_limits.clone())
|
(self.token_info(), self.latest_rate_limits.clone())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
|
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
|
||||||
match &mut self.token_info {
|
self.history.set_token_usage_full(context_window);
|
||||||
Some(info) => info.fill_to_context_window(context_window),
|
|
||||||
None => {
|
|
||||||
self.token_info = Some(TokenUsageInfo::full_context_window(context_window));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Pending input/approval moved to TurnState.
|
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user