diff --git a/src/artifacts/beets.rs b/src/artifacts/beets.rs index bbf8572..79b2b29 100644 --- a/src/artifacts/beets.rs +++ b/src/artifacts/beets.rs @@ -1,49 +1,279 @@ -use std::process::Stdio; +use std::{collections::HashSet, fmt::Display, process::Stdio}; +use chrono::{DateTime, NaiveDate, NaiveTime, Utc}; use tokio::process::Command; use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize, de::Visitor}; use uuid::Uuid; use crate::artifacts::{Artifact, ArtifactBuilder, Contents, Merge, SourceID, Track, tools::{DataSource, ToolDescription}}; - -#[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)] -pub struct BeatsQueryArgs { - pub artist: Option, - pub album: Option, - pub genre: Option, - pub title: Option, - pub year: Option, - pub label: Option +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub struct DateQuery { + year: u32, + month_and_day: Option<(u32, u32)>, + hour: Option, + minute: Option } +impl Display for DateQuery { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.year)?; + if let Some((month, day)) = self.month_and_day { + write!(f, "-{}-{}", month, day)?; + } + + match (self.hour, self.minute) { + (Some(hour), None) => write!(f, "T{}", hour), + (Some(hour), Some(minute)) => write!(f, "T{}:{}", hour, minute), + (_, _) => Ok(()) + } + } +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub enum RelativeDateUnit { + Day, + Week, + Month, + Year +} + +impl Display for RelativeDateUnit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Day => f.write_str("d"), + Self::Week => f.write_str("w"), + Self::Month => f.write_str("m"), + Self::Year => f.write_str("y") + } + } +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub enum RelativeDateQuery { + Earlier(u32, RelativeDateUnit), + Later(u32, RelativeDateUnit) +} + +impl Display for RelativeDateQuery { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Earlier(q, unit) => write!(f, "-{}{}", q, unit), + Self::Later(q, unit) => write!(f, "+{}{}", q, unit) + } + } +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub enum DateRangeQuery { + Before(DateQuery), + After(DateQuery), + #[serde(untagged)] + Between(DateQuery, DateQuery), + #[serde(untagged)] + Exact(DateQuery) +} + +impl Display for DateRangeQuery { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Exact(q) => write!(f, "{}", q), + Self::Before(q) => write!(f, "..{}", q), + Self::After(q) => write!(f, "{}..", q), + Self::Between(q, z) => write!(f, "{}..{}", q, z), + } + } +} + +/// Matches the numeric value of a field against another number or range of numbers +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub enum NumericRangeQuery { + /// Matches all numbers equal to or greater than the provided value + Above(u32), + /// Matches all numbers equal to or less than the provided value + Below(u32), + #[serde(untagged)] + /// Matches all numbers within the given range, inclusive + Between(u32, u32), + /// Matches exactly this number + #[serde(untagged)] + Exact(u32) +} + +impl Display for NumericRangeQuery { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Exact(v) => write!(f, "{}", v), + Self::Above(v) => write!(f, "{}..", v), + Self::Below(v) => write!(f, "..{}", v), + Self::Between(v, z) => write!(f, "{}..{}", v, z) + } + } +} + +/// Match the string value of a field in various case-insensitive ways +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +pub enum BeetsStringFieldMatch { + /// Matches a field against a regular expression + RegularExpression(String), + /// Performs full-text comparison + ExactMatch(String), + /// WIll match against a partial text comparison; if the provided value is somewhere within the field, it will match + #[serde(untagged)] + Substring(String), +} + +impl Display for BeetsStringFieldMatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::RegularExpression(q) => write!(f, ":{}", q), + Self::Substring(q) => f.write_str(q), + Self::ExactMatch(q) => write!(f, "=~{}", q), + } + } +} + +/// Matches the date value of a field +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +#[serde(untagged)] +pub enum BeetsDateFieldMatch { + /// Represents a known date in time, eg, "24/06/26 at 16:17" + Absolute(DateRangeQuery), + /// Represents a relative date in time, eg "three weeks ago" + Relative(RelativeDateQuery) +} + +impl Display for BeetsDateFieldMatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Absolute(q) => write!(f, "{}", q), + Self::Relative(q) => write!(f, "{}", q), + } + } +} + +/// A single field +#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)] +#[serde(tag = "field", content = "value")] +pub enum BeetsFieldQuery { + Artist(BeetsStringFieldMatch), + Album(BeetsStringFieldMatch), + Genre(BeetsStringFieldMatch), + Title(BeetsStringFieldMatch), + Label(BeetsStringFieldMatch), + + ReleaseDate(BeetsDateFieldMatch), + /// The last time this track was played + LastPlayed(BeetsDateFieldMatch), + /// The date that this track was first imported + Added(BeetsDateFieldMatch), + + /// The "rating" of a track on a 5 point scale, with higher points meaning higher energy + Rating(NumericRangeQuery), + + /// The number of times the track has been played + Playcount(NumericRangeQuery), + + #[serde(skip)] + MixxxPlaylist(String) +} + +impl Display for BeetsFieldQuery { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Artist(q) => write!(f, "artist:{}", q), + Self::Album(q) => write!(f, "album:{}", q), + Self::Genre(q) => write!(f, "genres:{}", q), + Self::Title(q) => write!(f, "title:{}", q), + Self::Label(q) => write!(f, "label:{}", q), + Self::ReleaseDate(q) => write!(f, "date:{}", q), + Self::LastPlayed(q) => write!(f, "mixxx_last_played:{}", q), + Self::Added(q) => write!(f, "added:{}", q), + Self::Rating(q) => write!(f, "mixxx_rating:{}", q), + Self::Playcount(q) => write!(f, "mixxx_times_played:{}", q), + Self::MixxxPlaylist(q) => write!(f, "mixxx_playlist:{}", q) + } + } +} + +/// A list of search parameters to query against the beets database. +#[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)] +pub struct BeetsQueryArgs { + pub search_fields: Vec +} + +/// A list of beets queries to run in parallel #[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)] pub struct BeatsQueryMultiArgs { - args: Vec + args: Vec } #[derive(Debug, Default, Deserialize)] struct BeetsTrack { + title: String, album: String, artist: String, - genres: Option>, + genres: Option>, label: Option, - title: String, - year: u32, + date: Option, + mixxx_last_played: Option, + added: Option, + mixxx_rating: Option, + mixxx_times_played: Option, + mb_trackid: Option } +#[derive(Debug)] +struct BeetsTimestamp(NaiveDate); + +struct BeetsTimestampVisitor; + +impl<'de> Visitor<'de> for BeetsTimestampVisitor { + type Value = BeetsTimestamp; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a beets date format string %Y-%m-%d") + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + log::debug!("visit str"); + Ok(BeetsTimestamp(NaiveDate::parse_from_str(v, "%Y-%m-%d").unwrap().into())) + } +} + +impl<'de> Deserialize<'de> for BeetsTimestamp { + + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de> { + deserializer.deserialize_str(BeetsTimestampVisitor) + } +} + +impl From for DateTime { + fn from(val: BeetsTimestamp) -> Self { + val.0.and_time(NaiveTime::default()).and_utc() + } +} + impl From for Artifact { fn from(val: BeetsTrack) -> Self { let track_data = Track { title: val.title, - label: val.label, - year: Some(val.year), - genres: val.genres.unwrap_or_default(), album: Some(val.album), artist: Some(val.artist), + genres: val.genres.unwrap_or_default(), + label: val.label, + release_date: val.date.map(BeetsTimestamp::into), + last_played: val.mixxx_last_played.map(BeetsTimestamp::into), + added: val.added.map(BeetsTimestamp::into), + rating: val.mixxx_rating, bpm: None, }; let builder = ArtifactBuilder::new(SourceID::Beets) @@ -89,62 +319,43 @@ impl BeetsDB { Ok(ret) } - async fn query_single(&self, args: &BeatsQueryArgs) -> Result, BeetsError> { + async fn query_single(&self, args: &BeetsQueryArgs) -> Result, BeetsError> { let mut beets_cmd = Command::new("beet"); - beets_cmd.args(["export", "-f", "json", "-i", "title,label,year,genres,album,artist,mb_trackid"]); - let mut valid = false; - if let Some(ref artist) = args.artist { - beets_cmd.arg(format!("artist:{}", artist)); - valid = true; - } - if let Some(ref genre) = args.genre { - beets_cmd.arg(format!("genre:{}", genre)); - valid = true; - } - if let Some(ref album) = args.album { - beets_cmd.arg(format!("album:{}", album)); - valid = true; - } - if let Some(ref title) = args.title { - beets_cmd.arg(format!("title:{}", title)); - valid = true; - } - if let Some(year) = args.year { - beets_cmd.arg(format!("year:{}", year)); - valid = true; - } - if let Some(ref label) = args.label { - beets_cmd.arg(format!("label:{}", label)); - valid = true; - } + beets_cmd.args(["export", "-f", "json", "-i", "title,artist,album,genres,label,date,mixxx_last_played,added,mixxx_rating,mixxx_times_played,mb_trackid"]); - if !valid { + if args.search_fields.is_empty() { log::warn!("Tried to execute an empty beets query"); return Err(BeetsError::EmptyQuery) } + for field in &args.search_fields { + beets_cmd.arg(format!("{}", field)); + } + log::debug!("Executing beets: {:?}", beets_cmd); let output = beets_cmd.stdout(Stdio::piped()).stderr(Stdio::null()).spawn()?.wait_with_output().await?; - let track = serde_json::from_str::>(str::from_utf8(&output.stdout).unwrap())?; - Ok(track.into_iter().map(|t| { t.into()}).collect()) + log::debug!("Output: {}", str::from_utf8(&output.stdout).unwrap()); + let track = serde_json::from_str::>(str::from_utf8(&output.stdout).unwrap()).unwrap(); + Ok(track.into_iter().map(BeetsTrack::into).collect()) } } impl DataSource for BeetsDB { - type Args = BeatsQueryMultiArgs; + type Args = BeetsQueryArgs; type Error = BeetsError; async fn synchronize(&self, artifact: &mut Artifact) -> Result, Self::Error> { if let Contents::Track(ref mut target_track) = artifact.contents { - let args = BeatsQueryArgs { - title: Some(target_track.title.clone()), - artist: target_track.artist.clone(), - album: target_track.album.clone(), - ..Default::default() + let args = BeetsQueryArgs { + search_fields: vec![ + BeetsFieldQuery::Title(BeetsStringFieldMatch::Substring(target_track.title.clone())), + BeetsFieldQuery::Artist(BeetsStringFieldMatch::Substring(target_track.artist.clone().unwrap_or_default())), + BeetsFieldQuery::Album(BeetsStringFieldMatch::Substring(target_track.album.clone().unwrap_or_default())) + ] }; - let results = self.query(&BeatsQueryMultiArgs { args: vec![args] }).await?; + let results = self.query(&args).await?; if let Some(first) = results.first() { artifact.merge(first.clone()); @@ -157,7 +368,7 @@ impl DataSource for BeetsDB { } fn query(&self, args: &Self::Args) -> impl Future, Self::Error>> { - self.query_multi(args) + self.query_single(args) } diff --git a/src/prediction/mod.rs b/src/prediction/mod.rs index 415b7f7..4ac3698 100644 --- a/src/prediction/mod.rs +++ b/src/prediction/mod.rs @@ -1,6 +1,6 @@ use std::{collections::HashMap, fmt::Debug, sync::Arc}; -use async_openai::types::chat::{ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, }; +use async_openai::types::chat::{ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestUserMessageArgs, }; use chrono::{DateTime, Utc}; use futures::stream::FuturesUnordered; use schemars::{JsonSchema, schema_for}; @@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Serializer, Value, ser::CompactFormatter}; use tokio::sync::{Mutex, mpsc::{self, UnboundedReceiver, UnboundedSender}}; -use crate::{SaveData, artifacts::{Contents, Track, archive::Archive, mixxx::{MixxxDB, MixxxQuery}, tools::DataSource}, prediction::{character::{Character, CharacterControl, CharacterOutput, character_task}, toolbox::{ArchiveToolbox, StageToolbox}}, scene::{Scene, StageDirection, conversation::{ConversationEntry, Speaker}}, sfx::SfxControl}; +use crate::{SaveData, artifacts::{Contents, Track, archive::Archive, beets::{BeetsDB, BeetsFieldQuery, BeetsQueryArgs}, mixxx::{MixxxDB, MixxxQuery}, tools::DataSource}, prediction::{character::{Character, CharacterControl, CharacterOutput, character_task}, toolbox::{ArchiveToolbox, StageToolbox}}, scene::{Scene, StageDirection, conversation::{ConversationEntry, Speaker}}, sfx::SfxControl}; use tokio_stream::StreamExt; pub mod character; @@ -220,8 +220,12 @@ impl Conversation { self.insert(ConversationEntry::ShipComputerCommand(command)).await; }, PredictionAction::SetPlaylist(playlist_name) => { - let args = MixxxQuery { playlist_name }; - match MixxxDB.query(&args).await { + let args = BeetsQueryArgs { + search_fields: vec![ + BeetsFieldQuery::MixxxPlaylist(playlist_name.clone()) + ] + }; + match BeetsDB.query(&args).await { Err(err) => log::info!("Failed to load mixxx playlist: {:?}.", err), Ok(playlist) => { self.current_playlist = vec![]; @@ -232,7 +236,7 @@ impl Conversation { self.current_playlist.push(as_track.clone()); } } - self.direction.playlist = args.playlist_name; + self.direction.playlist = playlist_name; log::info!("Mixxx playlist reloaded."); drop(archive);