gemini_adk_rs/telemetry/
mod.rs1use async_trait::async_trait;
9
10use gemini_genai_rs::prelude::FunctionCall;
11
12use crate::context::InvocationContext;
13use crate::error::{AgentError, ToolError};
14use crate::llm::{LlmRequest, LlmResponse};
15use crate::middleware::Middleware;
16
17pub mod logging;
18pub mod metrics;
19pub mod setup;
20pub mod spans;
21
22pub use setup::TelemetrySetup;
23
24pub struct TelemetryMiddleware {
31 agent_name: String,
32}
33
34impl TelemetryMiddleware {
35 pub fn new(agent_name: impl Into<String>) -> Self {
37 Self {
38 agent_name: agent_name.into(),
39 }
40 }
41
42 pub fn agent_name(&self) -> &str {
44 &self.agent_name
45 }
46}
47
48#[async_trait]
49impl Middleware for TelemetryMiddleware {
50 fn name(&self) -> &str {
51 "telemetry"
52 }
53
54 async fn before_agent(&self, _ctx: &InvocationContext) -> Result<(), AgentError> {
55 metrics::record_agent_started(&self.agent_name);
56 logging::log_agent_started(&self.agent_name, 0);
57 Ok(())
58 }
59
60 async fn after_agent(&self, _ctx: &InvocationContext) -> Result<(), AgentError> {
61 Ok(())
63 }
64
65 async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
66 metrics::record_agent_tool_dispatched(&self.agent_name, &call.name);
67 logging::log_tool_dispatch(&self.agent_name, &call.name, "function");
68 Ok(())
69 }
70
71 async fn after_tool(
72 &self,
73 call: &FunctionCall,
74 _result: &serde_json::Value,
75 ) -> Result<(), AgentError> {
76 logging::log_tool_result(&self.agent_name, &call.name, true, 0.0);
77 Ok(())
78 }
79
80 async fn on_tool_error(&self, call: &FunctionCall, _err: &ToolError) -> Result<(), AgentError> {
81 logging::log_tool_result(&self.agent_name, &call.name, false, 0.0);
82 Ok(())
83 }
84
85 async fn on_error(&self, err: &AgentError) -> Result<(), AgentError> {
86 metrics::record_agent_error(&self.agent_name, &err.to_string());
87 logging::log_agent_error(&self.agent_name, &err.to_string());
88 Ok(())
89 }
90
91 async fn before_model(&self, _request: &LlmRequest) -> Result<Option<LlmResponse>, AgentError> {
92 logging::log_tool_dispatch(&self.agent_name, "llm", "model_call");
93 Ok(None)
94 }
95
96 async fn after_model(
97 &self,
98 _request: &LlmRequest,
99 _response: &LlmResponse,
100 ) -> Result<Option<LlmResponse>, AgentError> {
101 if let Some(_usage) = &_response.usage {
102 metrics::record_agent_completed(&self.agent_name, 0.0); }
104 Ok(None)
105 }
106}
107
108#[cfg(test)]
109mod tests {
110 use super::*;
111
112 #[tokio::test]
113 async fn telemetry_middleware_hooks_dont_panic() {
114 let mw = TelemetryMiddleware::new("test_agent");
115 assert_eq!(mw.name(), "telemetry");
116 assert_eq!(mw.agent_name(), "test_agent");
117
118 let call = FunctionCall {
119 name: "my_tool".to_string(),
120 args: serde_json::json!({}),
121 id: None,
122 };
123 let result = serde_json::json!({"ok": true});
124 let tool_err = ToolError::ExecutionFailed("boom".to_string());
125 let agent_err = AgentError::Other("oops".to_string());
126
127 assert!(mw.before_tool(&call).await.is_ok());
129 assert!(mw.after_tool(&call, &result).await.is_ok());
130 assert!(mw.on_tool_error(&call, &tool_err).await.is_ok());
131 assert!(mw.on_error(&agent_err).await.is_ok());
132 }
133}