feat(ai): record usage after each OpenRouter extraction
Wire AiUsageRepository into AppState and add a fire-and-forget record_ai_usage helper. All extraction route handlers now capture the usage tuple and record it in the background.
This commit is contained in:
parent
6271b0e5e5
commit
0a93df845f
6 changed files with 86 additions and 19 deletions
|
|
@ -118,15 +118,15 @@ define_delete_handler!(
|
||||||
render_roaster_list_fragment
|
render_roaster_list_fragment
|
||||||
);
|
);
|
||||||
|
|
||||||
#[tracing::instrument(skip(state, _auth_user, headers, payload))]
|
#[tracing::instrument(skip(state, auth_user, headers, payload))]
|
||||||
pub(crate) async fn extract_roaster(
|
pub(crate) async fn extract_roaster(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
_auth_user: AuthenticatedUser,
|
auth_user: AuthenticatedUser,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
payload: FlexiblePayload<ExtractionInput>,
|
payload: FlexiblePayload<ExtractionInput>,
|
||||||
) -> Result<Response, ApiError> {
|
) -> Result<Response, ApiError> {
|
||||||
let (input, _) = payload.into_parts();
|
let (input, _) = payload.into_parts();
|
||||||
let result = ai::extract_roaster(
|
let (result, usage) = ai::extract_roaster(
|
||||||
&state.http_client,
|
&state.http_client,
|
||||||
&state.openrouter_api_key,
|
&state.openrouter_api_key,
|
||||||
&state.openrouter_model,
|
&state.openrouter_model,
|
||||||
|
|
@ -135,6 +135,14 @@ pub(crate) async fn extract_roaster(
|
||||||
.await
|
.await
|
||||||
.map_err(ApiError::from)?;
|
.map_err(ApiError::from)?;
|
||||||
|
|
||||||
|
super::support::record_ai_usage(
|
||||||
|
state.ai_usage_repo.clone(),
|
||||||
|
auth_user.0.id,
|
||||||
|
&state.openrouter_model,
|
||||||
|
"extract-roaster",
|
||||||
|
usage,
|
||||||
|
);
|
||||||
|
|
||||||
if is_datastar_request(&headers) {
|
if is_datastar_request(&headers) {
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
let signals = vec![
|
let signals = vec![
|
||||||
|
|
|
||||||
|
|
@ -257,15 +257,15 @@ impl TastingNotesInput {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tracing::instrument(skip(state, _auth_user, headers, payload))]
|
#[tracing::instrument(skip(state, auth_user, headers, payload))]
|
||||||
pub(crate) async fn extract_roast_info(
|
pub(crate) async fn extract_roast_info(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
_auth_user: AuthenticatedUser,
|
auth_user: AuthenticatedUser,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
payload: FlexiblePayload<ExtractionInput>,
|
payload: FlexiblePayload<ExtractionInput>,
|
||||||
) -> Result<Response, ApiError> {
|
) -> Result<Response, ApiError> {
|
||||||
let (input, _) = payload.into_parts();
|
let (input, _) = payload.into_parts();
|
||||||
let result = ai::extract_roast(
|
let (result, usage) = ai::extract_roast(
|
||||||
&state.http_client,
|
&state.http_client,
|
||||||
&state.openrouter_api_key,
|
&state.openrouter_api_key,
|
||||||
&state.openrouter_model,
|
&state.openrouter_model,
|
||||||
|
|
@ -274,6 +274,14 @@ pub(crate) async fn extract_roast_info(
|
||||||
.await
|
.await
|
||||||
.map_err(ApiError::from)?;
|
.map_err(ApiError::from)?;
|
||||||
|
|
||||||
|
super::support::record_ai_usage(
|
||||||
|
state.ai_usage_repo.clone(),
|
||||||
|
auth_user.0.id,
|
||||||
|
&state.openrouter_model,
|
||||||
|
"extract-roast",
|
||||||
|
usage,
|
||||||
|
);
|
||||||
|
|
||||||
if is_datastar_request(&headers) {
|
if is_datastar_request(&headers) {
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,17 +14,17 @@ use crate::domain::errors::RepositoryError;
|
||||||
use crate::domain::roasters::NewRoaster;
|
use crate::domain::roasters::NewRoaster;
|
||||||
use crate::domain::roasts::NewRoast;
|
use crate::domain::roasts::NewRoast;
|
||||||
use crate::domain::timeline::{NewTimelineEvent, TimelineEventDetail};
|
use crate::domain::timeline::{NewTimelineEvent, TimelineEventDetail};
|
||||||
use crate::infrastructure::ai::{self, ExtractionInput};
|
use crate::infrastructure::ai::{self, ExtractionInput, Usage};
|
||||||
|
|
||||||
#[tracing::instrument(skip(state, _auth_user, headers, payload))]
|
#[tracing::instrument(skip(state, auth_user, headers, payload))]
|
||||||
pub(crate) async fn extract_bag_scan(
|
pub(crate) async fn extract_bag_scan(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
_auth_user: AuthenticatedUser,
|
auth_user: AuthenticatedUser,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
payload: FlexiblePayload<ExtractionInput>,
|
payload: FlexiblePayload<ExtractionInput>,
|
||||||
) -> Result<Response, ApiError> {
|
) -> Result<Response, ApiError> {
|
||||||
let (input, _) = payload.into_parts();
|
let (input, _) = payload.into_parts();
|
||||||
let result = ai::extract_bag_scan(
|
let (result, usage) = ai::extract_bag_scan(
|
||||||
&state.http_client,
|
&state.http_client,
|
||||||
&state.openrouter_api_key,
|
&state.openrouter_api_key,
|
||||||
&state.openrouter_model,
|
&state.openrouter_model,
|
||||||
|
|
@ -33,6 +33,14 @@ pub(crate) async fn extract_bag_scan(
|
||||||
.await
|
.await
|
||||||
.map_err(ApiError::from)?;
|
.map_err(ApiError::from)?;
|
||||||
|
|
||||||
|
super::support::record_ai_usage(
|
||||||
|
state.ai_usage_repo.clone(),
|
||||||
|
auth_user.0.id,
|
||||||
|
&state.openrouter_model,
|
||||||
|
"extract-bag-scan",
|
||||||
|
usage,
|
||||||
|
);
|
||||||
|
|
||||||
if is_datastar_request(&headers) {
|
if is_datastar_request(&headers) {
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
|
||||||
|
|
@ -130,15 +138,16 @@ struct ScanResult {
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Populate a `BagScanSubmission` from AI extraction when image/prompt is provided.
|
/// Populate a `BagScanSubmission` from AI extraction when image/prompt is provided.
|
||||||
|
/// Returns the usage data so the caller can record it.
|
||||||
async fn extract_into_submission(
|
async fn extract_into_submission(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
submission: &mut BagScanSubmission,
|
submission: &mut BagScanSubmission,
|
||||||
) -> Result<(), ApiError> {
|
) -> Result<Option<Usage>, ApiError> {
|
||||||
let input = ExtractionInput {
|
let input = ExtractionInput {
|
||||||
image: submission.image.take(),
|
image: submission.image.take(),
|
||||||
prompt: submission.prompt.take(),
|
prompt: submission.prompt.take(),
|
||||||
};
|
};
|
||||||
let result = ai::extract_bag_scan(
|
let (result, usage) = ai::extract_bag_scan(
|
||||||
&state.http_client,
|
&state.http_client,
|
||||||
&state.openrouter_api_key,
|
&state.openrouter_api_key,
|
||||||
&state.openrouter_model,
|
&state.openrouter_model,
|
||||||
|
|
@ -178,14 +187,14 @@ async fn extract_into_submission(
|
||||||
submission.tasting_notes = TastingNotesInput::Text(notes.join(", "));
|
submission.tasting_notes = TastingNotesInput::Text(notes.join(", "));
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(usage)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(clippy::too_many_lines)]
|
#[allow(clippy::too_many_lines)]
|
||||||
#[tracing::instrument(skip(state, _auth_user, headers, payload))]
|
#[tracing::instrument(skip(state, auth_user, headers, payload))]
|
||||||
pub(crate) async fn submit_scan(
|
pub(crate) async fn submit_scan(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
_auth_user: AuthenticatedUser,
|
auth_user: AuthenticatedUser,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
payload: FlexiblePayload<BagScanSubmission>,
|
payload: FlexiblePayload<BagScanSubmission>,
|
||||||
) -> Result<Response, ApiError> {
|
) -> Result<Response, ApiError> {
|
||||||
|
|
@ -196,7 +205,14 @@ pub(crate) async fn submit_scan(
|
||||||
|| submission.prompt.as_deref().is_some_and(|s| !s.is_empty());
|
|| submission.prompt.as_deref().is_some_and(|s| !s.is_empty());
|
||||||
|
|
||||||
if has_raw_input {
|
if has_raw_input {
|
||||||
extract_into_submission(&state, &mut submission).await?;
|
let usage = extract_into_submission(&state, &mut submission).await?;
|
||||||
|
super::support::record_ai_usage(
|
||||||
|
state.ai_usage_repo.clone(),
|
||||||
|
auth_user.0.id,
|
||||||
|
&state.openrouter_model,
|
||||||
|
"extract-bag-scan",
|
||||||
|
usage,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Build and normalize the roaster
|
// Build and normalize the roaster
|
||||||
|
|
|
||||||
|
|
@ -236,6 +236,31 @@ pub(super) async fn load_cafe_options(state: &AppState) -> Result<Vec<CafeOption
|
||||||
Ok(cafes.into_iter().map(CafeOptionView::from).collect())
|
Ok(cafes.into_iter().map(CafeOptionView::from).collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Record AI usage in the background. Failures are logged but do not affect the response.
|
||||||
|
pub fn record_ai_usage(
|
||||||
|
repo: std::sync::Arc<dyn crate::domain::repositories::AiUsageRepository>,
|
||||||
|
user_id: crate::domain::ids::UserId,
|
||||||
|
model: &str,
|
||||||
|
endpoint: &str,
|
||||||
|
usage: Option<crate::infrastructure::ai::Usage>,
|
||||||
|
) {
|
||||||
|
let Some(usage) = usage else { return };
|
||||||
|
let new_usage = crate::domain::ai_usage::NewAiUsage {
|
||||||
|
user_id,
|
||||||
|
model: model.to_string(),
|
||||||
|
endpoint: endpoint.to_string(),
|
||||||
|
prompt_tokens: usage.prompt_tokens,
|
||||||
|
completion_tokens: usage.completion_tokens,
|
||||||
|
total_tokens: usage.total_tokens,
|
||||||
|
cost: usage.cost,
|
||||||
|
};
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(err) = repo.insert(new_usage).await {
|
||||||
|
tracing::warn!(error = %err, "failed to record AI usage");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
pub fn is_datastar_request(headers: &HeaderMap) -> bool {
|
pub fn is_datastar_request(headers: &HeaderMap) -> bool {
|
||||||
headers
|
headers
|
||||||
.get("datastar-request")
|
.get("datastar-request")
|
||||||
|
|
|
||||||
|
|
@ -12,13 +12,14 @@ use webauthn_rs::prelude::*;
|
||||||
use crate::application::routes::app_router;
|
use crate::application::routes::app_router;
|
||||||
use crate::domain::registration_tokens::NewRegistrationToken;
|
use crate::domain::registration_tokens::NewRegistrationToken;
|
||||||
use crate::domain::repositories::{
|
use crate::domain::repositories::{
|
||||||
BagRepository, BrewRepository, CafeRepository, CupRepository, GearRepository,
|
AiUsageRepository, BagRepository, BrewRepository, CafeRepository, CupRepository,
|
||||||
PasskeyCredentialRepository, RegistrationTokenRepository, RoastRepository, RoasterRepository,
|
GearRepository, PasskeyCredentialRepository, RegistrationTokenRepository, RoastRepository,
|
||||||
SessionRepository, TimelineEventRepository, TokenRepository, UserRepository,
|
RoasterRepository, SessionRepository, TimelineEventRepository, TokenRepository, UserRepository,
|
||||||
};
|
};
|
||||||
use crate::infrastructure::auth::{generate_session_token, hash_token};
|
use crate::infrastructure::auth::{generate_session_token, hash_token};
|
||||||
use crate::infrastructure::backup::BackupService;
|
use crate::infrastructure::backup::BackupService;
|
||||||
use crate::infrastructure::database::Database;
|
use crate::infrastructure::database::Database;
|
||||||
|
use crate::infrastructure::repositories::ai_usage::SqlAiUsageRepository;
|
||||||
use crate::infrastructure::repositories::bags::SqlBagRepository;
|
use crate::infrastructure::repositories::bags::SqlBagRepository;
|
||||||
use crate::infrastructure::repositories::brews::SqlBrewRepository;
|
use crate::infrastructure::repositories::brews::SqlBrewRepository;
|
||||||
use crate::infrastructure::repositories::cafes::SqlCafeRepository;
|
use crate::infrastructure::repositories::cafes::SqlCafeRepository;
|
||||||
|
|
@ -59,6 +60,7 @@ pub struct AppState {
|
||||||
pub session_repo: Arc<dyn SessionRepository>,
|
pub session_repo: Arc<dyn SessionRepository>,
|
||||||
pub passkey_repo: Arc<dyn PasskeyCredentialRepository>,
|
pub passkey_repo: Arc<dyn PasskeyCredentialRepository>,
|
||||||
pub registration_token_repo: Arc<dyn RegistrationTokenRepository>,
|
pub registration_token_repo: Arc<dyn RegistrationTokenRepository>,
|
||||||
|
pub ai_usage_repo: Arc<dyn AiUsageRepository>,
|
||||||
pub webauthn: Arc<Webauthn>,
|
pub webauthn: Arc<Webauthn>,
|
||||||
pub challenge_store: Arc<ChallengeStore>,
|
pub challenge_store: Arc<ChallengeStore>,
|
||||||
pub http_client: reqwest::Client,
|
pub http_client: reqwest::Client,
|
||||||
|
|
@ -101,6 +103,8 @@ pub async fn serve(config: ServerConfig) -> anyhow::Result<()> {
|
||||||
Arc::new(SqlPasskeyCredentialRepository::new(database.clone_pool()));
|
Arc::new(SqlPasskeyCredentialRepository::new(database.clone_pool()));
|
||||||
let registration_token_repo: Arc<dyn RegistrationTokenRepository> =
|
let registration_token_repo: Arc<dyn RegistrationTokenRepository> =
|
||||||
Arc::new(SqlRegistrationTokenRepository::new(database.clone_pool()));
|
Arc::new(SqlRegistrationTokenRepository::new(database.clone_pool()));
|
||||||
|
let ai_usage_repo: Arc<dyn AiUsageRepository> =
|
||||||
|
Arc::new(SqlAiUsageRepository::new(database.clone_pool()));
|
||||||
|
|
||||||
let backup_service = Arc::new(BackupService::new(database.clone_pool()));
|
let backup_service = Arc::new(BackupService::new(database.clone_pool()));
|
||||||
let challenge_store = Arc::new(ChallengeStore::new());
|
let challenge_store = Arc::new(ChallengeStore::new());
|
||||||
|
|
@ -122,6 +126,7 @@ pub async fn serve(config: ServerConfig) -> anyhow::Result<()> {
|
||||||
session_repo,
|
session_repo,
|
||||||
passkey_repo,
|
passkey_repo,
|
||||||
registration_token_repo,
|
registration_token_repo,
|
||||||
|
ai_usage_repo,
|
||||||
webauthn,
|
webauthn,
|
||||||
challenge_store,
|
challenge_store,
|
||||||
http_client: reqwest::Client::new(),
|
http_client: reqwest::Client::new(),
|
||||||
|
|
|
||||||
|
|
@ -164,6 +164,11 @@ async fn spawn_app_inner(
|
||||||
session_repo,
|
session_repo,
|
||||||
passkey_repo,
|
passkey_repo,
|
||||||
registration_token_repo,
|
registration_token_repo,
|
||||||
|
ai_usage_repo: Arc::new(
|
||||||
|
brewlog::infrastructure::repositories::ai_usage::SqlAiUsageRepository::new(
|
||||||
|
_database.clone_pool(),
|
||||||
|
),
|
||||||
|
),
|
||||||
webauthn: test_webauthn(),
|
webauthn: test_webauthn(),
|
||||||
challenge_store: Arc::new(ChallengeStore::new()),
|
challenge_store: Arc::new(ChallengeStore::new()),
|
||||||
http_client: reqwest::Client::new(),
|
http_client: reqwest::Client::new(),
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue