gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify workspace to path-based storage, remove legacy code

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files
src/config.rs
+6 -132
View File
@@ -60,52 +60,10 @@ impl DatabaseConfig {
}
}
/// LLM provider configuration.
/// LLM provider configuration (NEAR AI only).
#[derive(Debug, Clone)]
pub struct LlmConfig {
pub provider: LlmProvider,
pub openai: Option<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>,
pub nearai: NearAiConfig,
}
/// NEAR AI chat-api configuration.
@@ -121,66 +79,16 @@ pub struct NearAiConfig {
impl LlmConfig {
fn from_env() -> Result<Self, ConfigError> {
let provider: LlmProvider = optional_env("LLM_PROVIDER")?
.map(|s| s.parse())
.transpose()?
.unwrap_or(LlmProvider::NearAi);
let session_token = required_env("NEARAI_SESSION_TOKEN")?;
let openai = if let Some(api_key) = optional_env("OPENAI_API_KEY")? {
Some(OpenAiConfig {
api_key: SecretString::from(api_key),
model: optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4-turbo".to_string()),
base_url: optional_env("OPENAI_BASE_URL")?,
})
} else {
None
};
let anthropic = if let Some(api_key) = optional_env("ANTHROPIC_API_KEY")? {
Some(AnthropicConfig {
api_key: SecretString::from(api_key),
model: optional_env("ANTHROPIC_MODEL")?
.unwrap_or_else(|| "claude-3-opus-20240229".to_string()),
base_url: optional_env("ANTHROPIC_BASE_URL")?,
})
} else {
None
};
let nearai = if let Some(session_token) = optional_env("NEARAI_SESSION_TOKEN")? {
Some(NearAiConfig {
Ok(Self {
nearai: NearAiConfig {
session_token: SecretString::from(session_token),
model: optional_env("NEARAI_MODEL")?
.unwrap_or_else(|| "claude-3-5-sonnet-20241022".to_string()),
base_url: optional_env("NEARAI_BASE_URL")?
.unwrap_or_else(|| "https://api.near.ai".to_string()),
})
} else {
None
};
// Validate that the selected provider has configuration
match provider {
LlmProvider::OpenAi if openai.is_none() => {
return Err(ConfigError::MissingEnvVar("OPENAI_API_KEY".to_string()));
}
LlmProvider::Anthropic if anthropic.is_none() => {
return Err(ConfigError::MissingEnvVar("ANTHROPIC_API_KEY".to_string()));
}
LlmProvider::NearAi if nearai.is_none() => {
return Err(ConfigError::MissingEnvVar(
"NEARAI_SESSION_TOKEN".to_string(),
));
}
// Provider has valid configuration
LlmProvider::OpenAi | LlmProvider::Anthropic | LlmProvider::NearAi => {}
}
Ok(Self {
provider,
openai,
anthropic,
nearai,
},
})
}
}
@@ -354,37 +262,3 @@ where
.transpose()
.map(|opt| opt.unwrap_or(default))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_llm_provider_parsing() {
assert_eq!(
"openai".parse::<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());
}
}
src/lib.rs
+1 -1
View File
@@ -63,5 +63,5 @@ pub mod prelude {
pub use crate::llm::LlmProvider;
pub use crate::safety::{SanitizedOutput, Sanitizer};
pub use crate::tools::{Tool, ToolOutput, ToolRegistry};
pub use crate::workspace::{DocType, MemoryDocument, Workspace};
pub use crate::workspace::{MemoryDocument, Workspace};
}
src/llm/anthropic.rs
-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()),
})
}
}
src/llm/mod.rs
+3 -31
View File
@@ -1,17 +1,12 @@
//! LLM integration for the agent.
//!
//! Provides a unified interface to different LLM providers (OpenAI, Anthropic, NEAR AI)
//! and implements reasoning capabilities for planning, tool selection, and evaluation.
//! Uses the NEAR AI chat-api as the unified LLM provider.
mod anthropic;
mod nearai;
mod openai;
mod provider;
mod reasoning;
pub use anthropic::AnthropicProvider;
pub use nearai::NearAiProvider;
pub use openai::OpenAiProvider;
pub use provider::{
ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
@@ -20,33 +15,10 @@ pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, ToolSelection};
use std::sync::Arc;
use crate::config::{LlmConfig, LlmProvider as LlmProviderType};
use crate::config::LlmConfig;
use crate::error::LlmError;
/// Create an LLM provider based on configuration.
pub fn create_llm_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
match config.provider {
LlmProviderType::OpenAi => {
let openai_config = config.openai.as_ref().ok_or_else(|| LlmError::AuthFailed {
provider: "openai".to_string(),
})?;
Ok(Arc::new(OpenAiProvider::new(openai_config.clone())))
}
LlmProviderType::Anthropic => {
let anthropic_config =
config
.anthropic
.as_ref()
.ok_or_else(|| LlmError::AuthFailed {
provider: "anthropic".to_string(),
})?;
Ok(Arc::new(AnthropicProvider::new(anthropic_config.clone())))
}
LlmProviderType::NearAi => {
let nearai_config = config.nearai.as_ref().ok_or_else(|| LlmError::AuthFailed {
provider: "nearai".to_string(),
})?;
Ok(Arc::new(NearAiProvider::new(nearai_config.clone())))
}
}
Ok(Arc::new(NearAiProvider::new(config.nearai.clone())))
}
src/llm/openai.rs
-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()),
})
}
}
src/tools/builtin/memory.rs
+168 -61
View File
@@ -2,7 +2,7 @@
//!
//! These tools allow the agent to:
//! - Search past memories, decisions, and context
//! - Write important information to long-term memory
//! - Read and write files in the workspace
//!
//! # Usage
//!
@@ -18,7 +18,7 @@ use async_trait::async_trait;
use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::workspace::Workspace;
use crate::workspace::{Workspace, paths};
/// Tool for searching workspace memory.
///
@@ -135,7 +135,8 @@ impl Tool for MemoryWriteTool {
fn description(&self) -> &str {
"Write to persistent memory. Use for important facts, decisions, preferences, \
or lessons learned that should be remembered across sessions. Use 'memory' target \
for curated long-term facts, 'daily_log' for timestamped session notes."
for curated long-term facts, 'daily_log' for timestamped session notes, or \
provide a custom path for arbitrary file creation."
}
fn parameters_schema(&self) -> serde_json::Value {
@@ -148,9 +149,13 @@ impl Tool for MemoryWriteTool {
},
"target": {
"type": "string",
"enum": ["memory", "daily_log"],
"description": "Where to write: 'memory' for long-term curated facts, 'daily_log' for timestamped session notes",
"description": "Where to write: 'memory' for MEMORY.md, 'daily_log' for today's log, or a path like 'projects/alpha/notes.md'",
"default": "daily_log"
},
"append": {
"type": "boolean",
"description": "If true, append to existing content. If false, replace entirely.",
"default": true
}
},
"required": ["content"]
@@ -182,30 +187,53 @@ impl Tool for MemoryWriteTool {
.and_then(|v| v.as_str())
.unwrap_or("daily_log");
match target {
let append = params
.get("append")
.and_then(|v| v.as_bool())
.unwrap_or(true);
let path = match target {
"memory" => {
self.workspace
.append_memory(content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
if append {
self.workspace
.append_memory(content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
} else {
self.workspace
.write(paths::MEMORY, content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
}
paths::MEMORY.to_string()
}
"daily_log" => {
self.workspace
.append_daily_log(content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
format!("daily/{}.md", chrono::Utc::now().format("%Y-%m-%d"))
}
_ => {
return Err(ToolError::InvalidParameters(format!(
"invalid target '{}', must be 'memory' or 'daily_log'",
target
)));
path => {
if append {
self.workspace
.append(path, content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
} else {
self.workspace
.write(path, content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
}
path.to_string()
}
}
};
let output = serde_json::json!({
"status": "written",
"target": target,
"path": path,
"append": append,
"content_length": content.len(),
});
@@ -217,10 +245,9 @@ impl Tool for MemoryWriteTool {
}
}
/// Tool for reading specific memory documents.
/// Tool for reading workspace files.
///
/// Use this to read the full content of identity files, heartbeat checklist,
/// or other specific documents.
/// Use this to read the full content of any file in the workspace.
pub struct MemoryReadTool {
workspace: Arc<Workspace>,
}
@@ -239,25 +266,20 @@ impl Tool for MemoryReadTool {
}
fn description(&self) -> &str {
"Read a specific memory document by type. Use this to read identity files, \
heartbeat checklist, or full memory document content."
"Read a file from the workspace. Use this to read identity files, \
heartbeat checklist, memory, daily logs, or any custom file."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"doc_type": {
"path": {
"type": "string",
"enum": ["memory", "daily_log", "identity", "soul", "agents", "user", "heartbeat"],
"description": "The type of document to read"
},
"title": {
"type": "string",
"description": "Optional title (required for daily_log, format: YYYY-MM-DD)"
"description": "Path to the file (e.g., 'MEMORY.md', 'daily/2024-01-15.md', 'projects/alpha/notes.md')"
}
},
"required": ["doc_type"]
"required": ["path"]
})
}
@@ -268,27 +290,19 @@ impl Tool for MemoryReadTool {
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let doc_type_str = params
.get("doc_type")
let path = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("missing 'doc_type' parameter".to_string())
})?;
let doc_type = crate::workspace::DocType::try_from(doc_type_str)
.map_err(|e| ToolError::InvalidParameters(e.to_string()))?;
let title = params.get("title").and_then(|v| v.as_str());
.ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".to_string()))?;
let doc = self
.workspace
.get_document(doc_type, title)
.read(path)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Read failed: {}", e)))?;
let output = serde_json::json!({
"doc_type": doc_type_str,
"title": doc.title,
"path": doc.path,
"content": doc.content,
"word_count": doc.word_count(),
"updated_at": doc.updated_at.to_rfc3339(),
@@ -302,16 +316,88 @@ impl Tool for MemoryReadTool {
}
}
/// Tool for listing workspace files.
///
/// Use this to explore the workspace structure.
pub struct MemoryListTool {
workspace: Arc<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)]
mod tests {
use super::*;
// Integration tests would require a database connection.
// Unit tests for parameter validation:
#[test]
fn test_memory_search_schema() {
let workspace = Arc::new(Workspace::new(
fn make_test_workspace() -> Arc<Workspace> {
Arc::new(Workspace::new(
"test_user",
deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new(
tokio_postgres::Config::new(),
@@ -319,7 +405,12 @@ mod tests {
))
.build()
.unwrap(),
));
))
}
#[test]
fn test_memory_search_schema() {
let workspace = make_test_workspace();
let tool = MemorySearchTool::new(workspace);
assert_eq!(tool.name(), "memory_search");
@@ -337,26 +428,42 @@ mod tests {
#[test]
fn test_memory_write_schema() {
let workspace = Arc::new(Workspace::new(
"test_user",
deadpool_postgres::Pool::builder(deadpool_postgres::Manager::new(
tokio_postgres::Config::new(),
tokio_postgres::NoTls,
))
.build()
.unwrap(),
));
let workspace = make_test_workspace();
let tool = MemoryWriteTool::new(workspace);
assert_eq!(tool.name(), "memory_write");
let schema = tool.parameters_schema();
assert!(schema["properties"]["content"].is_object());
assert!(schema["properties"]["target"].is_object());
assert!(schema["properties"]["append"].is_object());
}
#[test]
fn test_memory_read_schema() {
let workspace = make_test_workspace();
let tool = MemoryReadTool::new(workspace);
assert_eq!(tool.name(), "memory_read");
let schema = tool.parameters_schema();
assert!(schema["properties"]["path"].is_object());
assert!(
schema["properties"]["target"]["enum"]
schema["required"]
.as_array()
.unwrap()
.contains(&"memory".into())
.contains(&"path".into())
);
}
#[test]
fn test_memory_list_schema() {
let workspace = make_test_workspace();
let tool = MemoryListTool::new(workspace);
assert_eq!(tool.name(), "memory_list");
let schema = tool.parameters_schema();
assert!(schema["properties"]["directory"].is_object());
}
}
src/workspace/document.rs
+107 -122
View File
@@ -1,99 +1,32 @@
//! Memory document types.
//! Memory document types for the workspace.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::error::WorkspaceError;
/// Document type in the workspace.
/// Well-known document paths.
///
/// Each type represents a different kind of persistent memory:
/// - **Memory**: Long-term curated facts and decisions (MEMORY.md)
/// - **DailyLog**: Append-only daily notes (memory/YYYY-MM-DD.md)
/// - **Identity**: Agent name and personality
/// - **Soul**: Core values and behavior principles
/// - **Agents**: Behavior instructions
/// - **User**: User context (name, preferences)
/// - **Heartbeat**: Periodic checklist for proactive execution
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DocType {
/// Long-term curated memory (MEMORY.md equivalent).
Memory,
/// Daily append-only logs.
DailyLog,
/// These are conventional paths that have special meaning in the workspace.
/// Agents can create arbitrary paths beyond these.
pub mod paths {
/// Long-term curated memory.
pub const MEMORY: &str = "MEMORY.md";
/// Agent identity (name, nature, vibe).
Identity,
/// Core values and principles (SOUL.md).
Soul,
/// Behavior instructions (AGENTS.md).
Agents,
/// User context (USER.md).
User,
/// Periodic checklist (HEARTBEAT.md).
Heartbeat,
}
impl DocType {
/// Get the string representation.
pub fn as_str(&self) -> &'static str {
match self {
DocType::Memory => "memory",
DocType::DailyLog => "daily_log",
DocType::Identity => "identity",
DocType::Soul => "soul",
DocType::Agents => "agents",
DocType::User => "user",
DocType::Heartbeat => "heartbeat",
}
}
/// Check if this document type is a singleton (one per user/agent).
pub fn is_singleton(&self) -> bool {
match self {
DocType::Memory
| DocType::Identity
| DocType::Soul
| DocType::Agents
| DocType::User
| DocType::Heartbeat => true,
DocType::DailyLog => false,
}
}
/// Check if this document should be included in the system prompt.
pub fn is_identity_document(&self) -> bool {
matches!(
self,
DocType::Identity | DocType::Soul | DocType::Agents | DocType::User
)
}
}
impl TryFrom<&str> for DocType {
type Error = WorkspaceError;
fn try_from(s: &str) -> Result<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())
}
pub const IDENTITY: &str = "IDENTITY.md";
/// Core values and principles.
pub const SOUL: &str = "SOUL.md";
/// Behavior instructions.
pub const AGENTS: &str = "AGENTS.md";
/// User context (name, preferences).
pub const USER: &str = "USER.md";
/// Periodic checklist for heartbeat.
pub const HEARTBEAT: &str = "HEARTBEAT.md";
/// Root runbook/readme.
pub const README: &str = "README.md";
/// Daily logs directory.
pub const DAILY_DIR: &str = "daily/";
/// Context directory (for identity-related docs).
pub const CONTEXT_DIR: &str = "context/";
}
/// A memory document stored in the database.
@@ -105,10 +38,8 @@ pub struct MemoryDocument {
pub user_id: String,
/// Optional agent ID for multi-agent isolation.
pub agent_id: Option<Uuid>,
/// Document type.
pub doc_type: DocType,
/// Optional title (e.g., date for daily logs).
pub title: Option<String>,
/// File path within the workspace (e.g., "context/vision.md").
pub path: String,
/// Full document content.
pub content: String,
/// Creation timestamp.
@@ -120,20 +51,18 @@ pub struct MemoryDocument {
}
impl MemoryDocument {
/// Create a new document (not persisted yet).
/// Create a new document with a path.
pub fn new(
user_id: impl Into<String>,
agent_id: Option<Uuid>,
doc_type: DocType,
title: Option<String>,
path: impl Into<String>,
) -> Self {
let now = Utc::now();
Self {
id: Uuid::new_v4(),
user_id: user_id.into(),
agent_id,
doc_type,
title,
path: path.into(),
content: String::new(),
created_at: now,
updated_at: now,
@@ -141,6 +70,17 @@ impl MemoryDocument {
}
}
/// Get the file name from the path.
pub fn file_name(&self) -> &str {
self.path.rsplit('/').next().unwrap_or(&self.path)
}
/// Get the parent directory from the path.
pub fn parent_dir(&self) -> Option<&str> {
let idx = self.path.rfind('/')?;
Some(&self.path[..idx])
}
/// Check if the document is empty.
pub fn is_empty(&self) -> bool {
self.content.is_empty()
@@ -150,6 +90,34 @@ impl MemoryDocument {
pub fn word_count(&self) -> usize {
self.content.split_whitespace().count()
}
/// Check if this is a well-known identity document.
pub fn is_identity_document(&self) -> bool {
matches!(
self.path.as_str(),
paths::IDENTITY | paths::SOUL | paths::AGENTS | paths::USER
)
}
}
/// An entry in a workspace directory listing.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkspaceEntry {
/// Path relative to listing directory.
pub path: String,
/// True if this is a directory (has children).
pub is_directory: bool,
/// Last update timestamp (latest among children for directories).
pub updated_at: Option<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.
@@ -194,43 +162,60 @@ mod tests {
use super::*;
#[test]
fn test_doc_type_roundtrip() {
for doc_type in [
DocType::Memory,
DocType::DailyLog,
DocType::Identity,
DocType::Soul,
DocType::Agents,
DocType::User,
DocType::Heartbeat,
] {
let s = doc_type.as_str();
let parsed = DocType::try_from(s).unwrap();
assert_eq!(parsed, doc_type);
}
fn test_memory_document_new() {
let doc = MemoryDocument::new("user1", None, "context/vision.md");
assert_eq!(doc.user_id, "user1");
assert_eq!(doc.path, "context/vision.md");
assert!(doc.content.is_empty());
}
#[test]
fn test_singleton_types() {
assert!(DocType::Memory.is_singleton());
assert!(DocType::Heartbeat.is_singleton());
assert!(!DocType::DailyLog.is_singleton());
fn test_memory_document_file_name() {
let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md");
assert_eq!(doc.file_name(), "README.md");
}
#[test]
fn test_identity_documents() {
assert!(DocType::Soul.is_identity_document());
assert!(DocType::Agents.is_identity_document());
assert!(!DocType::Memory.is_identity_document());
assert!(!DocType::DailyLog.is_identity_document());
fn test_memory_document_parent_dir() {
let doc = MemoryDocument::new("user1", None, "projects/alpha/README.md");
assert_eq!(doc.parent_dir(), Some("projects/alpha"));
let root_doc = MemoryDocument::new("user1", None, "README.md");
assert_eq!(root_doc.parent_dir(), None);
}
#[test]
fn test_memory_document_word_count() {
let mut doc = MemoryDocument::new("user1", None, DocType::Memory, None);
let mut doc = MemoryDocument::new("user1", None, "MEMORY.md");
assert_eq!(doc.word_count(), 0);
doc.content = "Hello world, this is a test.".to_string();
assert_eq!(doc.word_count(), 6);
}
#[test]
fn test_is_identity_document() {
let identity = MemoryDocument::new("user1", None, paths::IDENTITY);
assert!(identity.is_identity_document());
let soul = MemoryDocument::new("user1", None, paths::SOUL);
assert!(soul.is_identity_document());
let memory = MemoryDocument::new("user1", None, paths::MEMORY);
assert!(!memory.is_identity_document());
let custom = MemoryDocument::new("user1", None, "projects/notes.md");
assert!(!custom.is_identity_document());
}
#[test]
fn test_workspace_entry_name() {
let entry = WorkspaceEntry {
path: "projects/alpha".to_string(),
is_directory: true,
updated_at: None,
content_preview: None,
};
assert_eq!(entry.name(), "alpha");
}
}
src/workspace/mod.rs
+202 -95
View File
@@ -1,37 +1,44 @@
//! Workspace and memory system (OpenClaw-inspired).
//!
//! The workspace provides persistent memory for agents:
//! - **MEMORY.md**: Long-term curated memory (facts, decisions, preferences)
//! - **Daily logs**: Append-only daily notes (raw context)
//! - **Identity files**: Agent personality and user context
//! - **HEARTBEAT.md**: Periodic checklist for proactive execution
//! The workspace provides persistent memory for agents with a flexible
//! filesystem-like structure. Agents can create arbitrary markdown file
//! hierarchies that get indexed for full-text and semantic search.
//!
//! Memory is searchable via hybrid search (FTS + semantic embeddings).
//!
//! # Architecture
//! # Filesystem-like API
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │ Workspace │
//! │ ┌────────────────┐ ┌────────────────┐ ┌──────────────┐ │
//! │ │ MemoryDocument │ │ MemoryChunk │ │ Search │ │
//! │ │ (full docs) │──│ (chunked) │──│ (FTS+vector) │ │
//! │ └────────────────┘ └────────────────┘ └──────────────┘ │
//! │ │ │ │ │
//! │ └───────────────────┴──────────────────┘ │
//! │ │ │
//! │ ┌──────┴──────┐ │
//! │ │ Repository │ │
//! │ │ (PostgreSQL)│ │
//! │ └─────────────┘ │
//! └─────────────────────────────────────────────────────────────┘
//! workspace/
//! ├── README.md <- Root runbook/index
//! ├── MEMORY.md <- Long-term curated memory
//! ├── HEARTBEAT.md <- Periodic checklist
//! ├── context/ <- Identity and context
//! │ ├── vision.md
//! │ └── priorities.md
//! ├── daily/ <- Daily logs
//! │ ├── 2024-01-15.md
//! │ └── 2024-01-16.md
//! ├── projects/ <- Arbitrary structure
//! │ └── alpha/
//! │ ├── README.md
//! │ └── notes.md
//! └── ...
//! ```
//!
//! # Key Operations
//!
//! - `read(path)` - Read a file
//! - `write(path, content)` - Create or update a file
//! - `append(path, content)` - Append to a file
//! - `list(dir)` - List directory contents
//! - `delete(path)` - Delete a file
//! - `search(query)` - Full-text + semantic search across all files
//!
//! # Key Patterns
//!
//! 1. **Memory is persistence**: If you want to remember something, write it
//! 2. **Two-tier memory**: Daily logs (raw) + MEMORY.md (curated)
//! 3. **Hybrid search**: Vector similarity + BM25 full-text via RRF
//! 2. **Flexible structure**: Create any directory/file hierarchy you need
//! 3. **Self-documenting**: Use README.md files to describe directory structure
//! 4. **Hybrid search**: Vector similarity + BM25 full-text via RRF
mod chunker;
mod document;
@@ -40,7 +47,7 @@ mod repository;
mod search;
pub use chunker::{ChunkConfig, chunk_document};
pub use document::{DocType, MemoryChunk, MemoryDocument};
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
pub use embeddings::{EmbeddingProvider, OpenAiEmbeddings};
pub use repository::Repository;
pub use search::{SearchConfig, SearchResult};
@@ -101,15 +108,127 @@ impl Workspace {
self.agent_id
}
// ==================== Document Access ====================
// ==================== File Operations ====================
/// Read a file by path.
///
/// Returns the document if it exists, or an error if not found.
///
/// # Example
/// ```ignore
/// let doc = workspace.read("context/vision.md").await?;
/// println!("{}", doc.content);
/// ```
pub async fn read(&self, path: &str) -> Result<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).
///
/// Creates it if it doesn't exist.
pub async fn memory(&self) -> Result<MemoryDocument, WorkspaceError> {
self.repo
.get_or_create_document(&self.user_id, self.agent_id, DocType::Memory, None)
.await
self.read_or_create(paths::MEMORY).await
}
/// Get today's daily log.
@@ -122,38 +241,23 @@ impl Workspace {
/// Get a daily log for a specific date.
pub async fn daily_log(&self, date: NaiveDate) -> Result<MemoryDocument, WorkspaceError> {
let title = date.format("%Y-%m-%d").to_string();
self.repo
.get_or_create_document(
&self.user_id,
self.agent_id,
DocType::DailyLog,
Some(&title),
)
.await
let path = format!("daily/{}.md", date.format("%Y-%m-%d"));
self.read_or_create(&path).await
}
/// Get the heartbeat checklist (HEARTBEAT.md).
pub async fn heartbeat_checklist(&self) -> Result<Option<String>, WorkspaceError> {
match self
.repo
.get_document(&self.user_id, self.agent_id, DocType::Heartbeat, None)
.await
{
match self.read(paths::HEARTBEAT).await {
Ok(doc) => Ok(Some(doc.content)),
Err(WorkspaceError::DocumentNotFound { .. }) => Ok(None),
Err(e) => Err(e),
}
}
/// Get a document by type.
pub async fn get_document(
&self,
doc_type: DocType,
title: Option<&str>,
) -> Result<MemoryDocument, WorkspaceError> {
/// Helper to read or create a file.
async fn read_or_create(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
self.repo
.get_document(&self.user_id, self.agent_id, doc_type, title)
.get_or_create_document_by_path(&self.user_id, self.agent_id, path)
.await
}
@@ -164,6 +268,7 @@ impl Workspace {
/// This is for important facts, decisions, and preferences worth
/// remembering long-term.
pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> {
// Use double newline for memory entries (semantic separation)
let doc = self.memory().await?;
let new_content = if doc.content.is_empty() {
entry.to_string()
@@ -179,34 +284,11 @@ impl Workspace {
///
/// Daily logs are raw, append-only notes for the current day.
pub async fn append_daily_log(&self, entry: &str) -> Result<(), WorkspaceError> {
let doc = self.today_log().await?;
let today = Utc::now().date_naive();
let path = format!("daily/{}.md", today.format("%Y-%m-%d"));
let timestamp = Utc::now().format("%H:%M:%S");
let timestamped_entry = format!("[{}] {}", timestamp, entry);
let new_content = if doc.content.is_empty() {
timestamped_entry
} else {
format!("{}\n{}", doc.content, timestamped_entry)
};
self.repo.update_document(doc.id, &new_content).await?;
self.reindex_document(doc.id).await?;
Ok(())
}
/// Update a document's content entirely.
pub async fn update_document(
&self,
doc_type: DocType,
title: Option<&str>,
content: &str,
) -> Result<(), WorkspaceError> {
let doc = self
.repo
.get_or_create_document(&self.user_id, self.agent_id, doc_type, title)
.await?;
self.repo.update_document(doc.id, content).await?;
self.reindex_document(doc.id).await?;
Ok(())
self.append(&path, &timestamped_entry).await
}
// ==================== System Prompt ====================
@@ -219,19 +301,15 @@ impl Workspace {
let mut parts = Vec::new();
// Load identity files in order of importance
let identity_types = [
(DocType::Agents, "## Agent Instructions"),
(DocType::Soul, "## Core Values"),
(DocType::User, "## User Context"),
(DocType::Identity, "## Identity"),
let identity_files = [
(paths::AGENTS, "## Agent Instructions"),
(paths::SOUL, "## Core Values"),
(paths::USER, "## User Context"),
(paths::IDENTITY, "## Identity"),
];
for (doc_type, header) in identity_types {
if let Ok(doc) = self
.repo
.get_document(&self.user_id, self.agent_id, doc_type, None)
.await
{
for (path, header) in identity_files {
if let Ok(doc) = self.read(path).await {
if !doc.content.is_empty() {
parts.push(format!("{}\n\n{}", header, doc.content));
}
@@ -372,21 +450,50 @@ impl Workspace {
}
}
/// Normalize a file path (remove leading/trailing slashes, collapse //).
fn normalize_path(path: &str) -> String {
let path = path.trim().trim_matches('/');
// Collapse multiple slashes
let mut result = String::new();
let mut last_was_slash = false;
for c in path.chars() {
if c == '/' {
if !last_was_slash {
result.push(c);
}
last_was_slash = true;
} else {
result.push(c);
last_was_slash = false;
}
}
result
}
/// Normalize a directory path (ensure no trailing slash for consistency).
fn normalize_directory(path: &str) -> String {
let path = normalize_path(path);
path.trim_end_matches('/').to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_doc_type_display() {
assert_eq!(DocType::Memory.as_str(), "memory");
assert_eq!(DocType::DailyLog.as_str(), "daily_log");
assert_eq!(DocType::Heartbeat.as_str(), "heartbeat");
fn test_normalize_path() {
assert_eq!(normalize_path("foo/bar"), "foo/bar");
assert_eq!(normalize_path("/foo/bar/"), "foo/bar");
assert_eq!(normalize_path("foo//bar"), "foo/bar");
assert_eq!(normalize_path(" /foo/ "), "foo");
assert_eq!(normalize_path("README.md"), "README.md");
}
#[test]
fn test_doc_type_parse() {
assert_eq!(DocType::try_from("memory").unwrap(), DocType::Memory);
assert_eq!(DocType::try_from("daily_log").unwrap(), DocType::DailyLog);
assert!(DocType::try_from("invalid").is_err());
fn test_normalize_directory() {
assert_eq!(normalize_directory("foo/bar/"), "foo/bar");
assert_eq!(normalize_directory("foo/bar"), "foo/bar");
assert_eq!(normalize_directory("/"), "");
assert_eq!(normalize_directory(""), "");
}
}
src/workspace/repository.rs
+131 -78
View File
@@ -4,14 +4,14 @@
//! - Documents in `memory_documents` table
//! - Chunks in `memory_chunks` table (with FTS and vector indexes)
use chrono::Utc;
use chrono::{DateTime, Utc};
use deadpool_postgres::Pool;
use pgvector::Vector;
use uuid::Uuid;
use crate::error::WorkspaceError;
use crate::workspace::document::{DocType, MemoryChunk, MemoryDocument};
use crate::workspace::document::{MemoryChunk, MemoryDocument, WorkspaceEntry};
use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
/// Database repository for workspace operations.
@@ -37,50 +37,34 @@ impl Repository {
// ==================== Document Operations ====================
/// Get a document by type and optional title.
pub async fn get_document(
/// Get a document by its path.
pub async fn get_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
doc_type: DocType,
title: Option<&str>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
let conn = self.conn().await?;
let row = if let Some(title) = title {
conn.query_opt(
let row = conn
.query_opt(
r#"
SELECT id, user_id, agent_id, doc_type, title, content,
SELECT id, user_id, agent_id, path, content,
created_at, updated_at, metadata
FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
AND doc_type = $3 AND title = $4
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND path = $3
"#,
&[&user_id, &agent_id, &doc_type.as_str(), &title],
&[&user_id, &agent_id, &path],
)
.await
} else {
conn.query_opt(
r#"
SELECT id, user_id, agent_id, doc_type, title, content,
created_at, updated_at, metadata
FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
AND doc_type = $3 AND title IS NULL
"#,
&[&user_id, &agent_id, &doc_type.as_str()],
)
.await
};
let row = row.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
match row {
Some(row) => Ok(self.row_to_document(&row)?),
Some(row) => Ok(self.row_to_document(&row)),
None => Err(WorkspaceError::DocumentNotFound {
doc_type: doc_type.to_string(),
doc_type: path.to_string(),
user_id: user_id.to_string(),
}),
}
@@ -93,7 +77,7 @@ impl Repository {
let row = conn
.query_opt(
r#"
SELECT id, user_id, agent_id, doc_type, title, content,
SELECT id, user_id, agent_id, path, content,
created_at, updated_at, metadata
FROM memory_documents WHERE id = $1
"#,
@@ -105,7 +89,7 @@ impl Repository {
})?;
match row {
Some(row) => Ok(self.row_to_document(&row)?),
Some(row) => Ok(self.row_to_document(&row)),
None => Err(WorkspaceError::DocumentNotFound {
doc_type: "unknown".to_string(),
user_id: "unknown".to_string(),
@@ -113,16 +97,15 @@ impl Repository {
}
}
/// Get or create a document.
pub async fn get_or_create_document(
/// Get or create a document by path.
pub async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
doc_type: DocType,
title: Option<&str>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
// Try to get existing document first
match self.get_document(user_id, agent_id, doc_type, title).await {
match self.get_document_by_path(user_id, agent_id, path).await {
Ok(doc) => return Ok(doc),
Err(WorkspaceError::DocumentNotFound { .. }) => {}
Err(e) => return Err(e),
@@ -132,14 +115,15 @@ impl Repository {
let conn = self.conn().await?;
let id = Uuid::new_v4();
let now = Utc::now();
let metadata = serde_json::json!({});
conn.execute(
r#"
INSERT INTO memory_documents (id, user_id, agent_id, doc_type, title, content, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, '', $6, $7)
ON CONFLICT (user_id, agent_id, doc_type, title) DO NOTHING
INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata, created_at, updated_at)
VALUES ($1, $2, $3, $4, '', $5, $6, $7)
ON CONFLICT (user_id, agent_id, path) DO NOTHING
"#,
&[&id, &user_id, &agent_id, &doc_type.as_str(), &title, &now, &now],
&[&id, &user_id, &agent_id, &path, &metadata, &now, &now],
)
.await
.map_err(|e| WorkspaceError::SearchFailed {
@@ -147,7 +131,7 @@ impl Repository {
})?;
// Fetch the document (might have been created by concurrent request)
self.get_document(user_id, agent_id, doc_type, title).await
self.get_document_by_path(user_id, agent_id, path).await
}
/// Update a document's content.
@@ -166,31 +150,108 @@ impl Repository {
Ok(())
}
/// List documents by type.
/// Delete a document by its path.
pub async fn delete_document_by_path(
&self,
user_id: &str,
agent_id: Option<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(
&self,
user_id: &str,
agent_id: Option<Uuid>,
doc_type: Option<DocType>,
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
let conn = self.conn().await?;
let rows = if let Some(dt) = doc_type {
conn.query(
let rows = conn
.query(
r#"
SELECT id, user_id, agent_id, doc_type, title, content,
created_at, updated_at, metadata
FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND doc_type = $3
ORDER BY updated_at DESC
"#,
&[&user_id, &agent_id, &dt.as_str()],
)
.await
} else {
conn.query(
r#"
SELECT id, user_id, agent_id, doc_type, title, content,
SELECT id, user_id, agent_id, path, content,
created_at, updated_at, metadata
FROM memory_documents
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
@@ -199,30 +260,24 @@ impl Repository {
&[&user_id, &agent_id],
)
.await
};
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
let rows = rows.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
rows.iter().map(|r| self.row_to_document(r)).collect()
Ok(rows.iter().map(|r| self.row_to_document(r)).collect())
}
fn row_to_document(&self, row: &tokio_postgres::Row) -> Result<MemoryDocument, WorkspaceError> {
let doc_type_str: String = row.get("doc_type");
let doc_type = DocType::try_from(doc_type_str.as_str())?;
Ok(MemoryDocument {
fn row_to_document(&self, row: &tokio_postgres::Row) -> MemoryDocument {
MemoryDocument {
id: row.get("id"),
user_id: row.get("user_id"),
agent_id: row.get("agent_id"),
doc_type,
title: row.get("title"),
path: row.get("path"),
content: row.get("content"),
created_at: row.get("created_at"),
updated_at: row.get("updated_at"),
metadata: row.get("metadata"),
})
}
}
// ==================== Chunk Operations ====================
@@ -374,7 +429,6 @@ impl Repository {
) -> Result<Vec<RankedResult>, WorkspaceError> {
let conn = self.conn().await?;
// Use plainto_tsquery for natural language queries
let rows = conn
.query(
r#"
@@ -401,7 +455,7 @@ impl Repository {
chunk_id: row.get("chunk_id"),
document_id: row.get("document_id"),
content: row.get("content"),
rank: (i + 1) as u32, // 1-based rank
rank: (i + 1) as u32,
})
.collect())
}
@@ -417,7 +471,6 @@ impl Repository {
let conn = self.conn().await?;
let embedding_vec = Vector::from(embedding.to_vec());
// Use cosine distance (<=>)
let rows = conn
.query(
r#"
@@ -444,7 +497,7 @@ impl Repository {
chunk_id: row.get("chunk_id"),
document_id: row.get("document_id"),
content: row.get("content"),
rank: (i + 1) as u32, // 1-based rank
rank: (i + 1) as u32,
})
.collect())
}