use rskit_hook::{Event, EventType};
use rskit_llm::types::{AssistantMessage, CompletionRequest, CompletionResponse};
use rskit_tool::ToolInput;
use rskit_tool::ToolResult;
pub fn on_tool_call_type() -> EventType {
EventType::new("on_tool_call")
}
pub fn pre_tool_call_type() -> EventType {
on_tool_call_type()
}
pub fn on_tool_result_type() -> EventType {
EventType::new("on_tool_result")
}
pub fn post_tool_call_type() -> EventType {
on_tool_result_type()
}
pub fn on_llm_call_type() -> EventType {
EventType::new("on_llm_call")
}
pub fn pre_llm_call_type() -> EventType {
on_llm_call_type()
}
pub fn on_llm_response_type() -> EventType {
EventType::new("on_llm_response")
}
pub fn post_llm_call_type() -> EventType {
on_llm_response_type()
}
pub fn on_error_type() -> EventType {
EventType::new("on_error")
}
pub fn on_turn_start_type() -> EventType {
EventType::new("on_turn_start")
}
pub fn turn_start_type() -> EventType {
on_turn_start_type()
}
pub fn on_turn_complete_type() -> EventType {
EventType::new("on_turn_complete")
}
pub fn turn_end_type() -> EventType {
on_turn_complete_type()
}
pub fn on_event_type() -> EventType {
EventType::new("on_event")
}
pub fn on_mcp_call_type() -> EventType {
EventType::new("on_mcp_call")
}
pub fn on_mcp_result_type() -> EventType {
EventType::new("on_mcp_result")
}
#[derive(Debug, Clone)]
pub struct OnEvent {
pub event: rskit_ai::StreamEventRef,
}
impl Event for OnEvent {
fn event_type(&self) -> EventType {
on_event_type()
}
}
#[derive(Debug, Clone)]
pub struct OnMCPCall {
pub method: String,
}
impl Event for OnMCPCall {
fn event_type(&self) -> EventType {
on_mcp_call_type()
}
}
#[derive(Debug, Clone)]
pub struct OnMCPResult {
pub method: String,
pub result: Option<serde_json::Value>,
pub error: Option<String>,
}
impl Event for OnMCPResult {
fn event_type(&self) -> EventType {
on_mcp_result_type()
}
}
#[derive(Debug, Clone)]
pub struct PreToolCall {
pub name: String,
pub input: ToolInput,
}
impl Event for PreToolCall {
fn event_type(&self) -> EventType {
pre_tool_call_type()
}
}
#[derive(Debug, Clone)]
pub struct PostToolCall {
pub name: String,
pub input: ToolInput,
pub result: Option<ToolResult>,
pub error: Option<String>,
}
impl Event for PostToolCall {
fn event_type(&self) -> EventType {
post_tool_call_type()
}
}
#[derive(Debug, Clone)]
pub struct PreLLMCall {
pub request: CompletionRequest,
}
impl Event for PreLLMCall {
fn event_type(&self) -> EventType {
pre_llm_call_type()
}
}
#[derive(Debug, Clone)]
pub struct PostLLMCall {
pub response: CompletionResponse,
pub error: Option<String>,
}
impl Event for PostLLMCall {
fn event_type(&self) -> EventType {
post_llm_call_type()
}
}
#[derive(Debug, Clone)]
pub struct OnError {
pub error: String,
pub source: String,
}
impl Event for OnError {
fn event_type(&self) -> EventType {
on_error_type()
}
}
#[derive(Debug, Clone)]
pub struct TurnStart {
pub turn: u32,
}
impl Event for TurnStart {
fn event_type(&self) -> EventType {
turn_start_type()
}
}
#[derive(Debug, Clone)]
pub struct TurnEnd {
pub turn: u32,
pub message: AssistantMessage,
}
impl Event for TurnEnd {
fn event_type(&self) -> EventType {
turn_end_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rskit_hook::EventType;
use rskit_tool::ToolInput;
#[test]
fn test_pre_tool_call_event() {
let event = PreToolCall {
name: "calculator".to_string(),
input: ToolInput::new(serde_json::json!({"x": 1})).unwrap(),
};
assert_eq!(event.event_type(), EventType::new("on_tool_call"));
assert_eq!(event.name, "calculator");
}
#[test]
fn test_post_tool_call_event() {
let event = PostToolCall {
name: "calculator".to_string(),
input: ToolInput::new(serde_json::json!({})).unwrap(),
result: None,
error: Some("timeout".to_string()),
};
assert_eq!(event.event_type(), EventType::new("on_tool_result"));
}
#[test]
fn test_pre_llm_call_event() {
let event = PreLLMCall {
request: CompletionRequest {
model: "test".to_string(),
messages: vec![],
max_tokens: None,
temperature: None,
stream: false,
tools: None,
tool_choice: None,
},
};
assert_eq!(event.event_type(), EventType::new("on_llm_call"));
}
#[test]
fn test_on_error_event() {
let event = OnError {
error: "boom".to_string(),
source: "agent".to_string(),
};
assert_eq!(event.event_type(), EventType::new("on_error"));
}
#[test]
fn test_turn_start_event() {
let event = TurnStart { turn: 0 };
assert_eq!(event.event_type(), EventType::new("on_turn_start"));
assert_eq!(event.turn, 0);
}
#[test]
fn test_turn_end_event() {
use rskit_llm::types;
let event = TurnEnd {
turn: 3,
message: AssistantMessage {
content: types::text_content("done"),
tool_calls: vec![],
usage: None,
},
};
assert_eq!(event.event_type(), EventType::new("on_turn_complete"));
}
#[test]
fn test_event_type_constants() {
assert_eq!(pre_tool_call_type(), EventType::new("on_tool_call"));
assert_eq!(post_tool_call_type(), EventType::new("on_tool_result"));
assert_eq!(pre_llm_call_type(), EventType::new("on_llm_call"));
assert_eq!(post_llm_call_type(), EventType::new("on_llm_response"));
assert_eq!(on_error_type(), EventType::new("on_error"));
assert_eq!(turn_start_type(), EventType::new("on_turn_start"));
assert_eq!(turn_end_type(), EventType::new("on_turn_complete"));
}
}