~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/asyncai

0
0
main
152 lines · 4.7 KB Raw History
use serde::{Deserialize, Serialize};

use crate::ids::{AgentId, FileId, RunId, ThreadId};
use crate::message::{AgentContent, Role, RunStatus};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EventType {
    Error,
    EndEvents,
    InternalReQueue,
    NewAgentCreated,
    ExistingAgentRetrieved,
    FileIngested,
    NewThreadCreated,
    NewThreadRunCreated,
    NewThreadMessage,
    ThreadMessageAdded,
    ThreadRunComplete,
    ThreadRunInProgress,
    ThreadRunFailed,
    ThreadRunEventWithUnknownStatus,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ErrorEvent {
    pub message: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AgentEventData {
    pub agent_id: AgentId,
    pub agent_name: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct FileIngestedEvent {
    pub agent_id: AgentId,
    pub file_id: FileId,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct NewThreadCreatedEvent {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct NewThreadRunCreatedEvent {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
    pub run_id: RunId,
    pub run_status: RunStatus,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct NewThreadMessageEvent {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
    pub message_role: Role,
    pub message_content_type: String,
    pub message_content: AgentContent,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ThreadMessageAddedEvent {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
    pub message_id: usize,
    pub message_role: Role,
    pub message_content: AgentContent,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ThreadRunEventData {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
    pub run_id: RunId,
    pub run_status: RunStatus,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ThreadRunFailedEvent {
    pub agent_id: AgentId,
    pub thread_id: ThreadId,
    pub run_id: RunId,
    pub run_status: RunStatus,
    pub last_error: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "event_type", content = "event_data", rename_all = "snake_case")]
pub enum AgiEvent {
    Error(ErrorEvent),
    EndEvents,
    InternalReQueue,
    NewAgentCreated(AgentEventData),
    ExistingAgentRetrieved(AgentEventData),
    FileIngested(FileIngestedEvent),
    NewThreadCreated(NewThreadCreatedEvent),
    NewThreadRunCreated(NewThreadRunCreatedEvent),
    NewThreadMessage(NewThreadMessageEvent),
    ThreadMessageAdded(ThreadMessageAddedEvent),
    ThreadRunComplete(ThreadRunEventData),
    ThreadRunInProgress(ThreadRunEventData),
    ThreadRunFailed(ThreadRunFailedEvent),
    ThreadRunEventWithUnknownStatus(ThreadRunEventData),
}

impl AgiEvent {
    pub fn event_type(&self) -> EventType {
        match self {
            Self::Error(_) => EventType::Error,
            Self::EndEvents => EventType::EndEvents,
            Self::InternalReQueue => EventType::InternalReQueue,
            Self::NewAgentCreated(_) => EventType::NewAgentCreated,
            Self::ExistingAgentRetrieved(_) => EventType::ExistingAgentRetrieved,
            Self::FileIngested(_) => EventType::FileIngested,
            Self::NewThreadCreated(_) => EventType::NewThreadCreated,
            Self::NewThreadRunCreated(_) => EventType::NewThreadRunCreated,
            Self::NewThreadMessage(_) => EventType::NewThreadMessage,
            Self::ThreadMessageAdded(_) => EventType::ThreadMessageAdded,
            Self::ThreadRunComplete(_) => EventType::ThreadRunComplete,
            Self::ThreadRunInProgress(_) => EventType::ThreadRunInProgress,
            Self::ThreadRunFailed(_) => EventType::ThreadRunFailed,
            Self::ThreadRunEventWithUnknownStatus(_) => EventType::ThreadRunEventWithUnknownStatus,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn events_use_tagged_wire_shape() {
        let event = AgiEvent::NewThreadCreated(NewThreadCreatedEvent {
            agent_id: "a1".to_string(),
            thread_id: "t1".to_string(),
        });
        let value = serde_json::to_value(&event).unwrap();
        assert_eq!(value["event_type"], "new_thread_created");
        assert_eq!(value["event_data"]["agent_id"], "a1");
        assert_eq!(value["event_data"]["thread_id"], "t1");
    }

    #[test]
    fn unit_event_has_no_payload() {
        let value = serde_json::to_value(AgiEvent::EndEvents).unwrap();
        assert_eq!(value, serde_json::json!({"event_type": "end_events"}));
    }
}