// AI-budsjettsjekk — sjekker ai_budget-metadata mot ai_usage_log aggregat. // // Samlings- og brukernoder kan ha budsjettgrense i metadata: // { "ai_budget": { "monthly_limit_usd": 50.0 } } // // Sjekken aggregerer denne månedens forbruk fra ai_usage_log og // sammenligner med grensen. Ved overskridelse returneres feil // og en work_item-node opprettes. // // Ref: docs/infra/ai_gateway.md § 6.3, oppgave 28.3 use sqlx::PgPool; use uuid::Uuid; /// Estimerer kostnad fra tokens (samme formel som metrics.rs). /// Input: $3/MTok, Output: $15/MTok (konservativt gjennomsnitt). pub fn estimate_cost_usd(prompt_tokens: i64, completion_tokens: i64) -> f64 { let input_cost = prompt_tokens as f64 * 3.0 / 1_000_000.0; let output_cost = completion_tokens as f64 * 15.0 / 1_000_000.0; ((input_cost + output_cost) * 100.0).round() / 100.0 } /// Resultat av en budsjettsjekk. pub struct BudgetStatus { pub current_cost: f64, pub limit: f64, pub node_title: String, } /// Sjekker ai_budget for en samlings-node. /// /// Returnerer Ok(()) hvis budsjettet ikke er overskredet eller ikke er satt. /// Returnerer Err med BudgetStatus hvis overskredet. pub async fn check_collection_budget( db: &PgPool, collection_id: Uuid, ) -> Result<(), BudgetStatus> { check_budget_for_node(db, collection_id, "collection_node_id").await } /// Sjekker ai_budget for en bruker-node. pub async fn check_user_budget( db: &PgPool, user_id: Uuid, ) -> Result<(), BudgetStatus> { check_budget_for_node(db, user_id, "requested_by").await } async fn check_budget_for_node( db: &PgPool, node_id: Uuid, column: &str, ) -> Result<(), BudgetStatus> { // Hent metadata fra noden let metadata: Option = match sqlx::query_scalar("SELECT metadata FROM nodes WHERE id = $1") .bind(node_id) .fetch_optional(db) .await { Ok(m) => m, Err(e) => { tracing::warn!(error = %e, "Kunne ikke hente metadata for budsjettsjekk"); return Ok(()); // Fail open — ikke blokker AI-kall ved DB-feil } }; let monthly_limit = metadata .as_ref() .and_then(|m| m.get("ai_budget")) .and_then(|b| b.get("monthly_limit_usd")) .and_then(|v| v.as_f64()); let limit = match monthly_limit { Some(l) => l, None => return Ok(()), // Ingen budsjettgrense satt }; // Aggreger denne månedens forbruk // Bruker dynamisk kolonne via to separate queries for sikkerhet (unngår SQL injection) let (total_prompt, total_completion): (i64, i64) = match column { "collection_node_id" => { match sqlx::query_as::<_, (i64, i64)>( r#"SELECT COALESCE(SUM(prompt_tokens)::BIGINT, 0), COALESCE(SUM(completion_tokens)::BIGINT, 0) FROM ai_usage_log WHERE collection_node_id = $1 AND created_at >= date_trunc('month', now())"#, ) .bind(node_id) .fetch_one(db) .await { Ok(r) => r, Err(e) => { tracing::warn!(error = %e, "Kunne ikke aggregere AI-forbruk"); return Ok(()); // Fail open } } } "requested_by" => { match sqlx::query_as::<_, (i64, i64)>( r#"SELECT COALESCE(SUM(prompt_tokens)::BIGINT, 0), COALESCE(SUM(completion_tokens)::BIGINT, 0) FROM ai_usage_log WHERE requested_by = $1 AND created_at >= date_trunc('month', now())"#, ) .bind(node_id) .fetch_one(db) .await { Ok(r) => r, Err(e) => { tracing::warn!(error = %e, "Kunne ikke aggregere bruker-AI-forbruk"); return Ok(()); // Fail open } } } _ => return Ok(()), }; let current_cost = estimate_cost_usd(total_prompt, total_completion); if current_cost >= limit { let node_title: String = sqlx::query_scalar("SELECT title FROM nodes WHERE id = $1") .bind(node_id) .fetch_optional(db) .await .ok() .flatten() .unwrap_or_else(|| node_id.to_string()); return Err(BudgetStatus { current_cost, limit, node_title, }); } tracing::debug!( node_id = %node_id, current_cost, limit, "Budsjettsjekk OK" ); Ok(()) } /// Oppretter work_item-node når budsjett er overskredet. pub async fn create_budget_work_item( db: &PgPool, collection_id: Uuid, created_by: Option, status: &BudgetStatus, ) { let work_item_id = Uuid::now_v7(); let title = format!("AI-budsjett overskredet: {}", status.node_title); let content = format!( "AI-budsjettet for \"{}\" er overskredet.\n\ Forbruk denne måneden: ${:.2}\n\ Budsjettgrense: ${:.2}\n\n\ AI-kall er blokkert inntil budsjettet økes eller ny måned starter.", status.node_title, status.current_cost, status.limit ); let creator = created_by.unwrap_or(collection_id); let metadata = serde_json::json!({ "work_item_type": "budget_exceeded", "collection_id": collection_id.to_string(), "current_cost_usd": status.current_cost, "monthly_limit_usd": status.limit, }); if let Err(e) = sqlx::query( r#"INSERT INTO nodes (id, node_kind, title, content, visibility, metadata, created_by) VALUES ($1, 'content', $2, $3, 'hidden'::visibility, $4, $5)"#, ) .bind(work_item_id) .bind(&title) .bind(&content) .bind(&metadata) .bind(creator) .execute(db) .await { tracing::warn!(error = %e, "Kunne ikke opprette work_item for budsjettoverskridelse"); return; } // tagged-edge: budget_exceeded let tag_edge_id = Uuid::now_v7(); if let Err(e) = sqlx::query( r#"INSERT INTO edges (id, source_id, target_id, edge_type, metadata, system, created_by) VALUES ($1, $2, $2, 'tagged', '{"tag": "budget_exceeded"}'::jsonb, true, $3)"#, ) .bind(tag_edge_id) .bind(work_item_id) .bind(creator) .execute(db) .await { tracing::warn!(error = %e, "Kunne ikke opprette tagged-edge for work_item"); } // belongs_to-edge til samlingen let belongs_edge_id = Uuid::now_v7(); if let Err(e) = sqlx::query( r#"INSERT INTO edges (id, source_id, target_id, edge_type, metadata, system, created_by) VALUES ($1, $2, $3, 'belongs_to', '{}'::jsonb, true, $4)"#, ) .bind(belongs_edge_id) .bind(work_item_id) .bind(collection_id) .bind(creator) .execute(db) .await { tracing::warn!(error = %e, "Kunne ikke opprette belongs_to-edge for work_item"); } tracing::info!( work_item_id = %work_item_id, collection_id = %collection_id, current_cost = status.current_cost, limit = status.limit, "Work item opprettet for budsjettoverskridelse" ); }