weavegraph 0.7.0

//! Event types emitted by workflow nodes and the framework.
use std::fmt;

use chrono::{DateTime, Utc};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Emitted when the event stream closes, signalling clean stream termination.
pub const STREAM_END_SCOPE: &str = "__weavegraph_stream_end__";

/// Emitted after each [`AppRunner::invoke_next`](crate::runtimes::AppRunner::invoke_next) call
/// so subscribers can separate logical inputs without treating the bus as closed.
pub const INVOCATION_END_SCOPE: &str = "__weavegraph_invocation_end__";

/// Scope for internal framework diagnostics; filter on this to separate telemetry
/// from user-emitted node events.
pub const DIAGNOSTIC_SCOPE: &str = "__weavegraph_diagnostic__";

/// An event that can be emitted by a workflow node or by the framework itself.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum Event {
    /// A structured event from a workflow node.
    Node(NodeEvent),
    /// A framework-internal diagnostic.
    Diagnostic(DiagnosticEvent),
    /// An LLM streaming chunk, final marker, or error.
    LLM(LLMStreamingEvent),
}

impl Event {
    /// Construct a node event with only a scope and message.
    pub fn node_message(scope: impl Into<String>, message: impl Into<String>) -> Self {
        Event::Node(NodeEvent::new(None, None, scope.into(), message.into()))
    }

    /// Construct a node event with node ID, step, scope, and message.
    pub fn node_message_with_meta(
        node_id: impl Into<String>,
        step: u64,
        scope: impl Into<String>,
        message: impl Into<String>,
    ) -> Self {
        Event::Node(NodeEvent::new(
            Some(node_id.into()),
            Some(step),
            scope.into(),
            message.into(),
        ))
    }

    /// Construct a node event with full metadata including runtime labels.
    pub fn node_message_with_metadata(
        node_id: impl Into<String>,
        step: u64,
        scope: impl Into<String>,
        message: impl Into<String>,
        metadata: FxHashMap<String, Value>,
    ) -> Self {
        Event::Node(
            NodeEvent::new(
                Some(node_id.into()),
                Some(step),
                scope.into(),
                message.into(),
            )
            .with_metadata(metadata),
        )
    }

    /// Construct a diagnostic event.
    pub fn diagnostic(scope: impl Into<String>, message: impl Into<String>) -> Self {
        Event::Diagnostic(DiagnosticEvent {
            scope: scope.into(),
            message: message.into(),
        })
    }

    /// Returns the scope label for this event.
    pub fn scope_label(&self) -> Option<&str> {
        match self {
            Event::Node(n) => Some(n.scope()),
            Event::Diagnostic(d) => Some(d.scope()),
            Event::LLM(l) => Some(l.scope().as_ref()),
        }
    }

    /// Returns the primary message text.
    pub fn message(&self) -> &str {
        match self {
            Event::Node(n) => n.message(),
            Event::Diagnostic(d) => d.message(),
            Event::LLM(l) => l.chunk(),
        }
    }

