kcode-k1-chat-state 0.8.0

Provider-free open-format chat scheduling state
Documentation
use crate::{ActorState, ProviderCall, ToolCallId};

fn tool_id(value: u8) -> ToolCallId {
    ToolCallId::new([value; 12], u64::from(value))
}

fn call(value: u8) -> ProviderCall {
    ProviderCall {
        tool_call_id: tool_id(value),
        name: "WebSearch".into(),
        arguments: "{}".into(),
    }
}

#[test]
fn tool_messages_do_not_trigger_but_v2_result_does() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(active.job, String::new(), vec![call(1)])
        .unwrap();

    state
        .accept_tool_message(tool_id(1), "first".into())
        .unwrap();
    let flushed = state.flush_active_arrivals(active.job).unwrap();
    assert_eq!(flushed.len(), 1);
    assert_eq!(flushed[0].box_type(), "Tool Message");

    state
        .accept_tool_message(tool_id(1), "second".into())
        .unwrap();
    state.complete_inference(active.job, String::new()).unwrap();
    assert!(state.quiet());
    assert!(state.begin_inference().unwrap().is_none());

    state
        .accept_async_return_v2(
            tool_id(1),
            Ok("answer".into()),
            "k1.web-search-result/v1".into(),
            "{}".into(),
        )
        .unwrap();
    let followup = state.begin_inference().unwrap().unwrap();
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_inference(followup.job, String::new())
        .unwrap();
    assert!(state.quiet());
}

#[test]
fn messages_and_result_stay_fifo_and_schedule_one_followup() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(active.job, String::new(), vec![call(2)])
        .unwrap();
    state.accept_tool_message(tool_id(2), "one".into()).unwrap();
    state.accept_tool_message(tool_id(2), "two".into()).unwrap();
    assert_eq!(state.flush_active_arrivals(active.job).unwrap().len(), 2);
    state
        .accept_async_return(tool_id(2), Ok("done".into()))
        .unwrap();
    state.complete_inference(active.job, String::new()).unwrap();

    let kinds = state
        .boxes()
        .iter()
        .map(|value| value.box_type())
        .collect::<Vec<_>>();
    assert_eq!(
        &kinds[kinds.len() - 3..],
        ["Tool Message", "Tool Message", "Tool Result"]
    );
    let followup = state.begin_inference().unwrap().unwrap();
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_inference(followup.job, String::new())
        .unwrap();
    assert!(state.quiet());
}

#[test]
fn retry_uses_a_fresh_job_and_preserves_the_frontier() {
    let mut state = ActorState::new(false);
    state.accept_user("ask".into()).unwrap();
    let first = state.begin_inference().unwrap().unwrap();
    state.stall_inference(first.job, "retry".into()).unwrap();
    assert_eq!(state.take_halt(), Some("retry".into()));
    state.restart().unwrap();
    let second = state.begin_inference().unwrap().unwrap();
    assert!(second.job > first.job);
    assert_eq!(second.frontier, first.frontier);
    state.complete_inference(second.job, String::new()).unwrap();
}