diff --git a/.env.example b/.env.example index e08f40d3..85eb9a97 100644 --- a/.env.example +++ b/.env.example @@ -2,25 +2,12 @@ DATABASE_URL=postgres://near_agent:password@localhost:5432/near_agent DATABASE_POOL_SIZE=10 -# LLM Providers -# Default is NEAR AI which provides a unified interface to all models - -# NEAR AI (recommended - unified API with user authentication) +# LLM Provider (NEAR AI) +# NEAR AI provides a unified interface to all models with user authentication NEARAI_SESSION_TOKEN=sess_... NEARAI_MODEL=claude-3-5-sonnet-20241022 NEARAI_BASE_URL=https://api.near.ai -# OpenAI (alternative) -OPENAI_API_KEY=sk-... -OPENAI_MODEL=gpt-4-turbo-preview - -# Anthropic (alternative) -ANTHROPIC_API_KEY=sk-ant-... -ANTHROPIC_MODEL=claude-3-opus-20240229 - -# Default LLM provider: nearai, openai, or anthropic -LLM_PROVIDER=nearai - # Channel Configuration # CLI is always enabled diff --git a/CLAUDE.md b/CLAUDE.md index 577e6133..8db5377c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -53,11 +53,9 @@ src/ │ ├── validator.rs # Input validation (length, encoding, patterns) │ └── policy.rs # PolicyRule system with severity/actions │ -├── llm/ # LLM integration +├── llm/ # LLM integration (NEAR AI only) │ ├── provider.rs # LlmProvider trait, message types -│ ├── nearai.rs # NEAR AI chat-api (default, unified interface) -│ ├── openai.rs # OpenAI API implementation -│ ├── anthropic.rs # Anthropic API implementation +│ ├── nearai.rs # NEAR AI chat-api implementation │ └── reasoning.rs # Planning, tool selection, evaluation │ ├── tools/ # Extensible tool system @@ -76,7 +74,7 @@ src/ │ ├── workspace/ # Persistent memory system (OpenClaw-inspired) │ ├── mod.rs # Workspace struct, memory operations -│ ├── document.rs # DocType enum, MemoryDocument, MemoryChunk +│ ├── document.rs # MemoryDocument, MemoryChunk, WorkspaceEntry │ ├── chunker.rs # Document chunking (800 tokens, 15% overlap) │ ├── embeddings.rs # EmbeddingProvider trait, OpenAI implementation │ ├── search.rs # Hybrid search with RRF algorithm @@ -164,22 +162,11 @@ Environment variables (see `.env.example`): ```bash DATABASE_URL=postgres://user:pass@localhost/near_agent -# LLM Provider (default: nearai) -LLM_PROVIDER=nearai # Options: nearai, openai, anthropic - -# NEAR AI (recommended - unified API with user auth) +# NEAR AI (required) NEARAI_SESSION_TOKEN=sess_... NEARAI_MODEL=claude-3-5-sonnet-20241022 NEARAI_BASE_URL=https://api.near.ai -# OpenAI (alternative) -OPENAI_API_KEY=sk-... -OPENAI_MODEL=gpt-4-turbo - -# Anthropic (alternative) -ANTHROPIC_API_KEY=sk-ant-... -ANTHROPIC_MODEL=claude-3-opus-20240229 - # Agent settings AGENT_NAME=near-agent MAX_PARALLEL_JOBS=5 @@ -187,7 +174,7 @@ MAX_PARALLEL_JOBS=5 ### NEAR AI Provider -The default provider uses the NEAR AI chat-api (`https://api.near.ai/v1/responses`) which provides: +Uses the NEAR AI chat-api (`https://api.near.ai/v1/responses`) which provides: - Unified access to multiple models (OpenAI, Anthropic, etc.) - User authentication via session tokens - Usage tracking and billing through NEAR AI @@ -196,9 +183,9 @@ Session tokens have the format `sess_xxx` (37 characters). They are authenticate ## Database -Migrations in `migrations/`. Tables: +Single migration in `migrations/V1__initial.sql`. Tables: -**V1 (initial):** +**Core:** - `conversations` - Multi-channel conversation tracking - `agent_jobs` - Job metadata and status - `job_actions` - Event-sourced tool executions @@ -206,8 +193,8 @@ Migrations in `migrations/`. Tables: - `llm_calls` - Cost tracking - `estimation_snapshots` - Learning data -**V2 (workspace/memory):** -- `memory_documents` - Full documents (MEMORY.md, daily logs, identity files) +**Workspace/Memory:** +- `memory_documents` - Flexible path-based files (e.g., "context/vision.md", "daily/2024-01-15.md") - `memory_chunks` - Chunked content with FTS (tsvector) and vector (pgvector) indexes - `heartbeat_state` - Periodic execution tracking @@ -292,38 +279,59 @@ RUST_LOG=near_agent=debug,tower_http=debug cargo run ## Workspace & Memory System -Inspired by [OpenClaw](https://github.com/openclaw/openclaw), the workspace provides persistent memory for agents. +Inspired by [OpenClaw](https://github.com/openclaw/openclaw), the workspace provides persistent memory for agents with a flexible filesystem-like structure. ### Key Principles 1. **"Memory is files, not RAM"** - If you want to remember something, write it explicitly -2. **Two-tier memory** - Daily logs (raw) + curated MEMORY.md (distilled wisdom) -3. **Hybrid search** - Combines FTS (keyword) + vector (semantic) via Reciprocal Rank Fusion +2. **Flexible structure** - Create any directory/file hierarchy you need +3. **Self-documenting** - Use README.md files to describe directory structure +4. **Hybrid search** - Combines FTS (keyword) + vector (semantic) via Reciprocal Rank Fusion -### Document Types +### Filesystem Structure -| Type | Purpose | Singleton | -|------|---------|-----------| -| `Memory` | Long-term curated facts (MEMORY.md) | Yes | -| `DailyLog` | Append-only daily notes (keyed by date) | No | -| `Identity` | Agent name, nature, vibe | Yes | -| `Soul` | Core values and principles | Yes | -| `Agents` | Behavior instructions | Yes | -| `User` | User context (name, preferences) | Yes | -| `Heartbeat` | Periodic checklist | Yes | +``` +workspace/ +├── README.md <- Root runbook/index +├── MEMORY.md <- Long-term curated memory +├── HEARTBEAT.md <- Periodic checklist +├── IDENTITY.md <- Agent name, nature, vibe +├── SOUL.md <- Core values +├── AGENTS.md <- Behavior instructions +├── USER.md <- User context +├── context/ <- Identity-related docs +│ ├── vision.md +│ └── priorities.md +├── daily/ <- Daily logs +│ ├── 2024-01-15.md +│ └── 2024-01-16.md +├── projects/ <- Arbitrary structure +│ └── alpha/ +│ ├── README.md +│ └── notes.md +└── ... +``` ### Using the Workspace ```rust -use crate::workspace::{Workspace, DocType, OpenAiEmbeddings}; +use crate::workspace::{Workspace, OpenAiEmbeddings, paths}; // Create workspace for a user let workspace = Workspace::new("user_123", pool) .with_embeddings(Arc::new(OpenAiEmbeddings::new(api_key))); -// Write to memory +// Read/write any path +let doc = workspace.read("projects/alpha/notes.md").await?; +workspace.write("context/priorities.md", "# Priorities\n\n1. Feature X").await?; +workspace.append("daily/2024-01-15.md", "Completed task X").await?; + +// Convenience methods for well-known files workspace.append_memory("User prefers dark mode").await?; -workspace.append_daily_log("Completed task X").await?; +workspace.append_daily_log("Session note").await?; + +// List directory contents +let entries = workspace.list("projects/").await?; // Search (hybrid FTS + vector) let results = workspace.search("dark mode preference", 5).await?; @@ -334,11 +342,12 @@ let prompt = workspace.system_prompt().await?; ### Memory Tools -Three tools for LLM use: +Four tools for LLM use: - **`memory_search`** - Hybrid search, MUST be called before answering questions about prior work -- **`memory_write`** - Write to memory or daily_log target -- **`memory_read`** - Read specific document by type +- **`memory_write`** - Write to any path (memory, daily_log, or custom paths) +- **`memory_read`** - Read any file by path +- **`memory_list`** - List directory contents ### Hybrid Search (RRF) diff --git a/migrations/V1__initial.sql b/migrations/V1__initial.sql index 21dba82d..ea04b9b8 100644 --- a/migrations/V1__initial.sql +++ b/migrations/V1__initial.sql @@ -1,7 +1,12 @@ -- NEAR Agent Database Schema --- V1: Initial schema +-- V1: Complete schema with workspace and memory system + +-- Enable pgvector extension for semantic search +-- NOTE: Requires pgvector to be installed on PostgreSQL server +CREATE EXTENSION IF NOT EXISTS vector; + +-- ==================== Conversations ==================== --- Conversations from various channels CREATE TABLE conversations ( id UUID PRIMARY KEY, channel TEXT NOT NULL, @@ -16,7 +21,6 @@ CREATE INDEX idx_conversations_channel ON conversations(channel); CREATE INDEX idx_conversations_user ON conversations(user_id); CREATE INDEX idx_conversations_last_activity ON conversations(last_activity); --- Messages in conversations CREATE TABLE conversation_messages ( id UUID PRIMARY KEY, conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, @@ -27,7 +31,8 @@ CREATE TABLE conversation_messages ( CREATE INDEX idx_conversation_messages_conversation ON conversation_messages(conversation_id); --- Jobs we've worked on +-- ==================== Agent Jobs ==================== + CREATE TABLE agent_jobs ( id UUID PRIMARY KEY, marketplace_job_id UUID, @@ -59,7 +64,6 @@ CREATE INDEX idx_agent_jobs_marketplace ON agent_jobs(marketplace_job_id); CREATE INDEX idx_agent_jobs_conversation ON agent_jobs(conversation_id); CREATE INDEX idx_agent_jobs_stuck ON agent_jobs(stuck_since) WHERE stuck_since IS NOT NULL; --- Actions taken during job execution (event sourcing) CREATE TABLE job_actions ( id UUID PRIMARY KEY, job_id UUID NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE, @@ -80,7 +84,8 @@ CREATE TABLE job_actions ( CREATE INDEX idx_job_actions_job_id ON job_actions(job_id); CREATE INDEX idx_job_actions_tool ON job_actions(tool_name); --- Dynamic tools built by the agent +-- ==================== Dynamic Tools ==================== + CREATE TABLE dynamic_tools ( id UUID PRIMARY KEY, name TEXT NOT NULL UNIQUE, @@ -100,7 +105,8 @@ CREATE TABLE dynamic_tools ( CREATE INDEX idx_dynamic_tools_status ON dynamic_tools(status); CREATE INDEX idx_dynamic_tools_name ON dynamic_tools(name); --- LLM calls for cost tracking +-- ==================== LLM Calls ==================== + CREATE TABLE llm_calls ( id UUID PRIMARY KEY, job_id UUID REFERENCES agent_jobs(id) ON DELETE CASCADE, @@ -118,7 +124,8 @@ CREATE INDEX idx_llm_calls_job ON llm_calls(job_id); CREATE INDEX idx_llm_calls_conversation ON llm_calls(conversation_id); CREATE INDEX idx_llm_calls_provider ON llm_calls(provider); --- Estimation history for continuous learning +-- ==================== Estimation ==================== + CREATE TABLE estimation_snapshots ( id UUID PRIMARY KEY, job_id UUID NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE, @@ -136,7 +143,8 @@ CREATE TABLE estimation_snapshots ( CREATE INDEX idx_estimation_category ON estimation_snapshots(category); CREATE INDEX idx_estimation_job ON estimation_snapshots(job_id); --- Self-repair history +-- ==================== Self Repair ==================== + CREATE TABLE repair_attempts ( id UUID PRIMARY KEY, target_type TEXT NOT NULL, @@ -150,3 +158,193 @@ CREATE TABLE repair_attempts ( CREATE INDEX idx_repair_attempts_target ON repair_attempts(target_type, target_id); CREATE INDEX idx_repair_attempts_created ON repair_attempts(created_at); + +-- ==================== Workspace: Memory Documents ==================== +-- Flexible filesystem-like structure for agent memory. +-- Agents can create arbitrary paths like: +-- "README.md", "context/vision.md", "daily/2024-01-15.md", "projects/alpha/notes.md" + +CREATE TABLE memory_documents ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + agent_id UUID, -- NULL = shared across all agents for this user + + -- File path within workspace (e.g., "context/vision.md") + path TEXT NOT NULL, + content TEXT NOT NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + metadata JSONB NOT NULL DEFAULT '{}', + + CONSTRAINT unique_path_per_user UNIQUE (user_id, agent_id, path) +); + +CREATE INDEX idx_memory_documents_user ON memory_documents(user_id); +CREATE INDEX idx_memory_documents_path ON memory_documents(user_id, path); +CREATE INDEX idx_memory_documents_path_prefix ON memory_documents(user_id, path text_pattern_ops); +CREATE INDEX idx_memory_documents_updated ON memory_documents(updated_at DESC); + +-- ==================== Workspace: Memory Chunks ==================== +-- Documents are chunked for hybrid search (FTS + vector) + +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_index INT NOT NULL, + content TEXT NOT NULL, + + -- Full-text search vector + content_tsv TSVECTOR GENERATED ALWAYS AS (to_tsvector('english', content)) STORED, + + -- Semantic search embedding (text-embedding-3-small = 1536 dims) + embedding VECTOR(1536), + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT unique_chunk_per_doc UNIQUE (document_id, chunk_index) +); + +CREATE INDEX idx_memory_chunks_tsv ON memory_chunks USING GIN(content_tsv); +CREATE INDEX idx_memory_chunks_embedding ON memory_chunks + USING hnsw(embedding vector_cosine_ops) + WITH (m = 16, ef_construction = 64); +CREATE INDEX idx_memory_chunks_document ON memory_chunks(document_id); + +-- ==================== Workspace: Heartbeat State ==================== + +CREATE TABLE heartbeat_state ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + agent_id UUID, + last_run TIMESTAMPTZ, + next_run TIMESTAMPTZ, + interval_seconds INT NOT NULL DEFAULT 1800, + enabled BOOLEAN NOT NULL DEFAULT true, + consecutive_failures INT NOT NULL DEFAULT 0, + last_checks JSONB NOT NULL DEFAULT '{}', + 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 ==================== + +CREATE OR REPLACE FUNCTION update_updated_at_column() +RETURNS TRIGGER AS $$ +BEGIN + NEW.updated_at = NOW(); + RETURN NEW; +END; +$$ language 'plpgsql'; + +CREATE TRIGGER update_memory_documents_updated_at + BEFORE UPDATE ON memory_documents + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); + +-- Function to list files in a directory (prefix match) +CREATE OR REPLACE FUNCTION list_workspace_files( + p_user_id TEXT, + p_agent_id UUID, + p_directory TEXT DEFAULT '' +) +RETURNS TABLE ( + path TEXT, + is_directory BOOLEAN, + updated_at TIMESTAMPTZ, + content_preview TEXT +) AS $$ +BEGIN + -- Normalize directory path (ensure trailing slash for non-root) + IF p_directory != '' AND NOT p_directory LIKE '%/' THEN + p_directory := p_directory || '/'; + END IF; + + RETURN QUERY + WITH files AS ( + SELECT + d.path, + d.updated_at, + LEFT(d.content, 200) as content_preview, + -- Extract the immediate child name + CASE + WHEN p_directory = '' THEN + CASE + WHEN position('/' in d.path) > 0 + THEN substring(d.path from 1 for position('/' in d.path) - 1) + ELSE d.path + END + ELSE + CASE + WHEN position('/' in substring(d.path from length(p_directory) + 1)) > 0 + THEN substring( + substring(d.path from length(p_directory) + 1) + from 1 + for position('/' in substring(d.path from length(p_directory) + 1)) - 1 + ) + ELSE substring(d.path from length(p_directory) + 1) + END + END as child_name + FROM memory_documents d + WHERE d.user_id = p_user_id + AND d.agent_id IS NOT DISTINCT FROM p_agent_id + AND (p_directory = '' OR d.path LIKE p_directory || '%') + ) + SELECT DISTINCT ON (f.child_name) + CASE + WHEN p_directory = '' THEN f.child_name + ELSE p_directory || f.child_name + END as path, + EXISTS ( + SELECT 1 FROM memory_documents d2 + WHERE d2.user_id = p_user_id + AND d2.agent_id IS NOT DISTINCT FROM p_agent_id + AND d2.path LIKE + CASE WHEN p_directory = '' THEN f.child_name ELSE p_directory || f.child_name END + || '/%' + ) as is_directory, + MAX(f.updated_at) as updated_at, + CASE + WHEN EXISTS ( + SELECT 1 FROM memory_documents d2 + WHERE d2.user_id = p_user_id + AND d2.agent_id IS NOT DISTINCT FROM p_agent_id + AND d2.path LIKE + CASE WHEN p_directory = '' THEN f.child_name ELSE p_directory || f.child_name END + || '/%' + ) THEN NULL + ELSE MAX(f.content_preview) + END as content_preview + FROM files f + WHERE f.child_name != '' AND f.child_name IS NOT NULL + GROUP BY f.child_name + ORDER BY f.child_name, is_directory DESC; +END; +$$ LANGUAGE plpgsql; + +-- ==================== Views ==================== + +CREATE VIEW memory_documents_summary AS +SELECT + d.id, + d.user_id, + d.path, + 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; + +CREATE VIEW chunks_pending_embedding AS +SELECT + c.id as chunk_id, + c.document_id, + d.user_id, + d.path, + 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/migrations/V2__workspace_memory.sql b/migrations/V2__workspace_memory.sql deleted file mode 100644 index 7c690ed5..00000000 --- a/migrations/V2__workspace_memory.sql +++ /dev/null @@ -1,159 +0,0 @@ --- 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/config.rs b/src/config.rs index 6fb6f614..1123ea28 100644 --- a/src/config.rs +++ b/src/config.rs @@ -60,52 +60,10 @@ impl DatabaseConfig { } } -/// LLM provider configuration. +/// LLM provider configuration (NEAR AI only). #[derive(Debug, Clone)] pub struct LlmConfig { - pub provider: LlmProvider, - pub openai: Option, - pub anthropic: Option, - pub nearai: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum LlmProvider { - OpenAi, - Anthropic, - NearAi, -} - -impl std::str::FromStr for LlmProvider { - type Err = ConfigError; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "openai" => Ok(Self::OpenAi), - "anthropic" => Ok(Self::Anthropic), - "nearai" | "near-ai" | "near_ai" => Ok(Self::NearAi), - _ => Err(ConfigError::InvalidValue { - key: "LLM_PROVIDER".to_string(), - message: format!( - "unknown provider: {s}, expected 'openai', 'anthropic', or 'nearai'" - ), - }), - } - } -} - -#[derive(Debug, Clone)] -pub struct OpenAiConfig { - pub api_key: SecretString, - pub model: String, - pub base_url: Option, -} - -#[derive(Debug, Clone)] -pub struct AnthropicConfig { - pub api_key: SecretString, - pub model: String, - pub base_url: Option, + pub nearai: NearAiConfig, } /// NEAR AI chat-api configuration. @@ -121,66 +79,16 @@ pub struct NearAiConfig { impl LlmConfig { fn from_env() -> Result { - let provider: LlmProvider = optional_env("LLM_PROVIDER")? - .map(|s| s.parse()) - .transpose()? - .unwrap_or(LlmProvider::NearAi); + let session_token = required_env("NEARAI_SESSION_TOKEN")?; - let openai = if let Some(api_key) = optional_env("OPENAI_API_KEY")? { - Some(OpenAiConfig { - api_key: SecretString::from(api_key), - model: optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4-turbo".to_string()), - base_url: optional_env("OPENAI_BASE_URL")?, - }) - } else { - None - }; - - let anthropic = if let Some(api_key) = optional_env("ANTHROPIC_API_KEY")? { - Some(AnthropicConfig { - api_key: SecretString::from(api_key), - model: optional_env("ANTHROPIC_MODEL")? - .unwrap_or_else(|| "claude-3-opus-20240229".to_string()), - base_url: optional_env("ANTHROPIC_BASE_URL")?, - }) - } else { - None - }; - - let nearai = if let Some(session_token) = optional_env("NEARAI_SESSION_TOKEN")? { - Some(NearAiConfig { + Ok(Self { + nearai: NearAiConfig { session_token: SecretString::from(session_token), model: optional_env("NEARAI_MODEL")? .unwrap_or_else(|| "claude-3-5-sonnet-20241022".to_string()), base_url: optional_env("NEARAI_BASE_URL")? .unwrap_or_else(|| "https://api.near.ai".to_string()), - }) - } else { - None - }; - - // Validate that the selected provider has configuration - match provider { - LlmProvider::OpenAi if openai.is_none() => { - return Err(ConfigError::MissingEnvVar("OPENAI_API_KEY".to_string())); - } - LlmProvider::Anthropic if anthropic.is_none() => { - return Err(ConfigError::MissingEnvVar("ANTHROPIC_API_KEY".to_string())); - } - LlmProvider::NearAi if nearai.is_none() => { - return Err(ConfigError::MissingEnvVar( - "NEARAI_SESSION_TOKEN".to_string(), - )); - } - // Provider has valid configuration - LlmProvider::OpenAi | LlmProvider::Anthropic | LlmProvider::NearAi => {} - } - - Ok(Self { - provider, - openai, - anthropic, - nearai, + }, }) } } @@ -354,37 +262,3 @@ where .transpose() .map(|opt| opt.unwrap_or(default)) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_llm_provider_parsing() { - assert_eq!( - "openai".parse::().unwrap(), - LlmProvider::OpenAi - ); - assert_eq!( - "anthropic".parse::().unwrap(), - LlmProvider::Anthropic - ); - assert_eq!( - "OpenAI".parse::().unwrap(), - LlmProvider::OpenAi - ); - assert_eq!( - "nearai".parse::().unwrap(), - LlmProvider::NearAi - ); - assert_eq!( - "near-ai".parse::().unwrap(), - LlmProvider::NearAi - ); - assert_eq!( - "near_ai".parse::().unwrap(), - LlmProvider::NearAi - ); - assert!("invalid".parse::().is_err()); - } -} diff --git a/src/lib.rs b/src/lib.rs index 990a9c87..b755e745 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -63,5 +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}; + pub use crate::workspace::{MemoryDocument, Workspace}; } diff --git a/src/llm/anthropic.rs b/src/llm/anthropic.rs deleted file mode 100644 index da9d4e56..00000000 --- a/src/llm/anthropic.rs +++ /dev/null @@ -1,348 +0,0 @@ -//! Anthropic LLM provider implementation. - -use async_trait::async_trait; -use reqwest::Client; -use rust_decimal::Decimal; -use rust_decimal_macros::dec; -use secrecy::ExposeSecret; -use serde::{Deserialize, Serialize}; - -use crate::config::AnthropicConfig; -use crate::error::LlmError; -use crate::llm::provider::{ - ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, - ToolCompletionRequest, ToolCompletionResponse, -}; - -/// Anthropic API provider. -pub struct AnthropicProvider { - client: Client, - config: AnthropicConfig, - base_url: String, -} - -impl AnthropicProvider { - /// Create a new Anthropic provider. - pub fn new(config: AnthropicConfig) -> Self { - let base_url = config - .base_url - .clone() - .unwrap_or_else(|| "https://api.anthropic.com/v1".to_string()); - - Self { - client: Client::new(), - config, - base_url, - } - } - - fn build_messages(&self, messages: &[ChatMessage]) -> (Option, Vec) { - let mut system_message = None; - let mut anthropic_messages = Vec::new(); - - for msg in messages { - match msg.role { - Role::System => { - // Anthropic uses a separate system parameter - system_message = Some(msg.content.clone()); - } - Role::User => { - anthropic_messages.push(AnthropicMessage { - role: "user".to_string(), - content: AnthropicContent::Text(msg.content.clone()), - }); - } - Role::Assistant => { - anthropic_messages.push(AnthropicMessage { - role: "assistant".to_string(), - content: AnthropicContent::Text(msg.content.clone()), - }); - } - Role::Tool => { - // Tool results in Anthropic format - anthropic_messages.push(AnthropicMessage { - role: "user".to_string(), - content: AnthropicContent::ToolResult { - tool_use_id: msg.tool_call_id.clone().unwrap_or_default(), - content: msg.content.clone(), - }, - }); - } - } - } - - (system_message, anthropic_messages) - } -} - -#[derive(Debug, Serialize)] -struct AnthropicRequest { - model: String, - messages: Vec, - max_tokens: u32, - #[serde(skip_serializing_if = "Option::is_none")] - system: Option, - #[serde(skip_serializing_if = "Option::is_none")] - temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tools: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - tool_choice: Option, -} - -#[derive(Debug, Serialize)] -struct AnthropicMessage { - role: String, - content: AnthropicContent, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -enum AnthropicContent { - Text(String), - #[serde(rename_all = "snake_case")] - ToolResult { - #[serde(rename = "type")] - tool_use_id: String, - content: String, - }, - Blocks(Vec), -} - -#[derive(Debug, Serialize, Deserialize)] -#[serde(tag = "type")] -enum AnthropicContentBlock { - #[serde(rename = "text")] - Text { text: String }, - #[serde(rename = "tool_use")] - ToolUse { - id: String, - name: String, - input: serde_json::Value, - }, - #[serde(rename = "tool_result")] - ToolResult { - tool_use_id: String, - content: String, - }, -} - -#[derive(Debug, Serialize)] -struct AnthropicTool { - name: String, - description: String, - input_schema: serde_json::Value, -} - -#[derive(Debug, Serialize)] -struct AnthropicToolChoice { - #[serde(rename = "type")] - choice_type: String, -} - -#[derive(Debug, Deserialize)] -struct AnthropicResponse { - content: Vec, - stop_reason: Option, - usage: AnthropicUsage, -} - -#[derive(Debug, Deserialize)] -struct AnthropicUsage { - input_tokens: u32, - output_tokens: u32, -} - -#[derive(Debug, Deserialize)] -struct AnthropicError { - error: AnthropicErrorDetail, -} - -#[derive(Debug, Deserialize)] -struct AnthropicErrorDetail { - message: String, - #[serde(rename = "type")] - error_type: String, -} - -fn parse_finish_reason(reason: Option<&str>) -> FinishReason { - match reason { - Some("end_turn") | Some("stop_sequence") => FinishReason::Stop, - Some("max_tokens") => FinishReason::Length, - Some("tool_use") => FinishReason::ToolUse, - _ => FinishReason::Unknown, - } -} - -#[async_trait] -impl LlmProvider for AnthropicProvider { - fn model_name(&self) -> &str { - &self.config.model - } - - fn cost_per_token(&self) -> (Decimal, Decimal) { - // Pricing for Claude models (per 1M tokens, converted to per token) - match self.config.model.as_str() { - m if m.contains("opus") => { - (dec!(0.000015), dec!(0.000075)) // $15/$75 per 1M - } - m if m.contains("sonnet") => { - (dec!(0.000003), dec!(0.000015)) // $3/$15 per 1M - } - m if m.contains("haiku") => { - (dec!(0.00000025), dec!(0.00000125)) // $0.25/$1.25 per 1M - } - _ => (dec!(0.000003), dec!(0.000015)), // Default to Sonnet pricing - } - } - - async fn complete(&self, request: CompletionRequest) -> Result { - let (system, messages) = self.build_messages(&request.messages); - - let anthropic_request = AnthropicRequest { - model: self.config.model.clone(), - messages, - max_tokens: request.max_tokens.unwrap_or(4096), - system, - temperature: request.temperature, - tools: None, - tool_choice: None, - }; - - let response = self - .client - .post(format!("{}/messages", self.base_url)) - .header("x-api-key", self.config.api_key.expose_secret()) - .header("anthropic-version", "2023-06-01") - .header("Content-Type", "application/json") - .json(&anthropic_request) - .send() - .await?; - - if !response.status().is_success() { - let error: AnthropicError = - response - .json() - .await - .map_err(|e| LlmError::InvalidResponse { - provider: "anthropic".to_string(), - reason: format!("Failed to parse error response: {}", e), - })?; - return Err(LlmError::RequestFailed { - provider: "anthropic".to_string(), - reason: error.error.message, - }); - } - - let anthropic_response: AnthropicResponse = response.json().await?; - - // Extract text content - let content = anthropic_response - .content - .iter() - .filter_map(|block| match block { - AnthropicContentBlock::Text { text } => Some(text.clone()), - _ => None, - }) - .collect::>() - .join("\n"); - - Ok(CompletionResponse { - content, - input_tokens: anthropic_response.usage.input_tokens, - output_tokens: anthropic_response.usage.output_tokens, - finish_reason: parse_finish_reason(anthropic_response.stop_reason.as_deref()), - }) - } - - async fn complete_with_tools( - &self, - request: ToolCompletionRequest, - ) -> Result { - let (system, messages) = self.build_messages(&request.messages); - - let tools: Vec = request - .tools - .iter() - .map(|t| AnthropicTool { - name: t.name.clone(), - description: t.description.clone(), - input_schema: t.parameters.clone(), - }) - .collect(); - - let tool_choice = request.tool_choice.as_ref().map(|c| AnthropicToolChoice { - choice_type: match c.as_str() { - "auto" => "auto".to_string(), - "required" => "any".to_string(), - "none" => "none".to_string(), - _ => "auto".to_string(), - }, - }); - - let anthropic_request = AnthropicRequest { - model: self.config.model.clone(), - messages, - max_tokens: request.max_tokens.unwrap_or(4096), - system, - temperature: None, - tools: Some(tools), - tool_choice, - }; - - let response = self - .client - .post(format!("{}/messages", self.base_url)) - .header("x-api-key", self.config.api_key.expose_secret()) - .header("anthropic-version", "2023-06-01") - .header("Content-Type", "application/json") - .json(&anthropic_request) - .send() - .await?; - - if !response.status().is_success() { - let error: AnthropicError = - response - .json() - .await - .map_err(|e| LlmError::InvalidResponse { - provider: "anthropic".to_string(), - reason: format!("Failed to parse error response: {}", e), - })?; - return Err(LlmError::RequestFailed { - provider: "anthropic".to_string(), - reason: error.error.message, - }); - } - - let anthropic_response: AnthropicResponse = response.json().await?; - - // Extract text and tool calls - let mut content = None; - let mut tool_calls = Vec::new(); - - for block in anthropic_response.content { - match block { - AnthropicContentBlock::Text { text } => { - content = Some(text); - } - AnthropicContentBlock::ToolUse { id, name, input } => { - tool_calls.push(ToolCall { - id, - name, - arguments: input, - }); - } - _ => {} - } - } - - Ok(ToolCompletionResponse { - content, - tool_calls, - input_tokens: anthropic_response.usage.input_tokens, - output_tokens: anthropic_response.usage.output_tokens, - finish_reason: parse_finish_reason(anthropic_response.stop_reason.as_deref()), - }) - } -} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index f3643f90..bde4a55a 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -1,17 +1,12 @@ //! LLM integration for the agent. //! -//! Provides a unified interface to different LLM providers (OpenAI, Anthropic, NEAR AI) -//! and implements reasoning capabilities for planning, tool selection, and evaluation. +//! Uses the NEAR AI chat-api as the unified LLM provider. -mod anthropic; mod nearai; -mod openai; mod provider; mod reasoning; -pub use anthropic::AnthropicProvider; pub use nearai::NearAiProvider; -pub use openai::OpenAiProvider; pub use provider::{ ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, @@ -20,33 +15,10 @@ pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, ToolSelection}; use std::sync::Arc; -use crate::config::{LlmConfig, LlmProvider as LlmProviderType}; +use crate::config::LlmConfig; use crate::error::LlmError; /// Create an LLM provider based on configuration. pub fn create_llm_provider(config: &LlmConfig) -> Result, LlmError> { - match config.provider { - LlmProviderType::OpenAi => { - let openai_config = config.openai.as_ref().ok_or_else(|| LlmError::AuthFailed { - provider: "openai".to_string(), - })?; - Ok(Arc::new(OpenAiProvider::new(openai_config.clone()))) - } - LlmProviderType::Anthropic => { - let anthropic_config = - config - .anthropic - .as_ref() - .ok_or_else(|| LlmError::AuthFailed { - provider: "anthropic".to_string(), - })?; - Ok(Arc::new(AnthropicProvider::new(anthropic_config.clone()))) - } - LlmProviderType::NearAi => { - let nearai_config = config.nearai.as_ref().ok_or_else(|| LlmError::AuthFailed { - provider: "nearai".to_string(), - })?; - Ok(Arc::new(NearAiProvider::new(nearai_config.clone()))) - } - } + Ok(Arc::new(NearAiProvider::new(config.nearai.clone()))) } diff --git a/src/llm/openai.rs b/src/llm/openai.rs deleted file mode 100644 index 89a02290..00000000 --- a/src/llm/openai.rs +++ /dev/null @@ -1,335 +0,0 @@ -//! OpenAI LLM provider implementation. - -use async_trait::async_trait; -use reqwest::Client; -use rust_decimal::Decimal; -use rust_decimal_macros::dec; -use secrecy::ExposeSecret; -use serde::{Deserialize, Serialize}; - -use crate::config::OpenAiConfig; -use crate::error::LlmError; -use crate::llm::provider::{ - ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, - ToolCompletionRequest, ToolCompletionResponse, -}; - -/// OpenAI API provider. -pub struct OpenAiProvider { - client: Client, - config: OpenAiConfig, - base_url: String, -} - -impl OpenAiProvider { - /// Create a new OpenAI provider. - pub fn new(config: OpenAiConfig) -> Self { - let base_url = config - .base_url - .clone() - .unwrap_or_else(|| "https://api.openai.com/v1".to_string()); - - Self { - client: Client::new(), - config, - base_url, - } - } - - fn build_messages(&self, messages: &[ChatMessage]) -> Vec { - messages - .iter() - .map(|m| OpenAiMessage { - role: match m.role { - Role::System => "system".to_string(), - Role::User => "user".to_string(), - Role::Assistant => "assistant".to_string(), - Role::Tool => "tool".to_string(), - }, - content: Some(m.content.clone()), - tool_call_id: m.tool_call_id.clone(), - name: m.name.clone(), - tool_calls: None, - }) - .collect() - } -} - -#[derive(Debug, Serialize)] -struct OpenAiRequest { - model: String, - messages: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - max_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tools: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - tool_choice: Option, -} - -#[derive(Debug, Serialize, Deserialize)] -struct OpenAiMessage { - role: String, - #[serde(skip_serializing_if = "Option::is_none")] - content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_call_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - name: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, -} - -#[derive(Debug, Serialize)] -struct OpenAiTool { - #[serde(rename = "type")] - tool_type: String, - function: OpenAiFunction, -} - -#[derive(Debug, Serialize)] -struct OpenAiFunction { - name: String, - description: String, - parameters: serde_json::Value, -} - -#[derive(Debug, Deserialize)] -struct OpenAiResponse { - choices: Vec, - usage: OpenAiUsage, -} - -#[derive(Debug, Deserialize)] -struct OpenAiChoice { - message: OpenAiResponseMessage, - finish_reason: Option, -} - -#[derive(Debug, Deserialize)] -struct OpenAiResponseMessage { - content: Option, - tool_calls: Option>, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct OpenAiToolCall { - id: String, - #[serde(rename = "type")] - call_type: String, - function: OpenAiFunctionCall, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct OpenAiFunctionCall { - name: String, - arguments: String, -} - -#[derive(Debug, Deserialize)] -struct OpenAiUsage { - prompt_tokens: u32, - completion_tokens: u32, -} - -#[derive(Debug, Deserialize)] -struct OpenAiError { - error: OpenAiErrorDetail, -} - -#[derive(Debug, Deserialize)] -struct OpenAiErrorDetail { - message: String, - #[serde(rename = "type")] - error_type: Option, -} - -fn parse_finish_reason(reason: Option<&str>) -> FinishReason { - match reason { - Some("stop") => FinishReason::Stop, - Some("length") => FinishReason::Length, - Some("tool_calls") => FinishReason::ToolUse, - Some("content_filter") => FinishReason::ContentFilter, - _ => FinishReason::Unknown, - } -} - -#[async_trait] -impl LlmProvider for OpenAiProvider { - fn model_name(&self) -> &str { - &self.config.model - } - - fn cost_per_token(&self) -> (Decimal, Decimal) { - // Pricing for GPT-4 Turbo (per 1M tokens, converted to per token) - // These are approximate and should be updated based on actual pricing - match self.config.model.as_str() { - m if m.contains("gpt-4-turbo") || m.contains("gpt-4o") => { - (dec!(0.00001), dec!(0.00003)) // $10/$30 per 1M - } - m if m.contains("gpt-4") => { - (dec!(0.00003), dec!(0.00006)) // $30/$60 per 1M - } - m if m.contains("gpt-3.5") => { - (dec!(0.0000005), dec!(0.0000015)) // $0.50/$1.50 per 1M - } - _ => (dec!(0.00001), dec!(0.00003)), // Default to GPT-4 Turbo pricing - } - } - - async fn complete(&self, request: CompletionRequest) -> Result { - let openai_request = OpenAiRequest { - model: self.config.model.clone(), - messages: self.build_messages(&request.messages), - max_tokens: request.max_tokens, - temperature: request.temperature, - tools: None, - tool_choice: None, - }; - - let response = self - .client - .post(format!("{}/chat/completions", self.base_url)) - .header( - "Authorization", - format!("Bearer {}", self.config.api_key.expose_secret()), - ) - .header("Content-Type", "application/json") - .json(&openai_request) - .send() - .await?; - - if !response.status().is_success() { - let error: OpenAiError = - response - .json() - .await - .map_err(|e| LlmError::InvalidResponse { - provider: "openai".to_string(), - reason: format!("Failed to parse error response: {}", e), - })?; - return Err(LlmError::RequestFailed { - provider: "openai".to_string(), - reason: error.error.message, - }); - } - - let openai_response: OpenAiResponse = response.json().await?; - - let choice = openai_response - .choices - .first() - .ok_or_else(|| LlmError::InvalidResponse { - provider: "openai".to_string(), - reason: "No choices in response".to_string(), - })?; - - Ok(CompletionResponse { - content: choice.message.content.clone().unwrap_or_default(), - input_tokens: openai_response.usage.prompt_tokens, - output_tokens: openai_response.usage.completion_tokens, - finish_reason: parse_finish_reason(choice.finish_reason.as_deref()), - }) - } - - async fn complete_with_tools( - &self, - request: ToolCompletionRequest, - ) -> Result { - let tools: Vec = request - .tools - .iter() - .map(|t| OpenAiTool { - tool_type: "function".to_string(), - function: OpenAiFunction { - name: t.name.clone(), - description: t.description.clone(), - parameters: t.parameters.clone(), - }, - }) - .collect(); - - let tool_choice = request.tool_choice.as_ref().map(|c| match c.as_str() { - "auto" => serde_json::json!("auto"), - "required" => serde_json::json!("required"), - "none" => serde_json::json!("none"), - _ => serde_json::json!("auto"), - }); - - let openai_request = OpenAiRequest { - model: self.config.model.clone(), - messages: self.build_messages(&request.messages), - max_tokens: request.max_tokens, - temperature: None, - tools: Some(tools), - tool_choice, - }; - - let response = self - .client - .post(format!("{}/chat/completions", self.base_url)) - .header( - "Authorization", - format!("Bearer {}", self.config.api_key.expose_secret()), - ) - .header("Content-Type", "application/json") - .json(&openai_request) - .send() - .await?; - - if !response.status().is_success() { - let error: OpenAiError = - response - .json() - .await - .map_err(|e| LlmError::InvalidResponse { - provider: "openai".to_string(), - reason: format!("Failed to parse error response: {}", e), - })?; - return Err(LlmError::RequestFailed { - provider: "openai".to_string(), - reason: error.error.message, - }); - } - - let openai_response: OpenAiResponse = response.json().await?; - - let choice = openai_response - .choices - .first() - .ok_or_else(|| LlmError::InvalidResponse { - provider: "openai".to_string(), - reason: "No choices in response".to_string(), - })?; - - let tool_calls: Vec = choice - .message - .tool_calls - .as_ref() - .map(|calls| { - calls - .iter() - .filter_map(|c| { - let args: serde_json::Value = - serde_json::from_str(&c.function.arguments).ok()?; - Some(ToolCall { - id: c.id.clone(), - name: c.function.name.clone(), - arguments: args, - }) - }) - .collect() - }) - .unwrap_or_default(); - - Ok(ToolCompletionResponse { - content: choice.message.content.clone(), - tool_calls, - input_tokens: openai_response.usage.prompt_tokens, - output_tokens: openai_response.usage.completion_tokens, - finish_reason: parse_finish_reason(choice.finish_reason.as_deref()), - }) - } -} diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs index 480bd7a7..db468d5f 100644 --- a/src/tools/builtin/memory.rs +++ b/src/tools/builtin/memory.rs @@ -2,7 +2,7 @@ //! //! These tools allow the agent to: //! - Search past memories, decisions, and context -//! - Write important information to long-term memory +//! - Read and write files in the workspace //! //! # Usage //! @@ -18,7 +18,7 @@ use async_trait::async_trait; use crate::context::JobContext; use crate::tools::tool::{Tool, ToolError, ToolOutput}; -use crate::workspace::Workspace; +use crate::workspace::{Workspace, paths}; /// Tool for searching workspace memory. /// @@ -135,7 +135,8 @@ impl Tool for MemoryWriteTool { 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." + for curated long-term facts, 'daily_log' for timestamped session notes, or \ + provide a custom path for arbitrary file creation." } fn parameters_schema(&self) -> serde_json::Value { @@ -148,9 +149,13 @@ impl Tool for MemoryWriteTool { }, "target": { "type": "string", - "enum": ["memory", "daily_log"], - "description": "Where to write: 'memory' for long-term curated facts, 'daily_log' for timestamped session notes", + "description": "Where to write: 'memory' for MEMORY.md, 'daily_log' for today's log, or a path like 'projects/alpha/notes.md'", "default": "daily_log" + }, + "append": { + "type": "boolean", + "description": "If true, append to existing content. If false, replace entirely.", + "default": true } }, "required": ["content"] @@ -182,30 +187,53 @@ impl Tool for MemoryWriteTool { .and_then(|v| v.as_str()) .unwrap_or("daily_log"); - match target { + let append = params + .get("append") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + + let path = match target { "memory" => { - self.workspace - .append_memory(content) - .await - .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + if append { + self.workspace + .append_memory(content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } else { + self.workspace + .write(paths::MEMORY, content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } + paths::MEMORY.to_string() } "daily_log" => { self.workspace .append_daily_log(content) .await .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + format!("daily/{}.md", chrono::Utc::now().format("%Y-%m-%d")) } - _ => { - return Err(ToolError::InvalidParameters(format!( - "invalid target '{}', must be 'memory' or 'daily_log'", - target - ))); + path => { + if append { + self.workspace + .append(path, content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } else { + self.workspace + .write(path, content) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; + } + path.to_string() } - } + }; let output = serde_json::json!({ "status": "written", - "target": target, + "path": path, + "append": append, "content_length": content.len(), }); @@ -217,10 +245,9 @@ impl Tool for MemoryWriteTool { } } -/// Tool for reading specific memory documents. +/// Tool for reading workspace files. /// -/// Use this to read the full content of identity files, heartbeat checklist, -/// or other specific documents. +/// Use this to read the full content of any file in the workspace. pub struct MemoryReadTool { workspace: Arc, } @@ -239,25 +266,20 @@ impl Tool for MemoryReadTool { } fn description(&self) -> &str { - "Read a specific memory document by type. Use this to read identity files, \ - heartbeat checklist, or full memory document content." + "Read a file from the workspace. Use this to read identity files, \ + heartbeat checklist, memory, daily logs, or any custom file." } fn parameters_schema(&self) -> serde_json::Value { serde_json::json!({ "type": "object", "properties": { - "doc_type": { + "path": { "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)" + "description": "Path to the file (e.g., 'MEMORY.md', 'daily/2024-01-15.md', 'projects/alpha/notes.md')" } }, - "required": ["doc_type"] + "required": ["path"] }) } @@ -268,27 +290,19 @@ impl Tool for MemoryReadTool { ) -> Result { let start = std::time::Instant::now(); - let doc_type_str = params - .get("doc_type") + let path = params + .get("path") .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()); + .ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".to_string()))?; let doc = self .workspace - .get_document(doc_type, title) + .read(path) .await .map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?; let output = serde_json::json!({ - "doc_type": doc_type_str, - "title": doc.title, + "path": doc.path, "content": doc.content, "word_count": doc.word_count(), "updated_at": doc.updated_at.to_rfc3339(), @@ -302,16 +316,88 @@ impl Tool for MemoryReadTool { } } +/// Tool for listing workspace files. +/// +/// Use this to explore the workspace structure. +pub struct MemoryListTool { + workspace: Arc, +} + +impl MemoryListTool { + /// Create a new memory list tool. + pub fn new(workspace: Arc) -> Self { + Self { workspace } + } +} + +#[async_trait] +impl Tool for MemoryListTool { + fn name(&self) -> &str { + "memory_list" + } + + fn description(&self) -> &str { + "List files and directories in the workspace. Use this to explore \ + the workspace structure and discover available files." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "directory": { + "type": "string", + "description": "Directory to list (empty string or '/' for root)", + "default": "" + } + } + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let directory = params + .get("directory") + .and_then(|v| v.as_str()) + .unwrap_or(""); + + let entries = self + .workspace + .list(directory) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("List failed: {}", e)))?; + + let output = serde_json::json!({ + "directory": directory, + "entries": entries.iter().map(|e| serde_json::json!({ + "path": e.path, + "name": e.name(), + "is_directory": e.is_directory, + "updated_at": e.updated_at.map(|t| t.to_rfc3339()), + "preview": e.content_preview, + })).collect::>(), + "count": entries.len(), + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false // Internal tool + } +} + #[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( + fn make_test_workspace() -> Arc { + Arc::new(Workspace::new( "test_user", deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( tokio_postgres::Config::new(), @@ -319,7 +405,12 @@ mod tests { )) .build() .unwrap(), - )); + )) + } + + #[test] + fn test_memory_search_schema() { + let workspace = make_test_workspace(); let tool = MemorySearchTool::new(workspace); assert_eq!(tool.name(), "memory_search"); @@ -337,26 +428,42 @@ mod tests { #[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 workspace = make_test_workspace(); 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"].is_object()); + assert!(schema["properties"]["append"].is_object()); + } + + #[test] + fn test_memory_read_schema() { + let workspace = make_test_workspace(); + let tool = MemoryReadTool::new(workspace); + + assert_eq!(tool.name(), "memory_read"); + + let schema = tool.parameters_schema(); + assert!(schema["properties"]["path"].is_object()); assert!( - schema["properties"]["target"]["enum"] + schema["required"] .as_array() .unwrap() - .contains(&"memory".into()) + .contains(&"path".into()) ); } + + #[test] + fn test_memory_list_schema() { + let workspace = make_test_workspace(); + let tool = MemoryListTool::new(workspace); + + assert_eq!(tool.name(), "memory_list"); + + let schema = tool.parameters_schema(); + assert!(schema["properties"]["directory"].is_object()); + } } diff --git a/src/workspace/document.rs b/src/workspace/document.rs index 53352afb..23dcd5b2 100644 --- a/src/workspace/document.rs +++ b/src/workspace/document.rs @@ -1,99 +1,32 @@ -//! Memory document types. +//! Memory document types for the workspace. use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::error::WorkspaceError; - -/// Document type in the workspace. +/// Well-known document paths. /// -/// 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, +/// These are conventional paths that have special meaning in the workspace. +/// Agents can create arbitrary paths beyond these. +pub mod paths { + /// Long-term curated memory. + pub const MEMORY: &str = "MEMORY.md"; /// 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()) - } + pub const IDENTITY: &str = "IDENTITY.md"; + /// Core values and principles. + pub const SOUL: &str = "SOUL.md"; + /// Behavior instructions. + pub const AGENTS: &str = "AGENTS.md"; + /// User context (name, preferences). + pub const USER: &str = "USER.md"; + /// Periodic checklist for heartbeat. + pub const HEARTBEAT: &str = "HEARTBEAT.md"; + /// Root runbook/readme. + pub const README: &str = "README.md"; + /// Daily logs directory. + pub const DAILY_DIR: &str = "daily/"; + /// Context directory (for identity-related docs). + pub const CONTEXT_DIR: &str = "context/"; } /// A memory document stored in the database. @@ -105,10 +38,8 @@ pub struct MemoryDocument { 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, + /// File path within the workspace (e.g., "context/vision.md"). + pub path: String, /// Full document content. pub content: String, /// Creation timestamp. @@ -120,20 +51,18 @@ pub struct MemoryDocument { } impl MemoryDocument { - /// Create a new document (not persisted yet). + /// Create a new document with a path. pub fn new( user_id: impl Into, agent_id: Option, - doc_type: DocType, - title: Option, + path: impl Into, ) -> Self { let now = Utc::now(); Self { id: Uuid::new_v4(), user_id: user_id.into(), agent_id, - doc_type, - title, + path: path.into(), content: String::new(), created_at: now, updated_at: now, @@ -141,6 +70,17 @@ impl MemoryDocument { } } + /// Get the file name from the path. + pub fn file_name(&self) -> &str { + self.path.rsplit('/').next().unwrap_or(&self.path) + } + + /// Get the parent directory from the path. + pub fn parent_dir(&self) -> Option<&str> { + let idx = self.path.rfind('/')?; + Some(&self.path[..idx]) + } + /// Check if the document is empty. pub fn is_empty(&self) -> bool { self.content.is_empty() @@ -150,6 +90,34 @@ impl MemoryDocument { pub fn word_count(&self) -> usize { self.content.split_whitespace().count() } + + /// Check if this is a well-known identity document. + pub fn is_identity_document(&self) -> bool { + matches!( + self.path.as_str(), + paths::IDENTITY | paths::SOUL | paths::AGENTS | paths::USER + ) + } +} + +/// An entry in a workspace directory listing. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WorkspaceEntry { + /// Path relative to listing directory. + pub path: String, + /// True if this is a directory (has children). + pub is_directory: bool, + /// Last update timestamp (latest among children for directories). + pub updated_at: Option>, + /// Preview of content (first ~200 chars, None for directories). + pub content_preview: Option, +} + +impl WorkspaceEntry { + /// Get the entry name (last path component). + pub fn name(&self) -> &str { + self.path.rsplit('/').next().unwrap_or(&self.path) + } } /// A chunk of a memory document for search indexing. @@ -194,43 +162,60 @@ 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); - } + fn test_memory_document_new() { + let doc = MemoryDocument::new("user1", None, "context/vision.md"); + assert_eq!(doc.user_id, "user1"); + assert_eq!(doc.path, "context/vision.md"); + assert!(doc.content.is_empty()); } #[test] - fn test_singleton_types() { - assert!(DocType::Memory.is_singleton()); - assert!(DocType::Heartbeat.is_singleton()); - assert!(!DocType::DailyLog.is_singleton()); + fn test_memory_document_file_name() { + let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md"); + assert_eq!(doc.file_name(), "README.md"); } #[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()); + fn test_memory_document_parent_dir() { + let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md"); + assert_eq!(doc.parent_dir(), Some("projects/alpha")); + + let root_doc = MemoryDocument::new("user1", None, "README.md"); + assert_eq!(root_doc.parent_dir(), None); } #[test] fn test_memory_document_word_count() { - let mut doc = MemoryDocument::new("user1", None, DocType::Memory, None); + let mut doc = MemoryDocument::new("user1", None, "MEMORY.md"); assert_eq!(doc.word_count(), 0); doc.content = "Hello world, this is a test.".to_string(); assert_eq!(doc.word_count(), 6); } + + #[test] + fn test_is_identity_document() { + let identity = MemoryDocument::new("user1", None, paths::IDENTITY); + assert!(identity.is_identity_document()); + + let soul = MemoryDocument::new("user1", None, paths::SOUL); + assert!(soul.is_identity_document()); + + let memory = MemoryDocument::new("user1", None, paths::MEMORY); + assert!(!memory.is_identity_document()); + + let custom = MemoryDocument::new("user1", None, "projects/notes.md"); + assert!(!custom.is_identity_document()); + } + + #[test] + fn test_workspace_entry_name() { + let entry = WorkspaceEntry { + path: "projects/alpha".to_string(), + is_directory: true, + updated_at: None, + content_preview: None, + }; + assert_eq!(entry.name(), "alpha"); + } } diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index e500d648..12a4721c 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -1,37 +1,44 @@ //! 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 +//! The workspace provides persistent memory for agents with a flexible +//! filesystem-like structure. Agents can create arbitrary markdown file +//! hierarchies that get indexed for full-text and semantic search. //! -//! Memory is searchable via hybrid search (FTS + semantic embeddings). -//! -//! # Architecture +//! # Filesystem-like API //! //! ```text -//! ┌─────────────────────────────────────────────────────────────┐ -//! │ Workspace │ -//! │ ┌────────────────┐ ┌────────────────┐ ┌──────────────┐ │ -//! │ │ MemoryDocument │ │ MemoryChunk │ │ Search │ │ -//! │ │ (full docs) │──│ (chunked) │──│ (FTS+vector) │ │ -//! │ └────────────────┘ └────────────────┘ └──────────────┘ │ -//! │ │ │ │ │ -//! │ └───────────────────┴──────────────────┘ │ -//! │ │ │ -//! │ ┌──────┴──────┐ │ -//! │ │ Repository │ │ -//! │ │ (PostgreSQL)│ │ -//! │ └─────────────┘ │ -//! └─────────────────────────────────────────────────────────────┘ +//! workspace/ +//! ├── README.md <- Root runbook/index +//! ├── MEMORY.md <- Long-term curated memory +//! ├── HEARTBEAT.md <- Periodic checklist +//! ├── context/ <- Identity and context +//! │ ├── vision.md +//! │ └── priorities.md +//! ├── daily/ <- Daily logs +//! │ ├── 2024-01-15.md +//! │ └── 2024-01-16.md +//! ├── projects/ <- Arbitrary structure +//! │ └── alpha/ +//! │ ├── README.md +//! │ └── notes.md +//! └── ... //! ``` //! +//! # Key Operations +//! +//! - `read(path)` - Read a file +//! - `write(path, content)` - Create or update a file +//! - `append(path, content)` - Append to a file +//! - `list(dir)` - List directory contents +//! - `delete(path)` - Delete a file +//! - `search(query)` - Full-text + semantic search across all files +//! //! # 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 +//! 2. **Flexible structure**: Create any directory/file hierarchy you need +//! 3. **Self-documenting**: Use README.md files to describe directory structure +//! 4. **Hybrid search**: Vector similarity + BM25 full-text via RRF mod chunker; mod document; @@ -40,7 +47,7 @@ mod repository; mod search; pub use chunker::{ChunkConfig, chunk_document}; -pub use document::{DocType, MemoryChunk, MemoryDocument}; +pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths}; pub use embeddings::{EmbeddingProvider, OpenAiEmbeddings}; pub use repository::Repository; pub use search::{SearchConfig, SearchResult}; @@ -101,15 +108,127 @@ impl Workspace { self.agent_id } - // ==================== Document Access ==================== + // ==================== File Operations ==================== + + /// Read a file by path. + /// + /// Returns the document if it exists, or an error if not found. + /// + /// # Example + /// ```ignore + /// let doc = workspace.read("context/vision.md").await?; + /// println!("{}", doc.content); + /// ``` + pub async fn read(&self, path: &str) -> Result { + let path = normalize_path(path); + self.repo + .get_document_by_path(&self.user_id, self.agent_id, &path) + .await + } + + /// Write (create or update) a file. + /// + /// Creates parent directories implicitly (they're virtual in the DB). + /// Re-indexes the document for search after writing. + /// + /// # Example + /// ```ignore + /// workspace.write("projects/alpha/README.md", "# Project Alpha\n\nDescription here.").await?; + /// ``` + pub async fn write(&self, path: &str, content: &str) -> Result { + let path = normalize_path(path); + let doc = self + .repo + .get_or_create_document_by_path(&self.user_id, self.agent_id, &path) + .await?; + self.repo.update_document(doc.id, content).await?; + self.reindex_document(doc.id).await?; + + // Return updated doc + self.repo.get_document_by_id(doc.id).await + } + + /// Append content to a file. + /// + /// Creates the file if it doesn't exist. + /// Adds a newline separator between existing and new content. + pub async fn append(&self, path: &str, content: &str) -> Result<(), WorkspaceError> { + let path = normalize_path(path); + let doc = self + .repo + .get_or_create_document_by_path(&self.user_id, self.agent_id, &path) + .await?; + + let new_content = if doc.content.is_empty() { + content.to_string() + } else { + format!("{}\n{}", doc.content, content) + }; + + self.repo.update_document(doc.id, &new_content).await?; + self.reindex_document(doc.id).await?; + Ok(()) + } + + /// Check if a file exists. + pub async fn exists(&self, path: &str) -> Result { + let path = normalize_path(path); + match self + .repo + .get_document_by_path(&self.user_id, self.agent_id, &path) + .await + { + Ok(_) => Ok(true), + Err(WorkspaceError::DocumentNotFound { .. }) => Ok(false), + Err(e) => Err(e), + } + } + + /// Delete a file. + /// + /// Also deletes associated chunks. + pub async fn delete(&self, path: &str) -> Result<(), WorkspaceError> { + let path = normalize_path(path); + self.repo + .delete_document_by_path(&self.user_id, self.agent_id, &path) + .await + } + + /// List files and directories in a path. + /// + /// Returns immediate children (not recursive). + /// Use empty string or "/" for root directory. + /// + /// # Example + /// ```ignore + /// let entries = workspace.list("projects/").await?; + /// for entry in entries { + /// if entry.is_directory { + /// println!("📁 {}/", entry.name()); + /// } else { + /// println!("📄 {}", entry.name()); + /// } + /// } + /// ``` + pub async fn list(&self, directory: &str) -> Result, WorkspaceError> { + let directory = normalize_directory(directory); + self.repo + .list_directory(&self.user_id, self.agent_id, &directory) + .await + } + + /// List all files recursively (flat list of all paths). + pub async fn list_all(&self) -> Result, WorkspaceError> { + self.repo.list_all_paths(&self.user_id, self.agent_id).await + } + + // ==================== Convenience Methods ==================== /// 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 + self.read_or_create(paths::MEMORY).await } /// Get today's daily log. @@ -122,38 +241,23 @@ impl Workspace { /// 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 + let path = format!("daily/{}.md", date.format("%Y-%m-%d")); + self.read_or_create(&path).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 - { + match self.read(paths::HEARTBEAT).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 { + /// Helper to read or create a file. + async fn read_or_create(&self, path: &str) -> Result { self.repo - .get_document(&self.user_id, self.agent_id, doc_type, title) + .get_or_create_document_by_path(&self.user_id, self.agent_id, path) .await } @@ -164,6 +268,7 @@ impl Workspace { /// This is for important facts, decisions, and preferences worth /// remembering long-term. pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> { + // Use double newline for memory entries (semantic separation) let doc = self.memory().await?; let new_content = if doc.content.is_empty() { entry.to_string() @@ -179,34 +284,11 @@ impl Workspace { /// /// 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 today = Utc::now().date_naive(); + let path = format!("daily/{}.md", today.format("%Y-%m-%d")); let timestamp = Utc::now().format("%H:%M:%S"); let timestamped_entry = format!("[{}] {}", timestamp, entry); - - 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(()) + self.append(&path, ×tamped_entry).await } // ==================== System Prompt ==================== @@ -219,19 +301,15 @@ impl Workspace { 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"), + let identity_files = [ + (paths::AGENTS, "## Agent Instructions"), + (paths::SOUL, "## Core Values"), + (paths::USER, "## User Context"), + (paths::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 - { + for (path, header) in identity_files { + if let Ok(doc) = self.read(path).await { if !doc.content.is_empty() { parts.push(format!("{}\n\n{}", header, doc.content)); } @@ -372,21 +450,50 @@ impl Workspace { } } +/// Normalize a file path (remove leading/trailing slashes, collapse //). +fn normalize_path(path: &str) -> String { + let path = path.trim().trim_matches('/'); + // Collapse multiple slashes + let mut result = String::new(); + let mut last_was_slash = false; + for c in path.chars() { + if c == '/' { + if !last_was_slash { + result.push(c); + } + last_was_slash = true; + } else { + result.push(c); + last_was_slash = false; + } + } + result +} + +/// Normalize a directory path (ensure no trailing slash for consistency). +fn normalize_directory(path: &str) -> String { + let path = normalize_path(path); + path.trim_end_matches('/').to_string() +} + #[cfg(test)] mod tests { use super::*; #[test] - fn test_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"); + fn test_normalize_path() { + assert_eq!(normalize_path("foo/bar"), "foo/bar"); + assert_eq!(normalize_path("/foo/bar/"), "foo/bar"); + assert_eq!(normalize_path("foo//bar"), "foo/bar"); + assert_eq!(normalize_path(" /foo/ "), "foo"); + assert_eq!(normalize_path("README.md"), "README.md"); } #[test] - fn test_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()); + fn test_normalize_directory() { + assert_eq!(normalize_directory("foo/bar/"), "foo/bar"); + assert_eq!(normalize_directory("foo/bar"), "foo/bar"); + assert_eq!(normalize_directory("/"), ""); + assert_eq!(normalize_directory(""), ""); } } diff --git a/src/workspace/repository.rs b/src/workspace/repository.rs index 61e9c912..e350a2e3 100644 --- a/src/workspace/repository.rs +++ b/src/workspace/repository.rs @@ -4,14 +4,14 @@ //! - Documents in `memory_documents` table //! - Chunks in `memory_chunks` table (with FTS and vector indexes) -use chrono::Utc; +use chrono::{DateTime, 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::document::{MemoryChunk, MemoryDocument, WorkspaceEntry}; use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion}; /// Database repository for workspace operations. @@ -37,50 +37,34 @@ impl Repository { // ==================== Document Operations ==================== - /// Get a document by type and optional title. - pub async fn get_document( + /// Get a document by its path. + pub async fn get_document_by_path( &self, user_id: &str, agent_id: Option, - doc_type: DocType, - title: Option<&str>, + path: &str, ) -> Result { let conn = self.conn().await?; - let row = if let Some(title) = title { - conn.query_opt( + let row = conn + .query_opt( r#" - SELECT id, user_id, agent_id, doc_type, title, content, + SELECT id, user_id, agent_id, path, 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 + WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND path = $3 "#, - &[&user_id, &agent_id, &doc_type.as_str(), &title], + &[&user_id, &agent_id, &path], ) .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), - })?; + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; match row { - Some(row) => Ok(self.row_to_document(&row)?), + Some(row) => Ok(self.row_to_document(&row)), None => Err(WorkspaceError::DocumentNotFound { - doc_type: doc_type.to_string(), + doc_type: path.to_string(), user_id: user_id.to_string(), }), } @@ -93,7 +77,7 @@ impl Repository { let row = conn .query_opt( r#" - SELECT id, user_id, agent_id, doc_type, title, content, + SELECT id, user_id, agent_id, path, content, created_at, updated_at, metadata FROM memory_documents WHERE id = $1 "#, @@ -105,7 +89,7 @@ impl Repository { })?; match row { - Some(row) => Ok(self.row_to_document(&row)?), + Some(row) => Ok(self.row_to_document(&row)), None => Err(WorkspaceError::DocumentNotFound { doc_type: "unknown".to_string(), user_id: "unknown".to_string(), @@ -113,16 +97,15 @@ impl Repository { } } - /// Get or create a document. - pub async fn get_or_create_document( + /// Get or create a document by path. + pub async fn get_or_create_document_by_path( &self, user_id: &str, agent_id: Option, - doc_type: DocType, - title: Option<&str>, + path: &str, ) -> Result { // Try to get existing document first - match self.get_document(user_id, agent_id, doc_type, title).await { + match self.get_document_by_path(user_id, agent_id, path).await { Ok(doc) => return Ok(doc), Err(WorkspaceError::DocumentNotFound { .. }) => {} Err(e) => return Err(e), @@ -132,14 +115,15 @@ impl Repository { let conn = self.conn().await?; let id = Uuid::new_v4(); let now = Utc::now(); + let metadata = serde_json::json!({}); 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 + INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata, created_at, updated_at) + VALUES ($1, $2, $3, $4, '', $5, $6, $7) + ON CONFLICT (user_id, agent_id, path) DO NOTHING "#, - &[&id, &user_id, &agent_id, &doc_type.as_str(), &title, &now, &now], + &[&id, &user_id, &agent_id, &path, &metadata, &now, &now], ) .await .map_err(|e| WorkspaceError::SearchFailed { @@ -147,7 +131,7 @@ impl Repository { })?; // Fetch the document (might have been created by concurrent request) - self.get_document(user_id, agent_id, doc_type, title).await + self.get_document_by_path(user_id, agent_id, path).await } /// Update a document's content. @@ -166,31 +150,108 @@ impl Repository { Ok(()) } - /// List documents by type. + /// Delete a document by its path. + pub async fn delete_document_by_path( + &self, + user_id: &str, + agent_id: Option, + path: &str, + ) -> Result<(), WorkspaceError> { + let conn = self.conn().await?; + + // First get the document to delete its chunks + let doc = self.get_document_by_path(user_id, agent_id, path).await?; + self.delete_chunks(doc.id).await?; + + // Delete the document + conn.execute( + r#" + DELETE FROM memory_documents + WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND path = $3 + "#, + &[&user_id, &agent_id, &path], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Delete failed: {}", e), + })?; + + Ok(()) + } + + /// List files and directories in a directory path. + /// + /// Returns immediate children (not recursive). + /// Empty string lists the root directory. + pub async fn list_directory( + &self, + user_id: &str, + agent_id: Option, + directory: &str, + ) -> Result, WorkspaceError> { + let conn = self.conn().await?; + + let rows = conn + .query( + "SELECT path, is_directory, updated_at, content_preview FROM list_workspace_files($1, $2, $3)", + &[&user_id, &agent_id, &directory], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("List directory failed: {}", e), + })?; + + Ok(rows + .iter() + .map(|row| { + let updated_at: Option> = row.get("updated_at"); + WorkspaceEntry { + path: row.get("path"), + is_directory: row.get("is_directory"), + updated_at, + content_preview: row.get("content_preview"), + } + }) + .collect()) + } + + /// List all file paths in the workspace (flat list). + pub async fn list_all_paths( + &self, + user_id: &str, + agent_id: Option, + ) -> Result, WorkspaceError> { + let conn = self.conn().await?; + + let rows = conn + .query( + r#" + SELECT path FROM memory_documents + WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 + ORDER BY path + "#, + &[&user_id, &agent_id], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("List paths failed: {}", e), + })?; + + Ok(rows.iter().map(|row| row.get("path")).collect()) + } + + /// List all documents for a user. 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( + let rows = 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, + SELECT id, user_id, agent_id, path, content, created_at, updated_at, metadata FROM memory_documents WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 @@ -199,30 +260,24 @@ impl Repository { &[&user_id, &agent_id], ) .await - }; + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; - let rows = rows.map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })?; - - rows.iter().map(|r| self.row_to_document(r)).collect() + Ok(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 { + fn row_to_document(&self, row: &tokio_postgres::Row) -> MemoryDocument { + MemoryDocument { id: row.get("id"), user_id: row.get("user_id"), agent_id: row.get("agent_id"), - doc_type, - title: row.get("title"), + path: row.get("path"), content: row.get("content"), created_at: row.get("created_at"), updated_at: row.get("updated_at"), metadata: row.get("metadata"), - }) + } } // ==================== Chunk Operations ==================== @@ -374,7 +429,6 @@ impl Repository { ) -> Result, WorkspaceError> { let conn = self.conn().await?; - // Use plainto_tsquery for natural language queries let rows = conn .query( r#" @@ -401,7 +455,7 @@ impl Repository { chunk_id: row.get("chunk_id"), document_id: row.get("document_id"), content: row.get("content"), - rank: (i + 1) as u32, // 1-based rank + rank: (i + 1) as u32, }) .collect()) } @@ -417,7 +471,6 @@ impl Repository { let conn = self.conn().await?; let embedding_vec = Vector::from(embedding.to_vec()); - // Use cosine distance (<=>) let rows = conn .query( r#" @@ -444,7 +497,7 @@ impl Repository { chunk_id: row.get("chunk_id"), document_id: row.get("document_id"), content: row.get("content"), - rank: (i + 1) as u32, // 1-based rank + rank: (i + 1) as u32, }) .collect()) }