core/rag/src/storage/memory.rs - main - alexefan136/flowstack

core/rag/src/storage/memory.rs
732 строки25 KB
Alexander Efanov
Обновление репозитория
15 июл 2026, 12:19
15 июл 2026, 12:1976704c6
//! In-memory vector store для unit тестов и MVP.
//!
//! # Использование
//!
//! Подходит для:
//! - Unit тестов (без внешних зависимостей)
//! - Локальной разработки
//! - MVP и прототипирования
//!
//! **Не использовать в production** — данные теряются при рестарте.

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

use async_trait::async_trait;
use chrono::{DateTime, Utc};

use super::traits::{SearchParams, StorageStats, VectorStore};
use crate::domain::{Chunk, DocumentSummary, DocumentType, Embedding, RetrievalResult, SearchFilters};
use crate::{Error, Result};

// ============================================================================
// MemoryStore
// ============================================================================

/// In-memory vector store.
///
/// Хранит все данные в `HashMap` за `RwLock`.
/// Потокобезопасный — можно использовать из нескольких async задач.
pub struct MemoryStore {
    /// chunk_id -> (Chunk, Embedding)
    chunks: Arc<RwLock<HashMap<String, (Chunk, Embedding)>>>,
    /// Размерность векторов
    vector_size: usize,
}

impl MemoryStore {
    /// Создать новый in-memory store с размерностью по умолчанию (1536).
    pub fn new() -> Self {
        Self {
            chunks: Arc::new(RwLock::new(HashMap::new())),
            vector_size: 1536,
        }
    }

    /// Создать in-memory store с указанной размерностью векторов.
    pub fn with_vector_size(vector_size: usize) -> Self {
        Self {
            chunks: Arc::new(RwLock::new(HashMap::new())),
            vector_size,
        }
    }

    /// Проверить что chunk проходит фильтр.
    fn matches_filter(chunk: &Chunk, filters: &SearchFilters) -> bool {
        // Workspace filter (must match)
        if let Some(ref workspace_id) = filters.workspace_id {
            let chunk_workspace = chunk
                .metadata
                .get("workspace_id")
                .and_then(|v| v.as_str());
            if chunk_workspace != Some(workspace_id.as_str()) {
                return false;
            }
        }

        // Document types filter (OR — must match any)
        if let Some(ref doc_types) = filters.doc_types {
            let chunk_doc_type = chunk
                .metadata
                .get("doc_type")
                .and_then(|v| v.as_str());
            if let Some(dt) = chunk_doc_type {
                if !doc_types.contains(&dt.to_string()) {
                    return false;
                }
            } else {
                return false;
            }
        }

        // Sources filter (must match any)
        if let Some(ref sources) = filters.sources {
            let chunk_source = chunk.metadata.get("source").and_then(|v| v.as_str());
            if let Some(src) = chunk_source {
                if !sources.contains(&src.to_string()) {
                    return false;
                }
            } else {
                return false;
            }
        }

        // Extra filters (AND — must match all)
        for (key, expected) in &filters.extra {
            if let Some(actual) = chunk.metadata.get(key) {
                if actual != expected {
                    return false;
                }
            } else {
                return false;
            }
        }

        true
    }

    /// Извлечь значение поля из metadata chunk'а как строку.
    fn get_metadata_str(chunk: &Chunk, key: &str) -> Option<String> {
        chunk.metadata.get(key).and_then(|v| v.as_str()).map(|s| s.to_string())
    }

    /// Извлечь дату из metadata chunk'а (RFC3339 строка → DateTime<Utc>).
    fn get_metadata_datetime(chunk: &Chunk, key: &str) -> Option<DateTime<Utc>> {
        chunk
            .metadata
            .get(key)
            .and_then(|v| v.as_str())
            .and_then(|s| DateTime::parse_from_rfc3339(s).ok())
            .map(|dt| dt.with_timezone(&Utc))
    }
}

impl Default for MemoryStore {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// VectorStore trait implementation
// ============================================================================

#[async_trait]
impl VectorStore for MemoryStore {
    async fn initialize(&self) -> Result<()> {
        Ok(())
    }

