Skip to main content

llmtrace_sdk/
lib.rs

1//! LLMTrace SDK for Rust Applications
2//!
3//! This SDK provides embeddable tracing capabilities for Rust applications using LLM APIs.
4
5use llmtrace_core::{LLMProvider, Result, TenantId, TraceSpan};
6use uuid::Uuid;
7
8/// Configuration for the LLMTrace SDK
9#[derive(Debug, Clone)]
10pub struct SDKConfig {
11    pub tenant_id: TenantId,
12    pub endpoint: Option<String>,
13    pub enable_async_upload: bool,
14    pub buffer_size: usize,
15}
16
17impl Default for SDKConfig {
18    fn default() -> Self {
19        Self {
20            tenant_id: TenantId::new(),
21            endpoint: None,
22            enable_async_upload: true,
23            buffer_size: 1000,
24        }
25    }
26}
27
28/// Main SDK client for tracing LLM calls
29pub struct LLMTracer {
30    config: SDKConfig,
31    // TODO: Add trace buffer and uploader
32}
33
34impl LLMTracer {
35    /// Create a new LLM tracer with the given configuration
36    pub fn new(config: SDKConfig) -> Self {
37        Self { config }
38    }
39
40    /// Create a new LLM tracer with default configuration
41    pub fn with_defaults() -> Self {
42        Self::new(SDKConfig::default())
43    }
44
45    /// Start a new trace span
46    pub async fn start_span(
47        &self,
48        operation_name: &str,
49        provider: LLMProvider,
50        model: &str,
51    ) -> TraceSpan {
52        TraceSpan::new(
53            Uuid::new_v4(), // trace_id
54            self.config.tenant_id,
55            operation_name.to_string(),
56            provider,
57            model.to_string(),
58            String::new(), // empty prompt initially
59        )
60    }
61
62    /// Trace an LLM call with automatic instrumentation
63    pub async fn trace_call<F, R>(&self, _operation: &str, call: F) -> Result<R>
64    where
65        F: FnOnce() -> Result<R>,
66    {
67        // TODO: Implement automatic tracing wrapper
68        // For now, just execute the function
69        call()
70    }
71
72    /// Upload buffered traces (placeholder for async implementation)
73    pub async fn flush(&self) -> Result<()> {
74        // TODO: Implement trace upload to backend
75        Ok(())
76    }
77}
78
79/// Helper macro for instrumenting LLM calls
80#[macro_export]
81macro_rules! trace_llm_call {
82    ($tracer:expr, $operation:expr, $block:expr) => {{
83        // TODO: Implement macro-based instrumentation
84        $block
85    }};
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_sdk_config_default() {
94        let config = SDKConfig::default();
95        assert!(config.enable_async_upload);
96        assert_eq!(config.buffer_size, 1000);
97    }
98
99    #[tokio::test]
100    async fn test_tracer_creation() {
101        let tracer = LLMTracer::with_defaults();
102        assert_eq!(tracer.config.buffer_size, 1000);
103    }
104
105    #[tokio::test]
106    async fn test_start_span() {
107        let tracer = LLMTracer::with_defaults();
108        let span = tracer
109            .start_span("test_operation", LLMProvider::OpenAI, "gpt-4")
110            .await;
111
112        assert_eq!(span.operation_name, "test_operation");
113        assert_eq!(span.provider, LLMProvider::OpenAI);
114        assert_eq!(span.model_name, "gpt-4");
115        assert_eq!(span.tenant_id, tracer.config.tenant_id);
116    }
117}