Use policy scoring weights in release search (#74)
This commit was merged in pull request #74.
This commit is contained in:
+8
-2
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"db_name": "SQLite",
|
"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": {
|
"describe": {
|
||||||
"columns": [
|
"columns": [
|
||||||
{
|
{
|
||||||
@@ -62,6 +62,11 @@
|
|||||||
"name": "source_weights!: String",
|
"name": "source_weights!: String",
|
||||||
"ordinal": 11,
|
"ordinal": 11,
|
||||||
"type_info": "Text"
|
"type_info": "Text"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "score_weights!: String",
|
||||||
|
"ordinal": 12,
|
||||||
|
"type_info": "Text"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"parameters": {
|
"parameters": {
|
||||||
@@ -79,8 +84,9 @@
|
|||||||
false,
|
false,
|
||||||
false,
|
false,
|
||||||
false,
|
false,
|
||||||
|
false,
|
||||||
false
|
false
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
"hash": "b78bef6253289702123c03afcab624d5d50df897af4a477a9cf3dac196b52ecc"
|
"hash": "6e988f7e8d8ac2167c589a35626b11ab88e1fdcc28c432da4c3debf1202eee7a"
|
||||||
}
|
}
|
||||||
+115
-44
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
|||||||
use std::time::UNIX_EPOCH;
|
use std::time::UNIX_EPOCH;
|
||||||
|
|
||||||
use arr_core::policy::{evaluate, Candidate};
|
use arr_core::policy::{evaluate, Candidate};
|
||||||
|
use arr_core::score::score;
|
||||||
use arr_core::{
|
use arr_core::{
|
||||||
DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution, Rule,
|
DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution, Rule,
|
||||||
ScoreWeights, SizeBand, Source, TitleOverrides, Verdict,
|
ScoreWeights, SizeBand, Source, TitleOverrides, Verdict,
|
||||||
@@ -80,6 +81,7 @@ struct PolicyRow {
|
|||||||
size_bands: String,
|
size_bands: String,
|
||||||
resolution_pref: String,
|
resolution_pref: String,
|
||||||
source_weights: String,
|
source_weights: String,
|
||||||
|
score_weights: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
@@ -102,6 +104,13 @@ struct SizeBandJson {
|
|||||||
penalty_points_per_gib_over: i32,
|
penalty_points_per_gib_over: i32,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct ScoreWeightsJson {
|
||||||
|
size_at_target: i32,
|
||||||
|
source_tier: i32,
|
||||||
|
seeder_doubling: i32,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct OverridesJson {
|
struct OverridesJson {
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
@@ -254,7 +263,7 @@ pub async fn releases(
|
|||||||
Query(query): Query<ReleasesQuery>,
|
Query(query): Query<ReleasesQuery>,
|
||||||
) -> Result<Json<Vec<ClassifiedRelease>>, ApiError> {
|
) -> Result<Json<Vec<ClassifiedRelease>>, ApiError> {
|
||||||
let database = state.database().ok_or(ApiError::Unavailable)?;
|
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)?;
|
.fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?;
|
||||||
|
|
||||||
let tmdb = tmdb_client(&state)?
|
let tmdb = tmdb_client(&state)?
|
||||||
@@ -290,6 +299,7 @@ pub async fn releases(
|
|||||||
size_bands: movie.size_bands,
|
size_bands: movie.size_bands,
|
||||||
resolution_pref: movie.resolution_pref,
|
resolution_pref: movie.resolution_pref,
|
||||||
source_weights: movie.source_weights,
|
source_weights: movie.source_weights,
|
||||||
|
score_weights: movie.score_weights,
|
||||||
})?;
|
})?;
|
||||||
let overrides: OverridesJson = serde_json::from_value(movie.overrides)
|
let overrides: OverridesJson = serde_json::from_value(movie.overrides)
|
||||||
.map_err(|error| ApiError::Database(error.to_string()))?;
|
.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 bands: BTreeMap<String, SizeBandJson> = json(&row.size_bands)?;
|
||||||
let resolutions: Vec<String> = json(&row.resolution_pref)?;
|
let resolutions: Vec<String> = json(&row.resolution_pref)?;
|
||||||
let weights: BTreeMap<String, i32> = json(&row.source_weights)?;
|
let weights: BTreeMap<String, i32> = json(&row.source_weights)?;
|
||||||
|
let score_weights: ScoreWeightsJson = json(&row.score_weights)?;
|
||||||
Ok(Policy {
|
Ok(Policy {
|
||||||
id: PolicyId(row.policy_id),
|
id: PolicyId(row.policy_id),
|
||||||
name: row.policy_name,
|
name: row.policy_name,
|
||||||
@@ -451,9 +462,11 @@ fn policy_from_row(row: PolicyRow) -> Result<Policy, ApiError> {
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.filter_map(|(source, weight)| source_value(&source).map(|source| (source, weight)))
|
.filter_map(|(source, weight)| source_value(&source).map(|source| (source, weight)))
|
||||||
.collect(),
|
.collect(),
|
||||||
// The row's score_weights column is not selected here yet; the seeded
|
score_weights: ScoreWeights {
|
||||||
// values match these defaults. Wiring the column through is #68.
|
size_at_target: score_weights.size_at_target,
|
||||||
score_weights: ScoreWeights::default(),
|
source_tier: score_weights.source_tier,
|
||||||
|
seeder_doubling: score_weights.seeder_doubling,
|
||||||
|
},
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -472,7 +485,17 @@ fn classify(
|
|||||||
release.size,
|
release.size,
|
||||||
);
|
);
|
||||||
let (verdict, rule) = verdict(&evaluation.verdict);
|
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 {
|
Ok(ClassifiedRelease {
|
||||||
indexer_id: release.indexer_id,
|
indexer_id: release.indexer_id,
|
||||||
guid: release.guid,
|
guid: release.guid,
|
||||||
@@ -495,44 +518,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>) {
|
fn verdict(verdict: &Verdict) -> (&'static str, Option<String>) {
|
||||||
match verdict {
|
match verdict {
|
||||||
Verdict::Eligible => ("eligible", None),
|
Verdict::Eligible => ("eligible", None),
|
||||||
@@ -729,6 +714,12 @@ mod tests {
|
|||||||
let (_dir, state, base) = application(&tmdb, &prowlarr).await;
|
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)")
|
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");
|
.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"))
|
let response = reqwest::get(format!("{base}/api/releases?movie_id=1"))
|
||||||
.await
|
.await
|
||||||
@@ -744,7 +735,87 @@ mod tests {
|
|||||||
.find(|release| release["verdict"] == "rejected")
|
.find(|release| release["verdict"] == "rejected")
|
||||||
.expect("rejected");
|
.expect("rejected");
|
||||||
assert_eq!(rejected["rule"], "source");
|
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]
|
#[test]
|
||||||
|
|||||||
Reference in New Issue
Block a user