add persistence to the workon, set timeout

py/packages/corpora/tasks/sync.py

@@ -15,8 +15,11 @@ def process_tarball(corpus_id: str, tarball: bytes) -> None:
         for member in tar.getmembers():
             if member.isfile():
                 file_content = (
-                    tar.extractfile(member).read().decode("utf-8", errors="replace")
+                    tar.extractfile(member)
+                    .read()
+                    .decode("utf-8", errors="replace")
                 )
+
                 checksum = compute_checksum(file_content)
 
                 corpus_file, _ = CorpusTextFile.objects.get_or_create(

rs/core/corpora_cli/src/commands/chat.rs

@@ -6,12 +6,6 @@ use termimad::MadSkin;
 
 #[derive(Args)]
 pub struct ChatArgs {
-    // Both of these should be optional
-    // #[arg(help = "name of persistent session")]
-    // pub persist: String,
-    // #[arg(short, long, help = "list available sessions")]
-    // pub list: bool,
-    // this is how you make them optional:
     #[arg(short, long, help = "name of persistent session")]
     pub persist: Option<String>,
     #[arg(short, long, help = "list available sessions")]

rs/core/corpora_cli/src/commands/workon.rs

@@ -10,6 +10,9 @@ use std::path::Path;
 pub struct WorkonArgs {
     #[arg(help = "Path to the file or directory")]
     pub path: String,
+
+    #[arg(short, long, help = "name of persistent session")]
+    pub persist: Option<String>,
 }
 
 /// The `workon` command operation
@@ -38,17 +41,45 @@ pub fn run(ctx: &Context, args: WorkonArgs) {
     ctx.success("Current file content:");
     ctx.dim(&current_file_content);
 
-    let mut messages: Vec<MessageSchema> = vec![];
-    if !current_file_content.is_empty() {
-        messages.push(MessageSchema {
-            role: "user".to_string(),
-            text: format!(
-                "The original content of `{}` was:\n```{}\n```",
-                relative_path.display(),
-                current_file_content
-            ),
-        });
-    }
+    let mut messages: Vec<MessageSchema> = if let Some(session_name) = &args.persist {
+        match ctx.history.load_session(session_name) {
+            Ok(mut existing_messages) => {
+                ctx.success(&format!("Loaded session: {}", session_name));
+                for message in &existing_messages {
+                    ctx.dim(&format!("{}: {}", message.role, message.text));
+                }
+                // add the original content of the file to the chat history
+                if !current_file_content.is_empty() {
+                    existing_messages.push(MessageSchema {
+                        role: "user".to_string(),
+                        text: format!(
+                            "The current content of `{}` is:\n```\n{}\n```",
+                            relative_path.display(),
+                            current_file_content
+                        ),
+                    });
+                }
+                existing_messages
+            }
+            Err(_) => {
+                ctx.warn(&format!("No existing session found for: {}", session_name));
+                Vec::new()
+            }
+        }
+    } else {
+        if !current_file_content.is_empty() {
+            vec![MessageSchema {
+                role: "user".to_string(),
+                text: format!(
+                    "The original content of `{}` was:\n```\n{}\n```",
+                    relative_path.display(),
+                    current_file_content
+                ),
+            }]
+        } else {
+            Vec::new()
+        }
+    };
 
     loop {
         let user_input = ctx
@@ -104,6 +135,15 @@ pub fn run(ctx: &Context, args: WorkonArgs) {
             text: revision.clone(),
         });
 
+        if let Some(session_name) = &args.persist {
+            if let Err(err) = ctx.history.save_session(session_name, &messages) {
+                ctx.error(&format!(
+                    "Failed to save session '{}': {:?}",
+                    session_name, err
+                ));
+            }
+        }
+
         if Confirm::with_theme(&ColorfulTheme::default())
             .with_prompt("Write file?")
             .interact()

rs/core/corpora_cli/src/context/mod.rs

@@ -2,13 +2,15 @@ pub mod auth;
 pub mod collector;
 pub mod config;
 
+use std::process::Command;
 use std::sync::Arc;
+use std::time::Duration;
 
 use console::{Style, Term};
-use std::process::Command;
-use tempfile::NamedTempFile;
 
 use indicatif::{ProgressBar, ProgressStyle};
+use reqwest::blocking::ClientBuilder;
+use tempfile::NamedTempFile;
 
 use crate::history::files::FileChatHistory;
 use crate::history::ChatHistory;
@@ -36,9 +38,15 @@ impl Context {
         let token = get_bearer_token(&corpora_config)
             .expect("Failed to authenticate and retrieve bearer token");
 
+        let client = ClientBuilder::new()
+            .timeout(Duration::from_secs(60))
+            .build()
+            .expect("Failed to build the reqwest client");
+
         // Configure the API client
         let api_config = Configuration {
             base_path: corpora_config.server.base_url.clone(),
+            client,
             bearer_access_token: Some(token),
             ..Default::default()
         };