    async fn upsert_chunks(
        &self,
        chunks: Vec<Chunk>,
        embeddings: Vec<Embedding>,
    ) -> Result<usize> {
        if chunks.len() != embeddings.len() {
            return Err(Error::InvalidInput {
                message: format!(
                    "Chunks ({}) and embeddings ({}) count mismatch",
                    chunks.len(),
                    embeddings.len()
                ),
                field: None,
            });
        }

        for emb in &embeddings {
            if emb.dimensions != self.vector_size {
                return Err(Error::DimensionMismatch {
                    expected: self.vector_size,
                    actual: emb.dimensions,
                });
            }
        }

        let count = chunks.len();
        if count == 0 {
            return Ok(0);
        }

        let mut storage = self.chunks.write().map_err(|_| Error::Storage {
            message: "Failed to acquire write lock".into(),
            source: None,
        })?;

        for (chunk, embedding) in chunks.into_iter().zip(embeddings.into_iter()) {
            storage.insert(chunk.id.clone(), (chunk, embedding));
        }

        Ok(count)
    }

    async fn search(
        &self,
        query_embedding: &Embedding,
        filters: &SearchFilters,
        params: SearchParams,
    ) -> Result<Vec<RetrievalResult>> {
        if query_embedding.dimensions != self.vector_size {
            return Err(Error::DimensionMismatch {
                expected: self.vector_size,
                actual: query_embedding.dimensions,
            });
        }

        let storage = self.chunks.read().map_err(|_| Error::Storage {
            message: "Failed to acquire read lock".into(),
            source: None,
        })?;

        let mut results: Vec<(Chunk, f32)> = storage
            .values()
            .filter(|(chunk, _)| Self::matches_filter(chunk, filters))
            .map(|(chunk, embedding)| {
                let score = query_embedding.cosine_similarity(embedding);
                (chunk.clone(), score)
            })
            .filter(|(_, score)| {
                params
                    .score_threshold
                    .map(|threshold| *score >= threshold)
                    .unwrap_or(true)
            })
            .collect();

        // Сортируем по score descending
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Берём top_k
        results.truncate(params.top_k);

        Ok(results
            .into_iter()
            .map(|(chunk, score)| RetrievalResult {
                chunk,
                score,
                rerank_score: None,
            })
            .collect())
    }

    async fn delete_by_document(&self, document_id: &str) -> Result<usize> {
        let mut storage = self.chunks.write().map_err(|_| Error::Storage {
            message: "Failed to acquire write lock".into(),
            source: None,
        })?;

        let before = storage.len();
        storage.retain(|_, (chunk, _)| chunk.document_id != document_id);
        Ok(before - storage.len())
    }

    async fn delete_by_workspace(&self, workspace_id: &str) -> Result<usize> {
        let mut storage = self.chunks.write().map_err(|_| Error::Storage {
            message: "Failed to acquire write lock".into(),
            source: None,
        })?;

        let before = storage.len();
        storage.retain(|_, (chunk, _)| {
            chunk
                .metadata
                .get("workspace_id")
                .and_then(|v| v.as_str())
                .map(|wid| wid != workspace_id)
                .unwrap_or(true)
        });
        Ok(before - storage.len())
    }

    async fn count(&self, workspace_id: Option<&str>) -> Result<usize> {
        let storage = self.chunks.read().map_err(|_| Error::Storage {
            message: "Failed to acquire read lock".into(),
            source: None,
        })?;

        if let Some(wid) = workspace_id {
            Ok(storage
                .values()
                .filter(|(chunk, _)| {
                    chunk
                        .metadata
                        .get("workspace_id")
                        .and_then(|v| v.as_str())
                        .map(|w| w == wid)
                        .unwrap_or(false)
                })
                .count())
        } else {
            Ok(storage.len())
        }
    }

    async fn stats(&self) -> Result<StorageStats> {
        let storage = self.chunks.read().map_err(|_| Error::Storage {
            message: "Failed to acquire read lock".into(),
            source: None,
        })?;

        let mut workspaces = HashSet::new();
        let mut documents = HashSet::new();

        for (chunk, _) in storage.values() {
            if let Some(wid) = chunk.metadata.get("workspace_id").and_then(|v| v.as_str()) {
                workspaces.insert(wid.to_string());
            }
            documents.insert(chunk.document_id.clone());
        }

        Ok(StorageStats {
            total_chunks: storage.len(),
            total_documents: documents.len(),
            total_workspaces: workspaces.len(),
            vector_size: self.vector_size,
            backend: "memory".to_string(),
        })
    }

    async fn health_check(&self) -> Result<bool> {
        Ok(true)
    }

