mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-09-02 17:49:20 +00:00
Add workspace and memory system (OpenClaw-inspired)
Implements persistent memory for agents with hybrid search: - Database-backed workspace with PostgreSQL (not filesystem) - Memory documents: MEMORY.md, daily logs, identity files - Chunked content with FTS (tsvector) + vector (pgvector) indexes - Reciprocal Rank Fusion (RRF) for hybrid search combining BM25 and semantic - Memory tools: memory_search, memory_write, memory_read - Proactive heartbeat system for periodic execution (30 min default) - OpenAI embeddings provider (text-embedding-3-small) Key patterns from OpenClaw: - "Memory is files, not RAM" - explicit persistence required - Two-tier memory: daily logs (raw) + curated MEMORY.md - Session isolation via user_id/agent_id scoping Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
8c38566378
commit
4e238e60ac
@@ -0,0 +1,451 @@
|
||||
//! Database repository for workspace persistence.
|
||||
//!
|
||||
//! All workspace data is stored in PostgreSQL:
|
||||
//! - Documents in `memory_documents` table
|
||||
//! - Chunks in `memory_chunks` table (with FTS and vector indexes)
|
||||
|
||||
use chrono::Utc;
|
||||
use deadpool_postgres::Pool;
|
||||
use pgvector::Vector;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
use crate::workspace::document::{DocType, MemoryChunk, MemoryDocument};
|
||||
use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
||||
|
||||
/// Database repository for workspace operations.
|
||||
pub struct Repository {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
impl Repository {
|
||||
/// Create a new repository with a connection pool.
|
||||
pub fn new(pool: Pool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// Get a connection from the pool.
|
||||
async fn conn(&self) -> Result<deadpool_postgres::Object, WorkspaceError> {
|
||||
self.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Failed to get connection: {}", e),
|
||||
})
|
||||
}
|
||||
|
||||
// ==================== Document Operations ====================
|
||||
|
||||
/// Get a document by type and optional title.
|
||||
pub async fn get_document(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: DocType,
|
||||
title: Option<&str>,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let row = if let Some(title) = title {
|
||||
conn.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
AND doc_type = $3 AND title = $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &doc_type.as_str(), &title],
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
conn.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
AND doc_type = $3 AND title IS NULL
|
||||
"#,
|
||||
&[&user_id, &agent_id, &doc_type.as_str()],
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
let row = row.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(self.row_to_document(&row)?),
|
||||
None => Err(WorkspaceError::DocumentNotFound {
|
||||
doc_type: doc_type.to_string(),
|
||||
user_id: user_id.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a document by ID.
|
||||
pub async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let row = conn
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents WHERE id = $1
|
||||
"#,
|
||||
&[&id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(self.row_to_document(&row)?),
|
||||
None => Err(WorkspaceError::DocumentNotFound {
|
||||
doc_type: "unknown".to_string(),
|
||||
user_id: "unknown".to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or create a document.
|
||||
pub async fn get_or_create_document(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: DocType,
|
||||
title: Option<&str>,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
// Try to get existing document first
|
||||
match self.get_document(user_id, agent_id, doc_type, title).await {
|
||||
Ok(doc) => return Ok(doc),
|
||||
Err(WorkspaceError::DocumentNotFound { .. }) => {}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
|
||||
// Create new document
|
||||
let conn = self.conn().await?;
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO memory_documents (id, user_id, agent_id, doc_type, title, content, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, '', $6, $7)
|
||||
ON CONFLICT (user_id, agent_id, doc_type, title) DO NOTHING
|
||||
"#,
|
||||
&[&id, &user_id, &agent_id, &doc_type.as_str(), &title, &now, &now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Insert failed: {}", e),
|
||||
})?;
|
||||
|
||||
// Fetch the document (might have been created by concurrent request)
|
||||
self.get_document(user_id, agent_id, doc_type, title).await
|
||||
}
|
||||
|
||||
/// Update a document's content.
|
||||
pub async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"UPDATE memory_documents SET content = $2, updated_at = NOW() WHERE id = $1",
|
||||
&[&id, &content],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Update failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List documents by type.
|
||||
pub async fn list_documents(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: Option<DocType>,
|
||||
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let rows = if let Some(dt) = doc_type {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND doc_type = $3
|
||||
ORDER BY updated_at DESC
|
||||
"#,
|
||||
&[&user_id, &agent_id, &dt.as_str()],
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
ORDER BY updated_at DESC
|
||||
"#,
|
||||
&[&user_id, &agent_id],
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
let rows = rows.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
rows.iter().map(|r| self.row_to_document(r)).collect()
|
||||
}
|
||||
|
||||
fn row_to_document(&self, row: &tokio_postgres::Row) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let doc_type_str: String = row.get("doc_type");
|
||||
let doc_type = DocType::try_from(doc_type_str.as_str())?;
|
||||
|
||||
Ok(MemoryDocument {
|
||||
id: row.get("id"),
|
||||
user_id: row.get("user_id"),
|
||||
agent_id: row.get("agent_id"),
|
||||
doc_type,
|
||||
title: row.get("title"),
|
||||
content: row.get("content"),
|
||||
created_at: row.get("created_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
metadata: row.get("metadata"),
|
||||
})
|
||||
}
|
||||
|
||||
// ==================== Chunk Operations ====================
|
||||
|
||||
/// Delete all chunks for a document.
|
||||
pub async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"DELETE FROM memory_chunks WHERE document_id = $1",
|
||||
&[&document_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::ChunkingFailed {
|
||||
reason: format!("Delete failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Insert a chunk.
|
||||
pub async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let id = Uuid::new_v4();
|
||||
|
||||
let embedding_vec = embedding.map(|e| Vector::from(e.to_vec()));
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO memory_chunks (id, document_id, chunk_index, content, embedding)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
"#,
|
||||
&[&id, &document_id, &chunk_index, &content, &embedding_vec],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::ChunkingFailed {
|
||||
reason: format!("Insert failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// Update a chunk's embedding.
|
||||
pub async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let embedding_vec = Vector::from(embedding.to_vec());
|
||||
|
||||
conn.execute(
|
||||
"UPDATE memory_chunks SET embedding = $2 WHERE id = $1",
|
||||
&[&chunk_id, &embedding_vec],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::EmbeddingFailed {
|
||||
reason: format!("Update failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get chunks without embeddings for backfilling.
|
||||
pub async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id, c.document_id, c.chunk_index, c.content, c.created_at
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.embedding IS NULL
|
||||
LIMIT $3
|
||||
"#,
|
||||
&[&user_id, &agent_id, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|row| MemoryChunk {
|
||||
id: row.get("id"),
|
||||
document_id: row.get("document_id"),
|
||||
chunk_index: row.get("chunk_index"),
|
||||
content: row.get("content"),
|
||||
embedding: None,
|
||||
created_at: row.get("created_at"),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
// ==================== Search Operations ====================
|
||||
|
||||
/// Perform hybrid search combining FTS and vector similarity.
|
||||
pub async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
let fts_results = if config.use_fts {
|
||||
self.fts_search(user_id, agent_id, query, config.pre_fusion_limit)
|
||||
.await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let vector_results = if config.use_vector && embedding.is_some() {
|
||||
self.vector_search(
|
||||
user_id,
|
||||
agent_id,
|
||||
embedding.unwrap(),
|
||||
config.pre_fusion_limit,
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
Ok(reciprocal_rank_fusion(fts_results, vector_results, config))
|
||||
}
|
||||
|
||||
/// Full-text search using PostgreSQL ts_rank_cd.
|
||||
async fn fts_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<RankedResult>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
// Use plainto_tsquery for natural language queries
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id as chunk_id, c.document_id, c.content,
|
||||
ts_rank_cd(c.content_tsv, plainto_tsquery('english', $3)) as rank
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.content_tsv @@ plainto_tsquery('english', $3)
|
||||
ORDER BY rank DESC
|
||||
LIMIT $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &query, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("FTS query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, row)| RankedResult {
|
||||
chunk_id: row.get("chunk_id"),
|
||||
document_id: row.get("document_id"),
|
||||
content: row.get("content"),
|
||||
rank: (i + 1) as u32, // 1-based rank
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Vector similarity search using pgvector cosine distance.
|
||||
async fn vector_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
embedding: &[f32],
|
||||
limit: usize,
|
||||
) -> Result<Vec<RankedResult>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let embedding_vec = Vector::from(embedding.to_vec());
|
||||
|
||||
// Use cosine distance (<=>)
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id as chunk_id, c.document_id, c.content,
|
||||
1 - (c.embedding <=> $3) as similarity
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.embedding IS NOT NULL
|
||||
ORDER BY c.embedding <=> $3
|
||||
LIMIT $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &embedding_vec, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Vector query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, row)| RankedResult {
|
||||
chunk_id: row.get("chunk_id"),
|
||||
document_id: row.get("document_id"),
|
||||
content: row.get("content"),
|
||||
rank: (i + 1) as u32, // 1-based rank
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user