yolo/arr
3
0
Fork 0
Code Issues 12 Pull Requests Actions Packages Projects Releases Wiki Activity

fix(api): use policy scoring weights

Browse Source
This commit is contained in:
Miguel Palhas
2026-08-22 21:45:21 +01:00
parent 8444b6ebe5
commit 19342d70e8
2 changed files with 70 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.sqlx/query-b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc.json → .sqlx/query-6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a.json
Generated
+8 -2
View File
@@ -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"
}
crates/arr-api/src/search.rs
+62 -44
View File
@@ -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, MovieOverrides, Policy, PolicyId, RequiredAudio,
Resolution, Rule, ScoreWeights, SizeBand, Source, 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<ReleasesQuery>,
) -> Result<Json<Vec<ClassifiedRelease>>, 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<Policy, ApiError> {
let bands: BTreeMap<String, SizeBandJson> = json(&row.size_bands)?;
let resolutions: Vec<String> = json(&row.resolution_pref)?;
let weights: BTreeMap<String, i32> = 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<Policy, ApiError> {
.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,13 @@ 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(),
)
.total;
Ok(ClassifiedRelease {
indexer_id: release.indexer_id,
guid: release.guid,
@@ -495,44 +514,6 @@ fn classify(
})
}
fn score(
policy: &Policy,
parsed: &arr_parse::NameClaims,
size: Option<u64>,
seeders: Option<u32>,
) -> 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<String>) {
match verdict {
Verdict::Eligible => ("eligible", None),
@@ -729,6 +710,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 +731,38 @@ 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]