    fn backend_name(&self) -> &'static str {
        "memory"
    }

    /// Получить список всех уникальных документов с агрегированной статистикой.
    ///
    /// Группирует чанки по `document_id` и возвращает `DocumentSummary`
    /// для каждого уникального документа.
    async fn list_documents(
        &self,
        workspace_id: Option<&str>,
    ) -> Result<Vec<DocumentSummary>> {
        let storage = self.chunks.read().map_err(|_| Error::Storage {
            message: "Failed to acquire read lock".into(),
            source: None,
        })?;

        // Агрегация по document_id
        let mut docs_map: HashMap<String, DocumentSummary> = HashMap::new();

        for (chunk, _) in storage.values() {
            // Фильтр по workspace (если задан)
            if let Some(ws) = workspace_id {
                let chunk_ws = chunk
                    .metadata
                    .get("workspace_id")
                    .and_then(|v| v.as_str());
                if chunk_ws != Some(ws) {
                    continue;
                }
            }

            let doc_id = chunk.document_id.clone();
            if doc_id.is_empty() {
                continue;
            }

            let chunk_len = chunk.content.chars().count();

            // Собираем payload-подобный HashMap для использования DocumentSummary методов
            let mut payload: HashMap<String, serde_json::Value> = HashMap::new();

            if let Some(title) = Self::get_metadata_str(chunk, "title") {
                payload.insert("title".to_string(), serde_json::json!(title));
            }
            if let Some(source) = Self::get_metadata_str(chunk, "source") {
                payload.insert("source".to_string(), serde_json::json!(source));
            }
            if let Some(doc_type) = Self::get_metadata_str(chunk, "doc_type") {
                payload.insert("doc_type".to_string(), serde_json::json!(doc_type));
            }
            if let Some(author) = Self::get_metadata_str(chunk, "author") {
                payload.insert("author".to_string(), serde_json::json!(author));
            }
            if let Some(language) = Self::get_metadata_str(chunk, "language") {
                payload.insert("language".to_string(), serde_json::json!(language));
            }
            if let Some(ws_id) = Self::get_metadata_str(chunk, "workspace_id") {
                payload.insert("workspace_id".to_string(), serde_json::json!(ws_id));
            }
            if let Some(created_at) = chunk
                .metadata
                .get("created_at")
                .and_then(|v| v.as_str())
            {
                payload.insert("created_at".to_string(), serde_json::json!(created_at));
            }

            // Агрегируем: либо создаём новый summary, либо обновляем существующий
            match docs_map.get_mut(&doc_id) {
                Some(summary) => summary.merge_chunk(&payload, chunk_len),
                None => {
                    let mut summary =
                        DocumentSummary::from_first_chunk(doc_id.clone(), &payload);
                    summary.total_chars = chunk_len;
                    docs_map.insert(doc_id, summary);
                }
            }
        }

        // Сортируем по created_at DESC (новые сверху)
        let mut documents: Vec<DocumentSummary> = docs_map.into_values().collect();
        documents.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        Ok(documents)
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::{Document, DocumentType};

    fn make_chunk(doc_id: &str, workspace: &str, content: &str, index: usize) -> Chunk {
        let doc = Document::new(
            content.to_string(),
            "test.md".into(),
            DocumentType::Markdown,
            workspace.to_string(),
        );
        let mut doc = doc;
        doc.id = doc_id.to_string();
        Chunk::from_document(&doc, content.to_string(), index, 0)
    }

    fn make_chunk_with_title(
        doc_id: &str,
        workspace: &str,
        content: &str,
        title: &str,
        index: usize,
    ) -> Chunk {
        let mut chunk = make_chunk(doc_id, workspace, content, index);
        chunk.metadata.insert(
            "title".to_string(),
            serde_json::Value::String(title.to_string()),
        );
        chunk.metadata.insert(
            "source".to_string(),
            serde_json::Value::String(format!("{}.md", title.to_lowercase())),
        );
        chunk
    }

    fn make_embedding(dims: usize) -> Embedding {
        let vector: Vec<f32> = (0..dims).map(|i| (i as f32) / (dims as f32)).collect();
        Embedding::new(vector, "test".to_string())
    }

    #[tokio::test]
    async fn test_upsert_and_count() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk("doc-1", "ws-1", "content 1", 0),
            make_chunk("doc-2", "ws-1", "content 2", 0),
        ];
        let embeddings = vec![make_embedding(1536), make_embedding(1536)];

