From b5413eef467f2d853e704be1df7eefba1e22aa9c Mon Sep 17 00:00:00 2001 From: "naps62-yolo (agent)" Date: Sat, 22 Aug 2026 22:12:17 +0100 Subject: [PATCH] Use policy scoring weights in release search (#74) --- ...11ab88e1fdcc28c432da4c3debf1202eee7a.json} | 10 +- crates/arr-api/src/search.rs | 159 +++++++++++++----- 2 files changed, 123 insertions(+), 46 deletions(-) rename .sqlx/{query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json => query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json} (83%) diff --git a/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json b/.sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json similarity index 83% rename from .sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json rename to .sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json index 04ae29f..f613aaf 100644 --- a/.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json +++ b/.sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json @@ -1,6 +1,6 @@ { "db_name": "SQLite", - "query": "SELECT m.title AS \"title!: String\", m.tmdb_id AS \"tmdb_id!: i64\", m.original_language, m.overrides AS \"overrides!: serde_json::Value\", p.id AS \"policy_id!: i64\", p.name AS \"policy_name!: String\", p.required_audio AS \"required_audio!: String\", p.dub_blacklist AS \"dub_blacklist!: String\", p.hdr_rules AS \"hdr_rules!: String\", p.size_bands AS \"size_bands!: String\", p.resolution_pref AS \"resolution_pref!: String\", p.source_weights AS \"source_weights!: String\" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?", + "query": "SELECT m.title AS \"title!: String\", m.tmdb_id AS \"tmdb_id!: i64\", m.original_language, m.overrides AS \"overrides!: serde_json::Value\", p.id AS \"policy_id!: i64\", p.name AS \"policy_name!: String\", p.required_audio AS \"required_audio!: String\", p.dub_blacklist AS \"dub_blacklist!: String\", p.hdr_rules AS \"hdr_rules!: String\", p.size_bands AS \"size_bands!: String\", p.resolution_pref AS \"resolution_pref!: String\", p.source_weights AS \"source_weights!: String\", p.score_weights AS \"score_weights!: String\" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?", "describe": { "columns": [ { @@ -62,6 +62,11 @@ "name": "source_weights!: String", "ordinal": 11, "type_info": "Text" + }, + { + "name": "score_weights!: String", + "ordinal": 12, + "type_info": "Text" } ], "parameters": { @@ -79,8 +84,9 @@ false, false, false, + false, false ] }, - "hash": "b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc" + "hash": "6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a" } diff --git a/crates/arr-api/src/search.rs b/crates/arr-api/src/search.rs index 8745d56..e681ecf 100644 --- a/crates/arr-api/src/search.rs +++ b/crates/arr-api/src/search.rs @@ -2,6 +2,7 @@ use std::collections::BTreeMap; use std::time::UNIX_EPOCH; use arr_core::policy::{evaluate, Candidate}; +use arr_core::score::score; use arr_core::{ DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution, Rule, ScoreWeights, SizeBand, Source, TitleOverrides, Verdict, @@ -80,6 +81,7 @@ struct PolicyRow { size_bands: String, resolution_pref: String, source_weights: String, + score_weights: String, } #[derive(Debug, Deserialize)] @@ -102,6 +104,13 @@ struct SizeBandJson { penalty_points_per_gib_over: i32, } +#[derive(Debug, Deserialize)] +struct ScoreWeightsJson { + size_at_target: i32, + source_tier: i32, + seeder_doubling: i32, +} + #[derive(Debug, Deserialize)] struct OverridesJson { #[serde(default)] @@ -254,7 +263,7 @@ pub async fn releases( Query(query): Query, ) -> Result>, ApiError> { let database = state.database().ok_or(ApiError::Unavailable)?; - let movie = sqlx::query!(r#"SELECT m.title AS "title!: String", m.tmdb_id AS "tmdb_id!: i64", m.original_language, m.overrides AS "overrides!: serde_json::Value", p.id AS "policy_id!: i64", p.name AS "policy_name!: String", p.required_audio AS "required_audio!: String", p.dub_blacklist AS "dub_blacklist!: String", p.hdr_rules AS "hdr_rules!: String", p.size_bands AS "size_bands!: String", p.resolution_pref AS "resolution_pref!: String", p.source_weights AS "source_weights!: String" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?"#, query.movie_id) + let movie = sqlx::query!(r#"SELECT m.title AS "title!: String", m.tmdb_id AS "tmdb_id!: i64", m.original_language, m.overrides AS "overrides!: serde_json::Value", p.id AS "policy_id!: i64", p.name AS "policy_name!: String", p.required_audio AS "required_audio!: String", p.dub_blacklist AS "dub_blacklist!: String", p.hdr_rules AS "hdr_rules!: String", p.size_bands AS "size_bands!: String", p.resolution_pref AS "resolution_pref!: String", p.source_weights AS "source_weights!: String", p.score_weights AS "score_weights!: String" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?"#, query.movie_id) .fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?; let tmdb = tmdb_client(&state)? @@ -290,6 +299,7 @@ pub async fn releases( size_bands: movie.size_bands, resolution_pref: movie.resolution_pref, source_weights: movie.source_weights, + score_weights: movie.score_weights, })?; let overrides: OverridesJson = serde_json::from_value(movie.overrides) .map_err(|error| ApiError::Database(error.to_string()))?; @@ -405,6 +415,7 @@ fn policy_from_row(row: PolicyRow) -> Result { let bands: BTreeMap = json(&row.size_bands)?; let resolutions: Vec = json(&row.resolution_pref)?; let weights: BTreeMap = json(&row.source_weights)?; + let score_weights: ScoreWeightsJson = json(&row.score_weights)?; Ok(Policy { id: PolicyId(row.policy_id), name: row.policy_name, @@ -451,9 +462,11 @@ fn policy_from_row(row: PolicyRow) -> Result { .into_iter() .filter_map(|(source, weight)| source_value(&source).map(|source| (source, weight))) .collect(), - // The row's score_weights column is not selected here yet; the seeded - // values match these defaults. Wiring the column through is #68. - score_weights: ScoreWeights::default(), + score_weights: ScoreWeights { + size_at_target: score_weights.size_at_target, + source_tier: score_weights.source_tier, + seeder_doubling: score_weights.seeder_doubling, + }, }) } @@ -472,7 +485,17 @@ fn classify( release.size, ); let (verdict, rule) = verdict(&evaluation.verdict); - let score = score(policy, &parsed, release.size, release.seeders); + let score = score( + policy, + Candidate::PreGrab(&parsed), + release.size.unwrap_or_default(), + release.seeders.unwrap_or_default(), + ); + let score = if release.size.is_some() { + score.total + } else { + score.source.saturating_add(score.seeders) + }; Ok(ClassifiedRelease { indexer_id: release.indexer_id, guid: release.guid, @@ -495,44 +518,6 @@ fn classify( }) } -fn score( - policy: &Policy, - parsed: &arr_parse::NameClaims, - size: Option, - seeders: Option, -) -> i64 { - let resolution = parsed.resolution.map(Resolution::from); - let resolution_score = resolution - .and_then(|value| { - policy - .resolution_preference - .iter() - .position(|candidate| *candidate == value) - }) - .map_or(0, |position| { - 10_000 - i64::try_from(position).unwrap_or(0) * 5_000 - }); - let size_score = resolution - .and_then(|value| policy.size_bands.get(&value)) - .zip(size) - .map_or(0, |(band, bytes)| { - if bytes < band.floor_bytes { - -i64::try_from((band.floor_bytes - bytes) / 100_000_000).unwrap_or(i64::MAX) - } else { - -i64::try_from(bytes.abs_diff(band.target_bytes) / 100_000_000).unwrap_or(i64::MAX) - } - }); - let source_score = i64::from( - parsed - .source - .map(Source::from) - .and_then(|source| policy.source_weights.get(&source).copied()) - .unwrap_or(0), - ) * 5; - let seeder_score = seeders.map_or(0, |count| i64::from(count.saturating_add(1).ilog2()) * 2); - resolution_score + size_score + source_score + seeder_score -} - fn verdict(verdict: &Verdict) -> (&'static str, Option) { match verdict { Verdict::Eligible => ("eligible", None), @@ -729,6 +714,12 @@ mod tests { let (_dir, state, base) = application(&tmdb, &prowlarr).await; sqlx::query("INSERT INTO movies (tmdb_id, title, year, original_language, root_id) VALUES (693134, 'Dune Part Two', 2024, 'en', 2)") .execute(state.database().expect("database").pool()).await.expect("movie"); + sqlx::query( + "UPDATE policies SET score_weights = '{\"size_at_target\":0,\"source_tier\":0,\"seeder_doubling\":0}'", + ) + .execute(state.database().expect("database").pool()) + .await + .expect("score weights"); let response = reqwest::get(format!("{base}/api/releases?movie_id=1")) .await @@ -744,7 +735,87 @@ mod tests { .find(|release| release["verdict"] == "rejected") .expect("rejected"); assert_eq!(rejected["rule"], "source"); - assert!(releases[0]["score"].as_i64().is_some()); + let eligible = releases + .iter() + .find(|release| release["guid"] == "good") + .expect("eligible"); + assert_eq!(eligible["score"], 0); + } + + #[test] + fn policy_row_uses_persisted_score_weights() { + let policy = policy_from_row(PolicyRow { + policy_id: 1, + policy_name: "test".into(), + required_audio: r#"{"require":"original_language"}"#.into(), + dub_blacklist: "[]".into(), + hdr_rules: "{}".into(), + size_bands: + r#"{"2160p":{"floor_gib":8,"target_gib":22,"penalty_points_per_gib_over":60}}"# + .into(), + resolution_pref: r#"["2160p"]"#.into(), + source_weights: r#"{"WEB-DL":2}"#.into(), + score_weights: r#"{"size_at_target":2000,"source_tier":7,"seeder_doubling":11}"#.into(), + }) + .expect("policy row"); + + assert_eq!( + policy.score_weights, + ScoreWeights { + size_at_target: 2000, + source_tier: 7, + seeder_doubling: 11, + } + ); + } + + #[test] + fn releases_without_sizes_skip_the_size_score() { + let policy = Policy { + id: PolicyId(1), + name: "test".into(), + required_audio: RequiredAudio::OriginalLanguage, + dub_blacklist: Vec::new(), + hdr_rules: HdrRules { + rejected_dolby_vision_profiles: Vec::new(), + }, + size_bands: std::collections::BTreeMap::from([( + Resolution::R2160p, + SizeBand { + floor_bytes: gib(8), + target_bytes: gib(22), + penalty_points_per_gib_over: 60, + }, + )]), + resolution_preference: vec![Resolution::R2160p], + source_weights: std::collections::BTreeMap::from([(Source::WebDl, 2)]), + score_weights: ScoreWeights::default(), + }; + let release = SearchRelease { + indexer_id: 1, + guid: "release".into(), + name: "Dune.Part.Two.2024.2160p.WEB-DL".into(), + size: None, + seeders: Some(8), + publish_date: None, + download_url: "https://tracker/release".into(), + }; + let parsed = arr_parse::parse(&release.name); + let core_score = score(&policy, Candidate::PreGrab(&parsed), 0, 8); + + let classified = classify( + release, + &policy, + &MovieOverrides::default(), + &Language::Other("en".into()), + ) + .expect("classified release"); + + assert_eq!( + classified.score, + core_score.source.saturating_add(core_score.seeders) + ); + assert_ne!(classified.score, core_score.total); } #[test]