ferrin-core 0.2.0

//! Fans telemetry events out to the configured integrations and to
//! `tracing`.

mod modalities;

use std::panic::AssertUnwindSafe;
use std::panic::catch_unwind;
use std::sync::Arc;

use ferrin_spec::BoxFuture;
use ferrin_tool::ToolError;
use futures_util::FutureExt;
use futures_util::future::join_all;

use super::AbortEvent;
use super::EmbedEndEvent;
use super::EmbedStartEvent;
use super::EndEvent;
use super::ErrorEvent;
use super::ModelCallContext;
use super::ModelCallEndEvent;
use super::ModelCallOutcome;
use super::ModelCallStartEvent;
use super::RerankEndEvent;
use super::RerankStartEvent;
use super::StartEvent;
use super::StepEndEvent;
use super::StepStartEvent;
use super::Telemetry;
use super::TelemetryOptions;
use super::ToolExecutionContext;
use super::ToolExecutionEndEvent;
use super::ToolExecutionStartEvent;
use super::ToolOutcome;
use crate::error::Error;

/// Dispatches events to every integration of a [`TelemetryOptions`].
#[derive(Clone, Debug)]
pub(crate) struct TelemetryDispatcher {
    options: Arc<TelemetryOptions>,
}

macro_rules! dispatch {
    ($name:ident, $event:ty) => {
        pub(crate) async fn $name(&self, event: &$event) {
            if !self.options.enabled {
                return;
            }
            self.dispatch(|integration| integration.$name(event)).await;
        }
    };
}

macro_rules! dispatch_with_context {
    ($name:ident, $event:ty) => {
        pub(crate) async fn $name(&self, event: &$event) {
            if !self.options.enabled {
                return;
            }
            let mut recorded = event.clone();
            if !self.options.include_runtime_context {
                recorded.runtime_context = None;
            }
            self.dispatch(|integration| integration.$name(&recorded))
                .await;
        }
    };
}

impl TelemetryDispatcher {
    async fn dispatch<'a>(
        &'a self,
        callback: impl Fn(&'a dyn Telemetry) -> BoxFuture<'a, ()> + Send + Sync,
    ) {
        let futures = self.options.integrations.iter().filter_map(|integration| {
            catch_unwind(AssertUnwindSafe(|| callback(integration.as_ref())))
                .ok()
                .map(|future| AssertUnwindSafe(future).catch_unwind())
        });
        let _ = join_all(futures).await;
    }

    pub(crate) fn new(options: TelemetryOptions) -> Self {
        Self {
            options: Arc::new(options),
        }
    }

    pub(crate) fn record_inputs(&self) -> bool {
        self.options.enabled && self.options.record_inputs
    }

    pub(crate) fn record_outputs(&self) -> bool {
        self.options.enabled && self.options.record_outputs
    }

    dispatch_with_context!(on_start, StartEvent);
    dispatch_with_context!(on_step_start, StepStartEvent);
    dispatch_with_context!(on_language_model_call_start, ModelCallStartEvent);

    dispatch_with_context!(on_tool_execution_start, ToolExecutionStartEvent);
    pub(crate) async fn on_tool_execution_end(&self, event: &ToolExecutionEndEvent) {
        if !self.options.enabled {
            return;
        }
        let mut recorded = event.clone();
        if !self.options.include_runtime_context {
            recorded.runtime_context = None;
        }
        if !self.record_outputs() {
            recorded.output = None;
            recorded.error = recorded
                .error
                .map(|_| crate::generate_text::ToolErrorInfo::text(super::redact::REDACTED));
        }
        self.dispatch(|integration| integration.on_tool_execution_end(&recorded))
            .await;
    }

    dispatch!(on_abort, AbortEvent);

    pub(crate) async fn on_language_model_call_end(&self, event: &ModelCallEndEvent) {
        if !self.options.enabled {
            return;
        }
        let mut recorded = event.clone();
        if !self.options.include_runtime_context {
            recorded.runtime_context = None;
        }
        if !(self.record_inputs() && self.record_outputs()) {
            recorded.warnings = super::redact::warnings(&recorded.warnings);
        }
        if !self.record_outputs() {
            recorded.content = None;
            recorded.response.body = None;
        }
        self.dispatch(|integration| integration.on_language_model_call_end(&recorded))
            .await;
    }

