add session boilerplate

This commit is contained in:
b1ek 2024-10-22 12:56:35 +10:00
parent fa9a151d85
commit e083420c5c
Signed by: blek
GPG Key ID: A622C22C9BC616B2
3 changed files with 119 additions and 4 deletions

View File

@ -9,19 +9,19 @@ use crate::{cache::Cache, config::Config};
use crate::{WARN, RED, RESET};
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatMessagePayload {
pub struct ChatMessagePayload {
pub role: String,
pub content: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatPayload {
pub struct ChatPayload {
pub model: String,
pub messages: Vec<ChatMessagePayload>
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatChunk {
pub struct ChatChunk {
pub role: Option<String>,
pub message: String,
pub created: u64,
@ -31,7 +31,7 @@ struct ChatChunk {
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ErrChatChunk {
pub struct ErrChatChunk {
pub action: String,
pub status: u64,
#[serde(rename = "type")]

View File

@ -11,6 +11,7 @@ use crate::api::{get_res, get_vqd, simulate_browser_reqs};
mod cache;
mod config;
mod api;
mod session;
pub const GREEN: &str = "\x1b[1;32m";
pub const RED: &str = "\x1b[1;31m";

114
src/session.rs Normal file
View File

@ -0,0 +1,114 @@
use std::fs;
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::time::Duration;
use chrono::{DateTime, Local};
use serde::{Deserialize, Serialize};
use crate::api::ChatMessagePayload;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
session_id: String,
ttl: DateTime<Local>,
messages: Vec<ChatMessagePayload>
}
impl Session {
pub fn terminal_session_id<T: From<String>>() -> T {
#[cfg(unix)]
fn inner() -> Result<String, Box<dyn std::error::Error>> {
let child = Command::new("tty")
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()?;
let child = child.wait_with_output()?;
if ! child.status.success() {
Err("")? // value ignored
}
Ok(String::from_utf8(child.stdout)?)
}
#[cfg(not(unix))]
fn inner() -> Result<String, ()> {
Err(())
}
inner()
.unwrap_or("default".into())
.replace(['/', '\\', ':'], "-")
.trim()
.trim_start_matches('-')
.to_string()
.into()
}
fn path_for_id<T: Into<String>>(id: T) -> PathBuf {
let id: String = id.into();
std::env::temp_dir().join(format!("hey-duck-session-{id}"))
}
fn path(&self) -> PathBuf {
Self::path_for_id(&self.session_id)
}
fn restore_with_id<T: Into<String>>(id: T) -> Option<Self> {
fn inner(path: PathBuf) -> Result<Session, Box<dyn std::error::Error>> {
if path.is_file() {
let data = fs::read_to_string(path)?;
Ok(serde_json::from_str(&data)?)
} else {
Err("")? // value ignored
}
}
match inner(Self::path_for_id(id)) {
Ok(session) => {
if session.is_expired() {
session.destroy().expect("Couldn't destroy expired session");
None
} else {
Some(session)
}
},
Err(_) => None
}
}
fn new<T: Into<String>>(id: T) -> Self {
Self {
session_id: id.into(),
ttl: Local::now() + Duration::from_secs(60 * 5),
messages: vec![]
}
}
fn save(&self) -> Result<(), Box<dyn std::error::Error>> {
println!("{:?}", self.path());
fs::write(self.path(), serde_json::to_string_pretty(self)?)?;
Ok(())
}
fn is_expired(&self) -> bool {
self.ttl < Local::now()
}
pub fn destroy(&self) -> Result<(), Box<dyn std::error::Error>> {
Ok(fs::remove_file(self.path())?)
}
pub fn create_or_restore() -> Self {
let session_id: String = Self::terminal_session_id();
match Self::restore_with_id(&session_id) {
Some(session) => session,
None => {
let session = Self::new(&session_id);
session.save().expect("Couldn't save new session");
session
}
}
}
}