    /// Serialises the event to a normalised JSON value.
    ///
    /// The returned object has the shape:
    /// ```json
    /// {
    ///   "type": "node" | "diagnostic" | "llm",
    ///   "scope": "<scope>",
    ///   "message": "<text>",
    ///   "timestamp": "<rfc3339>",
    ///   "metadata": { ... }
    /// }
    /// ```
    ///
    /// # Example
    ///
    /// ```
    /// use weavegraph::event_bus::Event;
    ///
    /// let event = Event::node_message_with_meta("router", 5, "routing", "Processing request");
    /// let json = event.to_json_value();
    ///
    /// assert_eq!(json["type"], "node");
    /// assert_eq!(json["scope"], "routing");
    /// assert_eq!(json["message"], "Processing request");
    /// assert_eq!(json["metadata"]["node_id"], "router");
    /// assert_eq!(json["metadata"]["step"], 5);
    /// ```
    pub fn to_json_value(&self) -> serde_json::Value {
        use serde_json::json;

        let (event_type, metadata, timestamp) = match self {
            Event::Node(n) => {
                let mut meta: serde_json::Map<String, Value> = n
                    .metadata()
                    .iter()
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect();
                if let Some(id) = n.node_id() {
                    meta.insert("node_id".to_owned(), json!(id));
                }
                if let Some(step) = n.step() {
                    meta.insert("step".to_owned(), json!(step));
                }
                ("node", Value::Object(meta), Utc::now())
            }
            Event::Diagnostic(_) => ("diagnostic", json!({}), Utc::now()),
            Event::LLM(l) => {
                let mut meta: serde_json::Map<String, Value> = l
                    .metadata()
                    .iter()
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect();
                if let Some(id) = l.session_id() {
                    meta.insert("session_id".to_owned(), json!(id));
                }
                if let Some(id) = l.node_id() {
                    meta.insert("node_id".to_owned(), json!(id));
                }
                if let Some(id) = l.stream_id() {
                    meta.insert("stream_id".to_owned(), json!(id));
                }
                meta.insert("is_final".to_owned(), json!(l.is_final()));
                ("llm", Value::Object(meta), l.timestamp())
            }
        };

        json!({
            "type": event_type,
            "scope": self.scope_label(),
            "message": self.message(),
            "timestamp": timestamp.to_rfc3339(),
            "metadata": metadata,
        })
    }

    /// Serialises the event to a compact JSON string.
    ///
    /// # Example
    ///
    /// ```
    /// use weavegraph::event_bus::Event;
    ///
    /// let event = Event::diagnostic("test", "message");
    /// let json_str = event.to_json_string().unwrap();
    /// assert!(json_str.contains("\"type\":\"diagnostic\""));
    /// ```
    pub fn to_json_string(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string(&self.to_json_value())
    }

    /// Serialises the event to an indented JSON string.
    ///
    /// # Example
    ///
    /// ```
    /// use weavegraph::event_bus::Event;
    ///
    /// let event = Event::node_message("test", "hello");
    /// let json_str = event.to_json_pretty().unwrap();
    /// assert!(json_str.contains("  \"type\": \"node\""));
    /// ```
    pub fn to_json_pretty(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string_pretty(&self.to_json_value())
    }
}

impl fmt::Display for Event {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Event::Node(n) => match (n.node_id(), n.step()) {
                (Some(id), Some(step)) => write!(f, "[{id}@{step}] {}", n.message()),
                (Some(id), None) => write!(f, "[{id}] {}", n.message()),
                (None, Some(step)) => write!(f, "[step {step}] {}", n.message()),
                (None, None) => write!(f, "{}", n.message()),
            },
            Event::Diagnostic(d) => write!(f, "{}", d.message()),
            Event::LLM(l) => {
                if let Some(id) = l.stream_id() {
                    write!(f, "[LLM {id}] {}", l.chunk())
                } else if let Some(id) = l.node_id() {
                    write!(f, "[LLM {id}] {}", l.chunk())
                } else {
                    write!(f, "{}", l.chunk())
                }
            }
        }
    }
}

/// A structured event emitted by a workflow node during execution.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct NodeEvent {
    node_id: Option<String>,
    step: Option<u64>,
    scope: String,
    message: String,
    #[serde(default)]
    metadata: FxHashMap<String, Value>,
}

impl NodeEvent {
    /// Create a new node event.
    pub fn new(node_id: Option<String>, step: Option<u64>, scope: String, message: String) -> Self {
        Self {
            node_id,
            step,
            scope,
            message,
            metadata: FxHashMap::default(),
        }
    }

    /// Returns the node ID, if set.
    pub fn node_id(&self) -> Option<&str> {
        self.node_id.as_deref()
    }

