artifacts: beets: rewrite the query API to be more comprehensive and structured

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<String>,
-    pub album: Option<String>,
-    pub genre: Option<String>,
-    pub title: Option<String>,
-    pub year: Option<u32>,
-    pub label: Option<String>
+#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)]
+pub struct DateQuery {
+    year: u32,
+    month_and_day: Option<(u32, u32)>,
+    hour: Option<u32>,
+    minute: Option<u32>
 }
 
+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<BeetsFieldQuery>
+}
+
+/// A list of beets queries to run in parallel
 #[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)]
 pub struct BeatsQueryMultiArgs {
-    args: Vec<BeatsQueryArgs>
+    args: Vec<BeetsQueryArgs>
 }
 
 #[derive(Debug, Default, Deserialize)]
 struct BeetsTrack {
+    title: String,
     album: String,
     artist: String,
-    genres: Option<Vec<String>>,
+    genres: Option<HashSet<String>>,
     label: Option<String>,
-    title: String,
-    year: u32,
+    date: Option<BeetsTimestamp>,
+    mixxx_last_played: Option<BeetsTimestamp>,
+    added: Option<BeetsTimestamp>,
+    mixxx_rating: Option<u32>,
+    mixxx_times_played: Option<u32>,
+
     mb_trackid: Option<String>
 }
 
+#[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<E>(self, v: &str) -> Result<Self::Value, E>
+    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<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de> {
+        deserializer.deserialize_str(BeetsTimestampVisitor)
+    }
+}
+
+impl From<BeetsTimestamp> for DateTime<Utc> {
+    fn from(val: BeetsTimestamp) -> Self {
+        val.0.and_time(NaiveTime::default()).and_utc()
+    }
+}
+
 impl From<BeetsTrack> 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<Vec<Artifact>, BeetsError> {
+    async fn query_single(&self, args: &BeetsQueryArgs) -> Result<Vec<Artifact>, 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::<Vec<BeetsTrack>>(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::<Vec<BeetsTrack>>(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<Vec<Artifact>, 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<Output = Result<Vec<Artifact>, Self::Error>> {
-        self.query_multi(args)
+        self.query_single(args)
     }
 
     

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);