agent_io/
observability.rs

1//! Observability module for agent_io with tracing integration.
2//!
3//! This module provides:
4//! - `observe` macro for tracing functions
5//! - Span helpers for manual instrumentation
6//!
7//! Uses the `tracing` crate for structured logging and tracing.
8//!
9//! # Example
10//!
11//! ```rust,ignore
12//! use agent_io::observe;
13//!
14//! #[observe(name = "my_function")]
15//! async fn my_function() -> Result<(), Box<dyn std::error::Error>> {
16//!     Ok(())
17//! }
18//! ```
19
20use tracing::{Span, info_span};
21use tracing_subscriber::EnvFilter;
22
23/// Span types for categorization
24#[derive(Debug, Clone, Copy)]
25pub enum SpanType {
26    Default,
27    Llm,
28    Tool,
29}
30
31impl SpanType {
32    pub fn as_str(&self) -> &'static str {
33        match self {
34            SpanType::Default => "DEFAULT",
35            SpanType::Llm => "LLM",
36            SpanType::Tool => "TOOL",
37        }
38    }
39}
40
41/// Extension trait for spans
42pub trait SpanExt {
43    /// Set the span output
44    fn set_output(&self, output: &str);
45
46    /// Mark the span as having an error
47    fn set_error(&self, error: &str);
48}
49
50impl SpanExt for Span {
51    fn set_output(&self, output: &str) {
52        self.record("output", output);
53    }
54
55    fn set_error(&self, error: &str) {
56        self.record("error", true);
57        self.record("error.message", error);
58    }
59}
60
61/// Check if tracing is enabled
62pub fn is_tracing_enabled() -> bool {
63    tracing::dispatcher::has_been_set()
64}
65
66/// Get observability status
67pub fn get_observability_status() -> ObservabilityStatus {
68    ObservabilityStatus {
69        tracing_enabled: is_tracing_enabled(),
70    }
71}
72
73/// Observability status information
74#[derive(Debug, Clone)]
75pub struct ObservabilityStatus {
76    pub tracing_enabled: bool,
77}
78
79/// Macro to create an observed function
80///
81/// # Example
82///
83/// ```rust,ignore
84/// #[observe(name = "my_function")]
85/// async fn my_function() -> Result<(), Error> {
86///     // function body
87///     Ok(())
88/// }
89/// ```
90#[macro_export]
91macro_rules! observe {
92    (name = $name:expr, $item:item) => {
93        #[tracing::instrument(skip_all, name = $name)]
94        $item
95    };
96    ($item:item) => {
97        #[tracing::instrument(skip_all)]
98        $item
99    };
100}
101
102/// Macro for debug-only observation
103#[macro_export]
104macro_rules! observe_debug {
105    (name = $name:expr, $item:item) => {
106        #[cfg(debug_assertions)]
107        #[tracing::instrument(skip_all, name = $name)]
108        $item
109        #[cfg(not(debug_assertions))]
110        $item
111    };
112    ($item:item) => {
113        #[cfg(debug_assertions)]
114        #[tracing::instrument(skip_all)]
115        $item
116        #[cfg(not(debug_assertions))]
117        $item
118    };
119}
120
121/// Initialize default tracing subscriber for development
122pub fn init_default_subscriber() {
123    let _ = tracing_subscriber::fmt()
124        .with_env_filter(
125            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
126        )
127        .with_target(false)
128        .try_init();
129}
130
131/// Create a span for LLM invocation
132pub fn llm_span(model: &str, provider: &str) -> Span {
133    info_span!(
134        target: "agent_io::llm",
135        "llm_invoke",
136        model = %model,
137        provider = %provider
138    )
139}
140
141/// Create a span for tool execution
142pub fn tool_span(name: &str) -> Span {
143    info_span!(
144        target: "agent_io::tool",
145        "tool_execute",
146        tool_name = %name
147    )
148}
149
150/// Create a span for agent step
151pub fn agent_span(step: usize) -> Span {
152    info_span!(
153        target: "agent_io::agent",
154        "agent_step",
155        step = step
156    )
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_span_types() {
165        assert_eq!(SpanType::Default.as_str(), "DEFAULT");
166        assert_eq!(SpanType::Llm.as_str(), "LLM");
167        assert_eq!(SpanType::Tool.as_str(), "TOOL");
168    }
169}
agent_io/observability.rs

agent_io/
observability.rs