    /// Returns the step number at which this event was emitted, if set.
    pub fn step(&self) -> Option<u64> {
        self.step
    }

    /// Returns the scope label.
    pub fn scope(&self) -> &str {
        &self.scope
    }

    /// Returns the message text.
    pub fn message(&self) -> &str {
        &self.message
    }

    /// Returns the metadata map.
    pub fn metadata(&self) -> &FxHashMap<String, Value> {
        &self.metadata
    }

    /// Replace the metadata map and return `self`.
    pub fn with_metadata(mut self, metadata: FxHashMap<String, Value>) -> Self {
        self.metadata = metadata;
        self
    }
}

/// A framework-internal diagnostic event, emitted outside normal node execution.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct DiagnosticEvent {
    scope: String,
    message: String,
}

impl DiagnosticEvent {
    /// Returns the scope label.
    pub fn scope(&self) -> &str {
        &self.scope
    }

    /// Returns the diagnostic message text.
    pub fn message(&self) -> &str {
        &self.message
    }
}

/// Scope discriminant for LLM streaming events.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum LLMStreamingEventScope {
    /// An in-progress streaming session (default).
    Streaming,
    /// A single text chunk within a streaming response.
    Chunk,
    /// The final chunk, marking end-of-stream.
    Final,
    /// An error event that terminates the stream.
    Error,
}

impl AsRef<str> for LLMStreamingEventScope {
    fn as_ref(&self) -> &str {
        match self {
            Self::Streaming => "stream",
            Self::Chunk => "chunk",
            Self::Final => STREAM_END_SCOPE,
            Self::Error => "error",
        }
    }
}

/// An LLM event carrying a response chunk, a final marker, or an error.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LLMStreamingEvent {
    session_id: Option<String>,
    node_id: Option<String>,
    stream_id: Option<String>,
    chunk: String,
    is_final: bool,
    scope: LLMStreamingEventScope,
    metadata: FxHashMap<String, Value>,
    timestamp: DateTime<Utc>,
}

impl LLMStreamingEvent {
    /// Return a builder for constructing a new event; `chunk` is the only required field.
    pub fn builder(chunk: impl Into<String>) -> LLMStreamingEventBuilder {
        LLMStreamingEventBuilder::new(chunk)
    }

    /// Create a partial-chunk event.
    pub fn chunk_event(
        session_id: Option<String>,
        node_id: Option<String>,
        stream_id: Option<String>,
        chunk: impl Into<String>,
        metadata: FxHashMap<String, Value>,
    ) -> Self {
        Self {
            session_id,
            node_id,
            stream_id,
            chunk: chunk.into(),
            is_final: false,
            scope: LLMStreamingEventScope::Chunk,
            metadata,
            timestamp: Utc::now(),
        }
    }

    /// Create a final-chunk event marking the end of the stream.
    pub fn final_event(
        session_id: Option<String>,
        node_id: Option<String>,
        stream_id: Option<String>,
        chunk: impl Into<String>,
        metadata: FxHashMap<String, Value>,
    ) -> Self {
        Self {
            session_id,
            node_id,
            stream_id,
            chunk: chunk.into(),
            is_final: true,
            scope: LLMStreamingEventScope::Final,
            metadata,
            timestamp: Utc::now(),
        }
    }

    /// Create an error event marking a failed streaming session.
    pub fn error_event(
        session_id: Option<String>,
        node_id: Option<String>,
        stream_id: Option<String>,
        error_message: impl Into<String>,
    ) -> Self {
        let metadata = [("severity".to_owned(), Value::String("error".to_owned()))]
            .into_iter()
            .collect();
        Self {
            session_id,
            node_id,
            stream_id,
            chunk: error_message.into(),
            is_final: true,
            scope: LLMStreamingEventScope::Error,
            metadata,
            timestamp: Utc::now(),
        }
    }

    /// Returns the session ID, if set.
    pub fn session_id(&self) -> Option<&str> {
        self.session_id.as_deref()
    }

