pe-tools 0.1.0

//! Routing conditions — standard functions for tool-call routing.
//!
//! Used with `StateGraph::add_conditional_edge()` to route between
//! the chat node and the tool node in the standard ReAct pattern.

use pe_core::message::Message;
use pe_core::state::CoreState;

/// Standard ReAct routing condition.
///
/// Returns `["tools"]` if the last AI message has tool_calls,
/// `["__end__"]` otherwise.
///
/// # Usage
///
/// ```ignore
/// graph.add_conditional_edge("chat", tools_condition::<MyState>);
/// ```
pub fn tools_condition<S: CoreState>(state: &S) -> Vec<String> {
    let has_tool_calls = state
        .messages()
        .iter()
        .rev()
        .find_map(|m| {
            if let Message::Ai(ai) = m {
                Some(!ai.tool_calls.is_empty())
            } else {
                None
            }
        })
        .unwrap_or(false);

    if has_tool_calls {
        vec!["tools".to_string()]
    } else {
        vec!["__end__".to_string()]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pe_core::message::{AiMessage, MessageContent, ToolCall};
    use pe_core::state::{ExecutionContext, StateUpdate};
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone, Serialize, Deserialize)]
    struct TestState {
        messages: Vec<Message>,
        iterations: u32,
        thread_id: String,
        context: ExecutionContext,
    }

    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    struct TestUpdate;
    impl StateUpdate for TestUpdate {}

    impl pe_core::state::State for TestState {
        type Update = TestUpdate;
        fn apply(&mut self, _update: Self::Update) {}
    }

    impl CoreState for TestState {
        fn messages(&self) -> &[Message] {
            &self.messages
        }
        fn messages_mut(&mut self) -> &mut Vec<Message> {
            &mut self.messages
        }
        fn iterations(&self) -> u32 {
            self.iterations
        }
        fn set_iterations(&mut self, n: u32) {
            self.iterations = n;
        }
        fn thread_id(&self) -> &str {
            &self.thread_id
        }
        fn context(&self) -> &ExecutionContext {
            &self.context
        }
        fn context_mut(&mut self) -> &mut ExecutionContext {
            &mut self.context
        }
    }

    fn make_state(messages: Vec<Message>) -> TestState {
        TestState {
            messages,
            iterations: 0,
            thread_id: "t1".into(),
            context: ExecutionContext::new("test"),
        }
    }

    #[test]
    fn routes_to_tools_when_tool_calls_present() {
        let state = make_state(vec![Message::Ai(AiMessage {
            content: MessageContent::Text("Let me search...".into()),
            tool_calls: vec![ToolCall {
                id: "tc_1".into(),
                name: "search".into(),
                args: serde_json::json!({"q": "test"}),
            }],
            invalid_tool_calls: vec![],
            usage_metadata: None,
            response_metadata: Default::default(),
            id: None,
        })]);

        assert_eq!(tools_condition(&state), vec!["tools"]);
    }

    #[test]
    fn routes_to_end_when_no_tool_calls() {
        let state = make_state(vec![Message::ai("Done!")]);
        assert_eq!(tools_condition(&state), vec!["__end__"]);
    }

    #[test]
    fn routes_to_end_with_no_messages() {
        let state = make_state(vec![]);
        assert_eq!(tools_condition(&state), vec!["__end__"]);
    }

    #[test]
    fn routes_to_end_when_only_human_messages() {
        let state = make_state(vec![Message::human("hello")]);
        assert_eq!(tools_condition(&state), vec!["__end__"]);
    }

    #[test]
    fn finds_tool_calls_in_last_ai_message_ignoring_earlier() {
        let state = make_state(vec![
            // Earlier AI message WITH tool calls
            Message::Ai(AiMessage {
                content: MessageContent::Text("searching...".into()),
                tool_calls: vec![ToolCall {
                    id: "tc_old".into(),
                    name: "search".into(),
                    args: serde_json::json!({}),
                }],
                invalid_tool_calls: vec![],
                usage_metadata: None,
                response_metadata: Default::default(),
                id: None,
            }),
            // Tool result
            Message::tool("result", "tc_old"),
            // Latest AI message WITHOUT tool calls
            Message::ai("Here's the answer."),
        ]);

        assert_eq!(tools_condition(&state), vec!["__end__"]);
    }
}