1use llmtrace_core::{LLMProvider, Result, TenantId, TraceSpan};
6use uuid::Uuid;
7
8#[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
28pub struct LLMTracer {
30 config: SDKConfig,
31 }
33
34impl LLMTracer {
35 pub fn new(config: SDKConfig) -> Self {
37 Self { config }
38 }
39
40 pub fn with_defaults() -> Self {
42 Self::new(SDKConfig::default())
43 }
44
45 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(), self.config.tenant_id,
55 operation_name.to_string(),
56 provider,
57 model.to_string(),
58 String::new(), )
60 }
61
62 pub async fn trace_call<F, R>(&self, _operation: &str, call: F) -> Result<R>
64 where
65 F: FnOnce() -> Result<R>,
66 {
67 call()
70 }
71
72 pub async fn flush(&self) -> Result<()> {
74 Ok(())
76 }
77}
78
79#[macro_export]
81macro_rules! trace_llm_call {
82 ($tracer:expr, $operation:expr, $block:expr) => {{
83 $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}