Simplify workspace to path-based storage, remove legacy code

- Consolidate all migrations into V1__initial.sql
- Replace DocType enum with flexible path-based file storage
- Add list_workspace_files SQL function for directory listing
- Update memory tools for path-based API (memory_read, memory_write,
  memory_search, memory_list)
- Remove unused OpenAI/Anthropic providers (NEAR AI only)
- Simplify config to remove multi-provider support
- Update CLAUDE.md documentation

Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-02 21:38:53 -08:00
co-authored by Claude Opus 4.5
parent f29892b3fb
commit 3718cfa767
13 changed files with 877 additions and 1427 deletions
+2 -15
View File
@@ -2,25 +2,12 @@
DATABASE_URL=postgres://near_agent:password@localhost:5432/near_agent DATABASE_URL=postgres://near_agent:password@localhost:5432/near_agent
DATABASE_POOL_SIZE=10 DATABASE_POOL_SIZE=10
# LLM Providers # LLM Provider (NEAR AI)
# Default is NEAR AI which provides a unified interface to all models # NEAR AI provides a unified interface to all models with user authentication
# NEAR AI (recommended - unified API with user authentication)
NEARAI_SESSION_TOKEN=sess_... NEARAI_SESSION_TOKEN=sess_...
NEARAI_MODEL=claude-3-5-sonnet-20241022 NEARAI_MODEL=claude-3-5-sonnet-20241022
NEARAI_BASE_URL=https://api.near.ai 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 # Channel Configuration
# CLI is always enabled # CLI is always enabled
+50 -41
View File
@@ -53,11 +53,9 @@ src/
│ ├── validator.rs # Input validation (length, encoding, patterns) │ ├── validator.rs # Input validation (length, encoding, patterns)
│ └── policy.rs # PolicyRule system with severity/actions │ └── policy.rs # PolicyRule system with severity/actions
├── llm/ # LLM integration ├── llm/ # LLM integration (NEAR AI only)
│ ├── provider.rs # LlmProvider trait, message types │ ├── provider.rs # LlmProvider trait, message types
│ ├── nearai.rs # NEAR AI chat-api (default, unified interface) │ ├── nearai.rs # NEAR AI chat-api implementation
│ ├── openai.rs # OpenAI API implementation
│ ├── anthropic.rs # Anthropic API implementation
│ └── reasoning.rs # Planning, tool selection, evaluation │ └── reasoning.rs # Planning, tool selection, evaluation
├── tools/ # Extensible tool system ├── tools/ # Extensible tool system
@@ -76,7 +74,7 @@ src/
├── workspace/ # Persistent memory system (OpenClaw-inspired) ├── workspace/ # Persistent memory system (OpenClaw-inspired)
│ ├── mod.rs # Workspace struct, memory operations │ ├── 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) │ ├── chunker.rs # Document chunking (800 tokens, 15% overlap)
│ ├── embeddings.rs # EmbeddingProvider trait, OpenAI implementation │ ├── embeddings.rs # EmbeddingProvider trait, OpenAI implementation
│ ├── search.rs # Hybrid search with RRF algorithm │ ├── search.rs # Hybrid search with RRF algorithm
@@ -164,22 +162,11 @@ Environment variables (see `.env.example`):
```bash ```bash
DATABASE_URL=postgres://user:pass@localhost/near_agent DATABASE_URL=postgres://user:pass@localhost/near_agent
# LLM Provider (default: nearai) # NEAR AI (required)
LLM_PROVIDER=nearai # Options: nearai, openai, anthropic
# NEAR AI (recommended - unified API with user auth)
NEARAI_SESSION_TOKEN=sess_... NEARAI_SESSION_TOKEN=sess_...
NEARAI_MODEL=claude-3-5-sonnet-20241022 NEARAI_MODEL=claude-3-5-sonnet-20241022
NEARAI_BASE_URL=https://api.near.ai 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 settings
AGENT_NAME=near-agent AGENT_NAME=near-agent
MAX_PARALLEL_JOBS=5 MAX_PARALLEL_JOBS=5
@@ -187,7 +174,7 @@ MAX_PARALLEL_JOBS=5
### NEAR AI Provider ### 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.) - Unified access to multiple models (OpenAI, Anthropic, etc.)
- User authentication via session tokens - User authentication via session tokens
- Usage tracking and billing through NEAR AI - Usage tracking and billing through NEAR AI
@@ -196,9 +183,9 @@ Session tokens have the format `sess_xxx` (37 characters). They are authenticate
## Database ## Database
Migrations in `migrations/`. Tables: Single migration in `migrations/V1__initial.sql`. Tables:
**V1 (initial):** **Core:**
- `conversations` - Multi-channel conversation tracking - `conversations` - Multi-channel conversation tracking
- `agent_jobs` - Job metadata and status - `agent_jobs` - Job metadata and status
- `job_actions` - Event-sourced tool executions - `job_actions` - Event-sourced tool executions
@@ -206,8 +193,8 @@ Migrations in `migrations/`. Tables:
- `llm_calls` - Cost tracking - `llm_calls` - Cost tracking
- `estimation_snapshots` - Learning data - `estimation_snapshots` - Learning data
**V2 (workspace/memory):** **Workspace/Memory:**
- `memory_documents` - Full documents (MEMORY.md, daily logs, identity files) - `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 - `memory_chunks` - Chunked content with FTS (tsvector) and vector (pgvector) indexes
- `heartbeat_state` - Periodic execution tracking - `heartbeat_state` - Periodic execution tracking
@@ -292,38 +279,59 @@ RUST_LOG=near_agent=debug,tower_http=debug cargo run
## Workspace & Memory System ## 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 ### Key Principles
1. **"Memory is files, not RAM"** - If you want to remember something, write it explicitly 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) 2. **Flexible structure** - Create any directory/file hierarchy you need
3. **Hybrid search** - Combines FTS (keyword) + vector (semantic) via Reciprocal Rank Fusion 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 | ```
|------|---------|-----------| workspace/
| `Memory` | Long-term curated facts (MEMORY.md) | Yes | ├── README.md <- Root runbook/index
| `DailyLog` | Append-only daily notes (keyed by date) | No | ├── MEMORY.md <- Long-term curated memory
| `Identity` | Agent name, nature, vibe | Yes | ├── HEARTBEAT.md <- Periodic checklist
| `Soul` | Core values and principles | Yes | ├── IDENTITY.md <- Agent name, nature, vibe
| `Agents` | Behavior instructions | Yes | ├── SOUL.md <- Core values
| `User` | User context (name, preferences) | Yes | ├── AGENTS.md <- Behavior instructions
| `Heartbeat` | Periodic checklist | Yes | ├── 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 ### Using the Workspace
```rust ```rust
use crate::workspace::{Workspace, DocType, OpenAiEmbeddings}; use crate::workspace::{Workspace, OpenAiEmbeddings, paths};
// Create workspace for a user // Create workspace for a user
let workspace = Workspace::new("user_123", pool) let workspace = Workspace::new("user_123", pool)
.with_embeddings(Arc::new(OpenAiEmbeddings::new(api_key))); .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_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) // Search (hybrid FTS + vector)
let results = workspace.search("dark mode preference", 5).await?; let results = workspace.search("dark mode preference", 5).await?;
@@ -334,11 +342,12 @@ let prompt = workspace.system_prompt().await?;
### Memory Tools ### 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_search`** - Hybrid search, MUST be called before answering questions about prior work
- **`memory_write`** - Write to memory or daily_log target - **`memory_write`** - Write to any path (memory, daily_log, or custom paths)
- **`memory_read`** - Read specific document by type - **`memory_read`** - Read any file by path
- **`memory_list`** - List directory contents
### Hybrid Search (RRF) ### Hybrid Search (RRF)
+207 -9
View File
@@ -1,7 +1,12 @@
-- NEAR Agent Database Schema -- 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 ( CREATE TABLE conversations (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
channel TEXT NOT NULL, 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_user ON conversations(user_id);
CREATE INDEX idx_conversations_last_activity ON conversations(last_activity); CREATE INDEX idx_conversations_last_activity ON conversations(last_activity);
-- Messages in conversations
CREATE TABLE conversation_messages ( CREATE TABLE conversation_messages (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, 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); CREATE INDEX idx_conversation_messages_conversation ON conversation_messages(conversation_id);
-- Jobs we've worked on -- ==================== Agent Jobs ====================
CREATE TABLE agent_jobs ( CREATE TABLE agent_jobs (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
marketplace_job_id UUID, 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_conversation ON agent_jobs(conversation_id);
CREATE INDEX idx_agent_jobs_stuck ON agent_jobs(stuck_since) WHERE stuck_since IS NOT NULL; 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 ( CREATE TABLE job_actions (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
job_id UUID NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE, 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_job_id ON job_actions(job_id);
CREATE INDEX idx_job_actions_tool ON job_actions(tool_name); CREATE INDEX idx_job_actions_tool ON job_actions(tool_name);
-- Dynamic tools built by the agent -- ==================== Dynamic Tools ====================
CREATE TABLE dynamic_tools ( CREATE TABLE dynamic_tools (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
name TEXT NOT NULL UNIQUE, 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_status ON dynamic_tools(status);
CREATE INDEX idx_dynamic_tools_name ON dynamic_tools(name); CREATE INDEX idx_dynamic_tools_name ON dynamic_tools(name);
-- LLM calls for cost tracking -- ==================== LLM Calls ====================
CREATE TABLE llm_calls ( CREATE TABLE llm_calls (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
job_id UUID REFERENCES agent_jobs(id) ON DELETE CASCADE, 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_conversation ON llm_calls(conversation_id);
CREATE INDEX idx_llm_calls_provider ON llm_calls(provider); CREATE INDEX idx_llm_calls_provider ON llm_calls(provider);
-- Estimation history for continuous learning -- ==================== Estimation ====================
CREATE TABLE estimation_snapshots ( CREATE TABLE estimation_snapshots (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
job_id UUID NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE, 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_category ON estimation_snapshots(category);
CREATE INDEX idx_estimation_job ON estimation_snapshots(job_id); CREATE INDEX idx_estimation_job ON estimation_snapshots(job_id);
-- Self-repair history -- ==================== Self Repair ====================
CREATE TABLE repair_attempts ( CREATE TABLE repair_attempts (
id UUID PRIMARY KEY, id UUID PRIMARY KEY,
target_type TEXT NOT NULL, 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_target ON repair_attempts(target_type, target_id);
CREATE INDEX idx_repair_attempts_created ON repair_attempts(created_at); 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;
-159
View File
@@ -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;
+6 -132
View File
@@ -60,52 +60,10 @@ impl DatabaseConfig {
} }
} }
/// LLM provider configuration. /// LLM provider configuration (NEAR AI only).
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct LlmConfig { pub struct LlmConfig {
pub provider: LlmProvider, pub nearai: NearAiConfig,
pub openai: Option<OpenAiConfig>,
pub anthropic: Option<AnthropicConfig>,
pub nearai: Option<NearAiConfig>,
}
#[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<Self, Self::Err> {
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<String>,
}
#[derive(Debug, Clone)]
pub struct AnthropicConfig {
pub api_key: SecretString,
pub model: String,
pub base_url: Option<String>,
} }
/// NEAR AI chat-api configuration. /// NEAR AI chat-api configuration.
@@ -121,66 +79,16 @@ pub struct NearAiConfig {
impl LlmConfig { impl LlmConfig {
fn from_env() -> Result<Self, ConfigError> { fn from_env() -> Result<Self, ConfigError> {
let provider: LlmProvider = optional_env("LLM_PROVIDER")? let session_token = required_env("NEARAI_SESSION_TOKEN")?;
.map(|s| s.parse())
.transpose()?
.unwrap_or(LlmProvider::NearAi);
let openai = if let Some(api_key) = optional_env("OPENAI_API_KEY")? { Ok(Self {
Some(OpenAiConfig { nearai: NearAiConfig {
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 {
session_token: SecretString::from(session_token), session_token: SecretString::from(session_token),
model: optional_env("NEARAI_MODEL")? model: optional_env("NEARAI_MODEL")?
.unwrap_or_else(|| "claude-3-5-sonnet-20241022".to_string()), .unwrap_or_else(|| "claude-3-5-sonnet-20241022".to_string()),
base_url: optional_env("NEARAI_BASE_URL")? base_url: optional_env("NEARAI_BASE_URL")?
.unwrap_or_else(|| "https://api.near.ai".to_string()), .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() .transpose()
.map(|opt| opt.unwrap_or(default)) .map(|opt| opt.unwrap_or(default))
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_llm_provider_parsing() {
assert_eq!(
"openai".parse::<LlmProvider>().unwrap(),
LlmProvider::OpenAi
);
assert_eq!(
"anthropic".parse::<LlmProvider>().unwrap(),
LlmProvider::Anthropic
);
assert_eq!(
"OpenAI".parse::<LlmProvider>().unwrap(),
LlmProvider::OpenAi
);
assert_eq!(
"nearai".parse::<LlmProvider>().unwrap(),
LlmProvider::NearAi
);
assert_eq!(
"near-ai".parse::<LlmProvider>().unwrap(),
LlmProvider::NearAi
);
assert_eq!(
"near_ai".parse::<LlmProvider>().unwrap(),
LlmProvider::NearAi
);
assert!("invalid".parse::<LlmProvider>().is_err());
}
}
+1 -1
View File
@@ -63,5 +63,5 @@ pub mod prelude {
pub use crate::llm::LlmProvider; pub use crate::llm::LlmProvider;
pub use crate::safety::{SanitizedOutput, Sanitizer}; pub use crate::safety::{SanitizedOutput, Sanitizer};
pub use crate::tools::{Tool, ToolOutput, ToolRegistry}; pub use crate::tools::{Tool, ToolOutput, ToolRegistry};
pub use crate::workspace::{DocType, MemoryDocument, Workspace}; pub use crate::workspace::{MemoryDocument, Workspace};
} }
-348
View File
@@ -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<String>, Vec<AnthropicMessage>) {
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<AnthropicMessage>,
max_tokens: u32,
#[serde(skip_serializing_if = "Option::is_none")]
system: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<AnthropicTool>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<AnthropicToolChoice>,
}
#[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<AnthropicContentBlock>),
}
#[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<AnthropicContentBlock>,
stop_reason: Option<String>,
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<CompletionResponse, LlmError> {
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::<Vec<_>>()
.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<ToolCompletionResponse, LlmError> {
let (system, messages) = self.build_messages(&request.messages);
let tools: Vec<AnthropicTool> = 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()),
})
}
}
+3 -31
View File
@@ -1,17 +1,12 @@
//! LLM integration for the agent. //! LLM integration for the agent.
//! //!
//! Provides a unified interface to different LLM providers (OpenAI, Anthropic, NEAR AI) //! Uses the NEAR AI chat-api as the unified LLM provider.
//! and implements reasoning capabilities for planning, tool selection, and evaluation.
mod anthropic;
mod nearai; mod nearai;
mod openai;
mod provider; mod provider;
mod reasoning; mod reasoning;
pub use anthropic::AnthropicProvider;
pub use nearai::NearAiProvider; pub use nearai::NearAiProvider;
pub use openai::OpenAiProvider;
pub use provider::{ pub use provider::{
ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall, ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
@@ -20,33 +15,10 @@ pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, ToolSelection};
use std::sync::Arc; use std::sync::Arc;
use crate::config::{LlmConfig, LlmProvider as LlmProviderType}; use crate::config::LlmConfig;
use crate::error::LlmError; use crate::error::LlmError;
/// Create an LLM provider based on configuration. /// Create an LLM provider based on configuration.
pub fn create_llm_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> { pub fn create_llm_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
match config.provider { Ok(Arc::new(NearAiProvider::new(config.nearai.clone())))
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())))
}
}
} }
-335
View File
@@ -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<OpenAiMessage> {
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<OpenAiMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<OpenAiTool>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<serde_json::Value>,
}
#[derive(Debug, Serialize, Deserialize)]
struct OpenAiMessage {
role: String,
#[serde(skip_serializing_if = "Option::is_none")]
content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<Vec<OpenAiToolCall>>,
}
#[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<OpenAiChoice>,
usage: OpenAiUsage,
}
#[derive(Debug, Deserialize)]
struct OpenAiChoice {
message: OpenAiResponseMessage,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct OpenAiResponseMessage {
content: Option<String>,
tool_calls: Option<Vec<OpenAiToolCall>>,
}
#[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<String>,
}
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<CompletionResponse, LlmError> {
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<ToolCompletionResponse, LlmError> {
let tools: Vec<OpenAiTool> = 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<ToolCall> = 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()),
})
}
}
+163 -56
View File
@@ -2,7 +2,7 @@
//! //!
//! These tools allow the agent to: //! These tools allow the agent to:
//! - Search past memories, decisions, and context //! - Search past memories, decisions, and context
//! - Write important information to long-term memory //! - Read and write files in the workspace
//! //!
//! # Usage //! # Usage
//! //!
@@ -18,7 +18,7 @@ use async_trait::async_trait;
use crate::context::JobContext; use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput}; use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::workspace::Workspace; use crate::workspace::{Workspace, paths};
/// Tool for searching workspace memory. /// Tool for searching workspace memory.
/// ///
@@ -135,7 +135,8 @@ impl Tool for MemoryWriteTool {
fn description(&self) -> &str { fn description(&self) -> &str {
"Write to persistent memory. Use for important facts, decisions, preferences, \ "Write to persistent memory. Use for important facts, decisions, preferences, \
or lessons learned that should be remembered across sessions. Use 'memory' target \ 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 { fn parameters_schema(&self) -> serde_json::Value {
@@ -148,9 +149,13 @@ impl Tool for MemoryWriteTool {
}, },
"target": { "target": {
"type": "string", "type": "string",
"enum": ["memory", "daily_log"], "description": "Where to write: 'memory' for MEMORY.md, 'daily_log' for today's log, or a path like 'projects/alpha/notes.md'",
"description": "Where to write: 'memory' for long-term curated facts, 'daily_log' for timestamped session notes",
"default": "daily_log" "default": "daily_log"
},
"append": {
"type": "boolean",
"description": "If true, append to existing content. If false, replace entirely.",
"default": true
} }
}, },
"required": ["content"] "required": ["content"]
@@ -182,30 +187,53 @@ impl Tool for MemoryWriteTool {
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.unwrap_or("daily_log"); .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" => { "memory" => {
if append {
self.workspace self.workspace
.append_memory(content) .append_memory(content)
.await .await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; .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" => { "daily_log" => {
self.workspace self.workspace
.append_daily_log(content) .append_daily_log(content)
.await .await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?; .map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
format!("daily/{}.md", chrono::Utc::now().format("%Y-%m-%d"))
} }
_ => { path => {
return Err(ToolError::InvalidParameters(format!( if append {
"invalid target '{}', must be 'memory' or 'daily_log'", self.workspace
target .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!({ let output = serde_json::json!({
"status": "written", "status": "written",
"target": target, "path": path,
"append": append,
"content_length": content.len(), "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, /// Use this to read the full content of any file in the workspace.
/// or other specific documents.
pub struct MemoryReadTool { pub struct MemoryReadTool {
workspace: Arc<Workspace>, workspace: Arc<Workspace>,
} }
@@ -239,25 +266,20 @@ impl Tool for MemoryReadTool {
} }
fn description(&self) -> &str { fn description(&self) -> &str {
"Read a specific memory document by type. Use this to read identity files, \ "Read a file from the workspace. Use this to read identity files, \
heartbeat checklist, or full memory document content." heartbeat checklist, memory, daily logs, or any custom file."
} }
fn parameters_schema(&self) -> serde_json::Value { fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({ serde_json::json!({
"type": "object", "type": "object",
"properties": { "properties": {
"doc_type": { "path": {
"type": "string", "type": "string",
"enum": ["memory", "daily_log", "identity", "soul", "agents", "user", "heartbeat"], "description": "Path to the file (e.g., 'MEMORY.md', 'daily/2024-01-15.md', 'projects/alpha/notes.md')"
"description": "The type of document to read"
},
"title": {
"type": "string",
"description": "Optional title (required for daily_log, format: YYYY-MM-DD)"
} }
}, },
"required": ["doc_type"] "required": ["path"]
}) })
} }
@@ -268,27 +290,19 @@ impl Tool for MemoryReadTool {
) -> Result<ToolOutput, ToolError> { ) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let doc_type_str = params let path = params
.get("doc_type") .get("path")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.ok_or_else(|| { .ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".to_string()))?;
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 let doc = self
.workspace .workspace
.get_document(doc_type, title) .read(path)
.await .await
.map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?; .map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?;
let output = serde_json::json!({ let output = serde_json::json!({
"doc_type": doc_type_str, "path": doc.path,
"title": doc.title,
"content": doc.content, "content": doc.content,
"word_count": doc.word_count(), "word_count": doc.word_count(),
"updated_at": doc.updated_at.to_rfc3339(), "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<Workspace>,
}
impl MemoryListTool {
/// Create a new memory list tool.
pub fn new(workspace: Arc<Workspace>) -> 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<ToolOutput, ToolError> {
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::<Vec<_>>(),
"count": entries.len(),
});
Ok(ToolOutput::success(output, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false // Internal tool
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
// Integration tests would require a database connection. fn make_test_workspace() -> Arc<Workspace> {
// Unit tests for parameter validation: Arc::new(Workspace::new(
#[test]
fn test_memory_search_schema() {
let workspace = Arc::new(Workspace::new(
"test_user", "test_user",
deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new( deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new(
tokio_postgres::Config::new(), tokio_postgres::Config::new(),
@@ -319,7 +405,12 @@ mod tests {
)) ))
.build() .build()
.unwrap(), .unwrap(),
)); ))
}
#[test]
fn test_memory_search_schema() {
let workspace = make_test_workspace();
let tool = MemorySearchTool::new(workspace); let tool = MemorySearchTool::new(workspace);
assert_eq!(tool.name(), "memory_search"); assert_eq!(tool.name(), "memory_search");
@@ -337,26 +428,42 @@ mod tests {
#[test] #[test]
fn test_memory_write_schema() { fn test_memory_write_schema() {
let workspace = Arc::new(Workspace::new( let workspace = make_test_workspace();
"test_user",
deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new(
tokio_postgres::Config::new(),
tokio_postgres::NoTls,
))
.build()
.unwrap(),
));
let tool = MemoryWriteTool::new(workspace); let tool = MemoryWriteTool::new(workspace);
assert_eq!(tool.name(), "memory_write"); assert_eq!(tool.name(), "memory_write");
let schema = tool.parameters_schema(); let schema = tool.parameters_schema();
assert!(schema["properties"]["content"].is_object()); 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!( assert!(
schema["properties"]["target"]["enum"] schema["required"]
.as_array() .as_array()
.unwrap() .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());
}
} }
+107 -122
View File
@@ -1,99 +1,32 @@
//! Memory document types. //! Memory document types for the workspace.
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use uuid::Uuid; use uuid::Uuid;
use crate::error::WorkspaceError; /// Well-known document paths.
/// Document type in the workspace.
/// ///
/// Each type represents a different kind of persistent memory: /// These are conventional paths that have special meaning in the workspace.
/// - **Memory**: Long-term curated facts and decisions (MEMORY.md) /// Agents can create arbitrary paths beyond these.
/// - **DailyLog**: Append-only daily notes (memory/YYYY-MM-DD.md) pub mod paths {
/// - **Identity**: Agent name and personality /// Long-term curated memory.
/// - **Soul**: Core values and behavior principles pub const MEMORY: &str = "MEMORY.md";
/// - **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). /// Agent identity (name, nature, vibe).
Identity, pub const IDENTITY: &str = "IDENTITY.md";
/// Core values and principles (SOUL.md). /// Core values and principles.
Soul, pub const SOUL: &str = "SOUL.md";
/// Behavior instructions (AGENTS.md). /// Behavior instructions.
Agents, pub const AGENTS: &str = "AGENTS.md";
/// User context (USER.md). /// User context (name, preferences).
User, pub const USER: &str = "USER.md";
/// Periodic checklist (HEARTBEAT.md). /// Periodic checklist for heartbeat.
Heartbeat, pub const HEARTBEAT: &str = "HEARTBEAT.md";
} /// Root runbook/readme.
pub const README: &str = "README.md";
impl DocType { /// Daily logs directory.
/// Get the string representation. pub const DAILY_DIR: &str = "daily/";
pub fn as_str(&self) -> &'static str { /// Context directory (for identity-related docs).
match self { pub const CONTEXT_DIR: &str = "context/";
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<Self, Self::Error> {
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. /// A memory document stored in the database.
@@ -105,10 +38,8 @@ pub struct MemoryDocument {
pub user_id: String, pub user_id: String,
/// Optional agent ID for multi-agent isolation. /// Optional agent ID for multi-agent isolation.
pub agent_id: Option<Uuid>, pub agent_id: Option<Uuid>,
/// Document type. /// File path within the workspace (e.g., "context/vision.md").
pub doc_type: DocType, pub path: String,
/// Optional title (e.g., date for daily logs).
pub title: Option<String>,
/// Full document content. /// Full document content.
pub content: String, pub content: String,
/// Creation timestamp. /// Creation timestamp.
@@ -120,20 +51,18 @@ pub struct MemoryDocument {
} }
impl MemoryDocument { impl MemoryDocument {
/// Create a new document (not persisted yet). /// Create a new document with a path.
pub fn new( pub fn new(
user_id: impl Into<String>, user_id: impl Into<String>,
agent_id: Option<Uuid>, agent_id: Option<Uuid>,
doc_type: DocType, path: impl Into<String>,
title: Option<String>,
) -> Self { ) -> Self {
let now = Utc::now(); let now = Utc::now();
Self { Self {
id: Uuid::new_v4(), id: Uuid::new_v4(),
user_id: user_id.into(), user_id: user_id.into(),
agent_id, agent_id,
doc_type, path: path.into(),
title,
content: String::new(), content: String::new(),
created_at: now, created_at: now,
updated_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. /// Check if the document is empty.
pub fn is_empty(&self) -> bool { pub fn is_empty(&self) -> bool {
self.content.is_empty() self.content.is_empty()
@@ -150,6 +90,34 @@ impl MemoryDocument {
pub fn word_count(&self) -> usize { pub fn word_count(&self) -> usize {
self.content.split_whitespace().count() 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<DateTime<Utc>>,
/// Preview of content (first ~200 chars, None for directories).
pub content_preview: Option<String>,
}
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. /// A chunk of a memory document for search indexing.
@@ -194,43 +162,60 @@ mod tests {
use super::*; use super::*;
#[test] #[test]
fn test_doc_type_roundtrip() { fn test_memory_document_new() {
for doc_type in [ let doc = MemoryDocument::new("user1", None, "context/vision.md");
DocType::Memory, assert_eq!(doc.user_id, "user1");
DocType::DailyLog, assert_eq!(doc.path, "context/vision.md");
DocType::Identity, assert!(doc.content.is_empty());
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] #[test]
fn test_singleton_types() { fn test_memory_document_file_name() {
assert!(DocType::Memory.is_singleton()); let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md");
assert!(DocType::Heartbeat.is_singleton()); assert_eq!(doc.file_name(), "README.md");
assert!(!DocType::DailyLog.is_singleton());
} }
#[test] #[test]
fn test_identity_documents() { fn test_memory_document_parent_dir() {
assert!(DocType::Soul.is_identity_document()); let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md");
assert!(DocType::Agents.is_identity_document()); assert_eq!(doc.parent_dir(), Some("projects/alpha"));
assert!(!DocType::Memory.is_identity_document());
assert!(!DocType::DailyLog.is_identity_document()); let root_doc = MemoryDocument::new("user1", None, "README.md");
assert_eq!(root_doc.parent_dir(), None);
} }
#[test] #[test]
fn test_memory_document_word_count() { 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); assert_eq!(doc.word_count(), 0);
doc.content = "Hello world, this is a test.".to_string(); doc.content = "Hello world, this is a test.".to_string();
assert_eq!(doc.word_count(), 6); 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");
}
} }
+202 -95
View File
@@ -1,37 +1,44 @@
//! Workspace and memory system (OpenClaw-inspired). //! Workspace and memory system (OpenClaw-inspired).
//! //!
//! The workspace provides persistent memory for agents: //! The workspace provides persistent memory for agents with a flexible
//! - **MEMORY.md**: Long-term curated memory (facts, decisions, preferences) //! filesystem-like structure. Agents can create arbitrary markdown file
//! - **Daily logs**: Append-only daily notes (raw context) //! hierarchies that get indexed for full-text and semantic search.
//! - **Identity files**: Agent personality and user context
//! - **HEARTBEAT.md**: Periodic checklist for proactive execution
//! //!
//! Memory is searchable via hybrid search (FTS + semantic embeddings). //! # Filesystem-like API
//!
//! # Architecture
//! //!
//! ```text //! ```text
//! ┌─────────────────────────────────────────────────────────────┐ //! workspace/
//! Workspace │ //! ├── README.md <- Root runbook/index
//! │ ┌────────────────┐ ┌────────────────┐ ┌──────────────┐ │ //! ├── MEMORY.md <- Long-term curated memory
//! │ │ MemoryDocument │ MemoryChunk │ Search │ │ //! ├── HEARTBEAT.md <- Periodic checklist
//! │ │ (full docs) │──│ (chunked) │──│ (FTS+vector) │ │ //! ├── context/ <- Identity and context
//! │ └────────────────┘ └────────────────┘ └──────────────┘ │ //! │ ├── vision.md
//! │ │ │ │ │ //! │ └── priorities.md
//! │ └───────────────────┴──────────────────┘ │ //! ├── daily/ <- Daily logs
//! │ │ │ //! │ ├── 2024-01-15.md
//! │ ┌──────┴──────┐ │ //! │ └── 2024-01-16.md
//! │ Repository │ │ //! ├── projects/ <- Arbitrary structure
//! │ │ (PostgreSQL)│ │ //! │ └── 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 //! # Key Patterns
//! //!
//! 1. **Memory is persistence**: If you want to remember something, write it //! 1. **Memory is persistence**: If you want to remember something, write it
//! 2. **Two-tier memory**: Daily logs (raw) + MEMORY.md (curated) //! 2. **Flexible structure**: Create any directory/file hierarchy you need
//! 3. **Hybrid search**: Vector similarity + BM25 full-text via RRF //! 3. **Self-documenting**: Use README.md files to describe directory structure
//! 4. **Hybrid search**: Vector similarity + BM25 full-text via RRF
mod chunker; mod chunker;
mod document; mod document;
@@ -40,7 +47,7 @@ mod repository;
mod search; mod search;
pub use chunker::{ChunkConfig, chunk_document}; 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 embeddings::{EmbeddingProvider, OpenAiEmbeddings};
pub use repository::Repository; pub use repository::Repository;
pub use search::{SearchConfig, SearchResult}; pub use search::{SearchConfig, SearchResult};
@@ -101,15 +108,127 @@ impl Workspace {
self.agent_id 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<MemoryDocument, WorkspaceError> {
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<MemoryDocument, WorkspaceError> {
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<bool, WorkspaceError> {
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<Vec<WorkspaceEntry>, 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<Vec<String>, 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). /// Get the main MEMORY.md document (long-term curated memory).
/// ///
/// Creates it if it doesn't exist. /// Creates it if it doesn't exist.
pub async fn memory(&self) -> Result<MemoryDocument, WorkspaceError> { pub async fn memory(&self) -> Result<MemoryDocument, WorkspaceError> {
self.repo self.read_or_create(paths::MEMORY).await
.get_or_create_document(&self.user_id, self.agent_id, DocType::Memory, None)
.await
} }
/// Get today's daily log. /// Get today's daily log.
@@ -122,38 +241,23 @@ impl Workspace {
/// Get a daily log for a specific date. /// Get a daily log for a specific date.
pub async fn daily_log(&self, date: NaiveDate) -> Result<MemoryDocument, WorkspaceError> { pub async fn daily_log(&self, date: NaiveDate) -> Result<MemoryDocument, WorkspaceError> {
let title = date.format("%Y-%m-%d").to_string(); let path = format!("daily/{}.md", date.format("%Y-%m-%d"));
self.repo self.read_or_create(&path).await
.get_or_create_document(
&self.user_id,
self.agent_id,
DocType::DailyLog,
Some(&title),
)
.await
} }
/// Get the heartbeat checklist (HEARTBEAT.md). /// Get the heartbeat checklist (HEARTBEAT.md).
pub async fn heartbeat_checklist(&self) -> Result<Option<String>, WorkspaceError> { pub async fn heartbeat_checklist(&self) -> Result<Option<String>, WorkspaceError> {
match self match self.read(paths::HEARTBEAT).await {
.repo
.get_document(&self.user_id, self.agent_id, DocType::Heartbeat, None)
.await
{
Ok(doc) => Ok(Some(doc.content)), Ok(doc) => Ok(Some(doc.content)),
Err(WorkspaceError::DocumentNotFound { .. }) => Ok(None), Err(WorkspaceError::DocumentNotFound { .. }) => Ok(None),
Err(e) => Err(e), Err(e) => Err(e),
} }
} }
/// Get a document by type. /// Helper to read or create a file.
pub async fn get_document( async fn read_or_create(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
&self,
doc_type: DocType,
title: Option<&str>,
) -> Result<MemoryDocument, WorkspaceError> {
self.repo 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 .await
} }
@@ -164,6 +268,7 @@ impl Workspace {
/// This is for important facts, decisions, and preferences worth /// This is for important facts, decisions, and preferences worth
/// remembering long-term. /// remembering long-term.
pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> { pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> {
// Use double newline for memory entries (semantic separation)
let doc = self.memory().await?; let doc = self.memory().await?;
let new_content = if doc.content.is_empty() { let new_content = if doc.content.is_empty() {
entry.to_string() entry.to_string()
@@ -179,34 +284,11 @@ impl Workspace {
/// ///
/// Daily logs are raw, append-only notes for the current day. /// Daily logs are raw, append-only notes for the current day.
pub async fn append_daily_log(&self, entry: &str) -> Result<(), WorkspaceError> { 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 timestamp = Utc::now().format("%H:%M:%S");
let timestamped_entry = format!("[{}] {}", timestamp, entry); let timestamped_entry = format!("[{}] {}", timestamp, entry);
self.append(&path, &timestamped_entry).await
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 ==================== // ==================== System Prompt ====================
@@ -219,19 +301,15 @@ impl Workspace {
let mut parts = Vec::new(); let mut parts = Vec::new();
// Load identity files in order of importance // Load identity files in order of importance
let identity_types = [ let identity_files = [
(DocType::Agents, "## Agent Instructions"), (paths::AGENTS, "## Agent Instructions"),
(DocType::Soul, "## Core Values"), (paths::SOUL, "## Core Values"),
(DocType::User, "## User Context"), (paths::USER, "## User Context"),
(DocType::Identity, "## Identity"), (paths::IDENTITY, "## Identity"),
]; ];
for (doc_type, header) in identity_types { for (path, header) in identity_files {
if let Ok(doc) = self if let Ok(doc) = self.read(path).await {
.repo
.get_document(&self.user_id, self.agent_id, doc_type, None)
.await
{
if !doc.content.is_empty() { if !doc.content.is_empty() {
parts.push(format!("{}\n\n{}", header, doc.content)); 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
#[test] #[test]
fn test_doc_type_display() { fn test_normalize_path() {
assert_eq!(DocType::Memory.as_str(), "memory"); assert_eq!(normalize_path("foo/bar"), "foo/bar");
assert_eq!(DocType::DailyLog.as_str(), "daily_log"); assert_eq!(normalize_path("/foo/bar/"), "foo/bar");
assert_eq!(DocType::Heartbeat.as_str(), "heartbeat"); assert_eq!(normalize_path("foo//bar"), "foo/bar");
assert_eq!(normalize_path(" /foo/ "), "foo");
assert_eq!(normalize_path("README.md"), "README.md");
} }
#[test] #[test]
fn test_doc_type_parse() { fn test_normalize_directory() {
assert_eq!(DocType::try_from("memory").unwrap(), DocType::Memory); assert_eq!(normalize_directory("foo/bar/"), "foo/bar");
assert_eq!(DocType::try_from("daily_log").unwrap(), DocType::DailyLog); assert_eq!(normalize_directory("foo/bar"), "foo/bar");
assert!(DocType::try_from("invalid").is_err()); assert_eq!(normalize_directory("/"), "");
assert_eq!(normalize_directory(""), "");
} }
} }
+127 -74
View File
@@ -4,14 +4,14 @@
//! - Documents in `memory_documents` table //! - Documents in `memory_documents` table
//! - Chunks in `memory_chunks` table (with FTS and vector indexes) //! - Chunks in `memory_chunks` table (with FTS and vector indexes)
use chrono::Utc; use chrono::{DateTime, Utc};
use deadpool_postgres::Pool; use deadpool_postgres::Pool;
use pgvector::Vector; use pgvector::Vector;
use uuid::Uuid; use uuid::Uuid;
use crate::error::WorkspaceError; 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}; use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
/// Database repository for workspace operations. /// Database repository for workspace operations.
@@ -37,50 +37,34 @@ impl Repository {
// ==================== Document Operations ==================== // ==================== Document Operations ====================
/// Get a document by type and optional title. /// Get a document by its path.
pub async fn get_document( pub async fn get_document_by_path(
&self, &self,
user_id: &str, user_id: &str,
agent_id: Option<Uuid>, agent_id: Option<Uuid>,
doc_type: DocType, path: &str,
title: Option<&str>,
) -> Result<MemoryDocument, WorkspaceError> { ) -> Result<MemoryDocument, WorkspaceError> {
let conn = self.conn().await?; let conn = self.conn().await?;
let row = if let Some(title) = title { let row = conn
conn.query_opt( .query_opt(
r#" r#"
SELECT id, user_id, agent_id, doc_type, title, content, SELECT id, user_id, agent_id, path, content,
created_at, updated_at, metadata created_at, updated_at, metadata
FROM memory_documents FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND path = $3
AND doc_type = $3 AND title = $4
"#, "#,
&[&user_id, &agent_id, &doc_type.as_str(), &title], &[&user_id, &agent_id, &path],
) )
.await .await
} else { .map_err(|e| WorkspaceError::SearchFailed {
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), reason: format!("Query failed: {}", e),
})?; })?;
match row { match row {
Some(row) => Ok(self.row_to_document(&row)?), Some(row) => Ok(self.row_to_document(&row)),
None => Err(WorkspaceError::DocumentNotFound { None => Err(WorkspaceError::DocumentNotFound {
doc_type: doc_type.to_string(), doc_type: path.to_string(),
user_id: user_id.to_string(), user_id: user_id.to_string(),
}), }),
} }
@@ -93,7 +77,7 @@ impl Repository {
let row = conn let row = conn
.query_opt( .query_opt(
r#" r#"
SELECT id, user_id, agent_id, doc_type, title, content, SELECT id, user_id, agent_id, path, content,
created_at, updated_at, metadata created_at, updated_at, metadata
FROM memory_documents WHERE id = $1 FROM memory_documents WHERE id = $1
"#, "#,
@@ -105,7 +89,7 @@ impl Repository {
})?; })?;
match row { match row {
Some(row) => Ok(self.row_to_document(&row)?), Some(row) => Ok(self.row_to_document(&row)),
None => Err(WorkspaceError::DocumentNotFound { None => Err(WorkspaceError::DocumentNotFound {
doc_type: "unknown".to_string(), doc_type: "unknown".to_string(),
user_id: "unknown".to_string(), user_id: "unknown".to_string(),
@@ -113,16 +97,15 @@ impl Repository {
} }
} }
/// Get or create a document. /// Get or create a document by path.
pub async fn get_or_create_document( pub async fn get_or_create_document_by_path(
&self, &self,
user_id: &str, user_id: &str,
agent_id: Option<Uuid>, agent_id: Option<Uuid>,
doc_type: DocType, path: &str,
title: Option<&str>,
) -> Result<MemoryDocument, WorkspaceError> { ) -> Result<MemoryDocument, WorkspaceError> {
// Try to get existing document first // 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), Ok(doc) => return Ok(doc),
Err(WorkspaceError::DocumentNotFound { .. }) => {} Err(WorkspaceError::DocumentNotFound { .. }) => {}
Err(e) => return Err(e), Err(e) => return Err(e),
@@ -132,14 +115,15 @@ impl Repository {
let conn = self.conn().await?; let conn = self.conn().await?;
let id = Uuid::new_v4(); let id = Uuid::new_v4();
let now = Utc::now(); let now = Utc::now();
let metadata = serde_json::json!({});
conn.execute( conn.execute(
r#" r#"
INSERT INTO memory_documents (id, user_id, agent_id, doc_type, title, content, created_at, updated_at) INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, '', $6, $7) VALUES ($1, $2, $3, $4, '', $5, $6, $7)
ON CONFLICT (user_id, agent_id, doc_type, title) DO NOTHING 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 .await
.map_err(|e| WorkspaceError::SearchFailed { .map_err(|e| WorkspaceError::SearchFailed {
@@ -147,7 +131,7 @@ impl Repository {
})?; })?;
// Fetch the document (might have been created by concurrent request) // 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. /// Update a document's content.
@@ -166,31 +150,108 @@ impl Repository {
Ok(()) 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<Uuid>,
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<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, 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<DateTime<Utc>> = 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<Uuid>,
) -> Result<Vec<String>, 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( pub async fn list_documents(
&self, &self,
user_id: &str, user_id: &str,
agent_id: Option<Uuid>, agent_id: Option<Uuid>,
doc_type: Option<DocType>,
) -> Result<Vec<MemoryDocument>, WorkspaceError> { ) -> Result<Vec<MemoryDocument>, WorkspaceError> {
let conn = self.conn().await?; let conn = self.conn().await?;
let rows = if let Some(dt) = doc_type { let rows = conn
conn.query( .query(
r#" 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
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 created_at, updated_at, metadata
FROM memory_documents FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
@@ -199,30 +260,24 @@ impl Repository {
&[&user_id, &agent_id], &[&user_id, &agent_id],
) )
.await .await
}; .map_err(|e| WorkspaceError::SearchFailed {
let rows = rows.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e), 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<MemoryDocument, WorkspaceError> { fn row_to_document(&self, row: &tokio_postgres::Row) -> MemoryDocument {
let doc_type_str: String = row.get("doc_type"); MemoryDocument {
let doc_type = DocType::try_from(doc_type_str.as_str())?;
Ok(MemoryDocument {
id: row.get("id"), id: row.get("id"),
user_id: row.get("user_id"), user_id: row.get("user_id"),
agent_id: row.get("agent_id"), agent_id: row.get("agent_id"),
doc_type, path: row.get("path"),
title: row.get("title"),
content: row.get("content"), content: row.get("content"),
created_at: row.get("created_at"), created_at: row.get("created_at"),
updated_at: row.get("updated_at"), updated_at: row.get("updated_at"),
metadata: row.get("metadata"), metadata: row.get("metadata"),
}) }
} }
// ==================== Chunk Operations ==================== // ==================== Chunk Operations ====================
@@ -374,7 +429,6 @@ impl Repository {
) -> Result<Vec<RankedResult>, WorkspaceError> { ) -> Result<Vec<RankedResult>, WorkspaceError> {
let conn = self.conn().await?; let conn = self.conn().await?;
// Use plainto_tsquery for natural language queries
let rows = conn let rows = conn
.query( .query(
r#" r#"
@@ -401,7 +455,7 @@ impl Repository {
chunk_id: row.get("chunk_id"), chunk_id: row.get("chunk_id"),
document_id: row.get("document_id"), document_id: row.get("document_id"),
content: row.get("content"), content: row.get("content"),
rank: (i + 1) as u32, // 1-based rank rank: (i + 1) as u32,
}) })
.collect()) .collect())
} }
@@ -417,7 +471,6 @@ impl Repository {
let conn = self.conn().await?; let conn = self.conn().await?;
let embedding_vec = Vector::from(embedding.to_vec()); let embedding_vec = Vector::from(embedding.to_vec());
// Use cosine distance (<=>)
let rows = conn let rows = conn
.query( .query(
r#" r#"
@@ -444,7 +497,7 @@ impl Repository {
chunk_id: row.get("chunk_id"), chunk_id: row.get("chunk_id"),
document_id: row.get("document_id"), document_id: row.get("document_id"),
content: row.get("content"), content: row.get("content"),
rank: (i + 1) as u32, // 1-based rank rank: (i + 1) as u32,
}) })
.collect()) .collect())
} }