diff --git a/Cargo.lock b/Cargo.lock index 34d6d22..d878224 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1379,6 +1379,7 @@ dependencies = [ "tarpc", "thiserror 2.0.18", "tokio", + "tracing", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 4b0121b..1cf550f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,12 @@ schemars = "1.2.1" thiserror = "2.0.18" +tracing = { version = "0.1.44", optional = true } + jsonwebtoken = "9.3.1" +[features] +tracing = ["dep:tracing"] + [dev-dependencies] serial_test = "3.4.0" diff --git a/README.md b/README.md index 3b491ec..304cea3 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,10 @@ async fn main() -> Result<(), Box> { --- +## Feature Flag + +**`tracing`**: Enables tracing for logging request/response traces. It requires a `RUST_LOG` environment variable to be set. + ## Contributing Pull requests and issues are welcome! Please open an issue to discuss your ideas or report bugs. diff --git a/src/agent.rs b/src/agent.rs index 03a2bfd..520910f 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -28,6 +28,10 @@ impl AgentServer { tarpc::serde_transport::tcp::listen(self.socket_addr, Json::default).await?; listener.config_mut().max_frame_length(usize::MAX); + #[cfg(feature = "tracing")] + tracing::info!("Listening on: {}", listener.local_addr()); + + #[cfg(not(feature = "tracing"))] println!("Listening on: {}", listener.local_addr()); listener @@ -73,6 +77,10 @@ pub(crate) trait AgentWorker { impl AgentWorker for AgentServer { /// Handles a user message by passing it to the completion provider and returning the response. + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "agent.message", skip(self, _context, user_message)) + )] async fn message( self, _context: ::tarpc::context::Context, diff --git a/src/jwt.rs b/src/jwt.rs index 5ea8824..2b8be14 100644 --- a/src/jwt.rs +++ b/src/jwt.rs @@ -5,6 +5,10 @@ use crate::error::Error; /// JWT token decoder /// Decodes a JWT token and returns the claims if valid. +#[cfg_attr( + feature = "tracing", + tracing::instrument(name = "jwt.decode", skip(token, hmac_secret)) +)] pub(crate) fn decode_jwt(token: &str, hmac_secret: &str) -> Result { let validation = jsonwebtoken::Validation::new(Algorithm::HS256); diff --git a/src/lib.rs b/src/lib.rs index 6ec7061..801e640 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,7 +28,11 @@ //! Ok(()) //!} //!``` - +//! +//! ## Feature Flags +//! +//! `tracing`: Enables `tracing` instrumentation for request/response traces. +//! mod agent; mod builder; pub mod error; diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index d58acca..4771050 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -89,6 +89,10 @@ impl OllamaAI { #[async_trait::async_trait] impl CompletionProvider for OllamaAI { + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "ollama.chat", skip(self, prompt)) + )] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response) diff --git a/src/providers/openai.rs b/src/providers/openai.rs index 71e17fe..0982fd4 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -104,6 +104,10 @@ impl OpenAI { #[async_trait::async_trait] impl CompletionProvider for OpenAI { + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "openai.chat", skip(self, prompt)) + )] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response)