neomemx 0.1.2

A high-performance memory library for AI agents with semantic search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! Engine builder for configuration

use crate::embeddings::EmbeddingBase;
use crate::engine::config::{EngineConfig, ExtractionConfig, GraphConfig};
use crate::error::NeomemxError;
use crate::error::Result;
use crate::extraction::{FactConsolidator, FactExtractor};
use crate::llm::LlmBase;
use crate::search::ResultReranker;
use crate::storage::{GraphBackend, HistoryStore};
use crate::vector_store::VectorStoreBase;
use std::sync::Arc;

/// Builder for engine configuration
pub struct EngineBuilder {
    llm_provider: Option<Arc<dyn LlmBase>>,
    embedding_provider: Option<Arc<dyn EmbeddingBase>>,
    vector_backend: Option<Arc<dyn VectorStoreBase>>,
    graph_backend: Option<Arc<dyn GraphBackend>>,
    history_store: Option<Arc<dyn HistoryStore>>,
    fact_extractor: Option<Arc<dyn FactExtractor>>,
    fact_consolidator: Option<Arc<dyn FactConsolidator>>,
    reranker: Option<Arc<dyn ResultReranker>>,
    extraction_config: ExtractionConfig,
    graph_config: GraphConfig,
}

impl EngineBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self {
            llm_provider: None,
            embedding_provider: None,
            vector_backend: None,
            graph_backend: None,
            history_store: None,
            fact_extractor: None,
            fact_consolidator: None,
            reranker: None,
            extraction_config: ExtractionConfig::default(),
            graph_config: GraphConfig::default(),
        }
    }

    /// Set LLM provider
    pub fn with_llm_provider(mut self, provider: Arc<dyn LlmBase>) -> Self {
        self.llm_provider = Some(provider);
        self
    }

    /// Set embedding provider
    pub fn with_embedding_provider(mut self, provider: Arc<dyn EmbeddingBase>) -> Self {
        self.embedding_provider = Some(provider);
        self
    }

    /// Set vector backend
    pub fn with_vector_backend(mut self, backend: Arc<dyn VectorStoreBase>) -> Self {
        self.vector_backend = Some(backend);
        self
    }

    /// Set graph backend (optional)
    pub fn with_graph_backend(mut self, backend: Arc<dyn GraphBackend>) -> Self {
        self.graph_backend = Some(backend);
        self
    }

    /// Set history store
    pub fn with_history_store(mut self, store: Arc<dyn HistoryStore>) -> Self {
        self.history_store = Some(store);
        self
    }

    /// Set fact extractor
    pub fn with_fact_extractor(mut self, extractor: Arc<dyn FactExtractor>) -> Self {
        self.fact_extractor = Some(extractor);
        self
    }

    /// Set fact consolidator
    pub fn with_fact_consolidator(mut self, consolidator: Arc<dyn FactConsolidator>) -> Self {
        self.fact_consolidator = Some(consolidator);
        self
    }

    /// Set reranker (optional)
    pub fn with_reranker(mut self, reranker: Option<Arc<dyn ResultReranker>>) -> Self {
        self.reranker = reranker;
        self
    }

    /// Enable/disable fact extraction
    pub fn enable_extraction(mut self, enabled: bool) -> Self {
        self.extraction_config.enabled = enabled;
        self
    }

    /// Set extraction configuration
    pub fn with_extraction_config(mut self, config: ExtractionConfig) -> Self {
        self.extraction_config = config;
        self
    }

    /// Enable/disable graph features
    pub fn enable_graph(mut self, enabled: bool) -> Self {
        self.graph_config.enabled = enabled;
        self
    }

    /// Set graph configuration
    pub fn with_graph_config(mut self, config: GraphConfig) -> Self {
        self.graph_config = config;
        self
    }

    /// Build the engine configuration
    pub fn build(self) -> Result<EngineConfig> {
        Ok(EngineConfig {
            llm_provider: self
                .llm_provider
                .ok_or_else(|| NeomemxError::config_missing("llm_provider"))?,
            embedding_provider: self
                .embedding_provider
                .ok_or_else(|| NeomemxError::config_missing("embedding_provider"))?,
            vector_backend: self
                .vector_backend
                .ok_or_else(|| NeomemxError::config_missing("vector_backend"))?,
            graph_backend: self.graph_backend,
            history_store: self
                .history_store
                .ok_or_else(|| NeomemxError::config_missing("history_store"))?,
            fact_extractor: self
                .fact_extractor
                .ok_or_else(|| NeomemxError::config_missing("fact_extractor"))?,
            fact_consolidator: self
                .fact_consolidator
                .ok_or_else(|| NeomemxError::config_missing("fact_consolidator"))?,
            reranker: self.reranker,
            extraction_config: self.extraction_config,
            graph_config: self.graph_config,
        })
    }
}

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