mod chunker;
mod document;
mod embeddings;
#[cfg(feature = "postgres")]
mod repository;
mod search;
pub use chunker::{ChunkConfig, chunk_document};
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
#[cfg(feature = "postgres")]
pub use repository::Repository;
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
use std::sync::Arc;
use chrono::{NaiveDate, Utc};
#[cfg(feature = "postgres")]
use deadpool_postgres::Pool;
use uuid::Uuid;
use crate::error::WorkspaceError;
enum WorkspaceStorage {
#[cfg(feature = "postgres")]
Repo(Repository),
Db(Arc<dyn crate::db::Database>),
}
impl WorkspaceStorage {
async fn get_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.get_document_by_path(user_id, agent_id, path).await,
Self::Db(db) => db.get_document_by_path(user_id, agent_id, path).await,
}
}
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.get_document_by_id(id).await,
Self::Db(db) => db.get_document_by_id(id).await,
}
}
async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => {
repo.get_or_create_document_by_path(user_id, agent_id, path)
.await
}
Self::Db(db) => {
db.get_or_create_document_by_path(user_id, agent_id, path)
.await
}
}
}
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.update_document(id, content).await,
Self::Db(db) => db.update_document(id, content).await,
}
}
async fn delete_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<(), WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.delete_document_by_path(user_id, agent_id, path).await,
Self::Db(db) => db.delete_document_by_path(user_id, agent_id, path).await,
}
}
async fn list_directory(
&self,
user_id: &str,
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.list_directory(user_id, agent_id, directory).await,
Self::Db(db) => db.list_directory(user_id, agent_id, directory).await,
}
}
async fn list_all_paths(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.list_all_paths(user_id, agent_id).await,
Self::Db(db) => db.list_all_paths(user_id, agent_id).await,
}
}
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.delete_chunks(document_id).await,
Self::Db(db) => db.delete_chunks(document_id).await,
}
}
async fn insert_chunk(
&self,
document_id: Uuid,
chunk_index: i32,
content: &str,
embedding: Option<&[f32]>,
) -> Result<Uuid, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => {
repo.insert_chunk(document_id, chunk_index, content, embedding)
.await
}
Self::Db(db) => {
db.insert_chunk(document_id, chunk_index, content, embedding)
.await
}
}
}
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => repo.update_chunk_embedding(chunk_id, embedding).await,
Self::Db(db) => db.update_chunk_embedding(chunk_id, embedding).await,
}
}
async fn get_chunks_without_embeddings(
&self,
user_id: &str,
agent_id: Option<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => {
repo.get_chunks_without_embeddings(user_id, agent_id, limit)
.await
}
Self::Db(db) => {
db.get_chunks_without_embeddings(user_id, agent_id, limit)
.await
}
}
}
async fn hybrid_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
query: &str,
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
match self {
#[cfg(feature = "postgres")]
Self::Repo(repo) => {
repo.hybrid_search(user_id, agent_id, query, embedding, config)
.await
}
Self::Db(db) => {
db.hybrid_search(user_id, agent_id, query, embedding, config)
.await
}
}
}
}
const HEARTBEAT_SEED: &str = "\
# Heartbeat Checklist
<!-- Keep this file empty to skip heartbeat API calls.
Add tasks below when you want the agent to check something periodically.
Example:
- [ ] Check for unread emails needing a reply
- [ ] Review today's calendar for upcoming meetings
- [ ] Check CI build status for main branch
-->";
pub struct Workspace {
user_id: String,
agent_id: Option<Uuid>,
storage: WorkspaceStorage,
embeddings: Option<Arc<dyn EmbeddingProvider>>,
}
impl Workspace {
#[cfg(feature = "postgres")]
pub fn new(user_id: impl Into<String>, pool: Pool) -> Self {
Self {
user_id: user_id.into(),
agent_id: None,
storage: WorkspaceStorage::Repo(Repository::new(pool)),
embeddings: None,
}
}
pub fn new_with_db(user_id: impl Into<String>, db: Arc<dyn crate::db::Database>) -> Self {
Self {
user_id: user_id.into(),
agent_id: None,
storage: WorkspaceStorage::Db(db),
embeddings: None,
}
}
pub fn with_agent(mut self, agent_id: Uuid) -> Self {
self.agent_id = Some(agent_id);
self
}
pub fn with_embeddings(mut self, provider: Arc<dyn EmbeddingProvider>) -> Self {
self.embeddings = Some(provider);
self
}
pub fn user_id(&self) -> &str {
&self.user_id
}
pub fn agent_id(&self) -> Option<Uuid> {
self.agent_id
}
pub async fn read(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
let path = normalize_path(path);
self.storage
.get_document_by_path(&self.user_id, self.agent_id, &path)
.await
}
pub async fn write(&self, path: &str, content: &str) -> Result<MemoryDocument, WorkspaceError> {
let path = normalize_path(path);
let doc = self
.storage
.get_or_create_document_by_path(&self.user_id, self.agent_id, &path)
.await?;
self.storage.update_document(doc.id, content).await?;
self.reindex_document(doc.id).await?;
self.storage.get_document_by_id(doc.id).await
}
pub async fn append(&self, path: &str, content: &str) -> Result<(), WorkspaceError> {
let path = normalize_path(path);
let doc = self
.storage
.get_or_create_document_by_path(&self.user_id, self.agent_id, &path)
.await?;
let new_content = if doc.content.is_empty() {
content.to_string()
} else {
format!("{}\n{}", doc.content, content)
};
self.storage.update_document(doc.id, &new_content).await?;
self.reindex_document(doc.id).await?;
Ok(())
}
pub async fn exists(&self, path: &str) -> Result<bool, WorkspaceError> {
let path = normalize_path(path);
match self
.storage
.get_document_by_path(&self.user_id, self.agent_id, &path)
.await
{
Ok(_) => Ok(true),
Err(WorkspaceError::DocumentNotFound { .. }) => Ok(false),
Err(e) => Err(e),
}
}
pub async fn delete(&self, path: &str) -> Result<(), WorkspaceError> {
let path = normalize_path(path);
self.storage
.delete_document_by_path(&self.user_id, self.agent_id, &path)
.await
}
pub async fn list(&self, directory: &str) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
let directory = normalize_directory(directory);
self.storage
.list_directory(&self.user_id, self.agent_id, &directory)
.await
}
pub async fn list_all(&self) -> Result<Vec<String>, WorkspaceError> {
self.storage
.list_all_paths(&self.user_id, self.agent_id)
.await
}
pub async fn memory(&self) -> Result<MemoryDocument, WorkspaceError> {
self.read_or_create(paths::MEMORY).await
}
pub async fn today_log(&self) -> Result<MemoryDocument, WorkspaceError> {
let today = Utc::now().date_naive();
self.daily_log(today).await
}
pub async fn daily_log(&self, date: NaiveDate) -> Result<MemoryDocument, WorkspaceError> {
let path = format!("daily/{}.md", date.format("%Y-%m-%d"));
self.read_or_create(&path).await
}
pub async fn heartbeat_checklist(&self) -> Result<Option<String>, WorkspaceError> {
match self.read(paths::HEARTBEAT).await {
Ok(doc) => Ok(Some(doc.content)),
Err(WorkspaceError::DocumentNotFound { .. }) => Ok(Some(HEARTBEAT_SEED.to_string())),
Err(e) => Err(e),
}
}
async fn read_or_create(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
self.storage
.get_or_create_document_by_path(&self.user_id, self.agent_id, path)
.await
}
pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> {
let doc = self.memory().await?;
let new_content = if doc.content.is_empty() {
entry.to_string()
} else {
format!("{}\n\n{}", doc.content, entry)
};
self.storage.update_document(doc.id, &new_content).await?;
self.reindex_document(doc.id).await?;
Ok(())
}
pub async fn append_daily_log(&self, entry: &str) -> Result<(), WorkspaceError> {
let today = Utc::now().date_naive();
let path = format!("daily/{}.md", today.format("%Y-%m-%d"));
let timestamp = Utc::now().format("%H:%M:%S");
let timestamped_entry = format!("[{}] {}", timestamp, entry);
self.append(&path, ×tamped_entry).await
}
pub async fn system_prompt(&self) -> Result<String, WorkspaceError> {
let mut parts = Vec::new();
let identity_files = [
(paths::AGENTS, "## Agent Instructions"),
(paths::SOUL, "## Core Values"),
(paths::USER, "## User Context"),
(paths::IDENTITY, "## Identity"),
];
for (path, header) in identity_files {
if let Ok(doc) = self.read(path).await
&& !doc.content.is_empty()
{
parts.push(format!("{}\n\n{}", header, doc.content));
}
}
let today = Utc::now().date_naive();
let yesterday = today.pred_opt().unwrap_or(today);
for date in [today, yesterday] {
if let Ok(doc) = self.daily_log(date).await
&& !doc.content.is_empty()
{
let header = if date == today {
"## Today's Notes"
} else {
"## Yesterday's Notes"
};
parts.push(format!("{}\n\n{}", header, doc.content));
}
}
Ok(parts.join("\n\n---\n\n"))
}
pub async fn search(
&self,
query: &str,
limit: usize,
) -> Result<Vec<SearchResult>, WorkspaceError> {
self.search_with_config(query, SearchConfig::default().with_limit(limit))
.await
}
pub async fn search_with_config(
&self,
query: &str,
config: SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
let embedding = if let Some(ref provider) = self.embeddings {
Some(
provider
.embed(query)
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: e.to_string(),
})?,
)
} else {
None
};
self.storage
.hybrid_search(
&self.user_id,
self.agent_id,
query,
embedding.as_deref(),
&config,
)
.await
}
async fn reindex_document(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
let doc = self.storage.get_document_by_id(document_id).await?;
let chunks = chunk_document(&doc.content, ChunkConfig::default());
self.storage.delete_chunks(document_id).await?;
for (index, content) in chunks.into_iter().enumerate() {
let embedding = if let Some(ref provider) = self.embeddings {
match provider.embed(&content).await {
Ok(emb) => Some(emb),
Err(e) => {
tracing::warn!("Failed to generate embedding: {}", e);
None
}
}
} else {
None
};
self.storage
.insert_chunk(document_id, index as i32, &content, embedding.as_deref())
.await?;
}
Ok(())
}
pub async fn seed_if_empty(&self) -> Result<usize, WorkspaceError> {
let seed_files: &[(&str, &str)] = &[
(
paths::README,
"# Workspace\n\n\
This is your agent's persistent memory. Files here are indexed for search\n\
and used to build the agent's context.\n\n\
## Structure\n\n\
- `MEMORY.md` - Long-term notes and facts worth remembering\n\
- `IDENTITY.md` - Agent name, nature, personality\n\
- `SOUL.md` - Core values and principles\n\
- `AGENTS.md` - Behavior instructions for the agent\n\
- `USER.md` - Information about you (the user)\n\
- `HEARTBEAT.md` - Periodic background task checklist\n\
- `daily/` - Automatic daily session logs\n\
- `context/` - Additional context documents\n\n\
Edit these files to shape how your agent thinks and acts.",
),
(
paths::MEMORY,
"# Memory\n\n\
Long-term notes, decisions, and facts worth remembering.\n\
The agent appends here during conversations.",
),
(
paths::IDENTITY,
"# Identity\n\n\
Name: IronClaw\n\
Nature: A secure personal AI assistant\n\n\
Edit this file to give your agent a custom name and personality.",
),
(
paths::SOUL,
"# Core Values\n\n\
- Protect user privacy and data security above all else\n\
- Be honest about limitations and uncertainty\n\
- Prefer action over lengthy deliberation\n\
- Ask for clarification rather than guessing on important decisions\n\
- Learn from mistakes and remember lessons",
),
(
paths::AGENTS,
"# Agent Instructions\n\n\
You are a personal AI assistant with access to tools and persistent memory.\n\n\
## Guidelines\n\n\
- Always search memory before answering questions about prior conversations\n\
- Write important facts and decisions to memory for future reference\n\
- Use the daily log for session-level notes\n\
- Be concise but thorough",
),
(
paths::USER,
"# User Context\n\n\
The agent will fill this in as it learns about you.\n\
You can also edit this directly to provide context upfront.",
),
(paths::HEARTBEAT, HEARTBEAT_SEED),
];
let mut count = 0;
for (path, content) in seed_files {
match self.read(path).await {
Ok(_) => continue,
Err(WorkspaceError::DocumentNotFound { .. }) => {}
Err(e) => {
tracing::warn!("Failed to check {}: {}", path, e);
continue;
}
}
if let Err(e) = self.write(path, content).await {
tracing::warn!("Failed to seed {}: {}", path, e);
} else {
count += 1;
}
}
if count > 0 {
tracing::info!("Seeded {} workspace files", count);
}
Ok(count)
}
pub async fn backfill_embeddings(&self) -> Result<usize, WorkspaceError> {
let Some(ref provider) = self.embeddings else {
return Ok(0);
};
let chunks = self
.storage
.get_chunks_without_embeddings(&self.user_id, self.agent_id, 100)
.await?;
let mut count = 0;
for chunk in chunks {
match provider.embed(&chunk.content).await {
Ok(embedding) => {
self.storage
.update_chunk_embedding(chunk.id, &embedding)
.await?;
count += 1;
}
Err(e) => {
tracing::warn!("Failed to embed chunk {}: {}", chunk.id, e);
}
}
}
Ok(count)
}
}
fn normalize_path(path: &str) -> String {
let path = path.trim().trim_matches('/');
let mut result = String::new();
let mut last_was_slash = false;
for c in path.chars() {
if c == '/' {
if !last_was_slash {
result.push(c);
}
last_was_slash = true;
} else {
result.push(c);
last_was_slash = false;
}
}
result
}
fn normalize_directory(path: &str) -> String {
let path = normalize_path(path);
path.trim_end_matches('/').to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_normalize_path() {
assert_eq!(normalize_path("foo/bar"), "foo/bar");
assert_eq!(normalize_path("/foo/bar/"), "foo/bar");
assert_eq!(normalize_path("foo//bar"), "foo/bar");
assert_eq!(normalize_path(" /foo/ "), "foo");
assert_eq!(normalize_path("README.md"), "README.md");
}
#[test]
fn test_normalize_directory() {
assert_eq!(normalize_directory("foo/bar/"), "foo/bar");
assert_eq!(normalize_directory("foo/bar"), "foo/bar");
assert_eq!(normalize_directory("/"), "");
assert_eq!(normalize_directory(""), "");
}
}