    /// Returns the node ID, if set.
    pub fn node_id(&self) -> Option<&str> {
        self.node_id.as_deref()
    }

    /// Returns the stream ID, if set.
    pub fn stream_id(&self) -> Option<&str> {
        self.stream_id.as_deref()
    }

    /// Returns the text chunk.
    pub fn chunk(&self) -> &str {
        &self.chunk
    }

    /// Returns `true` if this is the final event in the stream.
    pub fn is_final(&self) -> bool {
        self.is_final
    }

    /// Returns the scope discriminant.
    pub fn scope(&self) -> &LLMStreamingEventScope {
        &self.scope
    }

    /// Returns the metadata map.
    pub fn metadata(&self) -> &FxHashMap<String, Value> {
        &self.metadata
    }

    /// Returns the event creation timestamp.
    pub fn timestamp(&self) -> DateTime<Utc> {
        self.timestamp
    }

    /// Replace the metadata map and return `self`.
    pub fn with_metadata(mut self, metadata: FxHashMap<String, Value>) -> Self {
        self.metadata = metadata;
        self
    }

    /// Replace the timestamp and return `self`.
    pub fn with_timestamp(mut self, timestamp: DateTime<Utc>) -> Self {
        self.timestamp = timestamp;
        self
    }
}

/// Builder for [`LLMStreamingEvent`].
///
/// Obtain one via [`LLMStreamingEvent::builder`]. All fields except `chunk` are optional;
/// `timestamp` defaults to [`Utc::now`] when [`build`](Self::build) is called.
pub struct LLMStreamingEventBuilder {
    session_id: Option<String>,
    node_id: Option<String>,
    stream_id: Option<String>,
    chunk: String,
    is_final: bool,
    scope: LLMStreamingEventScope,
    metadata: FxHashMap<String, Value>,
    timestamp: Option<DateTime<Utc>>,
}

impl LLMStreamingEventBuilder {
    fn new(chunk: impl Into<String>) -> Self {
        Self {
            session_id: None,
            node_id: None,
            stream_id: None,
            chunk: chunk.into(),
            is_final: false,
            scope: LLMStreamingEventScope::Streaming,
            metadata: FxHashMap::default(),
            timestamp: None,
        }
    }

    /// Set the session ID.
    pub fn session_id(mut self, id: impl Into<String>) -> Self {
        self.session_id = Some(id.into());
        self
    }

    /// Set the node ID.
    pub fn node_id(mut self, id: impl Into<String>) -> Self {
        self.node_id = Some(id.into());
        self
    }

    /// Set the stream ID.
    pub fn stream_id(mut self, id: impl Into<String>) -> Self {
        self.stream_id = Some(id.into());
        self
    }

    /// Mark the event as final.
    pub fn is_final(mut self, v: bool) -> Self {
        self.is_final = v;
        self
    }

    /// Set the scope discriminant (defaults to [`LLMStreamingEventScope::Streaming`]).
    pub fn scope(mut self, s: LLMStreamingEventScope) -> Self {
        self.scope = s;
        self
    }

    /// Set the metadata map.
    pub fn metadata(mut self, m: FxHashMap<String, Value>) -> Self {
        self.metadata = m;
        self
    }

    /// Fix the creation timestamp; defaults to [`Utc::now`] if omitted.
    pub fn timestamp(mut self, ts: DateTime<Utc>) -> Self {
        self.timestamp = Some(ts);
        self
    }

    /// Consume the builder and return the event.
    pub fn build(self) -> LLMStreamingEvent {
        LLMStreamingEvent {
            session_id: self.session_id,
            node_id: self.node_id,
            stream_id: self.stream_id,
            chunk: self.chunk,
            is_final: self.is_final,
            scope: self.scope,
            metadata: self.metadata,
            timestamp: self.timestamp.unwrap_or_else(Utc::now),
        }
    }
}