use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use anyhow::Result;
#[derive(Debug, Serialize, Deserialize)]
pub struct LLMCallStat {
pub program_name: String,
pub model_name: String,
pub call_count: i64,
pub total_tokens: i64,
pub total_prompt_tokens: i64,
pub total_completion_tokens: i64,
pub avg_execution_time_ms: f64,
pub cache_hits: i64,
pub unique_users: i64,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ProgramUsageStat {
pub program_name: String,
pub usage_count: i64,
pub successful_calls: i64,
pub failed_calls: i64,
pub avg_execution_time_ms: f64,
pub unique_users: i64,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TokenUsageStat {
pub program_name: String,
pub model_name: String,
pub total_tokens: i64,
pub total_prompt_tokens: i64,
pub total_completion_tokens: i64,
pub total_cost_estimate: f64,
pub unique_users: i64,
}
pub struct AnalyticsEngine {
conn: Connection,
}
impl AnalyticsEngine {
pub fn new(db_path: &str) -> Result<Self> {
let conn = Connection::open(db_path)?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS llm_calls (
id TEXT PRIMARY KEY,
program_name TEXT,
model_name TEXT,
prompt_tokens INTEGER,
completion_tokens INTEGER,
total_tokens INTEGER,
execution_time_ms REAL,
cache_hit BOOLEAN,
user_id TEXT,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS program_usage (
id TEXT PRIMARY KEY,
program_name TEXT,
execution_time_ms REAL,
success BOOLEAN,
error_message TEXT,
user_id TEXT,
input_params TEXT,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS token_usage (
id TEXT PRIMARY KEY,
program_name TEXT,
model_name TEXT,
prompt_tokens INTEGER,
completion_tokens INTEGER,
total_tokens INTEGER,
user_id TEXT,
cost_estimate REAL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
);"
)?;
Ok(AnalyticsEngine { conn })
}
pub fn track_llm_call(
&self,
program_name: &str,
model_name: &str,
prompt_tokens: Option<u32>,
completion_tokens: Option<u32>,
total_tokens: Option<u32>,
execution_time_ms: u32,
cache_hit: bool,
user_id: &str,
) -> Result<()> {
let id = uuid::Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO llm_calls
(id, program_name, model_name, prompt_tokens, completion_tokens, total_tokens,
execution_time_ms, cache_hit, user_id)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
[
&id,
program_name,
model_name,
&prompt_tokens.unwrap_or(0).to_string(),
&completion_tokens.unwrap_or(0).to_string(),
&total_tokens.unwrap_or(0).to_string(),
&execution_time_ms.to_string(),
&cache_hit.to_string(),
user_id,
],
)?;
Ok(())
}
pub fn track_program_usage(
&self,
program_name: &str,
execution_time_ms: u32,
success: bool,
error_message: Option<&str>,
user_id: &str,
input_params: &str,
) -> Result<()> {
let id = uuid::Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO program_usage
(id, program_name, execution_time_ms, success, error_message, user_id, input_params)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
[
&id,
program_name,
&execution_time_ms.to_string(),
&success.to_string(),
&error_message.unwrap_or("").to_string(),
user_id,
input_params,
],
)?;
Ok(())
}
pub fn track_token_usage(
&self,
program_name: &str,
model_name: &str,
prompt_tokens: u32,
completion_tokens: u32,
total_tokens: u32,
user_id: &str,
cost_estimate: f64,
) -> Result<()> {
let id = uuid::Uuid::new_v4().to_string();
self.conn.execute(
"INSERT INTO token_usage
(id, program_name, model_name, prompt_tokens, completion_tokens, total_tokens,
user_id, cost_estimate)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
[
&id,
program_name,
model_name,
&prompt_tokens.to_string(),
&completion_tokens.to_string(),
&total_tokens.to_string(),
user_id,
&cost_estimate.to_string(),
],
)?;
Ok(())
}
pub fn get_llm_call_stats(
&self,
program_name: Option<&str>,
model_name: Option<&str>,
) -> Result<Vec<LLMCallStat>> {
let mut query = "SELECT
program_name,
model_name,
COUNT(*) as call_count,
SUM(total_tokens) as total_tokens,
SUM(prompt_tokens) as total_prompt_tokens,
SUM(completion_tokens) as total_completion_tokens,
AVG(execution_time_ms) as avg_execution_time_ms,
SUM(CASE WHEN cache_hit = 'true' THEN 1 ELSE 0 END) as cache_hits,
COUNT(DISTINCT user_id) as unique_users
FROM llm_calls".to_string();
let mut params: Vec<String> = Vec::new();
if program_name.is_some() || model_name.is_some() {
query.push_str(" WHERE ");
let mut conditions = Vec::new();
if let Some(pn) = program_name {
conditions.push("program_name = ?".to_string());
params.push(pn.to_string());
}
if let Some(mn) = model_name {
conditions.push("model_name = ?".to_string());
params.push(mn.to_string());
}
query.push_str(&conditions.join(" AND "));
}
query.push_str(" GROUP BY program_name, model_name");
let mut stmt = self.conn.prepare(&query)?;
let rows = stmt.query_map([], |row| {
Ok(LLMCallStat {
program_name: row.get(0)?,
model_name: row.get(1)?,
call_count: row.get(2)?,
total_tokens: row.get(3)?,
total_prompt_tokens: row.get(4)?,
total_completion_tokens: row.get(5)?,
avg_execution_time_ms: row.get(6)?,
cache_hits: row.get(7)?,
unique_users: row.get(8)?,
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
pub fn get_program_usage_stats(
&self,
program_name: Option<&str>,
) -> Result<Vec<ProgramUsageStat>> {
let mut query = "SELECT
program_name,
COUNT(*) as usage_count,
SUM(CASE WHEN success = 'true' THEN 1 ELSE 0 END) as successful_calls,
SUM(CASE WHEN success = 'false' THEN 1 ELSE 0 END) as failed_calls,
AVG(execution_time_ms) as avg_execution_time_ms,
COUNT(DISTINCT user_id) as unique_users
FROM program_usage".to_string();
let mut params: Vec<String> = Vec::new();
if let Some(pn) = program_name {
query.push_str(" WHERE program_name = ?");
params.push(pn.to_string());
}
query.push_str(" GROUP BY program_name");
let mut stmt = self.conn.prepare(&query)?;
let rows = stmt.query_map([], |row| {
Ok(ProgramUsageStat {
program_name: row.get(0)?,
usage_count: row.get(1)?,
successful_calls: row.get(2)?,
failed_calls: row.get(3)?,
avg_execution_time_ms: row.get(4)?,
unique_users: row.get(5)?,
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
pub fn get_token_usage_stats(
&self,
program_name: Option<&str>,
model_name: Option<&str>,
) -> Result<Vec<TokenUsageStat>> {
let mut query = "SELECT
program_name,
model_name,
SUM(total_tokens) as total_tokens,
SUM(prompt_tokens) as total_prompt_tokens,
SUM(completion_tokens) as total_completion_tokens,
SUM(cost_estimate) as total_cost_estimate,
COUNT(DISTINCT user_id) as unique_users
FROM token_usage".to_string();
let mut params: Vec<String> = Vec::new();
if program_name.is_some() || model_name.is_some() {
query.push_str(" WHERE ");
let mut conditions = Vec::new();
if let Some(pn) = program_name {
conditions.push("program_name = ?".to_string());
params.push(pn.to_string());
}
if let Some(mn) = model_name {
conditions.push("model_name = ?".to_string());
params.push(mn.to_string());
}
query.push_str(&conditions.join(" AND "));
}
query.push_str(" GROUP BY program_name, model_name");
let mut stmt = self.conn.prepare(&query)?;
let rows = stmt.query_map([], |row| {
Ok(TokenUsageStat {
program_name: row.get(0)?,
model_name: row.get(1)?,
total_tokens: row.get(2)?,
total_prompt_tokens: row.get(3)?,
total_completion_tokens: row.get(4)?,
total_cost_estimate: row.get(5)?,
unique_users: row.get(6)?,
})
})?;
let mut stats = Vec::new();
for row in rows {
stats.push(row?);
}
Ok(stats)
}
}