use serde_json::Value;
use ulid::Ulid;
use crate::error::Result;
use crate::memory::{Memory, MemoryType};
use crate::store::{Mnemo, RecallRequest, RecallResult};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Role {
User,
Assistant,
System,
}
impl Role {
pub fn as_str(&self) -> &'static str {
match self {
Role::User => "user",
Role::Assistant => "assistant",
Role::System => "system",
}
}
}
#[derive(Clone, Debug)]
pub struct Turn {
pub role: Role,
pub content: String,
pub vector: Vec<f32>,
}
impl Turn {
pub fn new(role: Role, content: impl Into<String>, vector: Vec<f32>) -> Self {
Self { role, content: content.into(), vector }
}
pub fn user(content: impl Into<String>, vector: Vec<f32>) -> Self {
Self::new(Role::User, content, vector)
}
pub fn assistant(content: impl Into<String>, vector: Vec<f32>) -> Self {
Self::new(Role::Assistant, content, vector)
}
pub fn system(content: impl Into<String>, vector: Vec<f32>) -> Self {
Self::new(Role::System, content, vector)
}
}
pub struct Session<'db> {
db: &'db mut Mnemo,
agent_id: String,
session_id: String,
turns: Vec<Ulid>,
}
impl<'db> Session<'db> {
pub(crate) fn new(db: &'db mut Mnemo, agent_id: String) -> Self {
Self {
db,
agent_id,
session_id: Ulid::new().to_string(),
turns: Vec::new(),
}
}
pub fn id(&self) -> &str {
&self.session_id
}
pub fn agent(&self) -> &str {
&self.agent_id
}
pub fn turn_ids(&self) -> &[Ulid] {
&self.turns
}
pub fn turn_count(&self) -> usize {
self.turns.len()
}
pub fn add_turn(&mut self, turn: Turn) -> Result<Ulid> {
let memory = Memory::new(turn.content, MemoryType::Working, turn.vector)
.with_agent(&self.agent_id)
.with_session(&self.session_id)
.with_meta("role", Value::String(turn.role.as_str().to_string()));
let id = self.db.remember(memory)?;
self.turns.push(id);
Ok(id)
}
pub fn recall(&mut self, mut request: RecallRequest) -> Result<Vec<RecallResult>> {
request.agent_id = Some(self.agent_id.clone());
self.db.recall(&request)
}
pub fn close(self) -> Result<usize> {
let Session { db, turns, .. } = self;
let mut promoted = 0;
for id in &turns {
let mut memory = match db.get(id) {
Ok(m) => m,
Err(_) => continue,
};
if memory.memory_type == MemoryType::Working {
memory.memory_type = MemoryType::Episodic;
db.remember(memory)?;
promoted += 1;
}
}
db.flush()?;
Ok(promoted)
}
pub fn discard(self) -> Result<usize> {
let Session { db, turns, .. } = self;
let mut removed = 0;
for id in &turns {
if db.delete(id).is_ok() {
removed += 1;
}
}
db.flush()?;
Ok(removed)
}
}