Skip to main content

ai_agents_observability/
export.rs

1use crate::config::{ExportFormat, RawEventsFormat};
2use crate::manager::ObservabilityManager;
3use std::path::{Path, PathBuf};
4
5/// Result returned after writing configured observability export files.
6#[derive(Debug, Clone, Default)]
7pub struct ExportResult {
8    /// Paths written by this export call.
9    pub paths: Vec<PathBuf>,
10}
11
12/// Writes report, aggregate, raw event, and Prometheus files requested by config.
13pub fn export_observability(manager: &ObservabilityManager) -> std::io::Result<ExportResult> {
14    let mut result = ExportResult::default();
15    let base = PathBuf::from(&manager.config().export.path);
16
17    if manager.config().export.write_report && manager.wants_format(ExportFormat::Json) {
18        let path = output_path(&base, "report.json", "json");
19        ensure_parent(&path)?;
20        let report = manager.generate_report();
21        let json = serde_json::to_string_pretty(&report)?;
22        std::fs::write(&path, json)?;
23        result.paths.push(path);
24    }
25
26    if manager.wants_format(ExportFormat::Csv) {
27        let path = output_path(&base, "aggregates.csv", "csv");
28        ensure_parent(&path)?;
29        std::fs::write(&path, render_csv(manager))?;
30        result.paths.push(path);
31    }
32
33    if manager.config().export.write_raw_events {
34        match manager.config().export.raw_events_format {
35            RawEventsFormat::Jsonl if manager.wants_format(ExportFormat::Jsonl) => {
36                let path = output_path(&base, "events.jsonl", "jsonl");
37                ensure_parent(&path)?;
38                std::fs::write(&path, render_jsonl(manager)?)?;
39                result.paths.push(path);
40            }
41            RawEventsFormat::Json if manager.wants_format(ExportFormat::Json) => {
42                let path = raw_json_output_path(&base, manager.config().export.write_report);
43                ensure_parent(&path)?;
44                std::fs::write(&path, render_json_events(manager)?)?;
45                result.paths.push(path);
46            }
47            _ => {}
48        }
49    }
50
51    if manager.wants_format(ExportFormat::Prometheus) {
52        let path = output_path(&base, "metrics.prom", "prom");
53        ensure_parent(&path)?;
54        std::fs::write(&path, manager.render_prometheus())?;
55        result.paths.push(path);
56    }
57
58    Ok(result)
59}
60
61fn output_path(base: &Path, default_file: &str, extension: &str) -> PathBuf {
62    if base.extension().and_then(|ext| ext.to_str()) == Some(extension) {
63        base.to_path_buf()
64    } else if base.extension().is_some() && !base.exists() {
65        base.to_path_buf()
66    } else {
67        base.join(default_file)
68    }
69}
70
71fn raw_json_output_path(base: &Path, report_enabled: bool) -> PathBuf {
72    if report_enabled && base.extension().and_then(|ext| ext.to_str()) == Some("json") {
73        base.parent()
74            .filter(|parent| !parent.as_os_str().is_empty())
75            .map(|parent| parent.join("events.json"))
76            .unwrap_or_else(|| PathBuf::from("events.json"))
77    } else {
78        output_path(base, "events.json", "json")
79    }
80}
81
82fn ensure_parent(path: &Path) -> std::io::Result<()> {
83    if let Some(parent) = path.parent() {
84        if !parent.as_os_str().is_empty() {
85            std::fs::create_dir_all(parent)?;
86        }
87    }
88    Ok(())
89}
90
91fn render_jsonl(manager: &ObservabilityManager) -> std::io::Result<String> {
92    let mut output = String::new();
93    for event in manager.raw_events() {
94        output.push_str(&serde_json::to_string(&event)?);
95        output.push('\n');
96    }
97    Ok(output)
98}
99
100fn render_json_events(manager: &ObservabilityManager) -> std::io::Result<String> {
101    serde_json::to_string_pretty(&manager.raw_events()).map_err(std::io::Error::other)
102}
103
104fn render_csv(manager: &ObservabilityManager) -> String {
105    let mut output = String::from(
106        "dimensions,count,errors,min_ms,max_ms,avg_ms,p50_ms,p90_ms,p95_ms,p99_ms,total_tokens,total_cost_usd\n",
107    );
108    for metric in manager.get_metrics() {
109        let dimensions = serde_json::to_string(&metric.dimensions).unwrap_or_default();
110        output.push_str(&format!(
111            "\"{}\",{},{},{},{},{:.3},{},{},{},{},{},{:.8}\n",
112            dimensions.replace('"', "\"\""),
113            metric.count,
114            metric.errors,
115            metric.latency.min_ms,
116            metric.latency.max_ms,
117            metric.latency.avg_ms,
118            metric.latency.p50_ms,
119            metric.latency.p90_ms,
120            metric.latency.p95_ms,
121            metric.latency.p99_ms,
122            metric.tokens.total_tokens,
123            metric.cost.total_usd
124        ));
125    }
126    output
127}