1use crate::error::MemoryError;
2use crate::tokenizer::TokenCounter;
3use serde::{Deserialize, Serialize};
4use std::path::PathBuf;
5use std::sync::Arc;
6use std::time::Duration;
7
8#[derive(Clone, Serialize, Deserialize)]
10pub struct MemoryConfig {
11 pub base_dir: PathBuf,
14
15 pub embedding: EmbeddingConfig,
17
18 pub search: SearchConfig,
20
21 pub chunking: ChunkingConfig,
23
24 pub pool: PoolConfig,
26
27 pub limits: MemoryLimits,
29
30 #[serde(skip)]
32 pub token_counter: Option<Arc<dyn TokenCounter>>,
33
34 #[cfg(feature = "hnsw")]
36 #[serde(skip)]
37 pub hnsw: crate::hnsw::HnswConfig,
38}
39
40impl std::fmt::Debug for MemoryConfig {
41 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42 let mut s = f.debug_struct("MemoryConfig");
43 s.field("base_dir", &self.base_dir)
44 .field("embedding", &self.embedding)
45 .field("search", &self.search)
46 .field("chunking", &self.chunking)
47 .field("pool", &self.pool)
48 .field("limits", &self.limits)
49 .field(
50 "token_counter",
51 &self.token_counter.as_ref().map(|_| "custom"),
52 );
53 #[cfg(feature = "hnsw")]
54 s.field("hnsw", &self.hnsw);
55 s.finish()
56 }
57}
58
59impl Default for MemoryConfig {
60 fn default() -> Self {
61 Self {
62 base_dir: PathBuf::from("memory"),
63 embedding: EmbeddingConfig::default(),
64 search: SearchConfig::default(),
65 chunking: ChunkingConfig::default(),
66 pool: PoolConfig::default(),
67 limits: MemoryLimits::default(),
68 token_counter: None,
69 #[cfg(feature = "hnsw")]
70 hnsw: crate::hnsw::HnswConfig::default(),
71 }
72 }
73}
74
75impl MemoryConfig {
76 pub fn normalize_and_validate(mut self) -> Result<Self, MemoryError> {
80 self.embedding.normalize_and_validate()?;
81 self.limits = self.limits.normalize_and_validate()?;
82 let timeout_cap_secs = self.limits.embedding_timeout.as_secs().max(1);
83 self.embedding.timeout_secs = self.embedding.timeout_secs.min(timeout_cap_secs);
84 self.search
85 .normalize_and_validate(self.embedding.dimensions)?;
86 self.chunking.normalize_and_validate()?;
87 self.pool.normalize_and_validate()?;
88 #[cfg(feature = "hnsw")]
89 {
90 self.hnsw.dimensions = self.embedding.dimensions;
91 }
92 Ok(self)
93 }
94}
95
96#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct EmbeddingConfig {
99 pub ollama_url: String,
103
104 pub model: String,
106
107 pub dimensions: usize,
109
110 pub batch_size: usize,
112
113 pub timeout_secs: u64,
115}
116
117impl Default for EmbeddingConfig {
118 fn default() -> Self {
119 Self {
120 ollama_url: "http://localhost:11434".to_string(),
121 model: "nomic-embed-text".to_string(),
122 dimensions: 768,
123 batch_size: 32,
124 timeout_secs: 30,
125 }
126 }
127}
128
129impl EmbeddingConfig {
130 fn normalize_and_validate(&mut self) -> Result<(), MemoryError> {
131 if self.dimensions == 0 {
132 return Err(MemoryError::InvalidConfig {
133 field: "embedding.dimensions",
134 reason: "dimensions must be at least 1".to_string(),
135 });
136 }
137 if self.batch_size == 0 {
138 self.batch_size = 1;
139 }
140 if self.timeout_secs == 0 {
141 self.timeout_secs = 1;
142 }
143 #[cfg(not(feature = "candle-embedder"))]
147 {
148 let parsed =
149 reqwest::Url::parse(&self.ollama_url).map_err(|_| MemoryError::InvalidConfig {
150 field: "embedding.ollama_url",
151 reason: "must be an absolute http:// or https:// URL".to_string(),
152 })?;
153 match parsed.scheme() {
154 "http" | "https" if parsed.host_str().is_some() => {}
155 _ => {
156 return Err(MemoryError::InvalidConfig {
157 field: "embedding.ollama_url",
158 reason: "must be an absolute http:// or https:// URL".to_string(),
159 })
160 }
161 }
162 }
163 #[cfg(feature = "candle-embedder")]
167 {
168 let _ = &self.ollama_url; }
170 Ok(())
171 }
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize)]
176pub struct SearchConfig {
177 pub bm25_weight: f64,
179
180 pub vector_weight: f64,
182
183 #[serde(default = "default_zero")]
186 pub sparse_weight: f64,
187
188 #[serde(default = "default_sparse_top_k")]
190 pub sparse_top_k: usize,
191
192 #[serde(default = "default_zero")]
194 pub sparse_min_score: f64,
195
196 #[serde(default)]
199 pub derive_sparse_from_dense: bool,
200
201 #[serde(default = "default_sparse_derive_top_k")]
203 pub sparse_derive_top_k: usize,
204
205 #[serde(default = "default_sparse_derive_min_weight")]
207 pub sparse_derive_min_weight: f32,
208
209 #[serde(default = "default_zero")]
212 pub late_interaction_weight: f64,
213
214 pub bm25_k1: f64,
217
218 pub bm25_b: f64,
221
222 pub namespace_weights: std::collections::HashMap<String, f64>,
225
226 pub rrf_k: f64,
228
229 pub candidate_pool_size: usize,
231
232 pub default_top_k: usize,
234
235 pub min_similarity: f64,
237
238 pub recency_half_life_days: Option<f64>,
243
244 pub recency_weight: f64,
248
249 pub rerank_from_f32: bool,
254
255 #[serde(default)]
258 pub derived_vector_backend: DerivedVectorBackendPolicy,
259
260 #[serde(default = "default_turbo_quant_bits")]
262 pub turbo_quant_bits: u8,
263
264 #[serde(default = "default_turbo_quant_projections")]
266 pub turbo_quant_projections: usize,
267
268 #[serde(default)]
270 pub turbo_quant_seed: u64,
271
272 #[serde(default = "default_true")]
274 pub turbo_quant_require_exact_rerank: bool,
275
276 #[serde(default = "default_candidate_dims")]
282 pub candidate_dims: Option<usize>,
283
284 #[serde(default)]
288 pub compress_results: bool,
289}
290
291#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
293#[serde(rename_all = "snake_case")]
294pub enum DerivedVectorBackendPolicy {
295 #[default]
297 Disabled,
298 TurboQuantCandidateOnly,
300 ProveKvPoolCandidateOnly,
306}
307
308const fn default_turbo_quant_bits() -> u8 {
309 8
310}
311
312const fn default_turbo_quant_projections() -> usize {
313 64
314}
315
316const fn default_true() -> bool {
317 true
318}
319
320const fn default_zero() -> f64 {
321 0.0
322}
323
324const fn default_sparse_top_k() -> usize {
325 50
326}
327
328const fn default_sparse_derive_top_k() -> usize {
329 128
330}
331
332const MAX_SEARCH_CANDIDATE_POOL_SIZE: usize = 2_000;
333const MAX_SEARCH_DEFAULT_TOP_K: usize = 200;
334const MAX_SPARSE_TOP_K: usize = 1_000;
335const MAX_SPARSE_DERIVE_TOP_K: usize = 1_000;
336
337const fn default_sparse_derive_min_weight() -> f32 {
338 0.01
339}
340
341const fn default_candidate_dims() -> Option<usize> {
342 None
343}
344
345impl Default for SearchConfig {
346 fn default() -> Self {
347 Self {
348 bm25_weight: 1.0,
349 vector_weight: 1.0,
350 sparse_weight: 0.0,
351 sparse_top_k: default_sparse_top_k(),
352 sparse_min_score: 0.0,
353 derive_sparse_from_dense: false,
354 sparse_derive_top_k: default_sparse_derive_top_k(),
355 sparse_derive_min_weight: default_sparse_derive_min_weight(),
356 late_interaction_weight: 0.15,
357 bm25_k1: 1.2,
358 bm25_b: 0.75,
359 namespace_weights: std::collections::HashMap::new(),
360 rrf_k: 60.0,
361 candidate_pool_size: 50,
362 default_top_k: 5,
363 min_similarity: 0.3,
364 recency_half_life_days: None,
365 recency_weight: 0.5,
366 rerank_from_f32: true,
367 derived_vector_backend: DerivedVectorBackendPolicy::Disabled,
368 turbo_quant_bits: default_turbo_quant_bits(),
369 turbo_quant_projections: default_turbo_quant_projections(),
370 turbo_quant_seed: 0,
371 turbo_quant_require_exact_rerank: true,
372 candidate_dims: default_candidate_dims(),
373 compress_results: false,
374 }
375 }
376}
377
378impl SearchConfig {
379 pub(crate) fn uses_turbo_quant_backend(&self) -> bool {
380 self.derived_vector_backend == DerivedVectorBackendPolicy::TurboQuantCandidateOnly
381 }
382
383 pub(crate) fn uses_provekv_pool_backend(&self) -> bool {
384 self.derived_vector_backend == DerivedVectorBackendPolicy::ProveKvPoolCandidateOnly
385 }
386
387 pub(crate) fn uses_derived_vector_backend(&self) -> bool {
388 self.uses_turbo_quant_backend() || self.uses_provekv_pool_backend()
389 }
390
391 fn normalize_and_validate(&mut self, embedding_dimensions: usize) -> Result<(), MemoryError> {
392 #[cfg(not(feature = "turbo-quant-codec"))]
393 let _ = embedding_dimensions;
394 self.candidate_pool_size = self
395 .candidate_pool_size
396 .clamp(1, MAX_SEARCH_CANDIDATE_POOL_SIZE);
397 self.default_top_k = self.default_top_k.clamp(1, MAX_SEARCH_DEFAULT_TOP_K);
398 self.candidate_pool_size = self.candidate_pool_size.max(self.default_top_k);
399 self.sparse_top_k = self.sparse_top_k.clamp(1, MAX_SPARSE_TOP_K);
400 self.sparse_derive_top_k = self.sparse_derive_top_k.clamp(1, MAX_SPARSE_DERIVE_TOP_K);
401 if !self.rrf_k.is_finite() || self.rrf_k <= 0.0 {
402 return Err(MemoryError::InvalidConfig {
403 field: "search.rrf_k",
404 reason: "rrf_k must be finite and > 0".to_string(),
405 });
406 }
407 if !self.bm25_weight.is_finite() || self.bm25_weight < 0.0 {
408 return Err(MemoryError::InvalidConfig {
409 field: "search.bm25_weight",
410 reason: "bm25_weight must be finite and >= 0".to_string(),
411 });
412 }
413 if !self.vector_weight.is_finite() || self.vector_weight < 0.0 {
414 return Err(MemoryError::InvalidConfig {
415 field: "search.vector_weight",
416 reason: "vector_weight must be finite and >= 0".to_string(),
417 });
418 }
419 if !self.sparse_weight.is_finite() || self.sparse_weight < 0.0 {
420 return Err(MemoryError::InvalidConfig {
421 field: "search.sparse_weight",
422 reason: "sparse_weight must be finite and >= 0".to_string(),
423 });
424 }
425 if !self.sparse_min_score.is_finite() {
426 return Err(MemoryError::InvalidConfig {
427 field: "search.sparse_min_score",
428 reason: "sparse_min_score must be finite".to_string(),
429 });
430 }
431 if !self.sparse_derive_min_weight.is_finite() || self.sparse_derive_min_weight < 0.0 {
432 return Err(MemoryError::InvalidConfig {
433 field: "search.sparse_derive_min_weight",
434 reason: "sparse_derive_min_weight must be finite and >= 0".to_string(),
435 });
436 }
437 if !self.recency_weight.is_finite() || self.recency_weight < 0.0 {
438 return Err(MemoryError::InvalidConfig {
439 field: "search.recency_weight",
440 reason: "recency_weight must be finite and >= 0".to_string(),
441 });
442 }
443 if !self.min_similarity.is_finite() || !(-1.0..=1.0).contains(&self.min_similarity) {
444 return Err(MemoryError::InvalidConfig {
445 field: "search.min_similarity",
446 reason: "min_similarity must be finite and within [-1.0, 1.0]".to_string(),
447 });
448 }
449 if matches!(self.recency_half_life_days, Some(v) if !v.is_finite()) {
450 return Err(MemoryError::InvalidConfig {
451 field: "search.recency_half_life_days",
452 reason: "recency_half_life_days must be finite".to_string(),
453 });
454 }
455 if matches!(self.recency_half_life_days, Some(v) if v <= 0.0) {
456 return Err(MemoryError::InvalidConfig {
457 field: "search.recency_half_life_days",
458 reason: "recency_half_life_days must be > 0 when enabled".to_string(),
459 });
460 }
461 if self.uses_turbo_quant_backend() {
462 #[cfg(not(feature = "turbo-quant-codec"))]
463 {
464 return Err(MemoryError::InvalidConfig {
465 field: "search.derived_vector_backend",
466 reason: "turbo_quant_candidate_only requires the turbo-quant-codec feature"
467 .to_string(),
468 });
469 }
470 #[cfg(feature = "turbo-quant-codec")]
471 {
472 if embedding_dimensions % 2 != 0 {
473 return Err(MemoryError::InvalidConfig {
474 field: "embedding.dimensions",
475 reason: "TurboQuant requires even embedding dimensions".to_string(),
476 });
477 }
478 if self.turbo_quant_projections == 0 {
479 return Err(MemoryError::InvalidConfig {
480 field: "search.turbo_quant_projections",
481 reason: "TurboQuant projections must be at least 1".to_string(),
482 });
483 }
484 if !(2..=16).contains(&self.turbo_quant_bits) {
485 return Err(MemoryError::InvalidConfig {
486 field: "search.turbo_quant_bits",
487 reason: "TurboQuant bits must be within 2..=16".to_string(),
488 });
489 }
490 }
491 }
492 if self.uses_derived_vector_backend() && !self.turbo_quant_require_exact_rerank {
493 return Err(MemoryError::InvalidConfig {
494 field: "search.turbo_quant_require_exact_rerank",
495 reason: "derived vector candidate backends require exact f32 rerank".to_string(),
496 });
497 }
498 Ok(())
499 }
500}
501
502#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
504#[serde(rename_all = "snake_case")]
505pub enum ChunkingStrategy {
506 #[default]
508 Plain,
509 Sentence,
511 Code,
514 Markdown,
516}
517
518#[derive(Debug, Clone, Serialize, Deserialize)]
520pub struct ChunkingConfig {
521 pub target_size: usize,
523
524 pub min_size: usize,
526
527 pub max_size: usize,
529
530 pub overlap: usize,
532
533 #[serde(default)]
536 pub strategy: ChunkingStrategy,
537}
538
539impl Default for ChunkingConfig {
540 fn default() -> Self {
541 Self {
542 target_size: 1000,
543 min_size: 100,
544 max_size: 2000,
545 overlap: 200,
546 strategy: ChunkingStrategy::default(),
547 }
548 }
549}
550
551impl ChunkingConfig {
552 fn normalize_and_validate(&mut self) -> Result<(), MemoryError> {
553 if self.min_size == 0 {
554 self.min_size = 1;
555 }
556 if self.max_size == 0 {
557 return Err(MemoryError::InvalidConfig {
558 field: "chunking.max_size",
559 reason: "max_size must be at least 1".to_string(),
560 });
561 }
562 if self.max_size < self.min_size {
563 return Err(MemoryError::InvalidConfig {
564 field: "chunking.max_size",
565 reason: "max_size must be >= min_size".to_string(),
566 });
567 }
568 if self.target_size < self.min_size {
569 self.target_size = self.min_size;
570 }
571 if self.target_size > self.max_size {
572 self.target_size = self.max_size;
573 }
574 if self.overlap >= self.min_size {
575 self.overlap = self.min_size.saturating_sub(1);
576 }
577 Ok(())
578 }
579}
580
581#[derive(Debug, Clone, Serialize, Deserialize)]
586pub struct PoolConfig {
587 pub busy_timeout_ms: u32,
590
591 pub wal_autocheckpoint: u32,
594
595 pub enable_wal: bool,
598
599 pub max_read_connections: usize,
604
605 pub reader_timeout_secs: u64,
608}
609
610impl Default for PoolConfig {
611 fn default() -> Self {
612 Self {
613 busy_timeout_ms: 5000,
614 wal_autocheckpoint: 1000,
615 enable_wal: true,
616 max_read_connections: 4,
617 reader_timeout_secs: 30,
618 }
619 }
620}
621
622impl PoolConfig {
623 fn normalize_and_validate(&mut self) -> Result<(), MemoryError> {
624 if self.busy_timeout_ms == 0 {
625 self.busy_timeout_ms = 1;
626 }
627 if self.wal_autocheckpoint == 0 {
628 self.wal_autocheckpoint = 1;
629 }
630 if self.max_read_connections == 0 {
631 return Err(MemoryError::InvalidConfig {
632 field: "pool.max_read_connections",
633 reason: "set pool.max_read_connections to at least 1".to_string(),
634 });
635 }
636 if self.reader_timeout_secs == 0 {
637 self.reader_timeout_secs = 1;
638 }
639 self.reader_timeout_secs = self.reader_timeout_secs.min(300);
640 Ok(())
641 }
642}
643
644#[derive(Debug, Clone, Serialize, Deserialize)]
649pub struct MemoryLimits {
650 pub max_facts_per_namespace: usize,
653
654 pub max_chunks_per_document: usize,
657
658 pub max_content_bytes: usize,
661
662 pub max_embedding_concurrency: usize,
666
667 pub max_db_size_bytes: u64,
670
671 #[serde(with = "duration_secs")]
674 pub embedding_timeout: Duration,
675}
676
677impl Default for MemoryLimits {
678 fn default() -> Self {
679 Self {
680 max_facts_per_namespace: 100_000,
681 max_chunks_per_document: 1_000,
682 max_content_bytes: 1_048_576,
683 max_embedding_concurrency: 8,
684 max_db_size_bytes: 0,
685 embedding_timeout: Duration::from_secs(30),
686 }
687 }
688}
689
690impl MemoryLimits {
691 pub fn normalize_and_validate(mut self) -> Result<Self, MemoryError> {
693 if self.max_facts_per_namespace == 0 {
694 return Err(MemoryError::InvalidConfig {
695 field: "limits.max_facts_per_namespace",
696 reason: "must be at least 1".to_string(),
697 });
698 }
699 if self.max_chunks_per_document == 0 {
700 return Err(MemoryError::InvalidConfig {
701 field: "limits.max_chunks_per_document",
702 reason: "must be at least 1".to_string(),
703 });
704 }
705 if self.max_content_bytes == 0 {
706 return Err(MemoryError::InvalidConfig {
707 field: "limits.max_content_bytes",
708 reason: "must be at least 1".to_string(),
709 });
710 }
711 if self.max_embedding_concurrency > 32 {
713 self.max_embedding_concurrency = 32;
714 }
715 if self.max_embedding_concurrency == 0 {
716 self.max_embedding_concurrency = 1;
717 }
718 if self.embedding_timeout.is_zero() {
719 self.embedding_timeout = Duration::from_secs(1);
720 }
721 Ok(self)
722 }
723
724 pub fn validated(self) -> Self {
729 self.normalize_and_validate().unwrap_or_else(|err| {
730 tracing::warn!(
731 error = %err,
732 "invalid MemoryLimits supplied to validated(); using defaults"
733 );
734 let defaults = Self::default();
736 Self {
737 max_facts_per_namespace: defaults.max_facts_per_namespace,
738 max_chunks_per_document: defaults.max_chunks_per_document,
739 max_content_bytes: defaults.max_content_bytes,
740 max_embedding_concurrency: defaults.max_embedding_concurrency.clamp(1, 32),
741 max_db_size_bytes: defaults.max_db_size_bytes,
742 embedding_timeout: if defaults.embedding_timeout.is_zero() {
743 std::time::Duration::from_secs(1)
744 } else {
745 defaults.embedding_timeout
746 },
747 }
748 })
749 }
750}
751
752mod duration_secs {
753 use serde::{Deserialize, Deserializer, Serializer};
754 use std::time::Duration;
755
756 pub fn serialize<S: Serializer>(d: &Duration, s: S) -> Result<S::Ok, S::Error> {
757 s.serialize_u64(d.as_secs())
758 }
759
760 pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Duration, D::Error> {
761 let secs = u64::deserialize(d)?;
762 Ok(Duration::from_secs(secs))
763 }
764}