diff --git a/migrations/0019_add_ai_usage.sql b/migrations/0019_add_ai_usage.sql new file mode 100644 index 0000000..410f1c1 --- /dev/null +++ b/migrations/0019_add_ai_usage.sql @@ -0,0 +1,13 @@ +CREATE TABLE ai_usage ( + id INTEGER PRIMARY KEY, + user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, + model TEXT NOT NULL, + endpoint TEXT NOT NULL, + prompt_tokens INTEGER NOT NULL, + completion_tokens INTEGER NOT NULL, + total_tokens INTEGER NOT NULL, + cost REAL NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')) +); + +CREATE INDEX idx_ai_usage_user_id ON ai_usage(user_id); diff --git a/src/domain/ai_usage.rs b/src/domain/ai_usage.rs new file mode 100644 index 0000000..a4c5d88 --- /dev/null +++ b/src/domain/ai_usage.rs @@ -0,0 +1,38 @@ +use chrono::{DateTime, Utc}; +use serde::Serialize; + +use crate::domain::ids::{AiUsageId, UserId}; + +#[derive(Debug, Clone, Serialize)] +pub struct AiUsage { + pub id: AiUsageId, + pub user_id: UserId, + pub model: String, + pub endpoint: String, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub total_tokens: i64, + pub cost: f64, + pub created_at: DateTime, +} + +#[derive(Debug, Clone)] +pub struct NewAiUsage { + pub user_id: UserId, + pub model: String, + pub endpoint: String, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub total_tokens: i64, + pub cost: f64, +} + +/// Aggregated usage totals for display. +#[derive(Debug, Clone, Default, Serialize)] +pub struct AiUsageSummary { + pub total_calls: i64, + pub total_prompt_tokens: i64, + pub total_completion_tokens: i64, + pub total_tokens: i64, + pub total_cost: f64, +} diff --git a/src/domain/ids.rs b/src/domain/ids.rs index 50ebb42..5ee4311 100644 --- a/src/domain/ids.rs +++ b/src/domain/ids.rs @@ -61,3 +61,4 @@ define_id!(CafeId); define_id!(CupId); define_id!(PasskeyCredentialId); define_id!(RegistrationTokenId); +define_id!(AiUsageId); diff --git a/src/domain/mod.rs b/src/domain/mod.rs index f53a80b..8066ad9 100644 --- a/src/domain/mod.rs +++ b/src/domain/mod.rs @@ -1,3 +1,4 @@ +pub mod ai_usage; pub mod bags; pub mod brews; pub mod cafes; diff --git a/src/domain/repositories.rs b/src/domain/repositories.rs index f2ef840..8e26d52 100644 --- a/src/domain/repositories.rs +++ b/src/domain/repositories.rs @@ -1,4 +1,5 @@ use super::RepositoryError; +use crate::domain::ai_usage::{AiUsage, AiUsageSummary, NewAiUsage}; use crate::domain::listing::{ListRequest, Page, SortDirection, SortKey}; use crate::domain::bags::{Bag, BagFilter, BagSortKey, BagWithRoast, NewBag, UpdateBag}; @@ -261,3 +262,9 @@ pub trait RegistrationTokenRepository: Send + Sync { user_id: UserId, ) -> Result<(), RepositoryError>; } + +#[async_trait] +pub trait AiUsageRepository: Send + Sync { + async fn insert(&self, usage: NewAiUsage) -> Result; + async fn summary_for_user(&self, user_id: UserId) -> Result; +} diff --git a/src/infrastructure/ai.rs b/src/infrastructure/ai.rs index 9be4e8e..b868ce6 100644 --- a/src/infrastructure/ai.rs +++ b/src/infrastructure/ai.rs @@ -50,6 +50,14 @@ Only include fields you can identify with confidence. Each tasting note must be // --- Public types --- +#[derive(Debug, Clone, Deserialize)] +pub struct Usage { + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub total_tokens: i64, + pub cost: f64, +} + #[derive(Debug, Deserialize)] pub struct ExtractionInput { pub image: Option, @@ -88,13 +96,14 @@ pub async fn extract_roaster( api_key: &str, model: &str, input: &ExtractionInput, -) -> Result { - let content = call_openrouter(client, api_key, model, ROASTER_PROMPT, input).await?; +) -> Result<(ExtractedRoaster, Option), AppError> { + let (content, usage) = call_openrouter(client, api_key, model, ROASTER_PROMPT, input).await?; let json = extract_json(&content); - serde_json::from_str(json).map_err(|e| { + let extracted = serde_json::from_str(json).map_err(|e| { AppError::unexpected(format!("Failed to parse AI response as roaster data: {e}")) - }) + })?; + Ok((extracted, usage)) } pub async fn extract_roast( @@ -102,13 +111,14 @@ pub async fn extract_roast( api_key: &str, model: &str, input: &ExtractionInput, -) -> Result { - let content = call_openrouter(client, api_key, model, ROAST_PROMPT, input).await?; +) -> Result<(ExtractedRoast, Option), AppError> { + let (content, usage) = call_openrouter(client, api_key, model, ROAST_PROMPT, input).await?; let json = extract_json(&content); - serde_json::from_str(json).map_err(|e| { + let extracted = serde_json::from_str(json).map_err(|e| { AppError::unexpected(format!("Failed to parse AI response as roast data: {e}")) - }) + })?; + Ok((extracted, usage)) } pub async fn extract_bag_scan( @@ -116,13 +126,14 @@ pub async fn extract_bag_scan( api_key: &str, model: &str, input: &ExtractionInput, -) -> Result { - let content = call_openrouter(client, api_key, model, SCAN_PROMPT, input).await?; +) -> Result<(ExtractedBagScan, Option), AppError> { + let (content, usage) = call_openrouter(client, api_key, model, SCAN_PROMPT, input).await?; let json = extract_json(&content); - serde_json::from_str(json).map_err(|e| { + let extracted = serde_json::from_str(json).map_err(|e| { AppError::unexpected(format!("Failed to parse AI response as bag scan data: {e}")) - }) + })?; + Ok((extracted, usage)) } // --- Internal helpers --- @@ -133,7 +144,7 @@ async fn call_openrouter( model: &str, system_prompt: &str, input: &ExtractionInput, -) -> Result { +) -> Result<(String, Option), AppError> { let has_image = input.image.as_ref().is_some_and(|s| !s.trim().is_empty()); let has_prompt = input.prompt.as_ref().is_some_and(|s| !s.trim().is_empty()); @@ -212,7 +223,7 @@ async fn call_openrouter( )); } - Ok(content) + Ok((content, chat_response.usage)) } /// Extract a JSON object from a model response that may contain markdown @@ -273,6 +284,7 @@ struct ImageUrlDetail { #[derive(Debug, Deserialize)] struct ChatResponse { choices: Vec, + usage: Option, } #[derive(Debug, Deserialize)] @@ -303,7 +315,13 @@ mod tests { }, "finish_reason": "stop" } - ] + ], + "usage": { + "prompt_tokens": 194, + "completion_tokens": 42, + "total_tokens": 236, + "cost": 0.0012 + } }"#; let response: ChatResponse = serde_json::from_str(json).unwrap(); @@ -315,6 +333,33 @@ mod tests { assert_eq!(roaster.country.as_deref(), Some("United Kingdom")); assert_eq!(roaster.city.as_deref(), Some("London")); assert!(roaster.homepage.is_none()); + + let usage = response.usage.unwrap(); + assert_eq!(usage.prompt_tokens, 194); + assert_eq!(usage.completion_tokens, 42); + assert_eq!(usage.total_tokens, 236); + assert!((usage.cost - 0.0012).abs() < f64::EPSILON); + } + + #[test] + fn parse_chat_response_without_usage() { + let json = r#"{ + "id": "gen-abc123", + "model": "openrouter/free", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "{\"name\": \"Square Mile\"}" + }, + "finish_reason": "stop" + } + ] + }"#; + + let response: ChatResponse = serde_json::from_str(json).unwrap(); + assert!(response.usage.is_none()); } #[test] diff --git a/src/infrastructure/repositories/ai_usage.rs b/src/infrastructure/repositories/ai_usage.rs new file mode 100644 index 0000000..0c17fef --- /dev/null +++ b/src/infrastructure/repositories/ai_usage.rs @@ -0,0 +1,109 @@ +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use sqlx::query_as; + +use crate::domain::RepositoryError; +use crate::domain::ai_usage::{AiUsage, AiUsageSummary, NewAiUsage}; +use crate::domain::ids::{AiUsageId, UserId}; +use crate::domain::repositories::AiUsageRepository; +use crate::infrastructure::database::DatabasePool; + +#[derive(Clone)] +pub struct SqlAiUsageRepository { + pool: DatabasePool, +} + +impl SqlAiUsageRepository { + pub fn new(pool: DatabasePool) -> Self { + Self { pool } + } + + fn to_domain(record: AiUsageRecord) -> AiUsage { + AiUsage { + id: AiUsageId::from(record.id), + user_id: UserId::from(record.user_id), + model: record.model, + endpoint: record.endpoint, + prompt_tokens: record.prompt_tokens, + completion_tokens: record.completion_tokens, + total_tokens: record.total_tokens, + cost: record.cost, + created_at: record.created_at, + } + } +} + +#[async_trait] +impl AiUsageRepository for SqlAiUsageRepository { + async fn insert(&self, usage: NewAiUsage) -> Result { + let query = r" + INSERT INTO ai_usage (user_id, model, endpoint, prompt_tokens, completion_tokens, total_tokens, cost) + VALUES (?, ?, ?, ?, ?, ?, ?) + RETURNING id, user_id, model, endpoint, prompt_tokens, completion_tokens, total_tokens, cost, created_at + "; + + let record = query_as::<_, AiUsageRecord>(query) + .bind(i64::from(usage.user_id)) + .bind(&usage.model) + .bind(&usage.endpoint) + .bind(usage.prompt_tokens) + .bind(usage.completion_tokens) + .bind(usage.total_tokens) + .bind(usage.cost) + .fetch_one(&self.pool) + .await + .map_err(|err| RepositoryError::unexpected(err.to_string()))?; + + Ok(Self::to_domain(record)) + } + + async fn summary_for_user(&self, user_id: UserId) -> Result { + let query = r" + SELECT + COALESCE(COUNT(*), 0) as total_calls, + COALESCE(SUM(prompt_tokens), 0) as total_prompt_tokens, + COALESCE(SUM(completion_tokens), 0) as total_completion_tokens, + COALESCE(SUM(total_tokens), 0) as total_tokens, + COALESCE(SUM(cost), 0.0) as total_cost + FROM ai_usage + WHERE user_id = ? + "; + + let record = query_as::<_, AiUsageSummaryRecord>(query) + .bind(i64::from(user_id)) + .fetch_one(&self.pool) + .await + .map_err(|err| RepositoryError::unexpected(err.to_string()))?; + + Ok(AiUsageSummary { + total_calls: record.total_calls, + total_prompt_tokens: record.total_prompt_tokens, + total_completion_tokens: record.total_completion_tokens, + total_tokens: record.total_tokens, + total_cost: record.total_cost, + }) + } +} + +#[derive(sqlx::FromRow)] +struct AiUsageRecord { + id: i64, + user_id: i64, + model: String, + endpoint: String, + prompt_tokens: i64, + completion_tokens: i64, + total_tokens: i64, + cost: f64, + created_at: DateTime, +} + +#[derive(sqlx::FromRow)] +#[allow(clippy::struct_field_names)] +struct AiUsageSummaryRecord { + total_calls: i64, + total_prompt_tokens: i64, + total_completion_tokens: i64, + total_tokens: i64, + total_cost: f64, +} diff --git a/src/infrastructure/repositories/mod.rs b/src/infrastructure/repositories/mod.rs index a825035..49f901a 100644 --- a/src/infrastructure/repositories/mod.rs +++ b/src/infrastructure/repositories/mod.rs @@ -1,3 +1,4 @@ +pub mod ai_usage; pub mod bags; pub mod brews; pub mod cafes;