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
|
||||
);
|
||||
|
||||
#[tracing::instrument(skip(state, _auth_user, headers, payload))]
|
||||
#[tracing::instrument(skip(state, auth_user, headers, payload))]
|
||||
pub(crate) async fn extract_roaster(
|
||||
State(state): State<AppState>,
|
||||
_auth_user: AuthenticatedUser,
|
||||
auth_user: AuthenticatedUser,
|
||||
headers: HeaderMap,
|
||||
payload: FlexiblePayload<ExtractionInput>,
|
||||
) -> Result<Response, ApiError> {
|
||||
let (input, _) = payload.into_parts();
|
||||
let result = ai::extract_roaster(
|
||||
let (result, usage) = ai::extract_roaster(
|
||||
&state.http_client,
|
||||
&state.openrouter_api_key,
|
||||
&state.openrouter_model,
|
||||
|
|
@ -135,6 +135,14 @@ pub(crate) async fn extract_roaster(
|
|||
.await
|
||||
.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) {
|
||||
use serde_json::Value;
|
||||
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(
|
||||
State(state): State<AppState>,
|
||||
_auth_user: AuthenticatedUser,
|
||||
auth_user: AuthenticatedUser,
|
||||
headers: HeaderMap,
|
||||
payload: FlexiblePayload<ExtractionInput>,
|
||||
) -> Result<Response, ApiError> {
|
||||
let (input, _) = payload.into_parts();
|
||||
let result = ai::extract_roast(
|
||||
let (result, usage) = ai::extract_roast(
|
||||
&state.http_client,
|
||||
&state.openrouter_api_key,
|
||||
&state.openrouter_model,
|
||||
|
|
@ -274,6 +274,14 @@ pub(crate) async fn extract_roast_info(
|
|||
.await
|
||||
.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) {
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
|
|||
|
|
@ -14,17 +14,17 @@ use crate::domain::errors::RepositoryError;
|
|||
use crate::domain::roasters::NewRoaster;
|
||||
use crate::domain::roasts::NewRoast;
|
||||
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(
|
||||
State(state): State<AppState>,
|
||||
_auth_user: AuthenticatedUser,
|
||||
auth_user: AuthenticatedUser,
|
||||
headers: HeaderMap,
|
||||
payload: FlexiblePayload<ExtractionInput>,
|
||||
) -> Result<Response, ApiError> {
|
||||
let (input, _) = payload.into_parts();
|
||||
let result = ai::extract_bag_scan(
|
||||
let (result, usage) = ai::extract_bag_scan(
|
||||
&state.http_client,
|
||||
&state.openrouter_api_key,
|
||||
&state.openrouter_model,
|
||||
|
|
@ -33,6 +33,14 @@ pub(crate) async fn extract_bag_scan(
|
|||
.await
|
||||
.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) {
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -130,15 +138,16 @@ struct ScanResult {
|
|||
}
|
||||
|
||||
/// 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(
|
||||
state: &AppState,
|
||||
submission: &mut BagScanSubmission,
|
||||
) -> Result<(), ApiError> {
|
||||
) -> Result<Option<Usage>, ApiError> {
|
||||
let input = ExtractionInput {
|
||||
image: submission.image.take(),
|
||||
prompt: submission.prompt.take(),
|
||||
};
|
||||
let result = ai::extract_bag_scan(
|
||||
let (result, usage) = ai::extract_bag_scan(
|
||||
&state.http_client,
|
||||
&state.openrouter_api_key,
|
||||
&state.openrouter_model,
|
||||
|
|
@ -178,14 +187,14 @@ async fn extract_into_submission(
|
|||
submission.tasting_notes = TastingNotesInput::Text(notes.join(", "));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(usage)
|
||||
}
|
||||
|
||||
#[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(
|
||||
State(state): State<AppState>,
|
||||
_auth_user: AuthenticatedUser,
|
||||
auth_user: AuthenticatedUser,
|
||||
headers: HeaderMap,
|
||||
payload: FlexiblePayload<BagScanSubmission>,
|
||||
) -> Result<Response, ApiError> {
|
||||
|
|
@ -196,7 +205,14 @@ pub(crate) async fn submit_scan(
|
|||
|| submission.prompt.as_deref().is_some_and(|s| !s.is_empty());
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -236,6 +236,31 @@ pub(super) async fn load_cafe_options(state: &AppState) -> Result<Vec<CafeOption
|
|||
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 {
|
||||
headers
|
||||
.get("datastar-request")
|
||||
|
|
|
|||
|
|
@ -12,13 +12,14 @@ use webauthn_rs::prelude::*;
|
|||
use crate::application::routes::app_router;
|
||||
use crate::domain::registration_tokens::NewRegistrationToken;
|
||||
use crate::domain::repositories::{
|
||||
BagRepository, BrewRepository, CafeRepository, CupRepository, GearRepository,
|
||||
PasskeyCredentialRepository, RegistrationTokenRepository, RoastRepository, RoasterRepository,
|
||||
SessionRepository, TimelineEventRepository, TokenRepository, UserRepository,
|
||||
AiUsageRepository, BagRepository, BrewRepository, CafeRepository, CupRepository,
|
||||
GearRepository, PasskeyCredentialRepository, RegistrationTokenRepository, RoastRepository,
|
||||
RoasterRepository, SessionRepository, TimelineEventRepository, TokenRepository, UserRepository,
|
||||
};
|
||||
use crate::infrastructure::auth::{generate_session_token, hash_token};
|
||||
use crate::infrastructure::backup::BackupService;
|
||||
use crate::infrastructure::database::Database;
|
||||
use crate::infrastructure::repositories::ai_usage::SqlAiUsageRepository;
|
||||
use crate::infrastructure::repositories::bags::SqlBagRepository;
|
||||
use crate::infrastructure::repositories::brews::SqlBrewRepository;
|
||||
use crate::infrastructure::repositories::cafes::SqlCafeRepository;
|
||||
|
|
@ -59,6 +60,7 @@ pub struct AppState {
|
|||
pub session_repo: Arc<dyn SessionRepository>,
|
||||
pub passkey_repo: Arc<dyn PasskeyCredentialRepository>,
|
||||
pub registration_token_repo: Arc<dyn RegistrationTokenRepository>,
|
||||
pub ai_usage_repo: Arc<dyn AiUsageRepository>,
|
||||
pub webauthn: Arc<Webauthn>,
|
||||
pub challenge_store: Arc<ChallengeStore>,
|
||||
pub http_client: reqwest::Client,
|
||||
|
|
@ -101,6 +103,8 @@ pub async fn serve(config: ServerConfig) -> anyhow::Result<()> {
|
|||
Arc::new(SqlPasskeyCredentialRepository::new(database.clone_pool()));
|
||||
let registration_token_repo: Arc<dyn RegistrationTokenRepository> =
|
||||
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 challenge_store = Arc::new(ChallengeStore::new());
|
||||
|
|
@ -122,6 +126,7 @@ pub async fn serve(config: ServerConfig) -> anyhow::Result<()> {
|
|||
session_repo,
|
||||
passkey_repo,
|
||||
registration_token_repo,
|
||||
ai_usage_repo,
|
||||
webauthn,
|
||||
challenge_store,
|
||||
http_client: reqwest::Client::new(),
|
||||
|
|
|
|||
|
|
@ -164,6 +164,11 @@ async fn spawn_app_inner(
|
|||
session_repo,
|
||||
passkey_repo,
|
||||
registration_token_repo,
|
||||
ai_usage_repo: Arc::new(
|
||||
brewlog::infrastructure::repositories::ai_usage::SqlAiUsageRepository::new(
|
||||
_database.clone_pool(),
|
||||
),
|
||||
),
|
||||
webauthn: test_webauthn(),
|
||||
challenge_store: Arc::new(ChallengeStore::new()),
|
||||
http_client: reqwest::Client::new(),
|
||||
|
|
|
|||
Loading…
Reference in a new issue