wabot_testing/

chat_bot.rs

//! Drive a real `ChatBot` in a test. Port of
//! `wabot-ts/src/testing/chatBotHarness.ts`.

use std::sync::Arc;

use parking_lot::Mutex;
use wabot_core::injection::Container;
use wabot_feature_chat_bot::{
    ChatAdapter, ChatBot, ChatBotError, ChatItem, ChatMemory, ChatMessage, FunctionCall,
};
use wabot_feature_mindset::{Mindset, MindsetOperator, MindsetTool, ToolDefinition, ToolError};

use crate::memory::TestChatMemory;
use crate::mock_adapter::MockChatAdapter;

/// Everything the bot did in response to one message.
#[derive(Debug, Clone, Default)]
pub struct ChatTurn {
    /// Messages delivered through the reply callback — what the user
    /// actually saw.
    pub replies: Vec<ChatMessage>,
    /// Tool calls executed during the turn, with their results.
    pub tool_calls: Vec<FunctionCall>,
    /// Every item recorded during the turn, in order.
    pub items: Vec<ChatItem>,
}

impl ChatTurn {
    /// The reply texts, which is what most assertions want.
    pub fn texts(&self) -> Vec<String> {
        self.replies.iter().filter_map(|m| m.text.clone()).collect()
    }

    /// The single reply text, when the turn produced exactly one.
    ///
    /// # Panics
    ///
    /// If there wasn't exactly one — a turn that replied twice, or not
    /// at all, is a different outcome and silently taking the first
    /// would hide it.
    pub fn text(&self) -> String {
        let texts = self.texts();
        assert_eq!(texts.len(), 1, "expected exactly one reply, got {texts:?}");
        texts.into_iter().next().unwrap()
    }

    pub fn called(&self, name: &str) -> bool {
        self.tool_calls.iter().any(|c| c.name == name)
    }
}

/// Runs the **real** chat stack — real [`MindsetOperator`], real system
/// prompt, real tool loop with argument validation — against an in-RAM
/// memory and a scriptable adapter.
///
/// The point is that a test exercises production code paths. Anything
/// the harness reimplemented would be a second implementation free to
/// drift from the one that ships.
///
/// ```ignore
/// let harness = ChatBotHarness::builder(Arc::new(MyMindset))
///     .tools(OrderTools::register_tools(&container))
///     .container(container)
///     .build();
///
/// harness.adapter().call_tool("read_order", json!({ "id": 7 }));
/// harness.adapter().reply("It shipped yesterday.");
///
/// let turn = harness.send("where is order 7?").await.unwrap();
/// assert_eq!(turn.text(), "It shipped yesterday.");
/// assert!(turn.called("read_order"));
/// ```
pub struct ChatBotHarness {
    adapter: Arc<MockChatAdapter>,
    memory: Arc<TestChatMemory>,
    operator: Arc<MindsetOperator>,
    bot: ChatBot,
}

impl ChatBotHarness {
    /// A harness with a mock adapter, an in-RAM memory and no tools.
    pub fn new(mindset: Arc<dyn Mindset>) -> Self {
        Self::builder(mindset).build()
    }

    pub fn builder(mindset: Arc<dyn Mindset>) -> ChatBotHarnessBuilder {
        ChatBotHarnessBuilder {
            mindset,
            tools: Vec::new(),
            container: None,
            adapter: None,
        }
    }

    /// The scripted adapter — queue turns and assert on requests.
    pub fn adapter(&self) -> &Arc<MockChatAdapter> {
        &self.adapter
    }

    /// The operator the bot uses, for asserting on the real prompt and
    /// tool schema.
    pub fn operator(&self) -> &Arc<MindsetOperator> {
        &self.operator
    }

    /// Send a human message and collect what the bot did.
    pub async fn send(&self, message: impl IntoChatMessage) -> Result<ChatTurn, ChatBotError> {
        let before = self.memory.len();
        let replies: Arc<Mutex<Vec<ChatMessage>>> = Arc::new(Mutex::new(Vec::new()));

        let sink = replies.clone();
        let reply: wabot_feature_chat_bot::BotReplyFn = Arc::new(move |message| {
            let sink = sink.clone();
            Box::pin(async move {
                sink.lock().push(message);
            })
        });

        self.bot
            .send_message(message.into_chat_message(), reply)
            .await?;

        let replies = std::mem::take(&mut *replies.lock());
        let items = self.memory.items_from(before);
        let tool_calls = items
            .iter()
            .filter_map(|item| match item {
                ChatItem::FunctionCall { function_call } => Some(function_call.clone()),
                _ => None,
            })
            .collect();

        Ok(ChatTurn {
            replies,
            tool_calls,
            items,
        })
    }

    /// Run one tool directly — real validation, real dispatch — without
    /// scripting a conversation around it. Returns the string the model
    /// would have received.
    pub async fn call_tool(
        &self,
        name: &str,
        arguments: impl crate::mock_adapter::ToArguments,
    ) -> Result<String, ToolError> {
        self.operator
            .call_function(name, &arguments.to_arguments())
            .await
    }

    /// The real system prompt this mindset produces.
    pub async fn system_prompt(&self) -> String {
        self.operator.system_prompt().await
    }

    /// The real tool schema the model would be given.
    pub fn tools(&self) -> Result<Vec<MindsetTool>, ToolError> {
        self.operator.tools()
    }

    /// Everything recorded across every turn.
    pub fn history(&self) -> Vec<ChatItem> {
        self.memory.all()
    }

    pub fn memory(&self) -> &Arc<TestChatMemory> {
        &self.memory
    }
}

pub struct ChatBotHarnessBuilder {
    mindset: Arc<dyn Mindset>,
    tools: Vec<ToolDefinition>,
    container: Option<Container>,
    adapter: Option<Arc<MockChatAdapter>>,
}

impl ChatBotHarnessBuilder {
    /// Tools the mindset can call, usually
    /// `MyTools::register_tools(&container)`.
    pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
        self.tools.extend(tools);
        self
    }

    /// The container the tools resolve from. Register a tool set's
    /// dependencies (a fake database, say) there before building.
    pub fn container(mut self, container: Container) -> Self {
        self.container = Some(container);
        self
    }

    /// Share an adapter — to script it before building, or to reuse
    /// one across harnesses.
    pub fn adapter(mut self, adapter: Arc<MockChatAdapter>) -> Self {
        self.adapter = Some(adapter);
        self
    }

    pub fn build(self) -> ChatBotHarness {
        let container = self.container.unwrap_or_default();
        let adapter = self.adapter.unwrap_or_else(MockChatAdapter::arc);
        let memory = Arc::new(TestChatMemory::new());

        let operator =
            Arc::new(MindsetOperator::new(container, self.mindset).with_module_tools(self.tools));
        let bot = ChatBot::new(
            memory.clone() as Arc<dyn ChatMemory>,
            adapter.clone() as Arc<dyn ChatAdapter>,
            operator.clone(),
        );

        ChatBotHarness {
            adapter,
            memory,
            operator,
            bot,
        }
    }
}

/// So `send("hola")` works as well as `send(ChatMessage { … })`.
pub trait IntoChatMessage {
    fn into_chat_message(self) -> ChatMessage;
}

impl IntoChatMessage for ChatMessage {
    fn into_chat_message(self) -> ChatMessage {
        self
    }
}

impl IntoChatMessage for &str {
    fn into_chat_message(self) -> ChatMessage {
        ChatMessage::text(self)
    }
}

impl IntoChatMessage for String {
    fn into_chat_message(self) -> ChatMessage {
        ChatMessage::text(self)
    }
}