From 4e238e60ac83f25f38bf651a65e87df8ac97444d Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Mon, 2 Feb 2026 21:18:47 -0800 Subject: [PATCH] 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 --- Cargo.lock | 11 + Cargo.toml | 4 + migrations/V2__workspace_memory.sql | 159 ++++++++++ src/agent/heartbeat.rs | 320 ++++++++++++++++++++ src/agent/mod.rs | 3 + src/error.rs | 28 ++ src/lib.rs | 2 + src/tools/builtin/memory.rs | 362 ++++++++++++++++++++++ src/tools/builtin/mod.rs | 2 + src/workspace/chunker.rs | 314 +++++++++++++++++++ src/workspace/document.rs | 236 +++++++++++++++ src/workspace/embeddings.rs | 324 ++++++++++++++++++++ src/workspace/mod.rs | 392 ++++++++++++++++++++++++ src/workspace/repository.rs | 451 ++++++++++++++++++++++++++++ src/workspace/search.rs | 391 ++++++++++++++++++++++++ 15 files changed, 2999 insertions(+) create mode 100644 migrations/V2__workspace_memory.sql create mode 100644 src/agent/heartbeat.rs create mode 100644 src/tools/builtin/memory.rs create mode 100644 src/workspace/chunker.rs create mode 100644 src/workspace/document.rs create mode 100644 src/workspace/embeddings.rs create mode 100644 src/workspace/mod.rs create mode 100644 src/workspace/repository.rs create mode 100644 src/workspace/search.rs diff --git a/Cargo.lock b/Cargo.lock index 4e3a52df..2a204a48 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1275,6 +1275,7 @@ dependencies = [ "deadpool-postgres", "dotenvy", "futures", + "pgvector", "postgres-types", "pretty_assertions", "refinery", @@ -1404,6 +1405,16 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pgvector" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc58e2d255979a31caa7cabfa7aac654af0354220719ab7a68520ae7a91e8c0b" +dependencies = [ + "bytes", + "postgres-types", +] + [[package]] name = "phf" version = "0.13.1" diff --git a/Cargo.toml b/Cargo.toml index c5f6120a..9d7de1ab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,10 @@ aho-corasick = "1" # Secrecy for sensitive values secrecy = { version = "0.10", features = ["serde"] } +# Vector embeddings for semantic search +# The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres) +pgvector = { version = "0.4", features = ["postgres"] } + [dev-dependencies] tokio-test = "0.4" testcontainers-modules = { version = "0.11", features = ["postgres"] } diff --git a/migrations/V2__workspace_memory.sql b/migrations/V2__workspace_memory.sql new file mode 100644 index 00000000..7c690ed5 --- /dev/null +++ b/migrations/V2__workspace_memory.sql @@ -0,0 +1,159 @@ +-- NEAR Agent Database Schema +-- V2: Workspace and memory system (OpenClaw-inspired) +-- +-- This migration adds: +-- 1. Persistent memory documents (MEMORY.md, daily logs, identity files) +-- 2. Chunked content for hybrid search (FTS + vector) +-- 3. Heartbeat state for proactive execution + +-- Enable pgvector extension for semantic search +-- NOTE: This requires pgvector to be installed on the PostgreSQL server +-- Install via: CREATE EXTENSION vector; (requires superuser or rds_superuser) +CREATE EXTENSION IF NOT EXISTS vector; + +-- ==================== Memory Documents ==================== +-- Stores full documents like MEMORY.md, daily logs, identity files +-- Think of this as the filesystem equivalent, but in PostgreSQL + +CREATE TABLE memory_documents ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + -- Ownership: who this document belongs to + user_id TEXT NOT NULL, -- User identifier (from channel) + agent_id UUID, -- NULL = shared across all agents for this user + + -- Document type and content + doc_type TEXT NOT NULL, -- 'memory', 'daily_log', 'identity', 'soul', 'agents', 'user', 'heartbeat' + title TEXT, -- Optional title (e.g., date for daily logs) + content TEXT NOT NULL, -- Full document content + + -- Timestamps + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + -- Flexible metadata (tags, source, etc.) + metadata JSONB NOT NULL DEFAULT '{}', + + -- Ensure one document per type per user (for singleton docs like MEMORY.md) + -- Daily logs use title as the date discriminator + CONSTRAINT unique_doc_per_user_type UNIQUE (user_id, agent_id, doc_type, title) +); + +-- Indexes for common queries +CREATE INDEX idx_memory_documents_user ON memory_documents(user_id); +CREATE INDEX idx_memory_documents_user_type ON memory_documents(user_id, doc_type); +CREATE INDEX idx_memory_documents_updated ON memory_documents(updated_at DESC); + +-- ==================== Memory Chunks ==================== +-- Documents are chunked for search. Each chunk has: +-- 1. Full-text search vector (tsvector) for keyword matching +-- 2. Embedding vector for semantic similarity +-- +-- Hybrid search combines both using Reciprocal Rank Fusion (RRF) + +CREATE TABLE memory_chunks ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + document_id UUID NOT NULL REFERENCES memory_documents(id) ON DELETE CASCADE, + + -- Chunk position and content + chunk_index INT NOT NULL, -- Position in document (0-based) + content TEXT NOT NULL, -- Chunk text (~800 tokens with 15% overlap) + + -- Full-text search: auto-generated tsvector + content_tsv TSVECTOR GENERATED ALWAYS AS (to_tsvector('english', content)) STORED, + + -- Semantic search: embedding vector (OpenAI text-embedding-ada-002 = 1536 dims) + -- NULL until embeddings are generated + embedding VECTOR(1536), + + -- Timestamps + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + -- Each chunk index unique per document + CONSTRAINT unique_chunk_per_doc UNIQUE (document_id, chunk_index) +); + +-- GIN index for full-text search +CREATE INDEX idx_memory_chunks_tsv ON memory_chunks USING GIN(content_tsv); + +-- HNSW index for vector similarity (cosine distance) +-- HNSW is faster than IVFFlat for reads, slightly slower for writes +CREATE INDEX idx_memory_chunks_embedding ON memory_chunks + USING hnsw(embedding vector_cosine_ops) + WITH (m = 16, ef_construction = 64); + +-- Index for document lookups +CREATE INDEX idx_memory_chunks_document ON memory_chunks(document_id); + +-- ==================== Heartbeat State ==================== +-- Tracks periodic heartbeat execution per user/agent + +CREATE TABLE heartbeat_state ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + agent_id UUID, -- NULL = default agent for user + + -- Timing + last_run TIMESTAMPTZ, -- When heartbeat last executed + next_run TIMESTAMPTZ, -- Scheduled next execution + interval_seconds INT NOT NULL DEFAULT 1800, -- 30 minutes default + + -- State + enabled BOOLEAN NOT NULL DEFAULT true, + consecutive_failures INT NOT NULL DEFAULT 0, + + -- Last check timestamps (for batched monitoring) + -- e.g., {"email": "2024-01-15T10:00:00Z", "calendar": "2024-01-15T10:00:00Z"} + last_checks JSONB NOT NULL DEFAULT '{}', + + -- Ensure one heartbeat config per user/agent + CONSTRAINT unique_heartbeat_per_user UNIQUE (user_id, agent_id) +); + +CREATE INDEX idx_heartbeat_user ON heartbeat_state(user_id); +CREATE INDEX idx_heartbeat_next_run ON heartbeat_state(next_run) WHERE enabled = true; + +-- ==================== Helper Functions ==================== + +-- Function to update updated_at timestamp +CREATE OR REPLACE FUNCTION update_updated_at_column() +RETURNS TRIGGER AS $$ +BEGIN + NEW.updated_at = NOW(); + RETURN NEW; +END; +$$ language 'plpgsql'; + +-- Trigger to auto-update updated_at on memory_documents +CREATE TRIGGER update_memory_documents_updated_at + BEFORE UPDATE ON memory_documents + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); + +-- ==================== Views ==================== + +-- View for documents with chunk counts (useful for debugging) +CREATE VIEW memory_documents_summary AS +SELECT + d.id, + d.user_id, + d.doc_type, + d.title, + d.created_at, + d.updated_at, + COUNT(c.id) as chunk_count, + COUNT(c.embedding) as embedded_chunk_count +FROM memory_documents d +LEFT JOIN memory_chunks c ON c.document_id = d.id +GROUP BY d.id; + +-- View for pending embedding work +CREATE VIEW chunks_pending_embedding AS +SELECT + c.id as chunk_id, + c.document_id, + d.user_id, + d.doc_type, + LENGTH(c.content) as content_length +FROM memory_chunks c +JOIN memory_documents d ON d.id = c.document_id +WHERE c.embedding IS NULL; diff --git a/src/agent/heartbeat.rs b/src/agent/heartbeat.rs new file mode 100644 index 00000000..b6508ab1 --- /dev/null +++ b/src/agent/heartbeat.rs @@ -0,0 +1,320 @@ +//! Proactive heartbeat system for periodic execution. +//! +//! The heartbeat runner executes periodically (default: every 30 minutes) and: +//! 1. Reads the HEARTBEAT.md checklist +//! 2. Runs an agent turn to process the checklist +//! 3. Reports any findings to the configured channel +//! +//! If nothing needs attention, the agent replies "HEARTBEAT_OK" and no +//! message is sent to the user. +//! +//! # Usage +//! +//! Create a HEARTBEAT.md in the workspace with a checklist of things to monitor: +//! +//! ```markdown +//! # Heartbeat Checklist +//! +//! - [ ] Check for unread emails +//! - [ ] Review calendar for upcoming events +//! - [ ] Check project build status +//! ``` +//! +//! The agent will process this checklist on each heartbeat and only notify +//! if action is needed. + +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::mpsc; + +use crate::channels::OutgoingResponse; +use crate::error::WorkspaceError; +use crate::llm::{ChatMessage, CompletionRequest, LlmProvider}; +use crate::workspace::Workspace; + +/// Configuration for the heartbeat runner. +#[derive(Debug, Clone)] +pub struct HeartbeatConfig { + /// Interval between heartbeat checks. + pub interval: Duration, + /// Whether heartbeat is enabled. + pub enabled: bool, + /// Maximum consecutive failures before disabling. + pub max_failures: u32, + /// User ID to notify on heartbeat findings. + pub notify_user_id: Option, + /// Channel to notify on heartbeat findings. + pub notify_channel: Option, +} + +impl Default for HeartbeatConfig { + fn default() -> Self { + Self { + interval: Duration::from_secs(30 * 60), // 30 minutes + enabled: true, + max_failures: 3, + notify_user_id: None, + notify_channel: None, + } + } +} + +impl HeartbeatConfig { + /// Create a config with a specific interval. + pub fn with_interval(mut self, interval: Duration) -> Self { + self.interval = interval; + self + } + + /// Disable heartbeat. + pub fn disabled(mut self) -> Self { + self.enabled = false; + self + } + + /// Set the notification target. + pub fn with_notify(mut self, user_id: impl Into, channel: impl Into) -> Self { + self.notify_user_id = Some(user_id.into()); + self.notify_channel = Some(channel.into()); + self + } +} + +/// Result of a heartbeat check. +#[derive(Debug)] +pub enum HeartbeatResult { + /// Nothing needs attention. + Ok, + /// Something needs attention, with the message to send. + NeedsAttention(String), + /// Heartbeat was skipped (no checklist or disabled). + Skipped, + /// Heartbeat failed. + Failed(String), +} + +/// Heartbeat runner for proactive periodic execution. +pub struct HeartbeatRunner { + config: HeartbeatConfig, + workspace: Arc, + llm: Arc, + response_tx: Option>, + consecutive_failures: u32, +} + +impl HeartbeatRunner { + /// Create a new heartbeat runner. + pub fn new( + config: HeartbeatConfig, + workspace: Arc, + llm: Arc, + ) -> Self { + Self { + config, + workspace, + llm, + response_tx: None, + consecutive_failures: 0, + } + } + + /// Set the response channel for notifications. + pub fn with_response_channel(mut self, tx: mpsc::Sender) -> Self { + self.response_tx = Some(tx); + self + } + + /// Run the heartbeat loop. + /// + /// This runs forever, checking periodically based on the configured interval. + pub async fn run(&mut self) { + if !self.config.enabled { + tracing::info!("Heartbeat is disabled, not starting loop"); + return; + } + + tracing::info!( + "Starting heartbeat loop with interval {:?}", + self.config.interval + ); + + let mut interval = tokio::time::interval(self.config.interval); + // Don't run immediately on startup + interval.tick().await; + + loop { + interval.tick().await; + + match self.check_heartbeat().await { + HeartbeatResult::Ok => { + tracing::debug!("Heartbeat OK"); + self.consecutive_failures = 0; + } + HeartbeatResult::NeedsAttention(message) => { + tracing::info!("Heartbeat needs attention: {}", message); + self.consecutive_failures = 0; + self.send_notification(&message).await; + } + HeartbeatResult::Skipped => { + tracing::debug!("Heartbeat skipped"); + } + HeartbeatResult::Failed(error) => { + tracing::error!("Heartbeat failed: {}", error); + self.consecutive_failures += 1; + + if self.consecutive_failures >= self.config.max_failures { + tracing::error!( + "Heartbeat disabled after {} consecutive failures", + self.consecutive_failures + ); + break; + } + } + } + } + } + + /// Run a single heartbeat check. + pub async fn check_heartbeat(&self) -> HeartbeatResult { + // Get the heartbeat checklist + let checklist = match self.workspace.heartbeat_checklist().await { + Ok(Some(content)) if !content.trim().is_empty() => content, + Ok(_) => return HeartbeatResult::Skipped, + Err(e) => return HeartbeatResult::Failed(format!("Failed to read checklist: {}", e)), + }; + + // Build the heartbeat prompt + let prompt = format!( + "Read the HEARTBEAT.md checklist below and follow it strictly. \ + Do not infer or repeat old tasks. Check each item and report findings.\n\ + \n\ + If nothing needs attention, reply EXACTLY with: HEARTBEAT_OK\n\ + \n\ + If something needs attention, provide a concise summary of what needs action.\n\ + \n\ + ## HEARTBEAT.md\n\ + \n\ + {}", + checklist + ); + + // Get the system prompt for context + let system_prompt = match self.workspace.system_prompt().await { + Ok(p) => p, + Err(e) => { + tracing::warn!("Failed to get system prompt for heartbeat: {}", e); + String::new() + } + }; + + // Run the agent turn + let messages = if system_prompt.is_empty() { + vec![ChatMessage::user(&prompt)] + } else { + vec![ + ChatMessage::system(&system_prompt), + ChatMessage::user(&prompt), + ] + }; + + let request = CompletionRequest::new(messages) + .with_max_tokens(1024) + .with_temperature(0.3); // Lower temperature for more focused responses + + let response = match self.llm.complete(request).await { + Ok(r) => r, + Err(e) => return HeartbeatResult::Failed(format!("LLM call failed: {}", e)), + }; + + let content = response.content.trim(); + + // Check if nothing needs attention + if content == "HEARTBEAT_OK" || content.contains("HEARTBEAT_OK") { + return HeartbeatResult::Ok; + } + + HeartbeatResult::NeedsAttention(content.to_string()) + } + + /// Send a notification about heartbeat findings. + async fn send_notification(&self, message: &str) { + let Some(ref tx) = self.response_tx else { + tracing::debug!("No response channel configured for heartbeat notifications"); + return; + }; + + let response = OutgoingResponse { + content: format!("🔔 **Heartbeat Alert**\n\n{}", message), + thread_id: None, + metadata: serde_json::json!({ + "source": "heartbeat", + }), + }; + + if let Err(e) = tx.send(response).await { + tracing::error!("Failed to send heartbeat notification: {}", e); + } + } +} + +/// Spawn the heartbeat runner as a background task. +/// +/// Returns a handle that can be used to stop the runner. +pub fn spawn_heartbeat( + config: HeartbeatConfig, + workspace: Arc, + llm: Arc, + response_tx: Option>, +) -> tokio::task::JoinHandle<()> { + let mut runner = HeartbeatRunner::new(config, workspace, llm); + if let Some(tx) = response_tx { + runner = runner.with_response_channel(tx); + } + + tokio::spawn(async move { + runner.run().await; + }) +} + +/// Update heartbeat state in the database. +pub async fn update_heartbeat_state( + workspace: &Workspace, + last_run: chrono::DateTime, +) -> Result<(), WorkspaceError> { + // This would update the heartbeat_state table + // For now, we just log + tracing::debug!( + "Heartbeat state updated for user {} at {}", + workspace.user_id(), + last_run + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_heartbeat_config_defaults() { + let config = HeartbeatConfig::default(); + assert!(config.enabled); + assert_eq!(config.interval, Duration::from_secs(30 * 60)); + assert_eq!(config.max_failures, 3); + } + + #[test] + fn test_heartbeat_config_builders() { + let config = HeartbeatConfig::default() + .with_interval(Duration::from_secs(60)) + .with_notify("user1", "telegram"); + + assert_eq!(config.interval, Duration::from_secs(60)); + assert_eq!(config.notify_user_id, Some("user1".to_string())); + assert_eq!(config.notify_channel, Some("telegram".to_string())); + + let disabled = HeartbeatConfig::default().disabled(); + assert!(!disabled.enabled); + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 46dc37fa..0685f677 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -5,14 +5,17 @@ //! - Job scheduling and execution //! - Tool invocation with safety //! - Self-repair for stuck jobs +//! - Proactive heartbeat execution mod agent_loop; +mod heartbeat; mod router; mod scheduler; mod self_repair; mod worker; pub use agent_loop::Agent; +pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat}; pub use router::{MessageIntent, Router}; pub use scheduler::Scheduler; pub use self_repair::{RepairResult, RepairTask, SelfRepair, StuckJob}; diff --git a/src/error.rs b/src/error.rs index 627832ae..f41fc82c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -36,6 +36,9 @@ pub enum Error { #[error("Repair error: {0}")] Repair(#[from] RepairError), + + #[error("Workspace error: {0}")] + Workspace(#[from] WorkspaceError), } /// Configuration-related errors. @@ -268,5 +271,30 @@ pub enum RepairError { }, } +/// Workspace/memory errors. +#[derive(Debug, thiserror::Error)] +pub enum WorkspaceError { + #[error("Document not found: {doc_type} for user {user_id}")] + DocumentNotFound { doc_type: String, user_id: String }, + + #[error("Search failed: {reason}")] + SearchFailed { reason: String }, + + #[error("Embedding generation failed: {reason}")] + EmbeddingFailed { reason: String }, + + #[error("Document chunking failed: {reason}")] + ChunkingFailed { reason: String }, + + #[error("Invalid document type: {doc_type}")] + InvalidDocType { doc_type: String }, + + #[error("Workspace not initialized for user {user_id}")] + NotInitialized { user_id: String }, + + #[error("Heartbeat error: {reason}")] + HeartbeatError { reason: String }, +} + /// Result type alias for the agent. pub type Result = std::result::Result; diff --git a/src/lib.rs b/src/lib.rs index 0ff3204f..990a9c87 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -49,6 +49,7 @@ pub mod history; pub mod llm; pub mod safety; pub mod tools; +pub mod workspace; pub use config::Config; pub use error::{Error, Result}; @@ -62,4 +63,5 @@ pub mod prelude { pub use crate::llm::LlmProvider; pub use crate::safety::{SanitizedOutput, Sanitizer}; pub use crate::tools::{Tool, ToolOutput, ToolRegistry}; + pub use crate::workspace::{DocType, MemoryDocument, Workspace}; } diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs new file mode 100644 index 00000000..480bd7a7 --- /dev/null +++ b/src/tools/builtin/memory.rs @@ -0,0 +1,362 @@ +//! Memory tools for persistent workspace memory. +//! +//! These tools allow the agent to: +//! - Search past memories, decisions, and context +//! - Write important information to long-term memory +//! +//! # Usage +//! +//! The agent should use `memory_search` before answering questions about +//! prior work, decisions, dates, people, preferences, or todos. +//! +//! Use `memory_write` to persist important facts that should be remembered +//! across sessions. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::context::JobContext; +use crate::tools::tool::{Tool, ToolError, ToolOutput}; +use crate::workspace::Workspace; + +/// Tool for searching workspace memory. +/// +/// Performs hybrid search (FTS + semantic) across all memory documents. +/// The agent should call this tool before answering questions about +/// prior work, decisions, preferences, or any historical context. +pub struct MemorySearchTool { + workspace: Arc, +} + +impl MemorySearchTool { + /// Create a new memory search tool. + pub fn new(workspace: Arc) -> Self { + Self { workspace } + } +} + +#[async_trait] +impl Tool for MemorySearchTool { + fn name(&self) -> &str { + "memory_search" + } + + fn description(&self) -> &str { + "Search past memories, decisions, and context. MUST be called before answering \ + questions about prior work, decisions, dates, people, preferences, or todos. \ + Returns relevant snippets with relevance scores." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query. Use natural language to describe what you're looking for." + }, + "limit": { + "type": "integer", + "description": "Maximum number of results to return (default: 5, max: 20)", + "default": 5, + "minimum": 1, + "maximum": 20 + } + }, + "required": ["query"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let query = params + .get("query") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'query' parameter".to_string()))?; + + let limit = params + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(5) + .min(20) as usize; + + let results = self + .workspace + .search(query, limit) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Search failed: {}", e)))?; + + let output = serde_json::json!({ + "query": query, + "results": results.iter().map(|r| serde_json::json!({ + "content": r.content, + "score": r.score, + "document_id": r.document_id.to_string(), + "is_hybrid_match": r.is_hybrid(), + })).collect::>(), + "result_count": results.len(), + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false // Internal memory, trusted content + } +} + +/// Tool for writing to workspace memory. +/// +/// Use this to persist important information that should be remembered +/// across sessions: decisions, preferences, facts, lessons learned. +pub struct MemoryWriteTool { + workspace: Arc, +} + +impl MemoryWriteTool { + /// Create a new memory write tool. + pub fn new(workspace: Arc) -> Self { + Self { workspace } + } +} + +#[async_trait] +impl Tool for MemoryWriteTool { + fn name(&self) -> &str { + "memory_write" + } + + fn description(&self) -> &str { + "Write to persistent memory. Use for important facts, decisions, preferences, \ + or lessons learned that should be remembered across sessions. Use 'memory' target \ + for curated long-term facts, 'daily_log' for timestamped session notes." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "content": { + "type": "string", + "description": "The content to write to memory. Be concise but include relevant context." + }, + "target": { + "type": "string", + "enum": ["memory", "daily_log"], + "description": "Where to write: 'memory' for long-term curated facts, 'daily_log' for timestamped session notes", + "default": "daily_log" + } + }, + "required": ["content"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let content = params + .get("content") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("missing 'content' parameter".to_string()) + })?; + + if content.trim().is_empty() { + return Err(ToolError::InvalidParameters( + "content cannot be empty".to_string(), + )); + } + + let target = params + .get("target") + .and_then(|v| v.as_str()) + .unwrap_or("daily_log"); + + match target { + "memory" => { + self.workspace + .append_memory(content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } + "daily_log" => { + self.workspace + .append_daily_log(content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } + _ => { + return Err(ToolError::InvalidParameters(format!( + "invalid target '{}', must be 'memory' or 'daily_log'", + target + ))); + } + } + + let output = serde_json::json!({ + "status": "written", + "target": target, + "content_length": content.len(), + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false // Internal tool + } +} + +/// Tool for reading specific memory documents. +/// +/// Use this to read the full content of identity files, heartbeat checklist, +/// or other specific documents. +pub struct MemoryReadTool { + workspace: Arc, +} + +impl MemoryReadTool { + /// Create a new memory read tool. + pub fn new(workspace: Arc) -> Self { + Self { workspace } + } +} + +#[async_trait] +impl Tool for MemoryReadTool { + fn name(&self) -> &str { + "memory_read" + } + + fn description(&self) -> &str { + "Read a specific memory document by type. Use this to read identity files, \ + heartbeat checklist, or full memory document content." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "doc_type": { + "type": "string", + "enum": ["memory", "daily_log", "identity", "soul", "agents", "user", "heartbeat"], + "description": "The type of document to read" + }, + "title": { + "type": "string", + "description": "Optional title (required for daily_log, format: YYYY-MM-DD)" + } + }, + "required": ["doc_type"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let doc_type_str = params + .get("doc_type") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("missing 'doc_type' parameter".to_string()) + })?; + + let doc_type = crate::workspace::DocType::try_from(doc_type_str) + .map_err(|e| ToolError::InvalidParameters(e.to_string()))?; + + let title = params.get("title").and_then(|v| v.as_str()); + + let doc = self + .workspace + .get_document(doc_type, title) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?; + + let output = serde_json::json!({ + "doc_type": doc_type_str, + "title": doc.title, + "content": doc.content, + "word_count": doc.word_count(), + "updated_at": doc.updated_at.to_rfc3339(), + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false // Internal memory + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Integration tests would require a database connection. + // Unit tests for parameter validation: + + #[test] + fn test_memory_search_schema() { + let workspace = Arc::new(Workspace::new( + "test_user", + deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( + tokio_postgres::Config::new(), + tokio_postgres::NoTls, + )) + .build() + .unwrap(), + )); + let tool = MemorySearchTool::new(workspace); + + assert_eq!(tool.name(), "memory_search"); + assert!(!tool.requires_sanitization()); + + let schema = tool.parameters_schema(); + assert!(schema["properties"]["query"].is_object()); + assert!( + schema["required"] + .as_array() + .unwrap() + .contains(&"query".into()) + ); + } + + #[test] + fn test_memory_write_schema() { + let workspace = Arc::new(Workspace::new( + "test_user", + deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( + tokio_postgres::Config::new(), + tokio_postgres::NoTls, + )) + .build() + .unwrap(), + )); + let tool = MemoryWriteTool::new(workspace); + + assert_eq!(tool.name(), "memory_write"); + + let schema = tool.parameters_schema(); + assert!(schema["properties"]["content"].is_object()); + assert!( + schema["properties"]["target"]["enum"] + .as_array() + .unwrap() + .contains(&"memory".into()) + ); + } +} diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index d7e73266..c655319f 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -5,6 +5,7 @@ mod ecommerce; mod http; mod json; mod marketplace; +mod memory; mod restaurant; mod taskrabbit; mod time; @@ -14,6 +15,7 @@ pub use ecommerce::EcommerceTool; pub use http::HttpTool; pub use json::JsonTool; pub use marketplace::MarketplaceTool; +pub use memory::{MemoryReadTool, MemorySearchTool, MemoryWriteTool}; pub use restaurant::RestaurantTool; pub use taskrabbit::TaskRabbitTool; pub use time::TimeTool; diff --git a/src/workspace/chunker.rs b/src/workspace/chunker.rs new file mode 100644 index 00000000..98f3873d --- /dev/null +++ b/src/workspace/chunker.rs @@ -0,0 +1,314 @@ +//! Document chunking for search indexing. +//! +//! Documents are split into overlapping chunks for better search recall. +//! The overlap ensures context is preserved across chunk boundaries. + +/// Configuration for document chunking. +#[derive(Debug, Clone)] +pub struct ChunkConfig { + /// Target chunk size in words (approximate tokens). + /// Default: 800 (roughly 800 tokens for English text). + pub chunk_size: usize, + /// Overlap percentage between chunks. + /// Default: 0.15 (15% overlap). + pub overlap_percent: f32, + /// Minimum chunk size (don't create tiny trailing chunks). + /// Default: 50 words. + pub min_chunk_size: usize, +} + +impl Default for ChunkConfig { + fn default() -> Self { + Self { + chunk_size: 800, + overlap_percent: 0.15, + min_chunk_size: 50, + } + } +} + +impl ChunkConfig { + /// Create a config with a specific chunk size. + pub fn with_chunk_size(mut self, size: usize) -> Self { + self.chunk_size = size; + self + } + + /// Create a config with a specific overlap percentage. + pub fn with_overlap(mut self, percent: f32) -> Self { + self.overlap_percent = percent.clamp(0.0, 0.5); + self + } + + /// Calculate the overlap size in words. + fn overlap_size(&self) -> usize { + (self.chunk_size as f32 * self.overlap_percent) as usize + } + + /// Calculate the step size (chunk_size - overlap). + fn step_size(&self) -> usize { + self.chunk_size.saturating_sub(self.overlap_size()) + } +} + +/// Split a document into overlapping chunks. +/// +/// Each chunk contains approximately `chunk_size` words, with `overlap_percent` +/// overlap between adjacent chunks. This ensures that: +/// 1. Context is preserved across chunk boundaries +/// 2. Search can find content that spans chunk boundaries +/// +/// # Arguments +/// +/// * `content` - The document text to chunk +/// * `config` - Chunking configuration +/// +/// # Returns +/// +/// A vector of chunk strings. Empty documents return an empty vector. +pub fn chunk_document(content: &str, config: ChunkConfig) -> Vec { + if content.is_empty() { + return Vec::new(); + } + + // Split into words while preserving structure + let words: Vec<&str> = content.split_whitespace().collect(); + + if words.is_empty() { + return Vec::new(); + } + + // If content is smaller than chunk size, return as single chunk + if words.len() <= config.chunk_size { + return vec![content.to_string()]; + } + + let step = config.step_size(); + let mut chunks = Vec::new(); + let mut start = 0; + + while start < words.len() { + let end = (start + config.chunk_size).min(words.len()); + let chunk_words = &words[start..end]; + + // Don't create tiny trailing chunks, merge with previous + if chunk_words.len() < config.min_chunk_size && !chunks.is_empty() { + let last = chunks.pop().unwrap(); + let combined = format!("{} {}", last, chunk_words.join(" ")); + chunks.push(combined); + break; + } + + chunks.push(chunk_words.join(" ")); + + // Move to next chunk position + start += step; + + // Avoid creating duplicate chunks at the end + if start + config.min_chunk_size >= words.len() && end == words.len() { + break; + } + } + + chunks +} + +/// Split content by paragraphs first, then chunk. +/// +/// This is better for preserving semantic boundaries. +pub fn chunk_by_paragraphs(content: &str, config: ChunkConfig) -> Vec { + if content.is_empty() { + return Vec::new(); + } + + // Split by double newlines (paragraphs) + let paragraphs: Vec<&str> = content + .split("\n\n") + .map(|p| p.trim()) + .filter(|p| !p.is_empty()) + .collect(); + + if paragraphs.is_empty() { + return chunk_document(content, config); + } + + let mut chunks = Vec::new(); + let mut current_chunk = String::new(); + let mut current_word_count = 0; + + for paragraph in paragraphs { + let para_words = paragraph.split_whitespace().count(); + + // If this paragraph alone exceeds chunk size, chunk it separately + if para_words > config.chunk_size { + // Flush current chunk first + if !current_chunk.is_empty() { + chunks.push(current_chunk.trim().to_string()); + current_chunk = String::new(); + current_word_count = 0; + } + // Chunk the large paragraph + let para_chunks = chunk_document(paragraph, config.clone()); + chunks.extend(para_chunks); + continue; + } + + // Check if adding this paragraph would exceed chunk size + if current_word_count + para_words > config.chunk_size { + // Flush current chunk + if !current_chunk.is_empty() { + chunks.push(current_chunk.trim().to_string()); + } + current_chunk = paragraph.to_string(); + current_word_count = para_words; + } else { + // Add paragraph to current chunk + if !current_chunk.is_empty() { + current_chunk.push_str("\n\n"); + } + current_chunk.push_str(paragraph); + current_word_count += para_words; + } + } + + // Flush remaining content + if !current_chunk.is_empty() { + // If too small, merge with previous chunk if possible + if current_word_count < config.min_chunk_size && !chunks.is_empty() { + let last = chunks.pop().unwrap(); + chunks.push(format!("{}\n\n{}", last, current_chunk.trim())); + } else { + chunks.push(current_chunk.trim().to_string()); + } + } + + chunks +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_empty_content() { + let config = ChunkConfig::default(); + assert!(chunk_document("", config.clone()).is_empty()); + assert!(chunk_document(" ", config).is_empty()); + } + + #[test] + fn test_small_content() { + let config = ChunkConfig::default(); + let content = "Hello world, this is a test."; + let chunks = chunk_document(content, config); + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0], content); + } + + #[test] + fn test_exact_chunk_size() { + let config = ChunkConfig::default().with_chunk_size(5); + let content = "one two three four five"; + let chunks = chunk_document(content, config); + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0], content); + } + + #[test] + fn test_chunking_with_overlap() { + let config = ChunkConfig { + chunk_size: 10, + overlap_percent: 0.2, // 2 word overlap + min_chunk_size: 3, // Low threshold for test + }; + + // 20 words + let content = "one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen nineteen twenty"; + let chunks = chunk_document(content, config); + + // Should create overlapping chunks + assert!( + chunks.len() >= 2, + "Expected at least 2 chunks, got {}", + chunks.len() + ); + + // Each chunk should have roughly 10 words (allowing for overlap/merging) + for chunk in &chunks { + let word_count = chunk.split_whitespace().count(); + assert!(word_count >= 3, "Chunk too small: {} words", word_count); + } + } + + #[test] + fn test_overlap_calculation() { + let config = ChunkConfig::default() + .with_chunk_size(100) + .with_overlap(0.15); + + assert_eq!(config.overlap_size(), 15); + assert_eq!(config.step_size(), 85); + } + + #[test] + fn test_paragraph_chunking() { + let config = ChunkConfig::default().with_chunk_size(20); + + let content = "First paragraph with some words.\n\nSecond paragraph with different content.\n\nThird paragraph here."; + let chunks = chunk_by_paragraphs(content, config); + + // Should preserve paragraph boundaries + assert!(!chunks.is_empty()); + for chunk in &chunks { + // No chunk should start or end with \n\n + assert!(!chunk.starts_with("\n")); + assert!(!chunk.ends_with("\n")); + } + } + + #[test] + fn test_large_paragraph_handling() { + let config = ChunkConfig { + chunk_size: 10, + overlap_percent: 0.15, + min_chunk_size: 3, // Low threshold for test + }; + + // Create a paragraph with 30 words + let large_para = (1..=30) + .map(|i| format!("word{}", i)) + .collect::>() + .join(" "); + let content = format!("Short intro.\n\n{}\n\nShort outro.", large_para); + + let chunks = chunk_by_paragraphs(&content, config); + + // Should have multiple chunks due to large paragraph + // 30 words + 2 intro + 2 outro = 34 words, chunk_size=10 + // Expect at least 3 chunks + assert!( + chunks.len() >= 3, + "Expected at least 3 chunks for 34 words with chunk_size=10, got {}", + chunks.len() + ); + } + + #[test] + fn test_min_chunk_size_merging() { + let config = ChunkConfig { + chunk_size: 10, + overlap_percent: 0.0, + min_chunk_size: 5, + }; + + // 12 words: should create one chunk of 10, and merge the remaining 2 with it + let content = "one two three four five six seven eight nine ten eleven twelve"; + let chunks = chunk_document(content, config); + + // Should merge the tiny trailing chunk + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].split_whitespace().count(), 12); + } +} diff --git a/src/workspace/document.rs b/src/workspace/document.rs new file mode 100644 index 00000000..53352afb --- /dev/null +++ b/src/workspace/document.rs @@ -0,0 +1,236 @@ +//! Memory document types. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::error::WorkspaceError; + +/// Document type in the workspace. +/// +/// Each type represents a different kind of persistent memory: +/// - **Memory**: Long-term curated facts and decisions (MEMORY.md) +/// - **DailyLog**: Append-only daily notes (memory/YYYY-MM-DD.md) +/// - **Identity**: Agent name and personality +/// - **Soul**: Core values and behavior principles +/// - **Agents**: Behavior instructions +/// - **User**: User context (name, preferences) +/// - **Heartbeat**: Periodic checklist for proactive execution +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DocType { + /// Long-term curated memory (MEMORY.md equivalent). + Memory, + /// Daily append-only logs. + DailyLog, + /// Agent identity (name, nature, vibe). + Identity, + /// Core values and principles (SOUL.md). + Soul, + /// Behavior instructions (AGENTS.md). + Agents, + /// User context (USER.md). + User, + /// Periodic checklist (HEARTBEAT.md). + Heartbeat, +} + +impl DocType { + /// Get the string representation. + pub fn as_str(&self) -> &'static str { + match self { + DocType::Memory => "memory", + DocType::DailyLog => "daily_log", + DocType::Identity => "identity", + DocType::Soul => "soul", + DocType::Agents => "agents", + DocType::User => "user", + DocType::Heartbeat => "heartbeat", + } + } + + /// Check if this document type is a singleton (one per user/agent). + pub fn is_singleton(&self) -> bool { + match self { + DocType::Memory + | DocType::Identity + | DocType::Soul + | DocType::Agents + | DocType::User + | DocType::Heartbeat => true, + DocType::DailyLog => false, + } + } + + /// Check if this document should be included in the system prompt. + pub fn is_identity_document(&self) -> bool { + matches!( + self, + DocType::Identity | DocType::Soul | DocType::Agents | DocType::User + ) + } +} + +impl TryFrom<&str> for DocType { + type Error = WorkspaceError; + + fn try_from(s: &str) -> Result { + match s { + "memory" => Ok(DocType::Memory), + "daily_log" => Ok(DocType::DailyLog), + "identity" => Ok(DocType::Identity), + "soul" => Ok(DocType::Soul), + "agents" => Ok(DocType::Agents), + "user" => Ok(DocType::User), + "heartbeat" => Ok(DocType::Heartbeat), + _ => Err(WorkspaceError::InvalidDocType { + doc_type: s.to_string(), + }), + } + } +} + +impl std::fmt::Display for DocType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +/// A memory document stored in the database. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryDocument { + /// Unique document ID. + pub id: Uuid, + /// User identifier. + pub user_id: String, + /// Optional agent ID for multi-agent isolation. + pub agent_id: Option, + /// Document type. + pub doc_type: DocType, + /// Optional title (e.g., date for daily logs). + pub title: Option, + /// Full document content. + pub content: String, + /// Creation timestamp. + pub created_at: DateTime, + /// Last update timestamp. + pub updated_at: DateTime, + /// Flexible metadata. + pub metadata: serde_json::Value, +} + +impl MemoryDocument { + /// Create a new document (not persisted yet). + pub fn new( + user_id: impl Into, + agent_id: Option, + doc_type: DocType, + title: Option, + ) -> Self { + let now = Utc::now(); + Self { + id: Uuid::new_v4(), + user_id: user_id.into(), + agent_id, + doc_type, + title, + content: String::new(), + created_at: now, + updated_at: now, + metadata: serde_json::Value::Object(serde_json::Map::new()), + } + } + + /// Check if the document is empty. + pub fn is_empty(&self) -> bool { + self.content.is_empty() + } + + /// Get word count. + pub fn word_count(&self) -> usize { + self.content.split_whitespace().count() + } +} + +/// A chunk of a memory document for search indexing. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryChunk { + /// Unique chunk ID. + pub id: Uuid, + /// Parent document ID. + pub document_id: Uuid, + /// Position in the document (0-based). + pub chunk_index: i32, + /// Chunk text content. + pub content: String, + /// Embedding vector (if generated). + pub embedding: Option>, + /// Creation timestamp. + pub created_at: DateTime, +} + +impl MemoryChunk { + /// Create a new chunk (not persisted yet). + pub fn new(document_id: Uuid, chunk_index: i32, content: impl Into) -> Self { + Self { + id: Uuid::new_v4(), + document_id, + chunk_index, + content: content.into(), + embedding: None, + created_at: Utc::now(), + } + } + + /// Set the embedding. + pub fn with_embedding(mut self, embedding: Vec) -> Self { + self.embedding = Some(embedding); + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_doc_type_roundtrip() { + for doc_type in [ + DocType::Memory, + DocType::DailyLog, + DocType::Identity, + DocType::Soul, + DocType::Agents, + DocType::User, + DocType::Heartbeat, + ] { + let s = doc_type.as_str(); + let parsed = DocType::try_from(s).unwrap(); + assert_eq!(parsed, doc_type); + } + } + + #[test] + fn test_singleton_types() { + assert!(DocType::Memory.is_singleton()); + assert!(DocType::Heartbeat.is_singleton()); + assert!(!DocType::DailyLog.is_singleton()); + } + + #[test] + fn test_identity_documents() { + assert!(DocType::Soul.is_identity_document()); + assert!(DocType::Agents.is_identity_document()); + assert!(!DocType::Memory.is_identity_document()); + assert!(!DocType::DailyLog.is_identity_document()); + } + + #[test] + fn test_memory_document_word_count() { + let mut doc = MemoryDocument::new("user1", None, DocType::Memory, None); + assert_eq!(doc.word_count(), 0); + + doc.content = "Hello world, this is a test.".to_string(); + assert_eq!(doc.word_count(), 6); + } +} diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs new file mode 100644 index 00000000..5c80d1ed --- /dev/null +++ b/src/workspace/embeddings.rs @@ -0,0 +1,324 @@ +//! Embedding providers for semantic search. +//! +//! Embeddings convert text into dense vectors that capture semantic meaning. +//! Similar concepts have similar vectors, enabling semantic search. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +/// Error type for embedding operations. +#[derive(Debug, thiserror::Error)] +pub enum EmbeddingError { + #[error("HTTP request failed: {0}")] + HttpError(String), + + #[error("Invalid response: {0}")] + InvalidResponse(String), + + #[error("Rate limited, retry after {retry_after:?}")] + RateLimited { + retry_after: Option, + }, + + #[error("Authentication failed")] + AuthFailed, + + #[error("Text too long: {length} > {max}")] + TextTooLong { length: usize, max: usize }, +} + +impl From for EmbeddingError { + fn from(e: reqwest::Error) -> Self { + EmbeddingError::HttpError(e.to_string()) + } +} + +/// Trait for embedding providers. +#[async_trait] +pub trait EmbeddingProvider: Send + Sync { + /// Get the embedding dimension. + fn dimension(&self) -> usize; + + /// Get the model name. + fn model_name(&self) -> &str; + + /// Maximum input length in characters. + fn max_input_length(&self) -> usize; + + /// Generate an embedding for a single text. + async fn embed(&self, text: &str) -> Result, EmbeddingError>; + + /// Generate embeddings for multiple texts (batched). + /// + /// Default implementation calls embed() for each text. + async fn embed_batch(&self, texts: &[String]) -> Result>, EmbeddingError> { + let mut embeddings = Vec::with_capacity(texts.len()); + for text in texts { + embeddings.push(self.embed(text).await?); + } + Ok(embeddings) + } +} + +/// OpenAI embedding provider using text-embedding-ada-002 or text-embedding-3-small. +pub struct OpenAiEmbeddings { + client: reqwest::Client, + api_key: String, + model: String, + dimension: usize, +} + +impl OpenAiEmbeddings { + /// Create a new OpenAI embedding provider with the default model. + /// + /// Uses text-embedding-3-small which has 1536 dimensions. + pub fn new(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: "text-embedding-3-small".to_string(), + dimension: 1536, + } + } + + /// Use text-embedding-ada-002 model. + pub fn ada_002(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: "text-embedding-ada-002".to_string(), + dimension: 1536, + } + } + + /// Use text-embedding-3-large model. + pub fn large(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: "text-embedding-3-large".to_string(), + dimension: 3072, + } + } + + /// Use a custom model with specified dimension. + pub fn with_model( + api_key: impl Into, + model: impl Into, + dimension: usize, + ) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: model.into(), + dimension, + } + } +} + +#[derive(Debug, Serialize)] +struct OpenAiEmbeddingRequest<'a> { + model: &'a str, + input: &'a [String], +} + +#[derive(Debug, Deserialize)] +struct OpenAiEmbeddingResponse { + data: Vec, +} + +#[derive(Debug, Deserialize)] +struct OpenAiEmbeddingData { + embedding: Vec, +} + +#[async_trait] +impl EmbeddingProvider for OpenAiEmbeddings { + fn dimension(&self) -> usize { + self.dimension + } + + fn model_name(&self) -> &str { + &self.model + } + + fn max_input_length(&self) -> usize { + // text-embedding-3-small/large: 8191 tokens (~32k chars) + // text-embedding-ada-002: 8191 tokens + 32_000 + } + + async fn embed(&self, text: &str) -> Result, EmbeddingError> { + if text.len() > self.max_input_length() { + return Err(EmbeddingError::TextTooLong { + length: text.len(), + max: self.max_input_length(), + }); + } + + let embeddings = self.embed_batch(&[text.to_string()]).await?; + embeddings + .into_iter() + .next() + .ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string())) + } + + async fn embed_batch(&self, texts: &[String]) -> Result>, EmbeddingError> { + if texts.is_empty() { + return Ok(Vec::new()); + } + + let request = OpenAiEmbeddingRequest { + model: &self.model, + input: texts, + }; + + let response = self + .client + .post("https://api.openai.com/v1/embeddings") + .header("Authorization", format!("Bearer {}", self.api_key)) + .json(&request) + .send() + .await?; + + let status = response.status(); + + if status == reqwest::StatusCode::UNAUTHORIZED { + return Err(EmbeddingError::AuthFailed); + } + + if status == reqwest::StatusCode::TOO_MANY_REQUESTS { + let retry_after = response + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + .map(std::time::Duration::from_secs); + return Err(EmbeddingError::RateLimited { retry_after }); + } + + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(EmbeddingError::HttpError(format!( + "Status {}: {}", + status, error_text + ))); + } + + let result: OpenAiEmbeddingResponse = response.json().await.map_err(|e| { + EmbeddingError::InvalidResponse(format!("Failed to parse response: {}", e)) + })?; + + Ok(result.data.into_iter().map(|d| d.embedding).collect()) + } +} + +/// A mock embedding provider for testing. +#[cfg(test)] +pub struct MockEmbeddings { + dimension: usize, +} + +#[cfg(test)] +impl MockEmbeddings { + pub fn new(dimension: usize) -> Self { + Self { dimension } + } +} + +#[cfg(test)] +#[async_trait] +impl EmbeddingProvider for MockEmbeddings { + fn dimension(&self) -> usize { + self.dimension + } + + fn model_name(&self) -> &str { + "mock-embedding" + } + + fn max_input_length(&self) -> usize { + 10_000 + } + + async fn embed(&self, text: &str) -> Result, EmbeddingError> { + // Generate a deterministic embedding based on text hash + use std::hash::{Hash, Hasher}; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + text.hash(&mut hasher); + let hash = hasher.finish(); + + let mut embedding = Vec::with_capacity(self.dimension); + let mut seed = hash; + for _ in 0..self.dimension { + // Simple LCG for deterministic random values + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + let value = (seed as f32 / u64::MAX as f32) * 2.0 - 1.0; + embedding.push(value); + } + + // Normalize to unit length + let magnitude: f32 = embedding.iter().map(|x| x * x).sum::().sqrt(); + if magnitude > 0.0 { + for x in &mut embedding { + *x /= magnitude; + } + } + + Ok(embedding) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_mock_embeddings() { + let provider = MockEmbeddings::new(128); + + let embedding = provider.embed("hello world").await.unwrap(); + assert_eq!(embedding.len(), 128); + + // Check normalization (should be unit vector) + let magnitude: f32 = embedding.iter().map(|x| x * x).sum::().sqrt(); + assert!((magnitude - 1.0).abs() < 0.001); + } + + #[tokio::test] + async fn test_mock_embeddings_deterministic() { + let provider = MockEmbeddings::new(64); + + let emb1 = provider.embed("test").await.unwrap(); + let emb2 = provider.embed("test").await.unwrap(); + + // Same input should produce same embedding + assert_eq!(emb1, emb2); + } + + #[tokio::test] + async fn test_mock_embeddings_batch() { + let provider = MockEmbeddings::new(64); + + let texts = vec!["hello".to_string(), "world".to_string()]; + let embeddings = provider.embed_batch(&texts).await.unwrap(); + + assert_eq!(embeddings.len(), 2); + assert_eq!(embeddings[0].len(), 64); + assert_eq!(embeddings[1].len(), 64); + + // Different texts should produce different embeddings + assert_ne!(embeddings[0], embeddings[1]); + } + + #[test] + fn test_openai_embeddings_config() { + let provider = OpenAiEmbeddings::new("test-key"); + assert_eq!(provider.dimension(), 1536); + assert_eq!(provider.model_name(), "text-embedding-3-small"); + + let provider = OpenAiEmbeddings::large("test-key"); + assert_eq!(provider.dimension(), 3072); + assert_eq!(provider.model_name(), "text-embedding-3-large"); + } +} diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs new file mode 100644 index 00000000..e500d648 --- /dev/null +++ b/src/workspace/mod.rs @@ -0,0 +1,392 @@ +//! Workspace and memory system (OpenClaw-inspired). +//! +//! The workspace provides persistent memory for agents: +//! - **MEMORY.md**: Long-term curated memory (facts, decisions, preferences) +//! - **Daily logs**: Append-only daily notes (raw context) +//! - **Identity files**: Agent personality and user context +//! - **HEARTBEAT.md**: Periodic checklist for proactive execution +//! +//! Memory is searchable via hybrid search (FTS + semantic embeddings). +//! +//! # Architecture +//! +//! ```text +//! ┌─────────────────────────────────────────────────────────────┐ +//! │ Workspace │ +//! │ ┌────────────────┐ ┌────────────────┐ ┌──────────────┐ │ +//! │ │ MemoryDocument │ │ MemoryChunk │ │ Search │ │ +//! │ │ (full docs) │──│ (chunked) │──│ (FTS+vector) │ │ +//! │ └────────────────┘ └────────────────┘ └──────────────┘ │ +//! │ │ │ │ │ +//! │ └───────────────────┴──────────────────┘ │ +//! │ │ │ +//! │ ┌──────┴──────┐ │ +//! │ │ Repository │ │ +//! │ │ (PostgreSQL)│ │ +//! │ └─────────────┘ │ +//! └─────────────────────────────────────────────────────────────┘ +//! ``` +//! +//! # Key Patterns +//! +//! 1. **Memory is persistence**: If you want to remember something, write it +//! 2. **Two-tier memory**: Daily logs (raw) + MEMORY.md (curated) +//! 3. **Hybrid search**: Vector similarity + BM25 full-text via RRF + +mod chunker; +mod document; +mod embeddings; +mod repository; +mod search; + +pub use chunker::{ChunkConfig, chunk_document}; +pub use document::{DocType, MemoryChunk, MemoryDocument}; +pub use embeddings::{EmbeddingProvider, OpenAiEmbeddings}; +pub use repository::Repository; +pub use search::{SearchConfig, SearchResult}; + +use std::sync::Arc; + +use chrono::{NaiveDate, Utc}; +use deadpool_postgres::Pool; +use uuid::Uuid; + +use crate::error::WorkspaceError; + +/// Workspace provides database-backed memory storage for an agent. +/// +/// Each workspace is scoped to a user (and optionally an agent). +/// Documents are persisted to PostgreSQL and indexed for search. +pub struct Workspace { + /// User identifier (from channel). + user_id: String, + /// Optional agent ID for multi-agent isolation. + agent_id: Option, + /// Database repository. + repo: Repository, + /// Embedding provider for semantic search. + embeddings: Option>, +} + +impl Workspace { + /// Create a new workspace for a user. + pub fn new(user_id: impl Into, pool: Pool) -> Self { + Self { + user_id: user_id.into(), + agent_id: None, + repo: Repository::new(pool), + embeddings: None, + } + } + + /// Create a workspace with a specific agent ID. + pub fn with_agent(mut self, agent_id: Uuid) -> Self { + self.agent_id = Some(agent_id); + self + } + + /// Set the embedding provider for semantic search. + pub fn with_embeddings(mut self, provider: Arc) -> Self { + self.embeddings = Some(provider); + self + } + + /// Get the user ID. + pub fn user_id(&self) -> &str { + &self.user_id + } + + /// Get the agent ID. + pub fn agent_id(&self) -> Option { + self.agent_id + } + + // ==================== Document Access ==================== + + /// Get the main MEMORY.md document (long-term curated memory). + /// + /// Creates it if it doesn't exist. + pub async fn memory(&self) -> Result { + self.repo + .get_or_create_document(&self.user_id, self.agent_id, DocType::Memory, None) + .await + } + + /// Get today's daily log. + /// + /// Daily logs are append-only and keyed by date. + pub async fn today_log(&self) -> Result { + let today = Utc::now().date_naive(); + self.daily_log(today).await + } + + /// Get a daily log for a specific date. + pub async fn daily_log(&self, date: NaiveDate) -> Result { + let title = date.format("%Y-%m-%d").to_string(); + self.repo + .get_or_create_document( + &self.user_id, + self.agent_id, + DocType::DailyLog, + Some(&title), + ) + .await + } + + /// Get the heartbeat checklist (HEARTBEAT.md). + pub async fn heartbeat_checklist(&self) -> Result, WorkspaceError> { + match self + .repo + .get_document(&self.user_id, self.agent_id, DocType::Heartbeat, None) + .await + { + Ok(doc) => Ok(Some(doc.content)), + Err(WorkspaceError::DocumentNotFound { .. }) => Ok(None), + Err(e) => Err(e), + } + } + + /// Get a document by type. + pub async fn get_document( + &self, + doc_type: DocType, + title: Option<&str>, + ) -> Result { + self.repo + .get_document(&self.user_id, self.agent_id, doc_type, title) + .await + } + + // ==================== Memory Operations ==================== + + /// Append an entry to the main MEMORY.md document. + /// + /// This is for important facts, decisions, and preferences worth + /// remembering long-term. + 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.repo.update_document(doc.id, &new_content).await?; + self.reindex_document(doc.id).await?; + Ok(()) + } + + /// Append an entry to today's daily log. + /// + /// Daily logs are raw, append-only notes for the current day. + pub async fn append_daily_log(&self, entry: &str) -> Result<(), WorkspaceError> { + let doc = self.today_log().await?; + let timestamp = Utc::now().format("%H:%M:%S"); + let timestamped_entry = format!("[{}] {}", timestamp, entry); + + let new_content = if doc.content.is_empty() { + timestamped_entry + } else { + format!("{}\n{}", doc.content, timestamped_entry) + }; + self.repo.update_document(doc.id, &new_content).await?; + self.reindex_document(doc.id).await?; + Ok(()) + } + + /// Update a document's content entirely. + pub async fn update_document( + &self, + doc_type: DocType, + title: Option<&str>, + content: &str, + ) -> Result<(), WorkspaceError> { + let doc = self + .repo + .get_or_create_document(&self.user_id, self.agent_id, doc_type, title) + .await?; + self.repo.update_document(doc.id, content).await?; + self.reindex_document(doc.id).await?; + Ok(()) + } + + // ==================== System Prompt ==================== + + /// Build the system prompt from identity files. + /// + /// Loads AGENTS.md, SOUL.md, USER.md, and IDENTITY.md to compose + /// the agent's system prompt. + pub async fn system_prompt(&self) -> Result { + let mut parts = Vec::new(); + + // Load identity files in order of importance + let identity_types = [ + (DocType::Agents, "## Agent Instructions"), + (DocType::Soul, "## Core Values"), + (DocType::User, "## User Context"), + (DocType::Identity, "## Identity"), + ]; + + for (doc_type, header) in identity_types { + if let Ok(doc) = self + .repo + .get_document(&self.user_id, self.agent_id, doc_type, None) + .await + { + if !doc.content.is_empty() { + parts.push(format!("{}\n\n{}", header, doc.content)); + } + } + } + + // Add today's memory context (last 2 days of daily logs) + 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 { + if !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")) + } + + // ==================== Search ==================== + + /// Hybrid search across all memory documents. + /// + /// Combines full-text search (BM25) with semantic search (vector similarity) + /// using Reciprocal Rank Fusion (RRF). + pub async fn search( + &self, + query: &str, + limit: usize, + ) -> Result, WorkspaceError> { + self.search_with_config(query, SearchConfig::default().with_limit(limit)) + .await + } + + /// Search with custom configuration. + pub async fn search_with_config( + &self, + query: &str, + config: SearchConfig, + ) -> Result, WorkspaceError> { + // Generate embedding for semantic search if provider available + 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.repo + .hybrid_search( + &self.user_id, + self.agent_id, + query, + embedding.as_deref(), + &config, + ) + .await + } + + // ==================== Indexing ==================== + + /// Re-index a document (chunk and generate embeddings). + async fn reindex_document(&self, document_id: Uuid) -> Result<(), WorkspaceError> { + // Get the document + let doc = self.repo.get_document_by_id(document_id).await?; + + // Chunk the content + let chunks = chunk_document(&doc.content, ChunkConfig::default()); + + // Delete old chunks + self.repo.delete_chunks(document_id).await?; + + // Insert new chunks + for (index, content) in chunks.into_iter().enumerate() { + // Generate embedding if provider available + 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.repo + .insert_chunk(document_id, index as i32, &content, embedding.as_deref()) + .await?; + } + + Ok(()) + } + + /// Generate embeddings for chunks that don't have them yet. + /// + /// This is useful for backfilling embeddings after enabling the provider. + pub async fn backfill_embeddings(&self) -> Result { + let Some(ref provider) = self.embeddings else { + return Ok(0); + }; + + let chunks = self + .repo + .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.repo + .update_chunk_embedding(chunk.id, &embedding) + .await?; + count += 1; + } + Err(e) => { + tracing::warn!("Failed to embed chunk {}: {}", chunk.id, e); + } + } + } + + Ok(count) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_doc_type_display() { + assert_eq!(DocType::Memory.as_str(), "memory"); + assert_eq!(DocType::DailyLog.as_str(), "daily_log"); + assert_eq!(DocType::Heartbeat.as_str(), "heartbeat"); + } + + #[test] + fn test_doc_type_parse() { + assert_eq!(DocType::try_from("memory").unwrap(), DocType::Memory); + assert_eq!(DocType::try_from("daily_log").unwrap(), DocType::DailyLog); + assert!(DocType::try_from("invalid").is_err()); + } +} diff --git a/src/workspace/repository.rs b/src/workspace/repository.rs new file mode 100644 index 00000000..61e9c912 --- /dev/null +++ b/src/workspace/repository.rs @@ -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 { + 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, + doc_type: DocType, + title: Option<&str>, + ) -> Result { + 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 { + 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, + doc_type: DocType, + title: Option<&str>, + ) -> Result { + // 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, + doc_type: Option, + ) -> Result, 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 { + 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 { + 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, + limit: usize, + ) -> Result, 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, + query: &str, + embedding: Option<&[f32]>, + config: &SearchConfig, + ) -> Result, 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, + query: &str, + limit: usize, + ) -> Result, 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, + embedding: &[f32], + limit: usize, + ) -> Result, 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()) + } +} diff --git a/src/workspace/search.rs b/src/workspace/search.rs new file mode 100644 index 00000000..7cee5539 --- /dev/null +++ b/src/workspace/search.rs @@ -0,0 +1,391 @@ +//! Hybrid search combining full-text and semantic search. +//! +//! Uses Reciprocal Rank Fusion (RRF) to combine results from: +//! 1. PostgreSQL full-text search (ts_rank_cd) +//! 2. pgvector cosine similarity search +//! +//! RRF formula: score = sum(1 / (k + rank)) for each retrieval method +//! This is robust to different score scales and produces better results +//! than simple score averaging. + +use std::collections::HashMap; + +use uuid::Uuid; + +/// Configuration for hybrid search. +#[derive(Debug, Clone)] +pub struct SearchConfig { + /// Maximum number of results to return. + pub limit: usize, + /// RRF constant (typically 60). Higher values favor top results more. + pub rrf_k: u32, + /// Whether to include FTS results. + pub use_fts: bool, + /// Whether to include vector results. + pub use_vector: bool, + /// Minimum score threshold (0.0-1.0). + pub min_score: f32, + /// Maximum results to fetch from each method before fusion. + pub pre_fusion_limit: usize, +} + +impl Default for SearchConfig { + fn default() -> Self { + Self { + limit: 10, + rrf_k: 60, + use_fts: true, + use_vector: true, + min_score: 0.0, + pre_fusion_limit: 50, + } + } +} + +impl SearchConfig { + /// Set the result limit. + pub fn with_limit(mut self, limit: usize) -> Self { + self.limit = limit; + self + } + + /// Set the RRF constant. + pub fn with_rrf_k(mut self, k: u32) -> Self { + self.rrf_k = k; + self + } + + /// Disable FTS (only use vector search). + pub fn vector_only(mut self) -> Self { + self.use_fts = false; + self.use_vector = true; + self + } + + /// Disable vector search (only use FTS). + pub fn fts_only(mut self) -> Self { + self.use_fts = true; + self.use_vector = false; + self + } + + /// Set minimum score threshold. + pub fn with_min_score(mut self, score: f32) -> Self { + self.min_score = score.clamp(0.0, 1.0); + self + } +} + +/// A search result with hybrid scoring. +#[derive(Debug, Clone)] +pub struct SearchResult { + /// Document ID containing this chunk. + pub document_id: Uuid, + /// Chunk ID. + pub chunk_id: Uuid, + /// Chunk content. + pub content: String, + /// Combined RRF score (0.0-1.0 normalized). + pub score: f32, + /// Rank in FTS results (1-based, None if not in FTS results). + pub fts_rank: Option, + /// Rank in vector results (1-based, None if not in vector results). + pub vector_rank: Option, +} + +impl SearchResult { + /// Check if this result came from FTS. + pub fn from_fts(&self) -> bool { + self.fts_rank.is_some() + } + + /// Check if this result came from vector search. + pub fn from_vector(&self) -> bool { + self.vector_rank.is_some() + } + + /// Check if this result came from both methods (hybrid match). + pub fn is_hybrid(&self) -> bool { + self.fts_rank.is_some() && self.vector_rank.is_some() + } +} + +/// Raw result from a single search method. +#[derive(Debug, Clone)] +pub struct RankedResult { + pub chunk_id: Uuid, + pub document_id: Uuid, + pub content: String, + pub rank: u32, // 1-based rank +} + +/// Reciprocal Rank Fusion algorithm. +/// +/// Combines ranked results from multiple retrieval methods using the formula: +/// score(d) = sum(1 / (k + rank(d))) for each method where d appears +/// +/// # Arguments +/// +/// * `fts_results` - Results from full-text search, ordered by relevance +/// * `vector_results` - Results from vector search, ordered by similarity +/// * `config` - Search configuration +/// +/// # Returns +/// +/// Combined results sorted by RRF score (descending). +pub fn reciprocal_rank_fusion( + fts_results: Vec, + vector_results: Vec, + config: &SearchConfig, +) -> Vec { + let k = config.rrf_k as f32; + + // Track scores and metadata for each chunk + struct ChunkInfo { + document_id: Uuid, + content: String, + score: f32, + fts_rank: Option, + vector_rank: Option, + } + + let mut chunk_scores: HashMap = HashMap::new(); + + // Process FTS results + for result in fts_results { + let rrf_score = 1.0 / (k + result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += rrf_score; + info.fts_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + content: result.content, + score: rrf_score, + fts_rank: Some(result.rank), + vector_rank: None, + }); + } + + // Process vector results + for result in vector_results { + let rrf_score = 1.0 / (k + result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += rrf_score; + info.vector_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + content: result.content, + score: rrf_score, + fts_rank: None, + vector_rank: Some(result.rank), + }); + } + + // Convert to SearchResult and sort by score + let mut results: Vec = chunk_scores + .into_iter() + .map(|(chunk_id, info)| SearchResult { + document_id: info.document_id, + chunk_id, + content: info.content, + score: info.score, + fts_rank: info.fts_rank, + vector_rank: info.vector_rank, + }) + .collect(); + + // Normalize scores to 0-1 range + if let Some(max_score) = results.iter().map(|r| r.score).reduce(f32::max) { + if max_score > 0.0 { + for result in &mut results { + result.score /= max_score; + } + } + } + + // Filter by minimum score + if config.min_score > 0.0 { + results.retain(|r| r.score >= config.min_score); + } + + // Sort by score descending + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + // Limit results + results.truncate(config.limit); + + results +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_result(chunk_id: Uuid, doc_id: Uuid, rank: u32) -> RankedResult { + RankedResult { + chunk_id, + document_id: doc_id, + content: format!("content for chunk {}", chunk_id), + rank, + } + } + + #[test] + fn test_rrf_single_method() { + let config = SearchConfig::default().with_limit(10); + + let chunk1 = Uuid::new_v4(); + let chunk2 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)]; + + let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config); + + assert_eq!(results.len(), 2); + // First result should have higher score + assert!(results[0].score > results[1].score); + // All should have FTS rank + assert!(results.iter().all(|r| r.fts_rank.is_some())); + assert!(results.iter().all(|r| r.vector_rank.is_none())); + } + + #[test] + fn test_rrf_hybrid_match_boosted() { + let config = SearchConfig::default().with_limit(10); + + let chunk1 = Uuid::new_v4(); // In both + let chunk2 = Uuid::new_v4(); // FTS only + let chunk3 = Uuid::new_v4(); // Vector only + let doc = Uuid::new_v4(); + + let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)]; + + let vector_results = vec![make_result(chunk1, doc, 1), make_result(chunk3, doc, 2)]; + + let results = reciprocal_rank_fusion(fts_results, vector_results, &config); + + assert_eq!(results.len(), 3); + + // chunk1 should be first (hybrid match) + assert_eq!(results[0].chunk_id, chunk1); + assert!(results[0].is_hybrid()); + assert!(results[0].score > results[1].score); + + // Other chunks should not be hybrid + assert!(!results[1].is_hybrid()); + assert!(!results[2].is_hybrid()); + } + + #[test] + fn test_rrf_score_normalization() { + let config = SearchConfig::default(); + + let chunk1 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts_results = vec![make_result(chunk1, doc, 1)]; + + let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config); + + // Single result should have normalized score of 1.0 + assert_eq!(results.len(), 1); + assert!((results[0].score - 1.0).abs() < 0.001); + } + + #[test] + fn test_rrf_min_score_filter() { + let config = SearchConfig::default().with_limit(10).with_min_score(0.5); + + let chunk1 = Uuid::new_v4(); + let chunk2 = Uuid::new_v4(); + let chunk3 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + // chunk1 has rank 1, chunk3 has rank 100 (low score) + let fts_results = vec![ + make_result(chunk1, doc, 1), + make_result(chunk2, doc, 50), + make_result(chunk3, doc, 100), + ]; + + let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config); + + // Low-scoring results should be filtered out + // All results should have score >= 0.5 + for result in &results { + assert!(result.score >= 0.5); + } + } + + #[test] + fn test_rrf_limit() { + let config = SearchConfig::default().with_limit(2); + + let doc = Uuid::new_v4(); + let fts_results: Vec<_> = (1..=5) + .map(|i| make_result(Uuid::new_v4(), doc, i)) + .collect(); + + let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config); + + assert_eq!(results.len(), 2); + } + + #[test] + fn test_rrf_k_parameter() { + // Higher k values make ranking differences less pronounced + let chunk1 = Uuid::new_v4(); + let chunk2 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)]; + + // Low k: rank 1 score = 1/(10+1) = 0.091, rank 2 = 1/(10+2) = 0.083 + let config_low_k = SearchConfig::default().with_rrf_k(10); + let results_low = reciprocal_rank_fusion(fts_results.clone(), Vec::new(), &config_low_k); + + // High k: rank 1 score = 1/(100+1) = 0.0099, rank 2 = 1/(100+2) = 0.0098 + let config_high_k = SearchConfig::default().with_rrf_k(100); + let results_high = reciprocal_rank_fusion(fts_results, Vec::new(), &config_high_k); + + // With low k, the score difference is larger (relatively) + let diff_low = results_low[0].score - results_low[1].score; + let diff_high = results_high[0].score - results_high[1].score; + + // Low k should have larger relative difference + assert!(diff_low > diff_high); + } + + #[test] + fn test_search_config_builders() { + let config = SearchConfig::default() + .with_limit(20) + .with_rrf_k(30) + .with_min_score(0.1); + + assert_eq!(config.limit, 20); + assert_eq!(config.rrf_k, 30); + assert!((config.min_score - 0.1).abs() < 0.001); + assert!(config.use_fts); + assert!(config.use_vector); + + let fts_only = SearchConfig::default().fts_only(); + assert!(fts_only.use_fts); + assert!(!fts_only.use_vector); + + let vector_only = SearchConfig::default().vector_only(); + assert!(!vector_only.use_fts); + assert!(vector_only.use_vector); + } +}