use super::traits::{Event, EventType};
use crate::InternalMessage;
use serde::{Deserialize, Serialize};
use std::time::{SystemTime, UNIX_EPOCH};
fn generate_id() -> String {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("evt_{:x}", now)
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as u64
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
pub model_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub provider: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageEvent {
pub event_id: String,
pub session_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub project_hash: Option<String>,
pub timestamp_ms: u64,
pub sequence: u32,
pub message: InternalMessage,
#[serde(skip_serializing_if = "Option::is_none")]
pub token_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model_info: Option<ModelInfo>,
}
impl MessageEvent {
pub fn new(session_id: impl Into<String>, sequence: u32, message: InternalMessage) -> Self {
Self {
event_id: generate_id(),
session_id: session_id.into(),
project_hash: None,
timestamp_ms: now_ms(),
sequence,
message,
token_count: None,
model_info: None,
}
}
pub fn user(session_id: impl Into<String>, sequence: u32, content: impl Into<String>) -> Self {
Self::new(session_id, sequence, InternalMessage::user(content))
}
pub fn assistant(
session_id: impl Into<String>,
sequence: u32,
content: impl Into<String>,
) -> Self {
Self::new(session_id, sequence, InternalMessage::assistant(content))
}
pub fn system(
session_id: impl Into<String>,
sequence: u32,
content: impl Into<String>,
) -> Self {
Self::new(session_id, sequence, InternalMessage::system(content))
}
pub fn with_project(mut self, project_hash: impl Into<String>) -> Self {
self.project_hash = Some(project_hash.into());
self
}
pub fn with_token_count(mut self, count: usize) -> Self {
self.token_count = Some(count);
self
}
pub fn with_model_info(mut self, model: impl Into<String>, provider: Option<String>) -> Self {
self.model_info = Some(ModelInfo {
model_name: model.into(),
provider,
});
self
}
pub fn with_event_id(mut self, event_id: impl Into<String>) -> Self {
self.event_id = event_id.into();
self
}
}
impl Event for MessageEvent {
fn event_id(&self) -> &str {
&self.event_id
}
fn event_type(&self) -> EventType {
EventType::Message
}
fn session_id(&self) -> &str {
&self.session_id
}
fn timestamp_ms(&self) -> u64 {
self.timestamp_ms
}
fn sequence(&self) -> u32 {
self.sequence
}
fn to_json(&self) -> serde_json::Value {
serde_json::to_value(self).unwrap()
}
}