Enhance state direction command with ship computer outputs, and report token burn on the UI

src/main.rs

@@ -13,7 +13,7 @@ use futures::{StreamExt, future::FutureExt};
 use ratatui::prelude::*;
 use tui_skeleton::{AnimationMode, SkeletonText};
 
-use crate::{audio::{AudioInputControl, start_audio_input}, prediction::{BandcampResult, PossibleResponse}, scene::{ConversationEntry, Scene, StageActions, StageDirection}, transcription::TranscriptionControl, tts::{TtsControl, start_tts}};
+use crate::{audio::{AudioInputControl, start_audio_input}, prediction::{BandcampResult, PossibleResponse}, scene::{ConversationEntry, Scene, Scenery, StageActions, StageDirection}, transcription::TranscriptionControl, tts::{TtsControl, start_tts}};
 
 mod scene;
 mod events;
@@ -94,24 +94,6 @@ enum FocusState {
     UserInput
 }
 
-#[derive(Debug)]
-enum BandcampError {
-    Json(serde_json::Error),
-    Api(bandcamp::Error)
-}
-
-impl From<serde_json::Error> for BandcampError {
-    fn from(value: serde_json::Error) -> Self {
-        BandcampError::Json(value)
-    }
-}
-
-impl From<bandcamp::Error> for BandcampError {
-    fn from(value: bandcamp::Error) -> Self {
-        BandcampError::Api(value)
-    }
-}
-
 impl App {
     fn new(prediction_request_sink: watch::Sender<StageActions>, audio: AudioInputControl, transcription: TranscriptionControl, tts: TtsControl, sys_message_sink: mpsc::Sender<String>, initial_direction: StageDirection) -> Self {
         Self {
@@ -292,6 +274,9 @@ impl App {
             Span::from(" | ").style(ratatui::style::Color::DarkGray),
             Span::from(format!("Time Remaining: {}", formatted_time)).style(time_style),
             Span::from(" | ").style(ratatui::style::Color::DarkGray),
+            Span::from(format!("{} artifacts recorded", self.scene.scenery().artifacts.len())).style(ratatui::style::Color::LightBlue),
+            Span::from(" | ").style(ratatui::style::Color::DarkGray),
+            Span::from(format!("{} tokens sacrificed", self.scene.tokens_consumed)).style(ratatui::style::Color::LightCyan),
         ]);
         frame.render_widget(status_line, area);
     }
@@ -411,15 +396,7 @@ impl App {
         let command = parts.next().unwrap();
         let arg = parts.next().unwrap_or("");
         match command {
-            "/bandcamp" => {
+            // FIXME: Need some new kind of /bandcamp command to force loading of specific urls
-                if let Err(err) = self.add_bandcamp_artifact(arg).await {
-                    self.sys_message_sink.send(format!("Failed to fetch artifact: {:?}", err)).await.unwrap();
-                } else {
-                    self.sys_message_sink.send(format!("Added Bandcamp artifact from {}", arg)).await.unwrap();
-                    self.next_actions.push(ConversationEntry::ShipComputer(format!("Incoming transmission from {}", arg)));
-                    self.regenerate_responses();
-                }
-            },
             "/episode" => {
                 if let Ok(episode_number) = arg.trim().parse::<u32>() {
                     self.direction.episode_number = episode_number;
@@ -438,21 +415,10 @@ impl App {
                     self.sys_message_sink.send("Invalid timer format. Use /timer [minutes]".into()).await.unwrap();
                 }
             },
-            "/clear" => {
-                match arg.trim() {
-                    "artifacts" => {
-                        self.direction.artifacts.clear();
-                        self.sys_message_sink.send("Cleared artifacts.".into()).await.unwrap();
-                    },
-                    _ => {
-                        self.sys_message_sink.send("Unknown clear command. Use /clear [artifacts]".into()).await.unwrap();
-                    }
-                }
-                return;
-            },
             "/narrative" => {
                 self.direction.narrative = arg.to_string();
                 self.sys_message_sink.send(format!("Updated stage direction: {}", self.direction.narrative)).await.unwrap();
+                self.regenerate_responses();
             },
             "/event" => {
                 self.next_actions.push(ConversationEntry::StageDirection(arg.to_string()));
@@ -463,7 +429,7 @@ impl App {
                 self.regenerate_responses();
             },
             _ => {
-                self.sys_message_sink.send("Unknown command. Available commands: /bandcamp [url], /episode [number], /narrative [text], /reset".into()).await.unwrap();
+                self.sys_message_sink.send("Unknown command. Available commands: /episode [number], /narrative [text], /event [text], /computer [text], /timer [minutes]".into()).await.unwrap();
             }
         }
     }
@@ -476,6 +442,7 @@ impl App {
                         KeyCode::Tab => {
                             self.focus_state = FocusState::UserInput;
                             self.conversation_state.select(None);
+                            self.reply_state.select_first();
                         },
                         KeyCode::PageUp => self.conversation_state.scroll_down_by(5),
                         KeyCode::PageDown => self.conversation_state.scroll_up_by(5),
@@ -489,6 +456,7 @@ impl App {
                                 self.tts.speak(text.clone()).await;
                                 self.focus_state = FocusState::UserInput;
                                 self.conversation_state.select(None);
+                                self.reply_state.select_first();
                             }
                         },
                         _ => ()
@@ -499,6 +467,7 @@ impl App {
                         KeyCode::Tab => {
                             self.focus_state = FocusState::Conversation;
                             self.conversation_state.select_first();
+                            self.reply_state.select(None);
                         },
                         KeyCode::Char('r') if key.modifiers.contains(KeyModifiers::CONTROL) => self.regenerate_responses(),
                         KeyCode::Char('x') if key.modifiers.contains(KeyModifiers::CONTROL) => {
@@ -536,16 +505,6 @@ impl App {
         }
     }
 
-    async fn add_bandcamp_artifact(&mut self, url: &str) -> Result<(), BandcampError> {
-        let album = bandcamp::album_from_url(url).await?;
-        let result: BandcampResult = album.into();
-        let json = serde_json::to_string(&result)?;
-        self.direction.artifacts.push(json);
-        self.sys_message_sink.send("Added bandcamp album".into()).await.unwrap();
-
-        Ok(())
-    }
-
     fn regenerate_responses(&mut self) {
         let actions = StageActions {
             direction: self.direction.clone(),
@@ -570,7 +529,8 @@ impl App {
 #[derive(Serialize, Deserialize, Debug, Default)]
 pub struct SaveData {
     pub direction: StageDirection,
-    pub messages: Vec<ChatCompletionRequestMessage>
+    pub messages: Vec<ChatCompletionRequestMessage>,
+    pub scenery: Scenery
 }
 
 impl SaveData {
@@ -619,7 +579,7 @@ async fn main() {
 
     let (audio_ctrl, mic_stream, tts_output) = start_audio_input(&sys_message_sink).await;
     let tts_ctrl = start_tts(tts_output).await;
-    let (prediction_request_in, mut prediction_out) = prediction::start_prediction(sys_message_src, saved_session.messages).await;
+    let (prediction_request_in, mut prediction_out) = prediction::start_prediction(sys_message_src, saved_session.messages, saved_session.scenery).await;
     let transcription_ctrl = transcription::start_transcription(mic_stream).await;
 
     let mut app = App::new(prediction_request_in, audio_ctrl, transcription_ctrl, tts_ctrl, sys_message_sink, saved_session.direction);

src/prediction.rs

@@ -3,10 +3,12 @@ use std::process::{Command, Stdio};
 use async_openai::{Client, config::OpenAIConfig, types::chat::{ChatCompletionMessageToolCalls, ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestToolMessageArgs, ChatCompletionTool, ChatCompletionTools, CreateChatCompletionRequestArgs, FinishReason, FunctionObjectArgs, ResponseFormat, ResponseFormatJsonSchema}};
 use bandcamp::SearchResultItem;
 use chrono::{DateTime, Utc};
+use color_eyre::eyre::eyre;
 use schemars::{JsonSchema, schema_for};
 use serde::{Deserialize, Serialize};
+use serde_json::{Serializer, ser::CompactFormatter};
 
-use crate::{SaveData, scene::{ConversationEntry, Scene, StageActions, StageDirection}};
+use crate::{SaveData, scene::{Artifact, ConversationEntry, Scene, Scenery, StageActions, StageDirection}};
 
 
 const SYSTEM_PROMPT: &str = include_str!("system-prompt.txt");
@@ -28,12 +30,20 @@ struct Session {
     conversation: Vec<ConversationEntry>,
     header_message: ChatCompletionRequestMessage,
     messages: Vec<ChatCompletionRequestMessage>,
-    reply_options: GeneratedResponses
+    reply_options: GeneratedResponses,
+    scenery: Scenery,
+    tokens_consumed: usize
 }
 
-#[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)]
+#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)]
+enum StageEvent {
+    ShipComputer(String),
+    StageDirection(String)
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)]
 struct StageEventArgs {
-    text: String,
+    event: StageEvent
 }
 
 #[derive(Debug, Default, Serialize, Deserialize, Clone, JsonSchema)]
@@ -81,7 +91,7 @@ struct ToolResults {
 }
 
 impl Session {
-    fn from_initial_messages(messages: Vec<ChatCompletionRequestMessage>) -> Self {
+    fn from_initial_messages(messages: Vec<ChatCompletionRequestMessage>, scenery: Scenery) -> Self {
         let mut conversation = vec![];
         for msg in &messages {
             if let Ok(conversation_msg) = msg.clone().try_into() {
@@ -94,7 +104,9 @@ impl Session {
             conversation,
             header_message: ChatCompletionRequestSystemMessageArgs::default().content(SYSTEM_PROMPT).build().unwrap().into(),
             messages,
-            reply_options: Default::default()
+            reply_options: Default::default(),
+            scenery,
+            tokens_consumed: 0
         }
     }
 
@@ -105,15 +117,12 @@ impl Session {
     }
 
     async fn tool_stage_event(&mut self, args: StageEventArgs) -> ToolResults {
+        let msg = match args.event {
+            StageEvent::ShipComputer(text) => ConversationEntry::ShipComputer(text),
+            StageEvent::StageDirection(text) => ConversationEntry::StageDirection(text)
+        };
         ToolResults {
-            messages: vec![ConversationEntry::StageDirection(args.text)],
+            messages: vec![msg],
-            ..Default::default()
-        }
-    }
-
-    async fn tool_computer_event(&mut self, args: StageEventArgs) -> ToolResults {
-        ToolResults {
-            messages: vec![ConversationEntry::ShipComputer(args.text)],
             ..Default::default()
         }
     }
@@ -127,20 +136,23 @@ impl Session {
                 match result {
                     SearchResultItem::Artist(data) => {
                         let result: BandcampResult = bandcamp::fetch_artist(data.artist_id).await.unwrap().into();
-                        json_results.push(result);
+                        json_results.push(Artifact::Bandcamp(result));
                     },
                     SearchResultItem::Album(data) => {
                         let result: BandcampResult = bandcamp::fetch_album(data.band_id, data.album_id).await.unwrap().into();
-                        json_results.push(result);
+                        json_results.push(Artifact::Bandcamp(result));
                     }
                     _ => ()
                 }
             }
         }
-        messages.push(ConversationEntry::ShipComputer(format!("Artifact scan for '{}' complete. {} results.", args.query, json_results.len()).into()));
+        let artifact_count = json_results.len();
+        messages.push(ConversationEntry::ShipComputer(format!("Relay scan for '{}' complete. {} artifacts added to the archive.", args.query, artifact_count).into()));
+
+        self.scenery.artifacts.append(&mut json_results);
 
         ToolResults {
-            result: Some(serde_json::to_string(&json_results).unwrap()),
+            result: Some(format!("{} artifacts were added to the archive.", artifact_count)),
             messages
         }
     }
@@ -164,29 +176,47 @@ impl Session {
         if let Some(year) = args.year {
             beets_cmd.arg(format!("year:{}", year));
         }
-        let result = if let Ok(output) = beets_cmd.stdout(Stdio::piped()).spawn().unwrap().wait_with_output() {
+        if let Ok(output) = beets_cmd.stdout(Stdio::piped()).spawn().unwrap().wait_with_output() {
             messages.push(ConversationEntry::ShipComputer(format!("Executing archive query {:?}", beets_cmd)));
-            Some(minify::json::minify(str::from_utf8(&output.stdout).unwrap()))
+            self.scenery.artifacts.push(Artifact::BeetsTrack(serde_json::from_str(str::from_utf8(&output.stdout).unwrap()).unwrap()));
         } else {
             messages.push(ConversationEntry::ShipComputer("Unable to execute query!".into()));
-            None
         };
 
         ToolResults {
-            result,
+            result: None,
             messages
         }
     }
 
-    async fn regenerate_options(&mut self, direction: &StageDirection) {
+    fn generate_conversation(&self, direction: &StageDirection) -> Vec<ChatCompletionRequestMessage> {
-        loop {
+        let mut json_buf = vec![];
-            let direction_message: ChatCompletionRequestMessage = ChatCompletionRequestSystemMessageArgs::default().content(serde_json::to_string(&direction).unwrap()).build().unwrap().into();
+        let mut ser = Serializer::with_formatter(&mut json_buf, CompactFormatter);
+        direction.serialize(&mut ser).unwrap();
+        let direction_message: ChatCompletionRequestMessage = ChatCompletionRequestSystemMessageArgs::default()
+            .content(String::from_utf8(json_buf).unwrap())
+            .build().unwrap().into();
+
+        let mut json_buf = vec![];
+        let mut ser = Serializer::with_formatter(&mut json_buf, CompactFormatter);
+        self.scenery.serialize(&mut ser).unwrap();
+        let scenery_message: ChatCompletionRequestMessage = ChatCompletionRequestSystemMessageArgs::default()
+            .content(String::from_utf8(json_buf).unwrap())
+            .build().unwrap().into();
         let mut full_conversation = vec![
             self.header_message.clone(),
-                direction_message
+            direction_message,
+            scenery_message,
         ];
         full_conversation.append(&mut self.messages.clone());
 
+        full_conversation
+    }
+
+    async fn regenerate_options(&mut self, direction: &StageDirection) {
+        loop {
+            let full_conversation = self.generate_conversation(direction);
+
             let tools = vec![
                 ChatCompletionTools::Function(ChatCompletionTool {
                     function: FunctionObjectArgs::default()
@@ -195,13 +225,6 @@ impl Session {
                         .parameters(schema_for!(StageEventArgs))
                         .build().unwrap()
                 }),
-                ChatCompletionTools::Function(ChatCompletionTool {
-                    function: FunctionObjectArgs::default()
-                        .name("log_ship_computer_message")
-                        .description("Inserts a message from the ship computer into the scene script")
-                        .parameters(schema_for!(StageEventArgs))
-                        .build().unwrap()
-                }),
                 ChatCompletionTools::Function(ChatCompletionTool {
                     function: FunctionObjectArgs::default()
                         .name("archive_query")
@@ -212,13 +235,13 @@ impl Session {
                 ChatCompletionTools::Function(ChatCompletionTool {
                     function: FunctionObjectArgs::default()
                         .name("bandcamp_artifact_scan")
-                        .description("Scans Bandcamp to find artifacts to use in the scene that match the given search parameters")
+                        .description("Scans Bandcamp to find artifacts to use in the scene that match the given search parameters. To find an artist, provide only the artist name. To find an album, provide the artist and the album.")
                         .parameters(schema_for!(BandcampQueryArgs))
                         .build().unwrap()
                 })
             ];
             let request = CreateChatCompletionRequestArgs::default()
-                .messages(full_conversation.clone())
+                .messages(full_conversation)
                 .model("gpt-5.4-mini")
                 .tools(tools)
                 .max_completion_tokens(1024u32)
@@ -233,9 +256,13 @@ impl Session {
                 .build().unwrap();
 
             let response = self.client.chat().create(request).await.unwrap_or_else(|err| {
-                panic!("{} {:?}", err, full_conversation);
+                panic!("OpenAI Panic: {}", err);
             });
 
+            if let Some(usage) = response.usage {
+                self.tokens_consumed += usage.total_tokens as usize;
+            }
+
             if let Some(message) = response.choices.first() {
 
                 match message.finish_reason {
@@ -263,7 +290,6 @@ impl Session {
                                 let args = call.function.arguments.as_str();
                                 let tool_result = match func_name {
                                     "log_stage_event" => self.tool_stage_event(serde_json::from_str(args).unwrap()).await,
-                                    "log_ship_computer_message" => self.tool_computer_event(serde_json::from_str(args).unwrap()).await,
                                     "bandcamp_artifact_scan" => self.tool_bandcamp_scan(serde_json::from_str(args).unwrap()).await,
                                     "archive_query" => self.tool_artifact_query(serde_json::from_str(args).unwrap()).await,
                                     _ => unreachable!()
@@ -303,7 +329,7 @@ impl Session {
     }
 
     fn as_scene(&self) -> Scene {
-        Scene::new(self.reply_options.clone(), self.conversation.clone())
+        Scene::new(self.reply_options.clone(), self.conversation.clone(), self.scenery.clone(), self.tokens_consumed)
     }
 
     fn insert_conversation(&mut self, entry: ConversationEntry) {
@@ -315,11 +341,11 @@ impl Session {
     }
 }
 
-pub async fn start_prediction(mut sys_message_src: tokio::sync::mpsc::Receiver<String>, initial_messages: Vec<ChatCompletionRequestMessage>) -> (tokio::sync::watch::Sender<StageActions>, tokio::sync::watch::Receiver<Scene>) {
+pub async fn start_prediction(mut sys_message_src: tokio::sync::mpsc::Receiver<String>, initial_messages: Vec<ChatCompletionRequestMessage>, scenery: Scenery) -> (tokio::sync::watch::Sender<StageActions>, tokio::sync::watch::Receiver<Scene>) {
     let (prediction_in, prediction_out) = tokio::sync::watch::channel(Scene::default());
     let (prediction_request_in, mut prediction_request_out) = tokio::sync::watch::channel(StageActions::default());
 
-    let mut session = Session::from_initial_messages(initial_messages);
+    let mut session = Session::from_initial_messages(initial_messages, scenery);
 
     // Send the initial scene to the UI, after we have loaded the session from the first messages.
     prediction_in.send(session.as_scene()).unwrap();
@@ -340,7 +366,8 @@ pub async fn start_prediction(mut sys_message_src: tokio::sync::mpsc::Receiver<S
 
                         let mut save_data = SaveData {
                             direction: next_cxt.direction,
-                            messages: session.messages.clone()
+                            messages: session.messages.clone(),
+                            scenery: session.scenery.clone()
                         };
 
                         save_data.save();

src/scene.rs

@@ -3,7 +3,7 @@ use chrono::Duration;
 use serde::{Deserialize, Serialize};
 use sqlite::OpenFlags;
 
-use crate::prediction::{GeneratedResponses, PossibleResponse};
+use crate::prediction::{BandcampResult, GeneratedResponses, PossibleResponse};
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum ConversationEntry {
@@ -21,7 +21,7 @@ impl TryInto<ChatCompletionRequestMessage> for ConversationEntry {
             ConversationEntry::Eva(text) => Ok(ChatCompletionRequestMessage::Assistant(ChatCompletionRequestAssistantMessage { content: Some(text.into()), ..Default::default()})),
             ConversationEntry::ShipComputer(text) => Ok(ChatCompletionRequestMessage::System(ChatCompletionRequestSystemMessage { content: text.into(), name: Some("ship-computer".into()), ..Default::default() })),
             ConversationEntry::StageDirection(text) => Ok(ChatCompletionRequestMessage::System(ChatCompletionRequestSystemMessage { content: text.into(), name: Some("stage-direction".into()), ..Default::default() })),
-            ConversationEntry::SystemMessage(_) => Err(())
+            _ => Err(())
         }
     }
     
@@ -53,10 +53,20 @@ pub struct StageDirection {
     pub episode_number: u32,
     pub time_remaining: Duration,
     pub narrative: String,
-    pub artifacts: Vec<String>,
     pub current_playlist: Vec<PlaylistEntry>
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum Artifact {
+    Bandcamp(BandcampResult),
+    BeetsTrack(serde_json::Value)
+}
+
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
+pub struct Scenery {
+    pub artifacts: Vec<Artifact>
+}
+
 #[derive(Debug)]
 pub enum MixxxError {
     Sql(sqlite::Error)
@@ -114,16 +124,24 @@ pub struct PlaylistEntry {
 pub struct Scene {
     reply_options: GeneratedResponses,
     conversation: Vec<ConversationEntry>,
+    pub tokens_consumed: usize,
+    scenery: Scenery
 }
 
 impl Scene {
-    pub fn new(reply_options: GeneratedResponses, conversation: Vec<ConversationEntry>) -> Self {
+    pub fn new(reply_options: GeneratedResponses, conversation: Vec<ConversationEntry>, scenery: Scenery, tokens_consumed: usize) -> Self {
         Self {
             reply_options,
             conversation,
+            scenery,
+            tokens_consumed
         }
     }
 
+    pub fn scenery(&self) -> &Scenery {
+        &self.scenery
+    }
+    
     pub fn conversation(&self) -> &Vec<ConversationEntry> {
         &self.conversation
     }