Skip to main content

ai_agents_observability/
config.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::path::{Path, PathBuf};
4
5use crate::ObservabilityError;
6
7/// Top-level YAML and Rust configuration for metrics, privacy, aggregation, cost, and export behavior.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ObservabilityConfig {
10    /// Enables all observability collection when true.
11    #[serde(default)]
12    pub enabled: bool,
13    /// Controls which latency events are recorded.
14    #[serde(default)]
15    pub latency: LatencyConfig,
16    /// Controls token counting and estimation behavior.
17    #[serde(default)]
18    pub tokens: TokenConfig,
19    /// Controls cost estimation and pricing lookup behavior.
20    #[serde(default)]
21    pub cost: CostConfig,
22    /// Controls how the language dimension is resolved from runtime context.
23    #[serde(default)]
24    pub language: LanguageConfig,
25    /// Controls the configured aggregate metrics table.
26    #[serde(default)]
27    pub aggregation: AggregationConfig,
28    /// Controls raw text retention, hashing, truncation, and redaction.
29    #[serde(default)]
30    pub privacy: PrivacyConfig,
31    /// Controls file export formats and paths.
32    #[serde(default)]
33    pub export: ExportConfig,
34    /// Controls event queue and raw event buffer limits.
35    #[serde(default)]
36    pub buffer: BufferConfig,
37}
38
39impl Default for ObservabilityConfig {
40    fn default() -> Self {
41        Self {
42            enabled: false,
43            latency: LatencyConfig::default(),
44            tokens: TokenConfig::default(),
45            cost: CostConfig::default(),
46            language: LanguageConfig::default(),
47            aggregation: AggregationConfig::default(),
48            privacy: PrivacyConfig::default(),
49            export: ExportConfig::default(),
50            buffer: BufferConfig::default(),
51        }
52    }
53}
54
55impl ObservabilityConfig {
56    /// Validates bounds that would otherwise make aggregation or buffering unusable.
57    pub fn validate(&self) -> Result<(), ObservabilityError> {
58        if self.aggregation.window_size == 0 {
59            return Err(ObservabilityError::Config(
60                "observability.aggregation.window_size must be greater than zero".to_string(),
61            ));
62        }
63        if self.buffer.event_buffer == 0 {
64            return Err(ObservabilityError::Config(
65                "observability.buffer.event_buffer must be greater than zero".to_string(),
66            ));
67        }
68        for percentile in &self.aggregation.percentiles {
69            if !(0.0..=1.0).contains(percentile) {
70                return Err(ObservabilityError::Config(format!(
71                    "observability.aggregation.percentiles value {} is outside 0.0..=1.0",
72                    percentile
73                )));
74            }
75        }
76        Ok(())
77    }
78
79    /// Loads cost.pricing_file and merges it with inline pricing.
80    pub fn with_pricing_file_loaded(
81        mut self,
82        base_dir: Option<&Path>,
83    ) -> Result<Self, ObservabilityError> {
84        let Some(path) = self.cost.pricing_file.clone() else {
85            return Ok(self);
86        };
87        let resolved = resolve_pricing_path(&path, base_dir);
88        let content = std::fs::read_to_string(&resolved).map_err(ObservabilityError::Io)?;
89        let mut file_pricing = parse_pricing_file(&resolved, &content)?;
90        let inline_pricing = std::mem::take(&mut self.cost.pricing);
91        for (key, value) in inline_pricing {
92            file_pricing.insert(key.to_lowercase(), value);
93        }
94        self.cost.pricing = file_pricing;
95        Ok(self)
96    }
97}
98
99/// Latency switches for categories that can produce duration events.
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct LatencyConfig {
102    #[serde(default = "default_true")]
103    pub track_llm: bool,
104    #[serde(default = "default_true")]
105    pub track_tools: bool,
106    #[serde(default = "default_true")]
107    pub track_skills: bool,
108    #[serde(default = "default_true")]
109    pub track_orchestration: bool,
110    #[serde(default = "default_true")]
111    pub track_hitl: bool,
112    #[serde(default)]
113    pub detailed_breakdown: bool,
114}
115
116impl Default for LatencyConfig {
117    fn default() -> Self {
118        Self {
119            track_llm: true,
120            track_tools: true,
121            track_skills: true,
122            track_orchestration: true,
123            track_hitl: true,
124            detailed_breakdown: false,
125        }
126    }
127}
128
129/// Token counting settings for provider usage and fallback estimation.
130#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct TokenConfig {
132    #[serde(default = "default_true")]
133    pub count_input: bool,
134    #[serde(default = "default_true")]
135    pub count_output: bool,
136    #[serde(default = "default_true")]
137    pub estimate_when_missing: bool,
138    #[serde(default)]
139    pub breakdown_by_component: bool,
140}
141
142impl Default for TokenConfig {
143    fn default() -> Self {
144        Self {
145            count_input: true,
146            count_output: true,
147            estimate_when_missing: true,
148            breakdown_by_component: false,
149        }
150    }
151}
152
153/// Cost estimation settings and model pricing sources.
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct CostConfig {
156    /// Enables cost estimation when token usage is available.
157    #[serde(default = "default_true")]
158    pub enabled: bool,
159    /// Inline pricing keyed by model or provider/model.
160    #[serde(default)]
161    pub pricing: HashMap<String, ModelPricing>,
162    /// Optional JSON or YAML pricing file loaded by the runtime builder or Rust helper.
163    #[serde(default)]
164    pub pricing_file: Option<String>,
165    /// Controls how unknown model prices are represented.
166    #[serde(default)]
167    pub unknown_price_policy: UnknownPricePolicy,
168}
169
170impl Default for CostConfig {
171    fn default() -> Self {
172        Self {
173            enabled: true,
174            pricing: HashMap::new(),
175            pricing_file: None,
176            unknown_price_policy: UnknownPricePolicy::Omit,
177        }
178    }
179}
180
181/// Per-thousand-token price for one model.
182#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
183pub struct ModelPricing {
184    /// Price for one thousand input tokens in USD.
185    pub input_per_1k: f64,
186    /// Price for one thousand output tokens in USD.
187    pub output_per_1k: f64,
188}
189
190/// Behavior when token usage exists but no configured price matches the model.
191#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)]
192#[serde(rename_all = "snake_case")]
193pub enum UnknownPricePolicy {
194    #[default]
195    Omit,
196    Zero,
197    Error,
198}
199
200/// Language dimension lookup rules for reports and aggregations.
201#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct LanguageConfig {
203    #[serde(default = "default_language_paths")]
204    pub paths: Vec<String>,
205    #[serde(default = "default_unknown")]
206    pub fallback: String,
207}
208
209impl Default for LanguageConfig {
210    fn default() -> Self {
211        Self {
212            paths: default_language_paths(),
213            fallback: default_unknown(),
214        }
215    }
216}
217
218/// Grouping and rolling-window settings for aggregate metrics.
219#[derive(Debug, Clone, Serialize, Deserialize)]
220pub struct AggregationConfig {
221    #[serde(default = "default_dimensions")]
222    pub dimensions: Vec<AggregationDimension>,
223    #[serde(default = "default_percentiles")]
224    pub percentiles: Vec<f64>,
225    #[serde(default = "default_window_size")]
226    pub window_size: usize,
227}
228
229impl Default for AggregationConfig {
230    fn default() -> Self {
231        Self {
232            dimensions: default_dimensions(),
233            percentiles: default_percentiles(),
234            window_size: default_window_size(),
235        }
236    }
237}
238
239/// Supported fields that can be used as aggregation dimensions.
240#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
241#[serde(rename_all = "snake_case")]
242pub enum AggregationDimension {
243    Agent,
244    Actor,
245    Model,
246    Provider,
247    Alias,
248    Purpose,
249    Language,
250    State,
251    Tool,
252    Skill,
253    OrchestrationPattern,
254    Status,
255    Custom(String),
256}
257
258impl AggregationDimension {
259    /// Returns the stable dimension key used in reports and CSV output.
260    pub fn key(&self) -> String {
261        match self {
262            Self::Agent => "agent".to_string(),
263            Self::Actor => "actor".to_string(),
264            Self::Model => "model".to_string(),
265            Self::Provider => "provider".to_string(),
266            Self::Alias => "alias".to_string(),
267            Self::Purpose => "purpose".to_string(),
268            Self::Language => "language".to_string(),
269            Self::State => "state".to_string(),
270            Self::Tool => "tool".to_string(),
271            Self::Skill => "skill".to_string(),
272            Self::OrchestrationPattern => "orchestration_pattern".to_string(),
273            Self::Status => "status".to_string(),
274            Self::Custom(name) => format!("custom:{}", name),
275        }
276    }
277}
278
279/// Privacy controls for raw payload retention, hashes, truncation, and redaction.
280#[derive(Debug, Clone, Serialize, Deserialize)]
281pub struct PrivacyConfig {
282    #[serde(default)]
283    pub include_prompts: bool,
284    #[serde(default)]
285    pub include_responses: bool,
286    #[serde(default)]
287    pub include_tool_args: bool,
288    #[serde(default)]
289    pub include_tool_outputs: bool,
290    #[serde(default)]
291    pub max_text_chars: usize,
292    #[serde(default = "default_true")]
293    pub hash_inputs: bool,
294    #[serde(default = "default_redact_keys")]
295    pub redact_keys: Vec<String>,
296    #[serde(default = "default_redact_paths")]
297    pub redact_paths: Vec<String>,
298}
299
300impl Default for PrivacyConfig {
301    fn default() -> Self {
302        Self {
303            include_prompts: false,
304            include_responses: false,
305            include_tool_args: false,
306            include_tool_outputs: false,
307            max_text_chars: 0,
308            hash_inputs: true,
309            redact_keys: default_redact_keys(),
310            redact_paths: default_redact_paths(),
311        }
312    }
313}
314
315/// File export settings for reports, aggregates, raw events, and Prometheus metrics.
316#[derive(Debug, Clone, Serialize, Deserialize)]
317pub struct ExportConfig {
318    #[serde(default = "default_export_formats")]
319    pub formats: Vec<ExportFormat>,
320    #[serde(default = "default_export_path")]
321    pub path: String,
322    #[serde(default)]
323    pub write_raw_events: bool,
324    #[serde(default = "default_true")]
325    pub write_report: bool,
326    #[serde(default)]
327    pub raw_events_format: RawEventsFormat,
328}
329
330impl Default for ExportConfig {
331    fn default() -> Self {
332        Self {
333            formats: default_export_formats(),
334            path: default_export_path(),
335            write_raw_events: false,
336            write_report: true,
337            raw_events_format: RawEventsFormat::Jsonl,
338        }
339    }
340}
341
342/// Export formats supported by ObservabilityManager::export.
343#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)]
344#[serde(rename_all = "snake_case")]
345pub enum ExportFormat {
346    #[default]
347    Json,
348    Csv,
349    Jsonl,
350    Prometheus,
351}
352
353/// Raw event file shape used when raw event export is enabled.
354#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)]
355#[serde(rename_all = "snake_case")]
356pub enum RawEventsFormat {
357    #[default]
358    Jsonl,
359    Json,
360}
361
362/// Bounded queue and raw event retention limits.
363#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct BufferConfig {
365    #[serde(default = "default_event_buffer")]
366    pub event_buffer: usize,
367    #[serde(default = "default_raw_event_limit")]
368    pub raw_event_limit: usize,
369    #[serde(default = "default_true")]
370    pub drop_on_full: bool,
371}
372
373impl Default for BufferConfig {
374    fn default() -> Self {
375        Self {
376            event_buffer: default_event_buffer(),
377            raw_event_limit: default_raw_event_limit(),
378            drop_on_full: true,
379        }
380    }
381}
382
383pub fn default_true() -> bool {
384    true
385}
386
387fn default_unknown() -> String {
388    "unknown".to_string()
389}
390
391fn default_language_paths() -> Vec<String> {
392    vec![
393        "detected_language".to_string(),
394        "input.language".to_string(),
395        "user.language".to_string(),
396        "context.user.language".to_string(),
397    ]
398}
399
400fn default_dimensions() -> Vec<AggregationDimension> {
401    vec![AggregationDimension::Model, AggregationDimension::Purpose]
402}
403
404fn default_percentiles() -> Vec<f64> {
405    vec![0.5, 0.9, 0.95, 0.99]
406}
407
408fn default_window_size() -> usize {
409    1000
410}
411
412fn default_redact_keys() -> Vec<String> {
413    vec![
414        "api_key".to_string(),
415        "authorization".to_string(),
416        "token".to_string(),
417        "password".to_string(),
418        "secret".to_string(),
419    ]
420}
421
422fn default_redact_paths() -> Vec<String> {
423    vec![
424        "actor_facts".to_string(),
425        "relationship_memory".to_string(),
426        "persona.secrets".to_string(),
427    ]
428}
429
430fn default_export_formats() -> Vec<ExportFormat> {
431    vec![ExportFormat::Json]
432}
433
434fn default_export_path() -> String {
435    "./observability_data/".to_string()
436}
437
438fn default_event_buffer() -> usize {
439    4096
440}
441
442fn default_raw_event_limit() -> usize {
443    10_000
444}
445
446fn resolve_pricing_path(path: &str, base_dir: Option<&Path>) -> PathBuf {
447    let path = PathBuf::from(path);
448    if path.is_absolute() {
449        path
450    } else if let Some(base_dir) = base_dir {
451        base_dir.join(path)
452    } else {
453        path
454    }
455}
456
457fn parse_pricing_file(
458    path: &Path,
459    content: &str,
460) -> Result<HashMap<String, ModelPricing>, ObservabilityError> {
461    let parsed: HashMap<String, ModelPricing> = match path.extension().and_then(|ext| ext.to_str())
462    {
463        Some("json") => serde_json::from_str(content).map_err(ObservabilityError::Serialization)?,
464        Some("yaml") | Some("yml") | None => serde_yaml::from_str(content).map_err(|error| {
465            ObservabilityError::Config(format!(
466                "failed to parse observability.cost.pricing_file '{}': {}",
467                path.display(),
468                error
469            ))
470        })?,
471        Some(other) => {
472            return Err(ObservabilityError::Config(format!(
473                "unsupported observability.cost.pricing_file extension '{}': {}",
474                other,
475                path.display()
476            )));
477        }
478    };
479    Ok(parsed
480        .into_iter()
481        .map(|(key, value)| (key.to_lowercase(), value))
482        .collect())
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488
489    #[test]
490    fn defaults_are_privacy_safe() {
491        let config = ObservabilityConfig::default();
492        assert!(!config.enabled);
493        assert!(!config.privacy.include_prompts);
494        assert!(!config.privacy.include_responses);
495        assert!(!config.privacy.include_tool_args);
496        assert_eq!(config.privacy.max_text_chars, 0);
497    }
498
499    #[test]
500    fn deserializes_minimal_enabled_config() {
501        let yaml = r#"
502observability:
503  enabled: true
504  aggregation:
505    dimensions: [agent, model, purpose, language]
506"#;
507        #[derive(Deserialize)]
508        struct Wrapper {
509            observability: ObservabilityConfig,
510        }
511        let parsed: Wrapper = serde_yaml::from_str(yaml).unwrap();
512        assert!(parsed.observability.enabled);
513        assert_eq!(parsed.observability.aggregation.dimensions.len(), 4);
514    }
515
516    #[test]
517    fn validation_rejects_bad_percentile() {
518        let mut config = ObservabilityConfig::default();
519        config.aggregation.percentiles = vec![1.2];
520        assert!(config.validate().is_err());
521    }
522
523    #[test]
524    fn pricing_file_loads_and_inline_overrides() {
525        let dir = std::env::temp_dir().join(format!(
526            "ai_agents_observability_pricing_{}",
527            uuid::Uuid::new_v4()
528        ));
529        std::fs::create_dir_all(&dir).unwrap();
530        std::fs::write(
531            dir.join("pricing.yaml"),
532            "openai/test:\n  input_per_1k: 0.1\n  output_per_1k: 0.2\nopenai/other:\n  input_per_1k: 1.0\n  output_per_1k: 2.0\n",
533        )
534        .unwrap();
535
536        let mut config = ObservabilityConfig::default();
537        config.cost.pricing_file = Some("pricing.yaml".to_string());
538        config.cost.pricing.insert(
539            "openai/test".to_string(),
540            ModelPricing {
541                input_per_1k: 0.3,
542                output_per_1k: 0.4,
543            },
544        );
545
546        let loaded = config.with_pricing_file_loaded(Some(&dir)).unwrap();
547        let overridden = loaded.cost.pricing.get("openai/test").unwrap();
548        assert_eq!(overridden.input_per_1k, 0.3);
549        assert_eq!(overridden.output_per_1k, 0.4);
550        assert!(loaded.cost.pricing.contains_key("openai/other"));
551
552        let _ = std::fs::remove_dir_all(dir);
553    }
554}