Skip to content

Commit

Permalink
tune history trimming (#810)
Browse files Browse the repository at this point in the history
* wip: tune history trimming

* update tests

* trim history after 3 exchanges
  • Loading branch information
oppiliappan authored Aug 3, 2023
1 parent 4673946 commit 17c872f
Showing 1 changed file with 26 additions and 11 deletions.
37 changes: 26 additions & 11 deletions server/bleep/src/webserver/answer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1154,9 +1154,13 @@ impl Agent {

/// The full history of messages, including intermediate function calls
fn history(&self) -> Result<Vec<llm_gateway::api::Message>> {
const ANSWER_MAX_HISTORY_SIZE: usize = 3;
let history = self
.exchanges
.iter()
.rev()
.take(ANSWER_MAX_HISTORY_SIZE)
.rev()
.try_fold(Vec::new(), |mut acc, e| -> Result<_> {
let query = e
.query()
Expand Down Expand Up @@ -1541,33 +1545,44 @@ fn trim_history(
mut history: Vec<llm_gateway::api::Message>,
) -> Result<Vec<llm_gateway::api::Message>> {
const HEADROOM: usize = 2048;
const HIDDEN: &str = "[HIDDEN]";
const HIDDEN_WITH_INSTRUCTION: &str = "[HIDDEN]\nCall a function. Do not answer.";

let mut tiktoken_msgs = history.iter().map(|m| m.into()).collect::<Vec<_>>();

while tiktoken_rs::get_chat_completion_max_tokens(ANSWER_MODEL, &tiktoken_msgs)? < HEADROOM {
let idx = history
let _ = history
.iter_mut()
.position(|m| match m {
.zip(tiktoken_msgs.iter_mut())
.position(|(m, tm)| match m {
llm_gateway::api::Message::PlainText {
role,
ref mut content,
} if (role == "user" || role == "assistant") && content != "[HIDDEN]" => {
*content = "[HIDDEN]".into();
true
} => {
if (role == "user") && content != HIDDEN_WITH_INSTRUCTION {
*content = HIDDEN_WITH_INSTRUCTION.into();
tm.content = HIDDEN_WITH_INSTRUCTION.into();
true
} else if role == "assistant" && content != HIDDEN {
*content = HIDDEN.into();
tm.content = HIDDEN.into();
true
} else {
false
}
}
llm_gateway::api::Message::FunctionReturn {
role: _,
name: _,
ref mut content,
} if content != "[HIDDEN]" => {
*content = "[HIDDEN]".into();
} if content != HIDDEN_WITH_INSTRUCTION => {
*content = HIDDEN_WITH_INSTRUCTION.into();
tm.content = HIDDEN_WITH_INSTRUCTION.into();
true
}
_ => false,
})
.ok_or_else(|| anyhow!("could not find message to trim"))?;

tiktoken_msgs[idx].content = "[HIDDEN]".into();
}

Ok(history)
Expand Down Expand Up @@ -1701,9 +1716,9 @@ mod tests {
trim_history(history).unwrap(),
vec![
llm_gateway::api::Message::system("foo"),
llm_gateway::api::Message::user("[HIDDEN]"),
llm_gateway::api::Message::user("[HIDDEN]\nCall a function. Do not answer."),
llm_gateway::api::Message::assistant("[HIDDEN]"),
llm_gateway::api::Message::user("[HIDDEN]"),
llm_gateway::api::Message::user("[HIDDEN]\nCall a function. Do not answer."),
llm_gateway::api::Message::assistant("quux"),
llm_gateway::api::Message::user("fred"),
llm_gateway::api::Message::assistant("thud"),
Expand Down

0 comments on commit 17c872f

Please sign in to comment.