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