mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
Add workspace and memory system (OpenClaw-inspired)
Implements persistent memory for agents with hybrid search: - Database-backed workspace with PostgreSQL (not filesystem) - Memory documents: MEMORY.md, daily logs, identity files - Chunked content with FTS (tsvector) + vector (pgvector) indexes - Reciprocal Rank Fusion (RRF) for hybrid search combining BM25 and semantic - Memory tools: memory_search, memory_write, memory_read - Proactive heartbeat system for periodic execution (30 min default) - OpenAI embeddings provider (text-embedding-3-small) Key patterns from OpenClaw: - "Memory is files, not RAM" - explicit persistence required - Two-tier memory: daily logs (raw) + curated MEMORY.md - Session isolation via user_id/agent_id scoping Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
8c38566378
commit
4e238e60ac
@@ -0,0 +1,314 @@
|
||||
//! Document chunking for search indexing.
|
||||
//!
|
||||
//! Documents are split into overlapping chunks for better search recall.
|
||||
//! The overlap ensures context is preserved across chunk boundaries.
|
||||
|
||||
/// Configuration for document chunking.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChunkConfig {
|
||||
/// Target chunk size in words (approximate tokens).
|
||||
/// Default: 800 (roughly 800 tokens for English text).
|
||||
pub chunk_size: usize,
|
||||
/// Overlap percentage between chunks.
|
||||
/// Default: 0.15 (15% overlap).
|
||||
pub overlap_percent: f32,
|
||||
/// Minimum chunk size (don't create tiny trailing chunks).
|
||||
/// Default: 50 words.
|
||||
pub min_chunk_size: usize,
|
||||
}
|
||||
|
||||
impl Default for ChunkConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
chunk_size: 800,
|
||||
overlap_percent: 0.15,
|
||||
min_chunk_size: 50,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ChunkConfig {
|
||||
/// Create a config with a specific chunk size.
|
||||
pub fn with_chunk_size(mut self, size: usize) -> Self {
|
||||
self.chunk_size = size;
|
||||
self
|
||||
}
|
||||
|
||||
/// Create a config with a specific overlap percentage.
|
||||
pub fn with_overlap(mut self, percent: f32) -> Self {
|
||||
self.overlap_percent = percent.clamp(0.0, 0.5);
|
||||
self
|
||||
}
|
||||
|
||||
/// Calculate the overlap size in words.
|
||||
fn overlap_size(&self) -> usize {
|
||||
(self.chunk_size as f32 * self.overlap_percent) as usize
|
||||
}
|
||||
|
||||
/// Calculate the step size (chunk_size - overlap).
|
||||
fn step_size(&self) -> usize {
|
||||
self.chunk_size.saturating_sub(self.overlap_size())
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a document into overlapping chunks.
|
||||
///
|
||||
/// Each chunk contains approximately `chunk_size` words, with `overlap_percent`
|
||||
/// overlap between adjacent chunks. This ensures that:
|
||||
/// 1. Context is preserved across chunk boundaries
|
||||
/// 2. Search can find content that spans chunk boundaries
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `content` - The document text to chunk
|
||||
/// * `config` - Chunking configuration
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A vector of chunk strings. Empty documents return an empty vector.
|
||||
pub fn chunk_document(content: &str, config: ChunkConfig) -> Vec<String> {
|
||||
if content.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Split into words while preserving structure
|
||||
let words: Vec<&str> = content.split_whitespace().collect();
|
||||
|
||||
if words.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// If content is smaller than chunk size, return as single chunk
|
||||
if words.len() <= config.chunk_size {
|
||||
return vec![content.to_string()];
|
||||
}
|
||||
|
||||
let step = config.step_size();
|
||||
let mut chunks = Vec::new();
|
||||
let mut start = 0;
|
||||
|
||||
while start < words.len() {
|
||||
let end = (start + config.chunk_size).min(words.len());
|
||||
let chunk_words = &words[start..end];
|
||||
|
||||
// Don't create tiny trailing chunks, merge with previous
|
||||
if chunk_words.len() < config.min_chunk_size && !chunks.is_empty() {
|
||||
let last = chunks.pop().unwrap();
|
||||
let combined = format!("{} {}", last, chunk_words.join(" "));
|
||||
chunks.push(combined);
|
||||
break;
|
||||
}
|
||||
|
||||
chunks.push(chunk_words.join(" "));
|
||||
|
||||
// Move to next chunk position
|
||||
start += step;
|
||||
|
||||
// Avoid creating duplicate chunks at the end
|
||||
if start + config.min_chunk_size >= words.len() && end == words.len() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Split content by paragraphs first, then chunk.
|
||||
///
|
||||
/// This is better for preserving semantic boundaries.
|
||||
pub fn chunk_by_paragraphs(content: &str, config: ChunkConfig) -> Vec<String> {
|
||||
if content.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Split by double newlines (paragraphs)
|
||||
let paragraphs: Vec<&str> = content
|
||||
.split("\n\n")
|
||||
.map(|p| p.trim())
|
||||
.filter(|p| !p.is_empty())
|
||||
.collect();
|
||||
|
||||
if paragraphs.is_empty() {
|
||||
return chunk_document(content, config);
|
||||
}
|
||||
|
||||
let mut chunks = Vec::new();
|
||||
let mut current_chunk = String::new();
|
||||
let mut current_word_count = 0;
|
||||
|
||||
for paragraph in paragraphs {
|
||||
let para_words = paragraph.split_whitespace().count();
|
||||
|
||||
// If this paragraph alone exceeds chunk size, chunk it separately
|
||||
if para_words > config.chunk_size {
|
||||
// Flush current chunk first
|
||||
if !current_chunk.is_empty() {
|
||||
chunks.push(current_chunk.trim().to_string());
|
||||
current_chunk = String::new();
|
||||
current_word_count = 0;
|
||||
}
|
||||
// Chunk the large paragraph
|
||||
let para_chunks = chunk_document(paragraph, config.clone());
|
||||
chunks.extend(para_chunks);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if adding this paragraph would exceed chunk size
|
||||
if current_word_count + para_words > config.chunk_size {
|
||||
// Flush current chunk
|
||||
if !current_chunk.is_empty() {
|
||||
chunks.push(current_chunk.trim().to_string());
|
||||
}
|
||||
current_chunk = paragraph.to_string();
|
||||
current_word_count = para_words;
|
||||
} else {
|
||||
// Add paragraph to current chunk
|
||||
if !current_chunk.is_empty() {
|
||||
current_chunk.push_str("\n\n");
|
||||
}
|
||||
current_chunk.push_str(paragraph);
|
||||
current_word_count += para_words;
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining content
|
||||
if !current_chunk.is_empty() {
|
||||
// If too small, merge with previous chunk if possible
|
||||
if current_word_count < config.min_chunk_size && !chunks.is_empty() {
|
||||
let last = chunks.pop().unwrap();
|
||||
chunks.push(format!("{}\n\n{}", last, current_chunk.trim()));
|
||||
} else {
|
||||
chunks.push(current_chunk.trim().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_empty_content() {
|
||||
let config = ChunkConfig::default();
|
||||
assert!(chunk_document("", config.clone()).is_empty());
|
||||
assert!(chunk_document(" ", config).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_small_content() {
|
||||
let config = ChunkConfig::default();
|
||||
let content = "Hello world, this is a test.";
|
||||
let chunks = chunk_document(content, config);
|
||||
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0], content);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_chunk_size() {
|
||||
let config = ChunkConfig::default().with_chunk_size(5);
|
||||
let content = "one two three four five";
|
||||
let chunks = chunk_document(content, config);
|
||||
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0], content);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chunking_with_overlap() {
|
||||
let config = ChunkConfig {
|
||||
chunk_size: 10,
|
||||
overlap_percent: 0.2, // 2 word overlap
|
||||
min_chunk_size: 3, // Low threshold for test
|
||||
};
|
||||
|
||||
// 20 words
|
||||
let content = "one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen nineteen twenty";
|
||||
let chunks = chunk_document(content, config);
|
||||
|
||||
// Should create overlapping chunks
|
||||
assert!(
|
||||
chunks.len() >= 2,
|
||||
"Expected at least 2 chunks, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
|
||||
// Each chunk should have roughly 10 words (allowing for overlap/merging)
|
||||
for chunk in &chunks {
|
||||
let word_count = chunk.split_whitespace().count();
|
||||
assert!(word_count >= 3, "Chunk too small: {} words", word_count);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_overlap_calculation() {
|
||||
let config = ChunkConfig::default()
|
||||
.with_chunk_size(100)
|
||||
.with_overlap(0.15);
|
||||
|
||||
assert_eq!(config.overlap_size(), 15);
|
||||
assert_eq!(config.step_size(), 85);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_paragraph_chunking() {
|
||||
let config = ChunkConfig::default().with_chunk_size(20);
|
||||
|
||||
let content = "First paragraph with some words.\n\nSecond paragraph with different content.\n\nThird paragraph here.";
|
||||
let chunks = chunk_by_paragraphs(content, config);
|
||||
|
||||
// Should preserve paragraph boundaries
|
||||
assert!(!chunks.is_empty());
|
||||
for chunk in &chunks {
|
||||
// No chunk should start or end with \n\n
|
||||
assert!(!chunk.starts_with("\n"));
|
||||
assert!(!chunk.ends_with("\n"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_paragraph_handling() {
|
||||
let config = ChunkConfig {
|
||||
chunk_size: 10,
|
||||
overlap_percent: 0.15,
|
||||
min_chunk_size: 3, // Low threshold for test
|
||||
};
|
||||
|
||||
// Create a paragraph with 30 words
|
||||
let large_para = (1..=30)
|
||||
.map(|i| format!("word{}", i))
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
let content = format!("Short intro.\n\n{}\n\nShort outro.", large_para);
|
||||
|
||||
let chunks = chunk_by_paragraphs(&content, config);
|
||||
|
||||
// Should have multiple chunks due to large paragraph
|
||||
// 30 words + 2 intro + 2 outro = 34 words, chunk_size=10
|
||||
// Expect at least 3 chunks
|
||||
assert!(
|
||||
chunks.len() >= 3,
|
||||
"Expected at least 3 chunks for 34 words with chunk_size=10, got {}",
|
||||
chunks.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_min_chunk_size_merging() {
|
||||
let config = ChunkConfig {
|
||||
chunk_size: 10,
|
||||
overlap_percent: 0.0,
|
||||
min_chunk_size: 5,
|
||||
};
|
||||
|
||||
// 12 words: should create one chunk of 10, and merge the remaining 2 with it
|
||||
let content = "one two three four five six seven eight nine ten eleven twelve";
|
||||
let chunks = chunk_document(content, config);
|
||||
|
||||
// Should merge the tiny trailing chunk
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].split_whitespace().count(), 12);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
//! Memory document types.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
/// Document type in the workspace.
|
||||
///
|
||||
/// Each type represents a different kind of persistent memory:
|
||||
/// - **Memory**: Long-term curated facts and decisions (MEMORY.md)
|
||||
/// - **DailyLog**: Append-only daily notes (memory/YYYY-MM-DD.md)
|
||||
/// - **Identity**: Agent name and personality
|
||||
/// - **Soul**: Core values and behavior principles
|
||||
/// - **Agents**: Behavior instructions
|
||||
/// - **User**: User context (name, preferences)
|
||||
/// - **Heartbeat**: Periodic checklist for proactive execution
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum DocType {
|
||||
/// Long-term curated memory (MEMORY.md equivalent).
|
||||
Memory,
|
||||
/// Daily append-only logs.
|
||||
DailyLog,
|
||||
/// Agent identity (name, nature, vibe).
|
||||
Identity,
|
||||
/// Core values and principles (SOUL.md).
|
||||
Soul,
|
||||
/// Behavior instructions (AGENTS.md).
|
||||
Agents,
|
||||
/// User context (USER.md).
|
||||
User,
|
||||
/// Periodic checklist (HEARTBEAT.md).
|
||||
Heartbeat,
|
||||
}
|
||||
|
||||
impl DocType {
|
||||
/// Get the string representation.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
DocType::Memory => "memory",
|
||||
DocType::DailyLog => "daily_log",
|
||||
DocType::Identity => "identity",
|
||||
DocType::Soul => "soul",
|
||||
DocType::Agents => "agents",
|
||||
DocType::User => "user",
|
||||
DocType::Heartbeat => "heartbeat",
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this document type is a singleton (one per user/agent).
|
||||
pub fn is_singleton(&self) -> bool {
|
||||
match self {
|
||||
DocType::Memory
|
||||
| DocType::Identity
|
||||
| DocType::Soul
|
||||
| DocType::Agents
|
||||
| DocType::User
|
||||
| DocType::Heartbeat => true,
|
||||
DocType::DailyLog => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this document should be included in the system prompt.
|
||||
pub fn is_identity_document(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
DocType::Identity | DocType::Soul | DocType::Agents | DocType::User
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&str> for DocType {
|
||||
type Error = WorkspaceError;
|
||||
|
||||
fn try_from(s: &str) -> Result<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.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MemoryDocument {
|
||||
/// Unique document ID.
|
||||
pub id: Uuid,
|
||||
/// User identifier.
|
||||
pub user_id: String,
|
||||
/// Optional agent ID for multi-agent isolation.
|
||||
pub agent_id: Option<Uuid>,
|
||||
/// Document type.
|
||||
pub doc_type: DocType,
|
||||
/// Optional title (e.g., date for daily logs).
|
||||
pub title: Option<String>,
|
||||
/// Full document content.
|
||||
pub content: String,
|
||||
/// Creation timestamp.
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// Last update timestamp.
|
||||
pub updated_at: DateTime<Utc>,
|
||||
/// Flexible metadata.
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
impl MemoryDocument {
|
||||
/// Create a new document (not persisted yet).
|
||||
pub fn new(
|
||||
user_id: impl Into<String>,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: DocType,
|
||||
title: Option<String>,
|
||||
) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
id: Uuid::new_v4(),
|
||||
user_id: user_id.into(),
|
||||
agent_id,
|
||||
doc_type,
|
||||
title,
|
||||
content: String::new(),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
metadata: serde_json::Value::Object(serde_json::Map::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the document is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.content.is_empty()
|
||||
}
|
||||
|
||||
/// Get word count.
|
||||
pub fn word_count(&self) -> usize {
|
||||
self.content.split_whitespace().count()
|
||||
}
|
||||
}
|
||||
|
||||
/// A chunk of a memory document for search indexing.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MemoryChunk {
|
||||
/// Unique chunk ID.
|
||||
pub id: Uuid,
|
||||
/// Parent document ID.
|
||||
pub document_id: Uuid,
|
||||
/// Position in the document (0-based).
|
||||
pub chunk_index: i32,
|
||||
/// Chunk text content.
|
||||
pub content: String,
|
||||
/// Embedding vector (if generated).
|
||||
pub embedding: Option<Vec<f32>>,
|
||||
/// Creation timestamp.
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl MemoryChunk {
|
||||
/// Create a new chunk (not persisted yet).
|
||||
pub fn new(document_id: Uuid, chunk_index: i32, content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
id: Uuid::new_v4(),
|
||||
document_id,
|
||||
chunk_index,
|
||||
content: content.into(),
|
||||
embedding: None,
|
||||
created_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the embedding.
|
||||
pub fn with_embedding(mut self, embedding: Vec<f32>) -> Self {
|
||||
self.embedding = Some(embedding);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_doc_type_roundtrip() {
|
||||
for doc_type in [
|
||||
DocType::Memory,
|
||||
DocType::DailyLog,
|
||||
DocType::Identity,
|
||||
DocType::Soul,
|
||||
DocType::Agents,
|
||||
DocType::User,
|
||||
DocType::Heartbeat,
|
||||
] {
|
||||
let s = doc_type.as_str();
|
||||
let parsed = DocType::try_from(s).unwrap();
|
||||
assert_eq!(parsed, doc_type);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_singleton_types() {
|
||||
assert!(DocType::Memory.is_singleton());
|
||||
assert!(DocType::Heartbeat.is_singleton());
|
||||
assert!(!DocType::DailyLog.is_singleton());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_documents() {
|
||||
assert!(DocType::Soul.is_identity_document());
|
||||
assert!(DocType::Agents.is_identity_document());
|
||||
assert!(!DocType::Memory.is_identity_document());
|
||||
assert!(!DocType::DailyLog.is_identity_document());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_document_word_count() {
|
||||
let mut doc = MemoryDocument::new("user1", None, DocType::Memory, None);
|
||||
assert_eq!(doc.word_count(), 0);
|
||||
|
||||
doc.content = "Hello world, this is a test.".to_string();
|
||||
assert_eq!(doc.word_count(), 6);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,324 @@
|
||||
//! Embedding providers for semantic search.
|
||||
//!
|
||||
//! Embeddings convert text into dense vectors that capture semantic meaning.
|
||||
//! Similar concepts have similar vectors, enabling semantic search.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Error type for embedding operations.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum EmbeddingError {
|
||||
#[error("HTTP request failed: {0}")]
|
||||
HttpError(String),
|
||||
|
||||
#[error("Invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
|
||||
#[error("Rate limited, retry after {retry_after:?}")]
|
||||
RateLimited {
|
||||
retry_after: Option<std::time::Duration>,
|
||||
},
|
||||
|
||||
#[error("Authentication failed")]
|
||||
AuthFailed,
|
||||
|
||||
#[error("Text too long: {length} > {max}")]
|
||||
TextTooLong { length: usize, max: usize },
|
||||
}
|
||||
|
||||
impl From<reqwest::Error> for EmbeddingError {
|
||||
fn from(e: reqwest::Error) -> Self {
|
||||
EmbeddingError::HttpError(e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for embedding providers.
|
||||
#[async_trait]
|
||||
pub trait EmbeddingProvider: Send + Sync {
|
||||
/// Get the embedding dimension.
|
||||
fn dimension(&self) -> usize;
|
||||
|
||||
/// Get the model name.
|
||||
fn model_name(&self) -> &str;
|
||||
|
||||
/// Maximum input length in characters.
|
||||
fn max_input_length(&self) -> usize;
|
||||
|
||||
/// Generate an embedding for a single text.
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
|
||||
|
||||
/// Generate embeddings for multiple texts (batched).
|
||||
///
|
||||
/// Default implementation calls embed() for each text.
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
let mut embeddings = Vec::with_capacity(texts.len());
|
||||
for text in texts {
|
||||
embeddings.push(self.embed(text).await?);
|
||||
}
|
||||
Ok(embeddings)
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI embedding provider using text-embedding-ada-002 or text-embedding-3-small.
|
||||
pub struct OpenAiEmbeddings {
|
||||
client: reqwest::Client,
|
||||
api_key: String,
|
||||
model: String,
|
||||
dimension: usize,
|
||||
}
|
||||
|
||||
impl OpenAiEmbeddings {
|
||||
/// Create a new OpenAI embedding provider with the default model.
|
||||
///
|
||||
/// Uses text-embedding-3-small which has 1536 dimensions.
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
api_key: api_key.into(),
|
||||
model: "text-embedding-3-small".to_string(),
|
||||
dimension: 1536,
|
||||
}
|
||||
}
|
||||
|
||||
/// Use text-embedding-ada-002 model.
|
||||
pub fn ada_002(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
api_key: api_key.into(),
|
||||
model: "text-embedding-ada-002".to_string(),
|
||||
dimension: 1536,
|
||||
}
|
||||
}
|
||||
|
||||
/// Use text-embedding-3-large model.
|
||||
pub fn large(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
api_key: api_key.into(),
|
||||
model: "text-embedding-3-large".to_string(),
|
||||
dimension: 3072,
|
||||
}
|
||||
}
|
||||
|
||||
/// Use a custom model with specified dimension.
|
||||
pub fn with_model(
|
||||
api_key: impl Into<String>,
|
||||
model: impl Into<String>,
|
||||
dimension: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
api_key: api_key.into(),
|
||||
model: model.into(),
|
||||
dimension,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OpenAiEmbeddingRequest<'a> {
|
||||
model: &'a str,
|
||||
input: &'a [String],
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OpenAiEmbeddingResponse {
|
||||
data: Vec<OpenAiEmbeddingData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OpenAiEmbeddingData {
|
||||
embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for OpenAiEmbeddings {
|
||||
fn dimension(&self) -> usize {
|
||||
self.dimension
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
fn max_input_length(&self) -> usize {
|
||||
// text-embedding-3-small/large: 8191 tokens (~32k chars)
|
||||
// text-embedding-ada-002: 8191 tokens
|
||||
32_000
|
||||
}
|
||||
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
if text.len() > self.max_input_length() {
|
||||
return Err(EmbeddingError::TextTooLong {
|
||||
length: text.len(),
|
||||
max: self.max_input_length(),
|
||||
});
|
||||
}
|
||||
|
||||
let embeddings = self.embed_batch(&[text.to_string()]).await?;
|
||||
embeddings
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string()))
|
||||
}
|
||||
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let request = OpenAiEmbeddingRequest {
|
||||
model: &self.model,
|
||||
input: texts,
|
||||
};
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post("https://api.openai.com/v1/embeddings")
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.json(&request)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
return Err(EmbeddingError::AuthFailed);
|
||||
}
|
||||
|
||||
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
|
||||
let retry_after = response
|
||||
.headers()
|
||||
.get("retry-after")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.map(std::time::Duration::from_secs);
|
||||
return Err(EmbeddingError::RateLimited { retry_after });
|
||||
}
|
||||
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(EmbeddingError::HttpError(format!(
|
||||
"Status {}: {}",
|
||||
status, error_text
|
||||
)));
|
||||
}
|
||||
|
||||
let result: OpenAiEmbeddingResponse = response.json().await.map_err(|e| {
|
||||
EmbeddingError::InvalidResponse(format!("Failed to parse response: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(result.data.into_iter().map(|d| d.embedding).collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// A mock embedding provider for testing.
|
||||
#[cfg(test)]
|
||||
pub struct MockEmbeddings {
|
||||
dimension: usize,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl MockEmbeddings {
|
||||
pub fn new(dimension: usize) -> Self {
|
||||
Self { dimension }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for MockEmbeddings {
|
||||
fn dimension(&self) -> usize {
|
||||
self.dimension
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &str {
|
||||
"mock-embedding"
|
||||
}
|
||||
|
||||
fn max_input_length(&self) -> usize {
|
||||
10_000
|
||||
}
|
||||
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
// Generate a deterministic embedding based on text hash
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
text.hash(&mut hasher);
|
||||
let hash = hasher.finish();
|
||||
|
||||
let mut embedding = Vec::with_capacity(self.dimension);
|
||||
let mut seed = hash;
|
||||
for _ in 0..self.dimension {
|
||||
// Simple LCG for deterministic random values
|
||||
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
let value = (seed as f32 / u64::MAX as f32) * 2.0 - 1.0;
|
||||
embedding.push(value);
|
||||
}
|
||||
|
||||
// Normalize to unit length
|
||||
let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if magnitude > 0.0 {
|
||||
for x in &mut embedding {
|
||||
*x /= magnitude;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(embedding)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mock_embeddings() {
|
||||
let provider = MockEmbeddings::new(128);
|
||||
|
||||
let embedding = provider.embed("hello world").await.unwrap();
|
||||
assert_eq!(embedding.len(), 128);
|
||||
|
||||
// Check normalization (should be unit vector)
|
||||
let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
assert!((magnitude - 1.0).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mock_embeddings_deterministic() {
|
||||
let provider = MockEmbeddings::new(64);
|
||||
|
||||
let emb1 = provider.embed("test").await.unwrap();
|
||||
let emb2 = provider.embed("test").await.unwrap();
|
||||
|
||||
// Same input should produce same embedding
|
||||
assert_eq!(emb1, emb2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mock_embeddings_batch() {
|
||||
let provider = MockEmbeddings::new(64);
|
||||
|
||||
let texts = vec!["hello".to_string(), "world".to_string()];
|
||||
let embeddings = provider.embed_batch(&texts).await.unwrap();
|
||||
|
||||
assert_eq!(embeddings.len(), 2);
|
||||
assert_eq!(embeddings[0].len(), 64);
|
||||
assert_eq!(embeddings[1].len(), 64);
|
||||
|
||||
// Different texts should produce different embeddings
|
||||
assert_ne!(embeddings[0], embeddings[1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_openai_embeddings_config() {
|
||||
let provider = OpenAiEmbeddings::new("test-key");
|
||||
assert_eq!(provider.dimension(), 1536);
|
||||
assert_eq!(provider.model_name(), "text-embedding-3-small");
|
||||
|
||||
let provider = OpenAiEmbeddings::large("test-key");
|
||||
assert_eq!(provider.dimension(), 3072);
|
||||
assert_eq!(provider.model_name(), "text-embedding-3-large");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,392 @@
|
||||
//! Workspace and memory system (OpenClaw-inspired).
|
||||
//!
|
||||
//! The workspace provides persistent memory for agents:
|
||||
//! - **MEMORY.md**: Long-term curated memory (facts, decisions, preferences)
|
||||
//! - **Daily logs**: Append-only daily notes (raw context)
|
||||
//! - **Identity files**: Agent personality and user context
|
||||
//! - **HEARTBEAT.md**: Periodic checklist for proactive execution
|
||||
//!
|
||||
//! Memory is searchable via hybrid search (FTS + semantic embeddings).
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────────────────────┐
|
||||
//! │ Workspace │
|
||||
//! │ ┌────────────────┐ ┌────────────────┐ ┌──────────────┐ │
|
||||
//! │ │ MemoryDocument │ │ MemoryChunk │ │ Search │ │
|
||||
//! │ │ (full docs) │──│ (chunked) │──│ (FTS+vector) │ │
|
||||
//! │ └────────────────┘ └────────────────┘ └──────────────┘ │
|
||||
//! │ │ │ │ │
|
||||
//! │ └───────────────────┴──────────────────┘ │
|
||||
//! │ │ │
|
||||
//! │ ┌──────┴──────┐ │
|
||||
//! │ │ Repository │ │
|
||||
//! │ │ (PostgreSQL)│ │
|
||||
//! │ └─────────────┘ │
|
||||
//! └─────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! # Key Patterns
|
||||
//!
|
||||
//! 1. **Memory is persistence**: If you want to remember something, write it
|
||||
//! 2. **Two-tier memory**: Daily logs (raw) + MEMORY.md (curated)
|
||||
//! 3. **Hybrid search**: Vector similarity + BM25 full-text via RRF
|
||||
|
||||
mod chunker;
|
||||
mod document;
|
||||
mod embeddings;
|
||||
mod repository;
|
||||
mod search;
|
||||
|
||||
pub use chunker::{ChunkConfig, chunk_document};
|
||||
pub use document::{DocType, MemoryChunk, MemoryDocument};
|
||||
pub use embeddings::{EmbeddingProvider, OpenAiEmbeddings};
|
||||
pub use repository::Repository;
|
||||
pub use search::{SearchConfig, SearchResult};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{NaiveDate, Utc};
|
||||
use deadpool_postgres::Pool;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
/// Workspace provides database-backed memory storage for an agent.
|
||||
///
|
||||
/// Each workspace is scoped to a user (and optionally an agent).
|
||||
/// Documents are persisted to PostgreSQL and indexed for search.
|
||||
pub struct Workspace {
|
||||
/// User identifier (from channel).
|
||||
user_id: String,
|
||||
/// Optional agent ID for multi-agent isolation.
|
||||
agent_id: Option<Uuid>,
|
||||
/// Database repository.
|
||||
repo: Repository,
|
||||
/// Embedding provider for semantic search.
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
}
|
||||
|
||||
impl Workspace {
|
||||
/// Create a new workspace for a user.
|
||||
pub fn new(user_id: impl Into<String>, pool: Pool) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
agent_id: None,
|
||||
repo: Repository::new(pool),
|
||||
embeddings: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a workspace with a specific agent ID.
|
||||
pub fn with_agent(mut self, agent_id: Uuid) -> Self {
|
||||
self.agent_id = Some(agent_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the embedding provider for semantic search.
|
||||
pub fn with_embeddings(mut self, provider: Arc<dyn EmbeddingProvider>) -> Self {
|
||||
self.embeddings = Some(provider);
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the user ID.
|
||||
pub fn user_id(&self) -> &str {
|
||||
&self.user_id
|
||||
}
|
||||
|
||||
/// Get the agent ID.
|
||||
pub fn agent_id(&self) -> Option<Uuid> {
|
||||
self.agent_id
|
||||
}
|
||||
|
||||
// ==================== Document Access ====================
|
||||
|
||||
/// Get the main MEMORY.md document (long-term curated memory).
|
||||
///
|
||||
/// Creates it if it doesn't exist.
|
||||
pub async fn memory(&self) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
.get_or_create_document(&self.user_id, self.agent_id, DocType::Memory, None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Get today's daily log.
|
||||
///
|
||||
/// Daily logs are append-only and keyed by date.
|
||||
pub async fn today_log(&self) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let today = Utc::now().date_naive();
|
||||
self.daily_log(today).await
|
||||
}
|
||||
|
||||
/// Get a daily log for a specific date.
|
||||
pub async fn daily_log(&self, date: NaiveDate) -> Result<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
|
||||
}
|
||||
|
||||
/// 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
|
||||
{
|
||||
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> {
|
||||
self.repo
|
||||
.get_document(&self.user_id, self.agent_id, doc_type, title)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Memory Operations ====================
|
||||
|
||||
/// Append an entry to the main MEMORY.md document.
|
||||
///
|
||||
/// This is for important facts, decisions, and preferences worth
|
||||
/// remembering long-term.
|
||||
pub async fn append_memory(&self, entry: &str) -> Result<(), WorkspaceError> {
|
||||
let doc = self.memory().await?;
|
||||
let new_content = if doc.content.is_empty() {
|
||||
entry.to_string()
|
||||
} else {
|
||||
format!("{}\n\n{}", doc.content, entry)
|
||||
};
|
||||
self.repo.update_document(doc.id, &new_content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Append an entry to today's daily log.
|
||||
///
|
||||
/// Daily logs are raw, append-only notes for the current day.
|
||||
pub async fn append_daily_log(&self, entry: &str) -> Result<(), WorkspaceError> {
|
||||
let doc = self.today_log().await?;
|
||||
let timestamp = Utc::now().format("%H:%M:%S");
|
||||
let timestamped_entry = format!("[{}] {}", timestamp, entry);
|
||||
|
||||
let new_content = if doc.content.is_empty() {
|
||||
timestamped_entry
|
||||
} else {
|
||||
format!("{}\n{}", doc.content, timestamped_entry)
|
||||
};
|
||||
self.repo.update_document(doc.id, &new_content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update a document's content entirely.
|
||||
pub async fn update_document(
|
||||
&self,
|
||||
doc_type: DocType,
|
||||
title: Option<&str>,
|
||||
content: &str,
|
||||
) -> Result<(), WorkspaceError> {
|
||||
let doc = self
|
||||
.repo
|
||||
.get_or_create_document(&self.user_id, self.agent_id, doc_type, title)
|
||||
.await?;
|
||||
self.repo.update_document(doc.id, content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ==================== System Prompt ====================
|
||||
|
||||
/// Build the system prompt from identity files.
|
||||
///
|
||||
/// Loads AGENTS.md, SOUL.md, USER.md, and IDENTITY.md to compose
|
||||
/// the agent's system prompt.
|
||||
pub async fn system_prompt(&self) -> Result<String, WorkspaceError> {
|
||||
let mut parts = Vec::new();
|
||||
|
||||
// Load identity files in order of importance
|
||||
let identity_types = [
|
||||
(DocType::Agents, "## Agent Instructions"),
|
||||
(DocType::Soul, "## Core Values"),
|
||||
(DocType::User, "## User Context"),
|
||||
(DocType::Identity, "## Identity"),
|
||||
];
|
||||
|
||||
for (doc_type, header) in identity_types {
|
||||
if let Ok(doc) = self
|
||||
.repo
|
||||
.get_document(&self.user_id, self.agent_id, doc_type, None)
|
||||
.await
|
||||
{
|
||||
if !doc.content.is_empty() {
|
||||
parts.push(format!("{}\n\n{}", header, doc.content));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add today's memory context (last 2 days of daily logs)
|
||||
let today = Utc::now().date_naive();
|
||||
let yesterday = today.pred_opt().unwrap_or(today);
|
||||
|
||||
for date in [today, yesterday] {
|
||||
if let Ok(doc) = self.daily_log(date).await {
|
||||
if !doc.content.is_empty() {
|
||||
let header = if date == today {
|
||||
"## Today's Notes"
|
||||
} else {
|
||||
"## Yesterday's Notes"
|
||||
};
|
||||
parts.push(format!("{}\n\n{}", header, doc.content));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(parts.join("\n\n---\n\n"))
|
||||
}
|
||||
|
||||
// ==================== Search ====================
|
||||
|
||||
/// Hybrid search across all memory documents.
|
||||
///
|
||||
/// Combines full-text search (BM25) with semantic search (vector similarity)
|
||||
/// using Reciprocal Rank Fusion (RRF).
|
||||
pub async fn search(
|
||||
&self,
|
||||
query: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
self.search_with_config(query, SearchConfig::default().with_limit(limit))
|
||||
.await
|
||||
}
|
||||
|
||||
/// Search with custom configuration.
|
||||
pub async fn search_with_config(
|
||||
&self,
|
||||
query: &str,
|
||||
config: SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
// Generate embedding for semantic search if provider available
|
||||
let embedding = if let Some(ref provider) = self.embeddings {
|
||||
Some(
|
||||
provider
|
||||
.embed(query)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::EmbeddingFailed {
|
||||
reason: e.to_string(),
|
||||
})?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
self.repo
|
||||
.hybrid_search(
|
||||
&self.user_id,
|
||||
self.agent_id,
|
||||
query,
|
||||
embedding.as_deref(),
|
||||
&config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Indexing ====================
|
||||
|
||||
/// Re-index a document (chunk and generate embeddings).
|
||||
async fn reindex_document(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
// Get the document
|
||||
let doc = self.repo.get_document_by_id(document_id).await?;
|
||||
|
||||
// Chunk the content
|
||||
let chunks = chunk_document(&doc.content, ChunkConfig::default());
|
||||
|
||||
// Delete old chunks
|
||||
self.repo.delete_chunks(document_id).await?;
|
||||
|
||||
// Insert new chunks
|
||||
for (index, content) in chunks.into_iter().enumerate() {
|
||||
// Generate embedding if provider available
|
||||
let embedding = if let Some(ref provider) = self.embeddings {
|
||||
match provider.embed(&content).await {
|
||||
Ok(emb) => Some(emb),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to generate embedding: {}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
self.repo
|
||||
.insert_chunk(document_id, index as i32, &content, embedding.as_deref())
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Generate embeddings for chunks that don't have them yet.
|
||||
///
|
||||
/// This is useful for backfilling embeddings after enabling the provider.
|
||||
pub async fn backfill_embeddings(&self) -> Result<usize, WorkspaceError> {
|
||||
let Some(ref provider) = self.embeddings else {
|
||||
return Ok(0);
|
||||
};
|
||||
|
||||
let chunks = self
|
||||
.repo
|
||||
.get_chunks_without_embeddings(&self.user_id, self.agent_id, 100)
|
||||
.await?;
|
||||
|
||||
let mut count = 0;
|
||||
for chunk in chunks {
|
||||
match provider.embed(&chunk.content).await {
|
||||
Ok(embedding) => {
|
||||
self.repo
|
||||
.update_chunk_embedding(chunk.id, &embedding)
|
||||
.await?;
|
||||
count += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to embed chunk {}: {}", chunk.id, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_doc_type_display() {
|
||||
assert_eq!(DocType::Memory.as_str(), "memory");
|
||||
assert_eq!(DocType::DailyLog.as_str(), "daily_log");
|
||||
assert_eq!(DocType::Heartbeat.as_str(), "heartbeat");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_doc_type_parse() {
|
||||
assert_eq!(DocType::try_from("memory").unwrap(), DocType::Memory);
|
||||
assert_eq!(DocType::try_from("daily_log").unwrap(), DocType::DailyLog);
|
||||
assert!(DocType::try_from("invalid").is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,451 @@
|
||||
//! Database repository for workspace persistence.
|
||||
//!
|
||||
//! All workspace data is stored in PostgreSQL:
|
||||
//! - Documents in `memory_documents` table
|
||||
//! - Chunks in `memory_chunks` table (with FTS and vector indexes)
|
||||
|
||||
use chrono::Utc;
|
||||
use deadpool_postgres::Pool;
|
||||
use pgvector::Vector;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
use crate::workspace::document::{DocType, MemoryChunk, MemoryDocument};
|
||||
use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
||||
|
||||
/// Database repository for workspace operations.
|
||||
pub struct Repository {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
impl Repository {
|
||||
/// Create a new repository with a connection pool.
|
||||
pub fn new(pool: Pool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// Get a connection from the pool.
|
||||
async fn conn(&self) -> Result<deadpool_postgres::Object, WorkspaceError> {
|
||||
self.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Failed to get connection: {}", e),
|
||||
})
|
||||
}
|
||||
|
||||
// ==================== Document Operations ====================
|
||||
|
||||
/// Get a document by type and optional title.
|
||||
pub async fn get_document(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: DocType,
|
||||
title: Option<&str>,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let row = if let Some(title) = title {
|
||||
conn.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
AND doc_type = $3 AND title = $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &doc_type.as_str(), &title],
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
conn.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
AND doc_type = $3 AND title IS NULL
|
||||
"#,
|
||||
&[&user_id, &agent_id, &doc_type.as_str()],
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
let row = row.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(self.row_to_document(&row)?),
|
||||
None => Err(WorkspaceError::DocumentNotFound {
|
||||
doc_type: doc_type.to_string(),
|
||||
user_id: user_id.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a document by ID.
|
||||
pub async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let row = conn
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents WHERE id = $1
|
||||
"#,
|
||||
&[&id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
match row {
|
||||
Some(row) => Ok(self.row_to_document(&row)?),
|
||||
None => Err(WorkspaceError::DocumentNotFound {
|
||||
doc_type: "unknown".to_string(),
|
||||
user_id: "unknown".to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or create a document.
|
||||
pub async fn get_or_create_document(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: DocType,
|
||||
title: Option<&str>,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
// Try to get existing document first
|
||||
match self.get_document(user_id, agent_id, doc_type, title).await {
|
||||
Ok(doc) => return Ok(doc),
|
||||
Err(WorkspaceError::DocumentNotFound { .. }) => {}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
|
||||
// Create new document
|
||||
let conn = self.conn().await?;
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO memory_documents (id, user_id, agent_id, doc_type, title, content, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, '', $6, $7)
|
||||
ON CONFLICT (user_id, agent_id, doc_type, title) DO NOTHING
|
||||
"#,
|
||||
&[&id, &user_id, &agent_id, &doc_type.as_str(), &title, &now, &now],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Insert failed: {}", e),
|
||||
})?;
|
||||
|
||||
// Fetch the document (might have been created by concurrent request)
|
||||
self.get_document(user_id, agent_id, doc_type, title).await
|
||||
}
|
||||
|
||||
/// Update a document's content.
|
||||
pub async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"UPDATE memory_documents SET content = $2, updated_at = NOW() WHERE id = $1",
|
||||
&[&id, &content],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Update failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List documents by type.
|
||||
pub async fn list_documents(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
doc_type: Option<DocType>,
|
||||
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let rows = if let Some(dt) = doc_type {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2 AND doc_type = $3
|
||||
ORDER BY updated_at DESC
|
||||
"#,
|
||||
&[&user_id, &agent_id, &dt.as_str()],
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
conn.query(
|
||||
r#"
|
||||
SELECT id, user_id, agent_id, doc_type, title, content,
|
||||
created_at, updated_at, metadata
|
||||
FROM memory_documents
|
||||
WHERE user_id = $1 AND agent_id IS NOT DISTINCT FROM $2
|
||||
ORDER BY updated_at DESC
|
||||
"#,
|
||||
&[&user_id, &agent_id],
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
let rows = rows.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
rows.iter().map(|r| self.row_to_document(r)).collect()
|
||||
}
|
||||
|
||||
fn row_to_document(&self, row: &tokio_postgres::Row) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let doc_type_str: String = row.get("doc_type");
|
||||
let doc_type = DocType::try_from(doc_type_str.as_str())?;
|
||||
|
||||
Ok(MemoryDocument {
|
||||
id: row.get("id"),
|
||||
user_id: row.get("user_id"),
|
||||
agent_id: row.get("agent_id"),
|
||||
doc_type,
|
||||
title: row.get("title"),
|
||||
content: row.get("content"),
|
||||
created_at: row.get("created_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
metadata: row.get("metadata"),
|
||||
})
|
||||
}
|
||||
|
||||
// ==================== Chunk Operations ====================
|
||||
|
||||
/// Delete all chunks for a document.
|
||||
pub async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"DELETE FROM memory_chunks WHERE document_id = $1",
|
||||
&[&document_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::ChunkingFailed {
|
||||
reason: format!("Delete failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Insert a chunk.
|
||||
pub async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let id = Uuid::new_v4();
|
||||
|
||||
let embedding_vec = embedding.map(|e| Vector::from(e.to_vec()));
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO memory_chunks (id, document_id, chunk_index, content, embedding)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
"#,
|
||||
&[&id, &document_id, &chunk_index, &content, &embedding_vec],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::ChunkingFailed {
|
||||
reason: format!("Insert failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// Update a chunk's embedding.
|
||||
pub async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let embedding_vec = Vector::from(embedding.to_vec());
|
||||
|
||||
conn.execute(
|
||||
"UPDATE memory_chunks SET embedding = $2 WHERE id = $1",
|
||||
&[&chunk_id, &embedding_vec],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::EmbeddingFailed {
|
||||
reason: format!("Update failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get chunks without embeddings for backfilling.
|
||||
pub async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id, c.document_id, c.chunk_index, c.content, c.created_at
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.embedding IS NULL
|
||||
LIMIT $3
|
||||
"#,
|
||||
&[&user_id, &agent_id, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|row| MemoryChunk {
|
||||
id: row.get("id"),
|
||||
document_id: row.get("document_id"),
|
||||
chunk_index: row.get("chunk_index"),
|
||||
content: row.get("content"),
|
||||
embedding: None,
|
||||
created_at: row.get("created_at"),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
// ==================== Search Operations ====================
|
||||
|
||||
/// Perform hybrid search combining FTS and vector similarity.
|
||||
pub async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
let fts_results = if config.use_fts {
|
||||
self.fts_search(user_id, agent_id, query, config.pre_fusion_limit)
|
||||
.await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let vector_results = if config.use_vector && embedding.is_some() {
|
||||
self.vector_search(
|
||||
user_id,
|
||||
agent_id,
|
||||
embedding.unwrap(),
|
||||
config.pre_fusion_limit,
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
Ok(reciprocal_rank_fusion(fts_results, vector_results, config))
|
||||
}
|
||||
|
||||
/// Full-text search using PostgreSQL ts_rank_cd.
|
||||
async fn fts_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<RankedResult>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
// Use plainto_tsquery for natural language queries
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id as chunk_id, c.document_id, c.content,
|
||||
ts_rank_cd(c.content_tsv, plainto_tsquery('english', $3)) as rank
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.content_tsv @@ plainto_tsquery('english', $3)
|
||||
ORDER BY rank DESC
|
||||
LIMIT $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &query, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("FTS query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, row)| RankedResult {
|
||||
chunk_id: row.get("chunk_id"),
|
||||
document_id: row.get("document_id"),
|
||||
content: row.get("content"),
|
||||
rank: (i + 1) as u32, // 1-based rank
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Vector similarity search using pgvector cosine distance.
|
||||
async fn vector_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
embedding: &[f32],
|
||||
limit: usize,
|
||||
) -> Result<Vec<RankedResult>, WorkspaceError> {
|
||||
let conn = self.conn().await?;
|
||||
let embedding_vec = Vector::from(embedding.to_vec());
|
||||
|
||||
// Use cosine distance (<=>)
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT c.id as chunk_id, c.document_id, c.content,
|
||||
1 - (c.embedding <=> $3) as similarity
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE d.user_id = $1 AND d.agent_id IS NOT DISTINCT FROM $2
|
||||
AND c.embedding IS NOT NULL
|
||||
ORDER BY c.embedding <=> $3
|
||||
LIMIT $4
|
||||
"#,
|
||||
&[&user_id, &agent_id, &embedding_vec, &(limit as i64)],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::SearchFailed {
|
||||
reason: format!("Vector query failed: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, row)| RankedResult {
|
||||
chunk_id: row.get("chunk_id"),
|
||||
document_id: row.get("document_id"),
|
||||
content: row.get("content"),
|
||||
rank: (i + 1) as u32, // 1-based rank
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,391 @@
|
||||
//! Hybrid search combining full-text and semantic search.
|
||||
//!
|
||||
//! Uses Reciprocal Rank Fusion (RRF) to combine results from:
|
||||
//! 1. PostgreSQL full-text search (ts_rank_cd)
|
||||
//! 2. pgvector cosine similarity search
|
||||
//!
|
||||
//! RRF formula: score = sum(1 / (k + rank)) for each retrieval method
|
||||
//! This is robust to different score scales and produces better results
|
||||
//! than simple score averaging.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Configuration for hybrid search.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchConfig {
|
||||
/// Maximum number of results to return.
|
||||
pub limit: usize,
|
||||
/// RRF constant (typically 60). Higher values favor top results more.
|
||||
pub rrf_k: u32,
|
||||
/// Whether to include FTS results.
|
||||
pub use_fts: bool,
|
||||
/// Whether to include vector results.
|
||||
pub use_vector: bool,
|
||||
/// Minimum score threshold (0.0-1.0).
|
||||
pub min_score: f32,
|
||||
/// Maximum results to fetch from each method before fusion.
|
||||
pub pre_fusion_limit: usize,
|
||||
}
|
||||
|
||||
impl Default for SearchConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
limit: 10,
|
||||
rrf_k: 60,
|
||||
use_fts: true,
|
||||
use_vector: true,
|
||||
min_score: 0.0,
|
||||
pre_fusion_limit: 50,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchConfig {
|
||||
/// Set the result limit.
|
||||
pub fn with_limit(mut self, limit: usize) -> Self {
|
||||
self.limit = limit;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the RRF constant.
|
||||
pub fn with_rrf_k(mut self, k: u32) -> Self {
|
||||
self.rrf_k = k;
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable FTS (only use vector search).
|
||||
pub fn vector_only(mut self) -> Self {
|
||||
self.use_fts = false;
|
||||
self.use_vector = true;
|
||||
self
|
||||
}
|
||||
|
||||
/// Disable vector search (only use FTS).
|
||||
pub fn fts_only(mut self) -> Self {
|
||||
self.use_fts = true;
|
||||
self.use_vector = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set minimum score threshold.
|
||||
pub fn with_min_score(mut self, score: f32) -> Self {
|
||||
self.min_score = score.clamp(0.0, 1.0);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// A search result with hybrid scoring.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchResult {
|
||||
/// Document ID containing this chunk.
|
||||
pub document_id: Uuid,
|
||||
/// Chunk ID.
|
||||
pub chunk_id: Uuid,
|
||||
/// Chunk content.
|
||||
pub content: String,
|
||||
/// Combined RRF score (0.0-1.0 normalized).
|
||||
pub score: f32,
|
||||
/// Rank in FTS results (1-based, None if not in FTS results).
|
||||
pub fts_rank: Option<u32>,
|
||||
/// Rank in vector results (1-based, None if not in vector results).
|
||||
pub vector_rank: Option<u32>,
|
||||
}
|
||||
|
||||
impl SearchResult {
|
||||
/// Check if this result came from FTS.
|
||||
pub fn from_fts(&self) -> bool {
|
||||
self.fts_rank.is_some()
|
||||
}
|
||||
|
||||
/// Check if this result came from vector search.
|
||||
pub fn from_vector(&self) -> bool {
|
||||
self.vector_rank.is_some()
|
||||
}
|
||||
|
||||
/// Check if this result came from both methods (hybrid match).
|
||||
pub fn is_hybrid(&self) -> bool {
|
||||
self.fts_rank.is_some() && self.vector_rank.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw result from a single search method.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RankedResult {
|
||||
pub chunk_id: Uuid,
|
||||
pub document_id: Uuid,
|
||||
pub content: String,
|
||||
pub rank: u32, // 1-based rank
|
||||
}
|
||||
|
||||
/// Reciprocal Rank Fusion algorithm.
|
||||
///
|
||||
/// Combines ranked results from multiple retrieval methods using the formula:
|
||||
/// score(d) = sum(1 / (k + rank(d))) for each method where d appears
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `fts_results` - Results from full-text search, ordered by relevance
|
||||
/// * `vector_results` - Results from vector search, ordered by similarity
|
||||
/// * `config` - Search configuration
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// Combined results sorted by RRF score (descending).
|
||||
pub fn reciprocal_rank_fusion(
|
||||
fts_results: Vec<RankedResult>,
|
||||
vector_results: Vec<RankedResult>,
|
||||
config: &SearchConfig,
|
||||
) -> Vec<SearchResult> {
|
||||
let k = config.rrf_k as f32;
|
||||
|
||||
// Track scores and metadata for each chunk
|
||||
struct ChunkInfo {
|
||||
document_id: Uuid,
|
||||
content: String,
|
||||
score: f32,
|
||||
fts_rank: Option<u32>,
|
||||
vector_rank: Option<u32>,
|
||||
}
|
||||
|
||||
let mut chunk_scores: HashMap<Uuid, ChunkInfo> = HashMap::new();
|
||||
|
||||
// Process FTS results
|
||||
for result in fts_results {
|
||||
let rrf_score = 1.0 / (k + result.rank as f32);
|
||||
chunk_scores
|
||||
.entry(result.chunk_id)
|
||||
.and_modify(|info| {
|
||||
info.score += rrf_score;
|
||||
info.fts_rank = Some(result.rank);
|
||||
})
|
||||
.or_insert(ChunkInfo {
|
||||
document_id: result.document_id,
|
||||
content: result.content,
|
||||
score: rrf_score,
|
||||
fts_rank: Some(result.rank),
|
||||
vector_rank: None,
|
||||
});
|
||||
}
|
||||
|
||||
// Process vector results
|
||||
for result in vector_results {
|
||||
let rrf_score = 1.0 / (k + result.rank as f32);
|
||||
chunk_scores
|
||||
.entry(result.chunk_id)
|
||||
.and_modify(|info| {
|
||||
info.score += rrf_score;
|
||||
info.vector_rank = Some(result.rank);
|
||||
})
|
||||
.or_insert(ChunkInfo {
|
||||
document_id: result.document_id,
|
||||
content: result.content,
|
||||
score: rrf_score,
|
||||
fts_rank: None,
|
||||
vector_rank: Some(result.rank),
|
||||
});
|
||||
}
|
||||
|
||||
// Convert to SearchResult and sort by score
|
||||
let mut results: Vec<SearchResult> = chunk_scores
|
||||
.into_iter()
|
||||
.map(|(chunk_id, info)| SearchResult {
|
||||
document_id: info.document_id,
|
||||
chunk_id,
|
||||
content: info.content,
|
||||
score: info.score,
|
||||
fts_rank: info.fts_rank,
|
||||
vector_rank: info.vector_rank,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Normalize scores to 0-1 range
|
||||
if let Some(max_score) = results.iter().map(|r| r.score).reduce(f32::max) {
|
||||
if max_score > 0.0 {
|
||||
for result in &mut results {
|
||||
result.score /= max_score;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Filter by minimum score
|
||||
if config.min_score > 0.0 {
|
||||
results.retain(|r| r.score >= config.min_score);
|
||||
}
|
||||
|
||||
// Sort by score descending
|
||||
results.sort_by(|a, b| {
|
||||
b.score
|
||||
.partial_cmp(&a.score)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Limit results
|
||||
results.truncate(config.limit);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_result(chunk_id: Uuid, doc_id: Uuid, rank: u32) -> RankedResult {
|
||||
RankedResult {
|
||||
chunk_id,
|
||||
document_id: doc_id,
|
||||
content: format!("content for chunk {}", chunk_id),
|
||||
rank,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_single_method() {
|
||||
let config = SearchConfig::default().with_limit(10);
|
||||
|
||||
let chunk1 = Uuid::new_v4();
|
||||
let chunk2 = Uuid::new_v4();
|
||||
let doc = Uuid::new_v4();
|
||||
|
||||
let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)];
|
||||
|
||||
let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config);
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
// First result should have higher score
|
||||
assert!(results[0].score > results[1].score);
|
||||
// All should have FTS rank
|
||||
assert!(results.iter().all(|r| r.fts_rank.is_some()));
|
||||
assert!(results.iter().all(|r| r.vector_rank.is_none()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_hybrid_match_boosted() {
|
||||
let config = SearchConfig::default().with_limit(10);
|
||||
|
||||
let chunk1 = Uuid::new_v4(); // In both
|
||||
let chunk2 = Uuid::new_v4(); // FTS only
|
||||
let chunk3 = Uuid::new_v4(); // Vector only
|
||||
let doc = Uuid::new_v4();
|
||||
|
||||
let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)];
|
||||
|
||||
let vector_results = vec![make_result(chunk1, doc, 1), make_result(chunk3, doc, 2)];
|
||||
|
||||
let results = reciprocal_rank_fusion(fts_results, vector_results, &config);
|
||||
|
||||
assert_eq!(results.len(), 3);
|
||||
|
||||
// chunk1 should be first (hybrid match)
|
||||
assert_eq!(results[0].chunk_id, chunk1);
|
||||
assert!(results[0].is_hybrid());
|
||||
assert!(results[0].score > results[1].score);
|
||||
|
||||
// Other chunks should not be hybrid
|
||||
assert!(!results[1].is_hybrid());
|
||||
assert!(!results[2].is_hybrid());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_score_normalization() {
|
||||
let config = SearchConfig::default();
|
||||
|
||||
let chunk1 = Uuid::new_v4();
|
||||
let doc = Uuid::new_v4();
|
||||
|
||||
let fts_results = vec![make_result(chunk1, doc, 1)];
|
||||
|
||||
let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config);
|
||||
|
||||
// Single result should have normalized score of 1.0
|
||||
assert_eq!(results.len(), 1);
|
||||
assert!((results[0].score - 1.0).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_min_score_filter() {
|
||||
let config = SearchConfig::default().with_limit(10).with_min_score(0.5);
|
||||
|
||||
let chunk1 = Uuid::new_v4();
|
||||
let chunk2 = Uuid::new_v4();
|
||||
let chunk3 = Uuid::new_v4();
|
||||
let doc = Uuid::new_v4();
|
||||
|
||||
// chunk1 has rank 1, chunk3 has rank 100 (low score)
|
||||
let fts_results = vec![
|
||||
make_result(chunk1, doc, 1),
|
||||
make_result(chunk2, doc, 50),
|
||||
make_result(chunk3, doc, 100),
|
||||
];
|
||||
|
||||
let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config);
|
||||
|
||||
// Low-scoring results should be filtered out
|
||||
// All results should have score >= 0.5
|
||||
for result in &results {
|
||||
assert!(result.score >= 0.5);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_limit() {
|
||||
let config = SearchConfig::default().with_limit(2);
|
||||
|
||||
let doc = Uuid::new_v4();
|
||||
let fts_results: Vec<_> = (1..=5)
|
||||
.map(|i| make_result(Uuid::new_v4(), doc, i))
|
||||
.collect();
|
||||
|
||||
let results = reciprocal_rank_fusion(fts_results, Vec::new(), &config);
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_k_parameter() {
|
||||
// Higher k values make ranking differences less pronounced
|
||||
let chunk1 = Uuid::new_v4();
|
||||
let chunk2 = Uuid::new_v4();
|
||||
let doc = Uuid::new_v4();
|
||||
|
||||
let fts_results = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)];
|
||||
|
||||
// Low k: rank 1 score = 1/(10+1) = 0.091, rank 2 = 1/(10+2) = 0.083
|
||||
let config_low_k = SearchConfig::default().with_rrf_k(10);
|
||||
let results_low = reciprocal_rank_fusion(fts_results.clone(), Vec::new(), &config_low_k);
|
||||
|
||||
// High k: rank 1 score = 1/(100+1) = 0.0099, rank 2 = 1/(100+2) = 0.0098
|
||||
let config_high_k = SearchConfig::default().with_rrf_k(100);
|
||||
let results_high = reciprocal_rank_fusion(fts_results, Vec::new(), &config_high_k);
|
||||
|
||||
// With low k, the score difference is larger (relatively)
|
||||
let diff_low = results_low[0].score - results_low[1].score;
|
||||
let diff_high = results_high[0].score - results_high[1].score;
|
||||
|
||||
// Low k should have larger relative difference
|
||||
assert!(diff_low > diff_high);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_config_builders() {
|
||||
let config = SearchConfig::default()
|
||||
.with_limit(20)
|
||||
.with_rrf_k(30)
|
||||
.with_min_score(0.1);
|
||||
|
||||
assert_eq!(config.limit, 20);
|
||||
assert_eq!(config.rrf_k, 30);
|
||||
assert!((config.min_score - 0.1).abs() < 0.001);
|
||||
assert!(config.use_fts);
|
||||
assert!(config.use_vector);
|
||||
|
||||
let fts_only = SearchConfig::default().fts_only();
|
||||
assert!(fts_only.use_fts);
|
||||
assert!(!fts_only.use_vector);
|
||||
|
||||
let vector_only = SearchConfig::default().vector_only();
|
||||
assert!(!vector_only.use_fts);
|
||||
assert!(vector_only.use_vector);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user