"vscode:/vscode.git/clone" did not exist on "04c30254d4bd16b07958110165f59e9ac9fbf5c6"
Unverified Commit 1032076d authored by Ayush Agarwal's avatar Ayush Agarwal Committed by GitHub
Browse files

chore: added forced audit logging (#5552)


Signed-off-by: default avatarayushag <ayushag@nvidia.com>
parent 38fbb1db
......@@ -6,6 +6,7 @@ use std::sync::OnceLock;
#[derive(Clone, Copy)]
pub struct AuditPolicy {
pub enabled: bool,
pub force_logging: bool,
}
static POLICY: OnceLock<AuditPolicy> = OnceLock::new();
......@@ -14,6 +15,10 @@ static POLICY: OnceLock<AuditPolicy> = OnceLock::new();
pub fn init_from_env() -> AuditPolicy {
AuditPolicy {
enabled: std::env::var("DYN_AUDIT_SINKS").is_ok(),
force_logging: std::env::var("DYN_AUDIT_FORCE_LOGGING")
.ok()
.and_then(|v| v.parse::<bool>().ok())
.unwrap_or(false),
}
}
......
......@@ -56,7 +56,12 @@ impl AuditHandle {
}
pub fn create_handle(req: &NvCreateChatCompletionRequest, request_id: &str) -> Option<AuditHandle> {
if !config::policy().enabled || !req.inner.store.unwrap_or(false) {
let policy = config::policy();
if !policy.enabled {
return None;
}
// If force_logging is enabled, ignore the store flag
if !policy.force_logging && !req.inner.store.unwrap_or(false) {
return None;
}
let requested_streaming = req.inner.stream.unwrap_or(false);
......@@ -70,3 +75,40 @@ pub fn create_handle(req: &NvCreateChatCompletionRequest, request_id: &str) -> O
resp_full: None,
})
}
#[cfg(test)]
mod tests {
use super::*;
use temp_env::with_vars;
fn create_test_request(model: &str, store: bool) -> NvCreateChatCompletionRequest {
let json = serde_json::json!({
"model": model,
"messages": [{"role": "user", "content": "test"}],
"store": store
});
serde_json::from_value(json).expect("Failed to create test request")
}
/// Test that DYN_AUDIT_FORCE_LOGGING=true bypasses store=false
/// When force logging is enabled, audit handle should be created even when store=false
#[test]
fn test_force_logging_bypasses_store() {
with_vars(
[
("DYN_AUDIT_SINKS", Some("stderr")),
("DYN_AUDIT_FORCE_LOGGING", Some("true")),
],
|| {
// Create request with store=false
let request = create_test_request("test-model", false);
let handle = create_handle(&request, "test-id");
assert!(
handle.is_some(),
"When DYN_AUDIT_FORCE_LOGGING=true, handle should be created even with store=false"
);
},
);
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment