wabot-testing 0.1.0

//! A scriptable stand-in for an LLM. Port of
//! `wabot-ts/src/testing/MockChatAdapter.ts`.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use async_trait::async_trait;
use parking_lot::Mutex;
use wabot_feature_chat_bot::{
    ChatAdapter, ChatAdapterError, ChatAdapterRequest, ChatAdapterResponse, ChatItem, ChatMessage,
    FunctionCall, LanguageModelUsage,
};

/// One scripted turn: a fixed list of items, or a function of the
/// request when a test needs to answer differently depending on what
/// it was asked.
pub enum ScriptedTurn {
    Items(Vec<ChatItem>),
    #[allow(clippy::type_complexity)]
    Responder(Box<dyn Fn(&ChatAdapterRequest) -> Vec<ChatItem> + Send + Sync>),
}

/// What the adapter was asked, kept for assertions.
#[derive(Debug, Clone)]
pub struct RecordedRequest {
    pub system_prompt: String,
    pub tool_names: Vec<String>,
    pub models: Vec<String>,
    pub prev_items: Vec<ChatItem>,
}

impl RecordedRequest {
    fn of(request: &ChatAdapterRequest) -> Self {
        Self {
            system_prompt: request.system_prompt.clone(),
            tool_names: request.tools.iter().map(|t| t.name.clone()).collect(),
            models: request.models.iter().map(|m| m.model.clone()).collect(),
            prev_items: request.prev_items.clone(),
        }
    }
}

/// A deterministic [`ChatAdapter`]: script the model's turns, then
/// assert on what it was asked.
///
/// ```ignore
/// let adapter = MockChatAdapter::new();
/// adapter.call_tool("read_order", serde_json::json!({ "id": 7 }));
/// adapter.reply("Your order shipped.");
/// ```
///
/// **One queued turn is consumed per round-trip.** A tool-call turn
/// therefore needs another turn queued behind it, because the chat
/// loop calls the adapter again after running the tool — the mistake
/// is common enough that the exhaustion error says so.
///
/// Scripting takes `&self` so a test can queue turns after handing the
/// adapter to a harness.
pub struct MockChatAdapter {
    queue: Mutex<std::collections::VecDeque<ScriptedTurn>>,
    requests: Mutex<Vec<RecordedRequest>>,
    call_ids: AtomicUsize,
    fallback_reply: Mutex<Option<String>>,
}

impl Default for MockChatAdapter {
    fn default() -> Self {
        Self::new()
    }
}

impl MockChatAdapter {
    pub fn new() -> Self {
        Self {
            queue: Mutex::new(Default::default()),
            requests: Mutex::new(Vec::new()),
            call_ids: AtomicUsize::new(0),
            fallback_reply: Mutex::new(None),
        }
    }

    pub fn arc() -> Arc<Self> {
        Arc::new(Self::new())
    }

    /// Answer anything unscripted with `text` instead of failing —
    /// for tests that care about one turn and not about the rest.
    ///
    /// Off by default: a silent stand-in reply turns "the test didn't
    /// script enough" into a confusing assertion failure further down.
    pub fn with_fallback_reply(self, text: impl Into<String>) -> Self {
        *self.fallback_reply.lock() = Some(text.into());
        self
    }

    /// Queue a turn where the model answers with plain text.
    pub fn reply(&self, text: impl Into<String>) -> &Self {
        self.enqueue(ScriptedTurn::Items(vec![ChatItem::bot(ChatMessage::text(
            text,
        ))]))
    }

    /// Queue a turn where the model calls a tool. The tool itself runs
    /// for real — this only scripts the model's decision to call it.
    pub fn call_tool(&self, name: impl Into<String>, arguments: impl ToArguments) -> &Self {
        let id = self.call_ids.fetch_add(1, Ordering::SeqCst) + 1;
        self.enqueue(ScriptedTurn::Items(vec![ChatItem::call(FunctionCall {
            id: format!("mock-call-{id}"),
            name: name.into(),
            arguments: Some(arguments.to_arguments()),
            result: None,
            signature: None,
        })]))
    }

    /// Queue a raw turn.
    pub fn enqueue(&self, turn: ScriptedTurn) -> &Self {
        self.queue.lock().push_back(turn);
        self
    }

    /// Queue a turn computed from the request — for asserting on what
    /// the model was given *and* branching on it in one place.
    pub fn respond_with(
        &self,
        responder: impl Fn(&ChatAdapterRequest) -> Vec<ChatItem> + Send + Sync + 'static,
    ) -> &Self {
        self.enqueue(ScriptedTurn::Responder(Box::new(responder)))
    }

    /// Every request the adapter received, oldest first.
    pub fn requests(&self) -> Vec<RecordedRequest> {
        self.requests.lock().clone()
    }

    pub fn last_request(&self) -> Option<RecordedRequest> {
        self.requests.lock().last().cloned()
    }

    /// How many round-trips happened — the cheapest check that a tool
    /// loop ran as many times as expected, and no more.
    pub fn call_count(&self) -> usize {
        self.requests.lock().len()
    }

    /// Turns still queued. A test that ends with this above zero
    /// scripted more than it exercised, which usually means the
    /// conversation took a different path than it assumed.
    pub fn pending(&self) -> usize {
        self.queue.lock().len()
    }
}

#[async_trait]
impl ChatAdapter for MockChatAdapter {
    async fn next_items(
        &self,
        request: ChatAdapterRequest,
    ) -> Result<ChatAdapterResponse, ChatAdapterError> {
        self.requests.lock().push(RecordedRequest::of(&request));

        let turn = self.queue.lock().pop_front();
        let next_items = match turn {
            Some(ScriptedTurn::Items(items)) => items,
            Some(ScriptedTurn::Responder(responder)) => responder(&request),
            None => match self.fallback_reply.lock().clone() {
                Some(text) => vec![ChatItem::bot(ChatMessage::text(text))],
                None => {
                    return Err(ChatAdapterError::Other(
                        "MockChatAdapter: no scripted turn left. Queue one with reply() / \
                         call_tool() / enqueue(), or set with_fallback_reply(). Remember that \
                         after a call_tool() turn the chat loop calls the adapter again."
                            .into(),
                    ))
                }
            },
        };

        Ok(ChatAdapterResponse {
            next_items,
            usage: LanguageModelUsage {
                input_tokens: 1,
                output_tokens: 1,
                provider: Some("mock".into()),
                model: request.models.first().map(|m| m.model.clone()),
                ..Default::default()
            },
        })
    }
}

/// Tool arguments as either JSON or a raw string, so a test can write
/// the natural thing — and can also script the malformed JSON a real
/// model sometimes emits.
pub trait ToArguments {
    fn to_arguments(self) -> String;
}

impl ToArguments for serde_json::Value {
    fn to_arguments(self) -> String {
        self.to_string()
    }
}

impl ToArguments for &str {
    fn to_arguments(self) -> String {
        self.to_string()
    }
}

impl ToArguments for String {
    fn to_arguments(self) -> String {
        self
    }
}

/// `call_tool("x", NoArgs)` for a tool that takes none.
pub struct NoArgs;

impl ToArguments for NoArgs {
    fn to_arguments(self) -> String {
        "{}".into()
    }
}