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