85 lines
2.9 KiB
Rust
85 lines
2.9 KiB
Rust
mod common;
|
|
|
|
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
|
use agen::{Engine, EngineError, History, HistoryEntry, Item, Role};
|
|
use common::MockLlmClient;
|
|
|
|
fn completed_text_events(text: &str) -> Vec<Event> {
|
|
vec![
|
|
Event::text_block_start(0),
|
|
Event::text_delta(0, text),
|
|
Event::text_block_stop(0, None),
|
|
Event::Status(StatusEvent {
|
|
status: ResponseStatus::Completed,
|
|
}),
|
|
]
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_preserves_item_annotations_without_projecting_them() {
|
|
let client = MockLlmClient::new(completed_text_events("assistant reply"));
|
|
let engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client);
|
|
let mut history = History::<String>::new();
|
|
let mut next = 0usize;
|
|
let mut annotate = |item: &Item| {
|
|
next += 1;
|
|
let kind = match item {
|
|
Item::Message { role, .. } => match role {
|
|
Role::User => "user",
|
|
Role::Assistant => "assistant",
|
|
Role::System => "system",
|
|
},
|
|
Item::ToolCall { .. } => "tool_call",
|
|
Item::ToolResult { .. } => "tool_result",
|
|
Item::Reasoning { .. } => "reasoning",
|
|
};
|
|
Ok(format!("{next}:{kind}"))
|
|
};
|
|
|
|
let output = engine
|
|
.run_with_annotation(&mut history, "hello", &mut annotate)
|
|
.await;
|
|
|
|
assert!(matches!(output.result, agen::EngineRunExit::Finished));
|
|
assert_eq!(history.len(), 2);
|
|
assert_eq!(history.entries()[0].annotation, "1:user");
|
|
assert_eq!(history.entries()[1].annotation, "2:assistant");
|
|
assert_eq!(history.items_cloned().len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn append_failure_does_not_make_item_live() {
|
|
let client = MockLlmClient::new(vec![]);
|
|
let mut engine = Engine::<_, agen::state::Mutable, usize>::new_annotated(client);
|
|
let mut history = History::<usize>::new();
|
|
let mut fail = |_item: &Item| Err("commit failed".to_string());
|
|
|
|
let err = engine
|
|
.append_history_with(&mut history, [Item::user_message("uncommitted")], &mut fail)
|
|
.unwrap_err();
|
|
|
|
assert!(matches!(err, EngineError::HistoryAppend(message) if message == "commit failed"));
|
|
assert!(history.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn replacement_keeps_items_and_annotations_together() {
|
|
let mut history = History::from_entries(vec![
|
|
HistoryEntry::new(Item::user_message("old"), "old-ann".to_string()),
|
|
HistoryEntry::new(Item::user_message("second"), "second-ann".to_string()),
|
|
]);
|
|
|
|
history.truncate(1);
|
|
assert_eq!(history.entries()[0].item.as_text(), Some("old"));
|
|
assert_eq!(history.entries()[0].annotation, "old-ann");
|
|
|
|
let previous = history.replace_entries(vec![HistoryEntry::new(
|
|
Item::user_message("restored"),
|
|
"restored-ann".to_string(),
|
|
)]);
|
|
|
|
assert_eq!(previous.len(), 1);
|
|
assert_eq!(history.entries()[0].item.as_text(), Some("restored"));
|
|
assert_eq!(history.entries()[0].annotation, "restored-ann");
|
|
}
|