use crate::core::messages::{Message, Messages};
use crate::core::plugin::{AgentPlugin, PluginContext};
use async_trait::async_trait;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Default)]
pub struct History(pub Vec<Message>);
pub struct FileHistoryPlugin {
inner: Arc<RwLock<FileHistoryInner>>,
}
impl std::fmt::Debug for FileHistoryPlugin {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FileHistoryPlugin").finish_non_exhaustive()
}
}
#[derive(Debug)]
struct FileHistoryInner {
path: PathBuf,
messages: Messages,
}
impl FileHistoryPlugin {
pub fn new(path: impl AsRef<Path>) -> crate::error::Result<Self> {
let path = path.as_ref().to_path_buf();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
Ok(Self {
inner: Arc::new(RwLock::new(FileHistoryInner {
path,
messages: Vec::new(),
})),
})
}
pub async fn push(&self, msg: Message) -> crate::error::Result<()> {
let mut inner = self.inner.write().await;
inner.messages.push(msg);
let data = serde_json::to_string_pretty(&inner.messages)?;
fs::write(&inner.path, data).await?;
Ok(())
}
pub async fn load(&self) -> crate::error::Result<Messages> {
let inner = self.inner.read().await;
Ok(inner.messages.clone())
}
}
impl Clone for FileHistoryPlugin {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
#[async_trait]
impl AgentPlugin for FileHistoryPlugin {
fn name(&self) -> &'static str {
"file_history"
}
async fn init(&mut self, ctx: &mut PluginContext) {
let inner = self.inner.read().await;
ctx.insert(History(inner.messages.clone()));
}
async fn shutdown(&mut self, ctx: &mut PluginContext) {
if let Some(h) = ctx.get::<History>() {
let mut inner = self.inner.write().await;
inner.messages.clone_from(&h.0);
if let Ok(data) = serde_json::to_string_pretty(&inner.messages) {
let path = inner.path.clone();
let _ = fs::write(&path, data).await;
}
}
}
}
#[derive(Clone, Default)]
pub struct MemoryHistoryPlugin {
inner: Arc<RwLock<Messages>>,
}
impl std::fmt::Debug for MemoryHistoryPlugin {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MemoryHistoryPlugin")
.finish_non_exhaustive()
}
}
impl MemoryHistoryPlugin {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub async fn push(&self, msg: Message) {
self.inner.write().await.push(msg);
}
pub async fn messages(&self) -> Messages {
self.inner.read().await.clone()
}
}
#[async_trait]
impl AgentPlugin for MemoryHistoryPlugin {
fn name(&self) -> &'static str {
"memory_history"
}
async fn init(&mut self, ctx: &mut PluginContext) {
let msgs = self.inner.read().await.clone();
ctx.insert(History(msgs));
}
async fn shutdown(&mut self, ctx: &mut PluginContext) {
if let Some(h) = ctx.get::<History>() {
*self.inner.write().await = h.0.clone();
}
}
}