    pub(crate) async fn on_step_end(&self, event: &StepEndEvent) {
        if !self.options.enabled {
            return;
        }
        let recorded = StepEndEvent {
            call_id: event.call_id.clone(),
            step: Arc::new(self.recorded_step(&event.step)),
        };
        self.dispatch(|integration| integration.on_step_end(&recorded))
            .await;
    }

    pub(crate) async fn on_end(&self, event: &EndEvent) {
        if !self.options.enabled {
            return;
        }
        let recorded = EndEvent {
            runtime_context: self
                .options
                .include_runtime_context
                .then(|| event.runtime_context.clone())
                .flatten(),
            call_id: event.call_id.clone(),
            steps: event
                .steps
                .iter()
                .map(|step| self.recorded_step(step))
                .collect(),
            total_usage: event.total_usage.clone(),
            output_recorded: self
                .record_outputs()
                .then(|| event.output_recorded.clone())
                .flatten(),
        };
        self.dispatch(|integration| integration.on_end(&recorded))
            .await;
    }

    /// Filters the telemetry copy without changing application hooks or results.
    fn recorded_step(
        &self,
        step: &crate::generate_text::StepResult,
    ) -> crate::generate_text::StepResult {
        let mut recorded = step.clone();
        if !self.options.include_runtime_context {
            recorded.runtime_context = None;
        }
        if !self.options.include_tools_context {
            recorded.tools_context = None;
        }
        if !(self.record_inputs() && self.record_outputs()) {
            recorded.warnings = super::redact::warnings(&recorded.warnings);
        }
        if !self.record_inputs() {
            recorded.request.body = None;
            recorded.request.messages = None;
        }
        if !self.record_outputs() {
            recorded.content.clear();
            recorded.response.body = None;
            recorded.response.messages.clear();
            recorded.provider_metadata = None;
        }
        recorded
    }

    pub(crate) async fn on_error(&self, event: &ErrorEvent<'_>) {
        if !self.options.enabled {
            return;
        }
        let error = (!(self.record_inputs() && self.record_outputs()))
            .then(|| super::redact::redact_error(event.error, &self.options));
        let recorded = ErrorEvent {
            call_id: event.call_id,
            error: error.as_ref().unwrap_or(event.error),
            phase: event.phase,
        };
        self.dispatch(|integration| integration.on_error(&recorded))
            .await;
    }

    /// Wraps `call` with every integration's `execute_language_model_call`;
    /// the last integration becomes the outermost wrapper.
    pub(crate) fn execute_language_model_call<'a>(
        &'a self,
        ctx: &'a ModelCallContext,
        call: BoxFuture<'a, Result<ModelCallOutcome, Error>>,
    ) -> BoxFuture<'a, Result<ModelCallOutcome, Error>> {
        if !self.options.enabled {
            return call;
        }
        self.options
            .integrations
            .iter()
            .fold(call, |inner, integration| {
                integration.execute_language_model_call(ctx, inner)
            })
    }

    /// Wraps `call` with every integration's `execute_tool`.
    pub(crate) fn execute_tool<'a>(
        &'a self,
        ctx: &'a ToolExecutionContext,
        call: BoxFuture<'a, Result<ToolOutcome, ToolError>>,
    ) -> BoxFuture<'a, Result<ToolOutcome, ToolError>> {
        if !self.options.enabled {
            return call;
        }
        self.options
            .integrations
            .iter()
            .fold(call, |inner, integration| {
                integration.execute_tool(ctx, inner)
            })
    }
}

impl TelemetryDispatcher {
    dispatch!(on_embed_start, EmbedStartEvent);
    dispatch!(on_embed_end, EmbedEndEvent);
    dispatch!(on_rerank_start, RerankStartEvent);
    dispatch!(on_rerank_end, RerankEndEvent);
}

impl Default for TelemetryDispatcher {
    fn default() -> Self {
        Self::new(TelemetryOptions::default())
    }
}

impl dyn Telemetry {
    /// Convenience for tests: returns `true` when `self` is the same object.
    #[must_use]
    pub fn ptr_eq(this: &Arc<Self>, other: &Arc<Self>) -> bool {
        Arc::ptr_eq(this, other)
    }
}