        let count = store.upsert_chunks(chunks, embeddings).await.unwrap();
        assert_eq!(count, 2);

        let total = store.count(None).await.unwrap();
        assert_eq!(total, 2);

        let ws_count = store.count(Some("ws-1")).await.unwrap();
        assert_eq!(ws_count, 2);

        let ws2_count = store.count(Some("ws-2")).await.unwrap();
        assert_eq!(ws2_count, 0);
    }

    #[tokio::test]
    async fn test_search_with_filter() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk("doc-1", "ws-1", "Rust programming", 0),
            make_chunk("doc-2", "ws-2", "Python programming", 0),
        ];
        let embeddings = vec![make_embedding(1536), make_embedding(1536)];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        let query = make_embedding(1536);
        let filters = SearchFilters {
            workspace_id: Some("ws-1".into()),
            ..Default::default()
        };

        let results = store
            .search(&query, &filters, SearchParams::new(10))
            .await
            .unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].chunk.document_id, "doc-1");
    }

    #[tokio::test]
    async fn test_delete_by_document() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk("doc-1", "ws-1", "content 1", 0),
            make_chunk("doc-1", "ws-1", "content 2", 1),
            make_chunk("doc-2", "ws-1", "content 3", 0),
        ];
        let embeddings = vec![make_embedding(1536); 3];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        let deleted = store.delete_by_document("doc-1").await.unwrap();
        assert_eq!(deleted, 2);

        let count = store.count(None).await.unwrap();
        assert_eq!(count, 1);
    }

    #[tokio::test]
    async fn test_delete_by_workspace() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk("doc-1", "ws-1", "c1", 0),
            make_chunk("doc-2", "ws-1", "c2", 0),
            make_chunk("doc-3", "ws-2", "c3", 0),
        ];
        let embeddings = vec![make_embedding(1536); 3];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        let deleted = store.delete_by_workspace("ws-1").await.unwrap();
        assert_eq!(deleted, 2);

        let count = store.count(None).await.unwrap();
        assert_eq!(count, 1);
    }

    #[tokio::test]
    async fn test_dimension_mismatch() {
        let store = MemoryStore::with_vector_size(100);
        store.initialize().await.unwrap();

        let chunks = vec![make_chunk("doc-1", "ws-1", "c1", 0)];
        let embeddings = vec![make_embedding(200)];

        let result = store.upsert_chunks(chunks, embeddings).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_stats() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk("doc-1", "ws-1", "c1", 0),
            make_chunk("doc-2", "ws-1", "c2", 0),
            make_chunk("doc-3", "ws-2", "c3", 0),
        ];
        let embeddings = vec![make_embedding(1536); 3];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        let stats = store.stats().await.unwrap();
        assert_eq!(stats.total_chunks, 3);
        assert_eq!(stats.total_documents, 3);
        assert_eq!(stats.total_workspaces, 2);
        assert_eq!(stats.vector_size, 1536);
        assert_eq!(stats.backend, "memory");
    }

    #[tokio::test]
    async fn test_upsert_updates_existing() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunk = make_chunk("doc-1", "ws-1", "old content", 0);
        let embedding = make_embedding(1536);

        store.upsert_chunks(vec![chunk], vec![embedding]).await.unwrap();

        // Upsert again with new content, same ID
        let chunk2 = make_chunk("doc-1#0", "ws-1", "new content", 0);
        let embedding2 = make_embedding(1536);
        store
            .upsert_chunks(vec![chunk2], vec![embedding2])
            .await
            .unwrap();

        let count = store.count(None).await.unwrap();
        assert_eq!(count, 1);
    }

    #[tokio::test]
    async fn test_health_check() {
        let store = MemoryStore::new();
        assert!(store.health_check().await.unwrap());
    }

    #[tokio::test]
    async fn test_list_documents_empty() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let docs = store.list_documents(None).await.unwrap();
        assert!(docs.is_empty());

        let docs = store.list_documents(Some("ws-1")).await.unwrap();
        assert!(docs.is_empty());
    }

    #[tokio::test]
    async fn test_list_documents_aggregation() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        // 3 chunks одного документа + 1 chunk другого
        let chunks = vec![
            make_chunk_with_title("doc-1", "ws-1", "Chunk 1 of doc 1", "My Doc", 0),
            make_chunk_with_title("doc-1", "ws-1", "Chunk 2 of doc 1", "My Doc", 1),
            make_chunk_with_title("doc-1", "ws-1", "Chunk 3 of doc 1", "My Doc", 2),
            make_chunk_with_title("doc-2", "ws-1", "Only chunk of doc 2", "Other", 0),
        ];
        let embeddings = vec![make_embedding(1536); 4];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        let docs = store.list_documents(Some("ws-1")).await.unwrap();
        assert_eq!(docs.len(), 2);

        // Находим doc-1 и проверяем агрегацию
        let doc1 = docs.iter().find(|d| d.document_id == "doc-1").unwrap();
        assert_eq!(doc1.chunks_count, 3);
        assert_eq!(doc1.title, "My Doc");
        let expected_chars = "Chunk 1 of doc 1".len()
            + "Chunk 2 of doc 1".len()
            + "Chunk 3 of doc 1".len();
        assert_eq!(doc1.total_chars, expected_chars);

        // Находим doc-2
        let doc2 = docs.iter().find(|d| d.document_id == "doc-2").unwrap();
        assert_eq!(doc2.chunks_count, 1);
        assert_eq!(doc2.title, "Other");
    }

    #[tokio::test]
    async fn test_list_documents_workspace_filter() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let chunks = vec![
            make_chunk_with_title("doc-1", "ws-1", "c1", "Doc 1", 0),
            make_chunk_with_title("doc-2", "ws-1", "c2", "Doc 2", 0),
            make_chunk_with_title("doc-3", "ws-2", "c3", "Doc 3", 0),
        ];
        let embeddings = vec![make_embedding(1536); 3];

        store.upsert_chunks(chunks, embeddings).await.unwrap();

        // Все документы
        let all_docs = store.list_documents(None).await.unwrap();
        assert_eq!(all_docs.len(), 3);

        // Только ws-1
        let ws1_docs = store.list_documents(Some("ws-1")).await.unwrap();
        assert_eq!(ws1_docs.len(), 2);
        assert!(ws1_docs.iter().all(|d| d.workspace_id == "ws-1"));

        // Только ws-2
        let ws2_docs = store.list_documents(Some("ws-2")).await.unwrap();
        assert_eq!(ws2_docs.len(), 1);
        assert_eq!(ws2_docs[0].document_id, "doc-3");

        // Несуществующий workspace
        let ws3_docs = store.list_documents(Some("ws-3")).await.unwrap();
        assert!(ws3_docs.is_empty());
    }

    #[tokio::test]
    async fn test_list_documents_ignores_chunks_without_document_id() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let mut chunk_with_id = make_chunk("doc-1", "ws-1", "good", 0);
        let mut chunk_without_id = make_chunk("", "ws-1", "bad", 0);
        chunk_without_id.id = "orphan".to_string();
        chunk_without_id.document_id = "".to_string();

        let embeddings = vec![make_embedding(1536); 2];
        store
            .upsert_chunks(vec![chunk_with_id, chunk_without_id], embeddings)
            .await
            .unwrap();

        let docs = store.list_documents(None).await.unwrap();
        assert_eq!(docs.len(), 1);
        assert_eq!(docs[0].document_id, "doc-1");
    }

    #[tokio::test]
    async fn test_list_documents_doc_type_parsing() {
        let store = MemoryStore::new();
        store.initialize().await.unwrap();

        let mut chunk_pdf = make_chunk("doc-pdf", "ws-1", "pdf content", 0);
        chunk_pdf.metadata.insert(
            "doc_type".to_string(),
            serde_json::json!("pdf"),
        );

        let mut chunk_md = make_chunk("doc-md", "ws-1", "md content", 0);
        chunk_md.metadata.insert(
            "doc_type".to_string(),
            serde_json::json!("Markdown"),
        );

        let embeddings = vec![make_embedding(1536); 2];
        store
            .upsert_chunks(vec![chunk_pdf, chunk_md], embeddings)
            .await
            .unwrap();

        let docs = store.list_documents(None).await.unwrap();
        assert_eq!(docs.len(), 2);

        let pdf_doc = docs.iter().find(|d| d.document_id == "doc-pdf").unwrap();
        assert_eq!(pdf_doc.doc_type, DocumentType::Pdf);

        let md_doc = docs.iter().find(|d| d.document_id == "doc-md").unwrap();
        assert_eq!(md_doc.doc_type, DocumentType::Markdown);
    }
}