pub mod frame;
pub mod metrics;
pub mod runtime;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Instant;
use async_trait::async_trait;
use crate::error::Result;
use crate::observer::Observer;
use frame::{next_frame_id, Direction, Frame, FrameMeta, StartParams};
use metrics::MetricsData;
use runtime::EnvelopeSender;
pub use runtime::{run_processor, ProcessorRx, ProcessorTx};
#[derive(Debug, Clone)]
pub struct Envelope {
pub meta: FrameMeta,
pub frame: Frame,
pub direction: Direction,
}
impl Envelope {
pub fn new(frame: Frame, direction: Direction) -> Self {
Self {
meta: FrameMeta::new(&frame),
frame,
direction,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StopReason {
EndOfTask,
Stopped,
Cancelled,
}
#[derive(Clone)]
pub struct Clock {
base: Instant,
}
impl Clock {
pub fn new() -> Self {
Self {
base: Instant::now(),
}
}
pub fn now_ns(&self) -> i64 {
self.base.elapsed().as_nanos() as i64
}
}
impl Default for Clock {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct ProcessorSetup {
pub clock: Clock,
pub observer: Option<Observer>,
pub cancel: tokio_util::sync::CancellationToken,
pub enable_metrics: bool,
pub enable_usage_metrics: bool,
}
#[derive(Clone)]
pub struct Link {
pub(crate) next: Option<EnvelopeSender>,
pub(crate) prev: Option<EnvelopeSender>,
pub(crate) name: Arc<str>,
pub(crate) clock: Clock,
pub(crate) observer: Option<Observer>,
pub(crate) enable_metrics: bool,
pub(crate) enable_usage_metrics: bool,
pub(crate) ttfb_start: Arc<AtomicI64>,
pub(crate) processing_start: Arc<AtomicI64>,
}
impl Link {
pub fn name(&self) -> &str {
&self.name
}
pub async fn push(&self, meta: FrameMeta, frame: Frame, direction: Direction) {
let tx = match direction {
Direction::Downstream => self.next.as_ref(),
Direction::Upstream => self.prev.as_ref(),
};
let Some(tx) = tx else {
return;
};
if let Some(o) = &self.observer {
let dest = tx.name();
let ev = crate::observer::FramePushEvent {
source: &self.name,
destination: dest,
frame: &frame,
meta: &meta,
direction,
timestamp_ns: self.clock.now_ns(),
};
o.on_push(&ev).await;
}
let env = Envelope {
meta,
frame,
direction,
};
tx.send(env).await;
}
pub async fn push_down(&self, frame: Frame) {
let meta = FrameMeta::new(&frame);
self.push(meta, frame, Direction::Downstream).await;
}
pub async fn push_up(&self, frame: Frame) {
let meta = FrameMeta::new(&frame);
self.push(meta, frame, Direction::Upstream).await;
}
pub async fn push_error(&self, message: impl Into<String>, fatal: bool) {
let frame = Frame::Error {
message: message.into(),
fatal,
processor: Some(self.name.clone()),
};
self.push_up(frame).await;
}
pub async fn broadcast(&self, frame: Frame) {
let down_id = next_frame_id();
let up_id = next_frame_id();
let mut down_meta = FrameMeta::new(&frame);
down_meta.id = down_id;
down_meta.broadcast_sibling_id = Some(up_id);
let mut up_meta = FrameMeta::new(&frame);
up_meta.id = up_id;
up_meta.broadcast_sibling_id = Some(down_id);
self.push(down_meta, frame.clone(), Direction::Downstream)
.await;
self.push(up_meta, frame, Direction::Upstream).await;
}
pub fn start_ttfb(&self) {
if self.enable_metrics {
self.ttfb_start
.store(self.clock.now_ns(), Ordering::Relaxed);
}
}
pub async fn stop_ttfb(&self, model: Option<String>) {
if !self.enable_metrics {
return;
}
let start = self.ttfb_start.swap(0, Ordering::Relaxed);
if start == 0 {
return;
}
let seconds = (self.clock.now_ns() - start).max(0) as f64 / 1e9;
self.emit_metric(MetricsData::Ttfb {
processor: self.name.to_string(),
model,
seconds,
})
.await;
}
pub fn start_processing(&self) {
if self.enable_metrics {
self.processing_start
.store(self.clock.now_ns(), Ordering::Relaxed);
}
}
pub async fn stop_processing(&self, model: Option<String>) {
if !self.enable_metrics {
return;
}
let start = self.processing_start.swap(0, Ordering::Relaxed);
if start == 0 {
return;
}
let seconds = (self.clock.now_ns() - start).max(0) as f64 / 1e9;
self.emit_metric(MetricsData::Processing {
processor: self.name.to_string(),
model,
seconds,
})
.await;
}
pub async fn report_llm_usage(&self, model: Option<String>, tokens: metrics::LlmTokenUsage) {
if !self.enable_usage_metrics {
return;
}
self.emit_metric(MetricsData::LlmUsage {
processor: self.name.to_string(),
model,
tokens,
})
.await;
}
pub async fn report_tts_usage(&self, characters: u64) {
if !self.enable_usage_metrics {
return;
}
self.emit_metric(MetricsData::TtsUsage {
processor: self.name.to_string(),
characters,
})
.await;
}
async fn emit_metric(&self, data: MetricsData) {
self.push_down(Frame::Metrics(vec![data])).await;
}
}
#[async_trait]
pub trait FrameProcessor: Send + 'static {
fn name(&self) -> &str;
async fn start(&mut self, _setup: &ProcessorSetup, _params: &StartParams) -> Result<()> {
Ok(())
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
async fn stop(&mut self, _reason: StopReason) -> Result<()> {
Ok(())
}
fn can_generate_metrics(&self) -> bool {
false
}
#[doc(hidden)]
fn explode(&mut self) -> Option<Topology> {
None
}
}
pub fn into_topology(mut p: Box<dyn FrameProcessor>) -> Topology {
match p.explode() {
Some(topo) => topo,
None => Topology::Leaf(p),
}
}
pub enum Topology {
Leaf(Box<dyn FrameProcessor>),
Chain(Vec<Topology>),
Parallel(Vec<Topology>),
}