1use crate::config::{DerivedVectorBackendPolicy, SearchConfig};
4use crate::episodes;
5use crate::error::MemoryError;
6use crate::types::{
7 ExplainedResult, ScoreBreakdown, SearchContext, SearchResult, SearchSource, SearchSourceType,
8 VectorSearchReceiptV1,
9};
10use rusqlite::types::Value as SqlValue;
11use rusqlite::Connection;
12#[allow(unused_imports)]
18use rusqlite::OptionalExtension;
19use stack_ids::DigestBuilder;
20#[cfg(feature = "turbo-quant-codec")]
21use std::collections::BinaryHeap;
22use std::collections::{HashMap, HashSet};
23use std::sync::atomic::{AtomicUsize, Ordering};
24
25const VECTOR_SCAN_WARN_THRESHOLD: usize = 50_000;
27const VECTOR_SCAN_HARD_LIMIT: usize = 250_000;
29
30static VECTOR_SCAN_WARN_LIMIT: AtomicUsize = AtomicUsize::new(VECTOR_SCAN_WARN_THRESHOLD);
31static VECTOR_SCAN_BLOCK_LIMIT: AtomicUsize = AtomicUsize::new(VECTOR_SCAN_HARD_LIMIT);
32
33pub fn sanitize_fts_query(raw: &str) -> Option<String> {
43 let cleaned: String = raw
44 .chars()
45 .map(|c| {
46 if c.is_alphanumeric() || c.is_whitespace() || c == '_' {
47 c
48 } else {
49 ' '
50 }
51 })
52 .collect();
53
54 let tokens: Vec<&str> = cleaned
55 .split_whitespace()
56 .filter(|t| !matches!(t.to_uppercase().as_str(), "AND" | "OR" | "NOT" | "NEAR"))
57 .collect();
58
59 if tokens.is_empty() {
60 None
61 } else {
62 Some(
63 tokens
64 .into_iter()
65 .map(|token| format!("\"{}\"", token.replace('"', "\"\"")))
66 .collect::<Vec<_>>()
67 .join(" OR "),
68 )
69 }
70}
71
72pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Result<f32, MemoryError> {
74 if a.len() != b.len() {
75 return Err(MemoryError::EmbeddingDimensionMismatch {
76 expected: a.len(),
77 actual: b.len(),
78 });
79 }
80 if let Some((index, _)) = a.iter().enumerate().find(|(_, value)| !value.is_finite()) {
81 return Err(MemoryError::NonFiniteEmbeddingValue { index });
82 }
83 if let Some((index, _)) = b.iter().enumerate().find(|(_, value)| !value.is_finite()) {
84 return Err(MemoryError::NonFiniteEmbeddingValue { index });
85 }
86 let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
87 let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
88 let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
89 if norm_a == 0.0 || norm_b == 0.0 {
90 return Ok(0.0);
91 }
92 let similarity = dot / (norm_a * norm_b);
93 if !similarity.is_finite() {
94 return Err(MemoryError::Other(
95 "cosine similarity produced a non-finite score".to_string(),
96 ));
97 }
98 Ok(similarity)
99}
100
101fn days_since(timestamp: &str, evaluation_time: chrono::DateTime<chrono::Utc>) -> Option<f64> {
102 let dt = parse_search_timestamp(timestamp)?;
103 let duration = evaluation_time.naive_utc() - dt;
104 Some(duration.num_seconds() as f64 / 86_400.0)
105}
106
107fn parse_search_timestamp(timestamp: &str) -> Option<chrono::NaiveDateTime> {
108 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(timestamp, "%Y-%m-%d %H:%M:%S") {
109 return Some(dt);
110 }
111 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(timestamp, "%Y-%m-%d %H:%M:%S%.f") {
112 return Some(dt);
113 }
114 if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(timestamp) {
115 return Some(dt.naive_utc());
116 }
117 tracing::warn!(
118 timestamp,
119 "failed to parse search timestamp for recency scoring; recency contribution dropped"
120 );
121 None
122}
123
124fn recency_contribution(
125 config: &SearchConfig,
126 context: &SearchContext,
127 updated_at: Option<&str>,
128 best_rank: Option<usize>,
129) -> Option<f64> {
130 match (config.recency_half_life_days, updated_at) {
131 (Some(half_life), Some(ts)) if half_life > 0.0 => {
132 let age_days = days_since(ts, context.evaluation_time).map(|days| days.max(0.0))?;
133 let decay = 2.0_f64.powf(-age_days / half_life);
134 let rank = best_rank.unwrap_or(1).max(1) as f64;
135 Some(config.recency_weight * decay / (config.rrf_k + rank))
136 }
137 _ => None,
138 }
139}
140
141pub(crate) fn search_result_id(source: &SearchSource) -> String {
142 match source {
143 SearchSource::Fact { fact_id, .. } => format!("fact:{fact_id}"),
144 SearchSource::Chunk { chunk_id, .. } => format!("chunk:{chunk_id}"),
145 SearchSource::Message { message_id, .. } => format!("msg:{message_id}"),
146 SearchSource::Episode { episode_id, .. } => format!("episode:{episode_id}"),
147 SearchSource::Projection { projection_id, .. } => format!("projection:{projection_id}"),
148 }
149}
150
151pub fn source_dedup_key(source: &SearchSource) -> (u8, String) {
152 match source {
153 SearchSource::Fact { fact_id, .. } => (0, fact_id.clone()),
154 SearchSource::Chunk { chunk_id, .. } => (1, chunk_id.clone()),
155 SearchSource::Message {
156 message_id,
157 session_id,
158 ..
159 } => (2, format!("{session_id}:{message_id}")),
160 SearchSource::Episode { episode_id, .. } => (3, episode_id.clone()),
161 SearchSource::Projection { projection_id, .. } => (4, projection_id.clone()),
162 }
163}
164
165#[derive(Debug, Clone)]
167pub struct Bm25Hit {
168 pub id: String,
170 pub content: String,
172 pub source: SearchSource,
174 pub raw_score: f64,
176 pub updated_at: Option<String>,
178}
179
180#[derive(Debug, Clone)]
182pub struct VectorHit {
183 pub id: String,
185 pub content: String,
187 pub source: SearchSource,
189 pub similarity: f64,
191 pub updated_at: Option<String>,
193 pub source_rank: Option<usize>,
195 pub source_similarity: Option<f64>,
197 pub reranked_from_f32: bool,
199}
200
201#[allow(dead_code)]
202struct VectorRow {
203 id: String,
204 content: String,
205 blob: Vec<u8>,
206 updated_at: Option<String>,
207 source_type: SearchSourceType,
208 filter_namespace: Option<String>,
209 filter_session_id: Option<String>,
210 source: SearchSource,
211}
212
213struct RrfCandidate {
214 content: String,
215 source: SearchSource,
216 updated_at: Option<String>,
217 bm25_score: Option<f64>,
218 bm25_rank: Option<usize>,
219 vector_score: Option<f64>,
220 vector_rank: Option<usize>,
221 vector_source_rank: Option<usize>,
222 vector_source_score: Option<f64>,
223 vector_reranked_from_f32: bool,
224 late_interaction_rank: Option<usize>,
226 late_interaction_score: Option<f64>,
228}
229
230impl RrfCandidate {
231 fn explained(self, config: &SearchConfig, context: &SearchContext) -> ExplainedResult {
232 let bm25_contribution = self
233 .bm25_rank
234 .map(|rank| config.bm25_weight / (config.rrf_k + rank as f64));
235 let vector_contribution = self
236 .vector_rank
237 .map(|rank| config.vector_weight / (config.rrf_k + rank as f64));
238 let late_interaction_weight = config.late_interaction_weight;
241 let late_interaction_contribution = self
242 .late_interaction_rank
243 .map(|rank| late_interaction_weight / (config.rrf_k + rank as f64));
244 let best_rank = match (self.bm25_rank, self.vector_rank) {
245 (Some(a), Some(b)) => Some(a.min(b)),
246 (Some(a), None) | (None, Some(a)) => Some(a),
247 (None, None) => None,
248 };
249 let recency_score =
250 recency_contribution(config, context, self.updated_at.as_deref(), best_rank);
251 let rrf_score = bm25_contribution.unwrap_or(0.0)
252 + vector_contribution.unwrap_or(0.0)
253 + late_interaction_contribution.unwrap_or(0.0)
254 + recency_score.unwrap_or(0.0);
255
256 let breakdown = ScoreBreakdown {
257 rrf_score,
258 bm25_score: self.bm25_score,
259 vector_score: self.vector_score,
260 recency_score,
261 bm25_rank: self.bm25_rank,
262 vector_rank: self.vector_rank,
263 vector_source_rank: self.vector_source_rank,
264 vector_source_score: self.vector_source_score,
265 bm25_contribution,
266 vector_contribution,
267 vector_reranked_from_f32: self.vector_reranked_from_f32,
268 bm25_weight: config.bm25_weight,
269 vector_weight: config.vector_weight,
270 recency_weight: config.recency_half_life_days.map(|_| config.recency_weight),
271 rrf_k: config.rrf_k,
272 };
273
274 ExplainedResult {
275 result: SearchResult {
276 content: self.content,
277 source: self.source,
278 score: rrf_score,
279 bm25_rank: breakdown.bm25_rank,
280 vector_rank: breakdown.vector_rank,
281 cosine_similarity: breakdown.vector_score,
282 },
283 breakdown,
284 }
285 }
286}
287
288fn scan_vector_rows(
289 rows: impl Iterator<Item = Result<VectorRow, rusqlite::Error>>,
290 query_embedding: &[f32],
291 min_similarity: f64,
292 table_label: &str,
293) -> Result<(Vec<VectorHit>, usize), MemoryError> {
294 let expected_dims = query_embedding.len();
295 let mut hits = Vec::new();
296 let mut row_count = 0usize;
297 let warn_limit = VECTOR_SCAN_WARN_LIMIT.load(Ordering::Relaxed);
298 let hard_limit = VECTOR_SCAN_BLOCK_LIMIT.load(Ordering::Relaxed);
299
300 for row in rows {
301 let row = row?;
302 row_count += 1;
303 if warn_limit > 0 && row_count == warn_limit.saturating_add(1) {
304 tracing::warn!(
305 table = table_label,
306 count = row_count,
307 threshold = warn_limit,
308 "vector scan warning threshold exceeded"
309 );
310 }
311 if hard_limit > 0 && row_count > hard_limit {
312 return Err(MemoryError::VectorScanLimitExceeded {
313 table: table_label.to_string(),
314 scanned: row_count,
315 limit: hard_limit,
316 });
317 }
318
319 let stored_embedding = match crate::db::decode_f32_le(&row.blob, expected_dims) {
320 Ok(embedding) => embedding,
321 Err(error) => {
322 tracing::warn!(
323 error = %error,
324 table = table_label,
325 item = %row.id,
326 "Skipping row with invalid embedding blob"
327 );
328 continue;
329 }
330 };
331
332 if stored_embedding.len() != expected_dims {
333 tracing::warn!(
334 expected = expected_dims,
335 actual = stored_embedding.len(),
336 "Skipping {} with wrong embedding dimensions",
337 table_label
338 );
339 continue;
340 }
341
342 let similarity = cosine_similarity(query_embedding, &stored_embedding)? as f64;
343 if similarity >= min_similarity {
344 hits.push(VectorHit {
345 id: row.id,
346 content: row.content,
347 source: row.source,
348 similarity,
349 updated_at: row.updated_at,
350 source_rank: None,
351 source_similarity: None,
352 reranked_from_f32: false,
353 });
354 }
355 }
356
357 Ok((hits, row_count))
358}
359
360fn rank_vector_hits(mut hits: Vec<VectorHit>, pool_size: usize) -> Vec<VectorHit> {
361 hits.sort_by(|a, b| {
362 b.similarity.partial_cmp(&a.similarity).unwrap_or_else(|| {
363 if a.similarity.is_nan() {
364 std::cmp::Ordering::Greater
365 } else {
366 std::cmp::Ordering::Less
367 }
368 })
369 });
370
371 for (idx, hit) in hits.iter_mut().enumerate() {
372 hit.source_rank = Some(idx + 1);
373 hit.source_similarity = Some(hit.similarity);
374 }
375
376 hits.truncate(pool_size);
377 hits
378}
379
380pub(crate) fn bm25_search(
382 conn: &Connection,
383 sanitized_query: &str,
384 pool_size: usize,
385 namespaces: Option<&[&str]>,
386 source_types: Option<&[SearchSourceType]>,
387 session_ids: Option<&[&str]>,
388) -> Result<Vec<Bm25Hit>, MemoryError> {
389 let mut hits = Vec::new();
390
391 let search_facts = source_types
392 .map(|st| st.contains(&SearchSourceType::Facts))
393 .unwrap_or(true);
394 let search_chunks = source_types
395 .map(|st| st.contains(&SearchSourceType::Chunks))
396 .unwrap_or(true);
397 let search_messages = source_types
398 .map(|st| st.contains(&SearchSourceType::Messages))
399 .unwrap_or(false);
400 let search_episodes = source_types
401 .map(|st| st.contains(&SearchSourceType::Episodes))
402 .unwrap_or(true);
403
404 if search_facts {
405 let (ns_clause, ns_params) = build_filter_clause("f.namespace", namespaces, 3);
406 let sql = format!(
407 "SELECT fm.fact_id, f.content, f.namespace, bm25(facts_fts) AS score, f.updated_at
408 FROM facts_fts
409 JOIN facts_rowid_map fm ON facts_fts.rowid = fm.rowid
410 JOIN facts f ON f.id = fm.fact_id
411 WHERE facts_fts MATCH ?1 {}
412 ORDER BY score ASC
413 LIMIT ?2",
414 ns_clause
415 );
416
417 let mut params = vec![
418 SqlValue::Text(sanitized_query.to_string()),
419 SqlValue::Integer(pool_size as i64),
420 ];
421 params.extend(ns_params);
422
423 let mut stmt = conn.prepare(&sql)?;
424 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
425 let fact_id: String = row.get(0)?;
426 let content: String = row.get(1)?;
427 let namespace: String = row.get(2)?;
428 let raw_score: f64 = row.get(3)?;
429 let updated_at: Option<String> = row.get(4)?;
430 Ok(Bm25Hit {
431 id: format!("fact:{fact_id}"),
432 content,
433 source: SearchSource::Fact { fact_id, namespace },
434 raw_score,
435 updated_at,
436 })
437 })?;
438
439 for row in rows {
440 hits.push(row?);
441 }
442 }
443
444 if search_chunks {
445 let (ns_clause, ns_params) = build_filter_clause("d.namespace", namespaces, 3);
446 let sql = format!(
447 "SELECT cm.chunk_id, c.content, c.document_id, d.title, c.chunk_index,
448 bm25(chunks_fts) AS score, c.created_at
449 FROM chunks_fts
450 JOIN chunks_rowid_map cm ON chunks_fts.rowid = cm.rowid
451 JOIN chunks c ON c.id = cm.chunk_id
452 JOIN documents d ON d.id = c.document_id
453 WHERE chunks_fts MATCH ?1 {}
454 ORDER BY score ASC
455 LIMIT ?2",
456 ns_clause
457 );
458
459 let mut params = vec![
460 SqlValue::Text(sanitized_query.to_string()),
461 SqlValue::Integer(pool_size as i64),
462 ];
463 params.extend(ns_params);
464
465 let mut stmt = conn.prepare(&sql)?;
466 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
467 let chunk_id: String = row.get(0)?;
468 let content: String = row.get(1)?;
469 let document_id: String = row.get(2)?;
470 let document_title: String = row.get(3)?;
471 let chunk_index: i64 = row.get(4)?;
472 let raw_score: f64 = row.get(5)?;
473 let updated_at: Option<String> = row.get(6)?;
474 Ok(Bm25Hit {
475 id: format!("chunk:{chunk_id}"),
476 content,
477 source: SearchSource::Chunk {
478 chunk_id,
479 document_id,
480 document_title,
481 chunk_index: chunk_index as usize,
482 },
483 raw_score,
484 updated_at,
485 })
486 })?;
487
488 for row in rows {
489 hits.push(row?);
490 }
491 }
492
493 if search_messages {
494 let (sid_clause, sid_params) = build_filter_clause("m.session_id", session_ids, 3);
495 let sql = format!(
496 "SELECT mm.message_id, m.content, m.session_id, m.role,
497 bm25(messages_fts) AS score, m.created_at
498 FROM messages_fts
499 JOIN messages_rowid_map mm ON messages_fts.rowid = mm.rowid
500 JOIN messages m ON m.id = mm.message_id
501 WHERE messages_fts MATCH ?1 {}
502 ORDER BY score ASC
503 LIMIT ?2",
504 sid_clause
505 );
506
507 let mut params = vec![
508 SqlValue::Text(sanitized_query.to_string()),
509 SqlValue::Integer(pool_size as i64),
510 ];
511 params.extend(sid_params);
512
513 let mut stmt = conn.prepare(&sql)?;
514 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
515 let message_id: i64 = row.get(0)?;
516 let content: String = row.get(1)?;
517 let session_id: String = row.get(2)?;
518 let role: String = row.get(3)?;
519 let raw_score: f64 = row.get(4)?;
520 let updated_at: Option<String> = row.get(5)?;
521 Ok(Bm25Hit {
522 id: format!("msg:{message_id}"),
523 content,
524 source: SearchSource::Message {
525 message_id,
526 session_id,
527 role,
528 },
529 raw_score,
530 updated_at,
531 })
532 })?;
533
534 for row in rows {
535 hits.push(row?);
536 }
537 }
538
539 if search_episodes {
540 let (ns_clause, ns_params) = build_filter_clause("d.namespace", namespaces, 3);
541 let sql = format!(
542 "SELECT e.episode_id, e.document_id, e.search_text, e.effect_type, e.outcome,
543 bm25(episodes_fts) AS score, e.updated_at
544 FROM episodes_fts
545 JOIN episodes_rowid_map rm ON episodes_fts.rowid = rm.rowid
546 JOIN episodes e ON e.episode_id = rm.episode_id
547 JOIN documents d ON d.id = e.document_id
548 WHERE episodes_fts MATCH ?1 {}
549 ORDER BY score ASC
550 LIMIT ?2",
551 ns_clause
552 );
553
554 let mut params = vec![
555 SqlValue::Text(sanitized_query.to_string()),
556 SqlValue::Integer(pool_size as i64),
557 ];
558 params.extend(ns_params);
559
560 let mut stmt = conn.prepare(&sql)?;
561 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
562 let episode_id: String = row.get(0)?;
563 let document_id: String = row.get(1)?;
564 let content: String = row.get(2)?;
565 let effect_type: String = row.get(3)?;
566 let outcome: String = row.get(4)?;
567 let raw_score: f64 = row.get(5)?;
568 let updated_at: Option<String> = row.get(6)?;
569 Ok(Bm25Hit {
570 id: episodes::episode_item_key(&episode_id),
571 content,
572 source: SearchSource::Episode {
573 episode_id,
574 document_id,
575 effect_type,
576 outcome,
577 },
578 raw_score,
579 updated_at,
580 })
581 })?;
582
583 for row in rows {
584 hits.push(row?);
585 }
586 }
587
588 Ok(hits)
589}
590
591pub(crate) fn vector_search(
593 conn: &Connection,
594 query_embedding: &[f32],
595 pool_size: usize,
596 min_similarity: f64,
597 namespaces: Option<&[&str]>,
598 source_types: Option<&[SearchSourceType]>,
599 session_ids: Option<&[&str]>,
600) -> Result<Vec<VectorHit>, MemoryError> {
601 let mut hits = Vec::new();
602
603 let search_facts = source_types
604 .map(|st| st.contains(&SearchSourceType::Facts))
605 .unwrap_or(true);
606 let search_chunks = source_types
607 .map(|st| st.contains(&SearchSourceType::Chunks))
608 .unwrap_or(true);
609 let search_messages = source_types
610 .map(|st| st.contains(&SearchSourceType::Messages))
611 .unwrap_or(false);
612 let search_episodes = source_types
613 .map(|st| st.contains(&SearchSourceType::Episodes))
614 .unwrap_or(true);
615
616 if search_facts {
617 let (ns_clause, ns_params) = build_filter_clause("namespace", namespaces, 1);
618 let sql = format!(
619 "SELECT id, content, namespace, embedding, updated_at
620 FROM facts
621 WHERE embedding IS NOT NULL {}",
622 ns_clause
623 );
624
625 let mut stmt = conn.prepare(&sql)?;
626 let rows = stmt.query_map(rusqlite::params_from_iter(&ns_params), |row| {
627 let id: String = row.get(0)?;
628 let content: String = row.get(1)?;
629 let namespace: String = row.get(2)?;
630 let blob: Vec<u8> = row.get(3)?;
631 let updated_at: Option<String> = row.get(4)?;
632 Ok(VectorRow {
633 id: format!("fact:{id}"),
634 content,
635 blob,
636 updated_at,
637 source_type: SearchSourceType::Facts,
638 filter_namespace: Some(namespace.clone()),
639 filter_session_id: None,
640 source: SearchSource::Fact {
641 fact_id: id,
642 namespace,
643 },
644 })
645 })?;
646
647 let (fact_hits, fact_count) =
648 scan_vector_rows(rows, query_embedding, min_similarity, "fact")?;
649 hits.extend(fact_hits);
650
651 if vector_scan_warn_exceeded(fact_count) {
652 tracing::warn!(
653 count = fact_count,
654 "facts table exceeds vector scan threshold ({} rows)",
655 fact_count
656 );
657 }
658 }
659
660 if search_chunks {
661 let (ns_clause, ns_params) = build_filter_clause("d.namespace", namespaces, 1);
662 let sql = format!(
663 "SELECT c.id, c.content, c.document_id, d.title, c.chunk_index, c.embedding, c.created_at, d.namespace
664 FROM chunks c
665 JOIN documents d ON d.id = c.document_id
666 WHERE c.embedding IS NOT NULL {}",
667 ns_clause
668 );
669
670 let mut stmt = conn.prepare(&sql)?;
671 let rows = stmt.query_map(rusqlite::params_from_iter(&ns_params), |row| {
672 let id: String = row.get(0)?;
673 let content: String = row.get(1)?;
674 let document_id: String = row.get(2)?;
675 let document_title: String = row.get(3)?;
676 let chunk_index: i64 = row.get(4)?;
677 let blob: Vec<u8> = row.get(5)?;
678 let updated_at: Option<String> = row.get(6)?;
679 let namespace: String = row.get(7)?;
680 Ok(VectorRow {
681 id: format!("chunk:{id}"),
682 content,
683 blob,
684 updated_at,
685 source_type: SearchSourceType::Chunks,
686 filter_namespace: Some(namespace),
687 filter_session_id: None,
688 source: SearchSource::Chunk {
689 chunk_id: id,
690 document_id,
691 document_title,
692 chunk_index: chunk_index as usize,
693 },
694 })
695 })?;
696
697 let (chunk_hits, chunk_count) =
698 scan_vector_rows(rows, query_embedding, min_similarity, "chunk")?;
699 hits.extend(chunk_hits);
700
701 if vector_scan_warn_exceeded(chunk_count) {
702 tracing::warn!(
703 count = chunk_count,
704 "chunks table exceeds vector scan threshold ({} rows)",
705 chunk_count
706 );
707 }
708 }
709
710 if search_messages {
711 let (sid_clause, sid_params) = build_filter_clause("m.session_id", session_ids, 1);
712 let sql = format!(
713 "SELECT m.id, m.content, m.session_id, m.role, m.embedding, m.created_at
714 FROM messages m
715 WHERE m.embedding IS NOT NULL {}",
716 sid_clause
717 );
718
719 let mut stmt = conn.prepare(&sql)?;
720 let rows = stmt.query_map(rusqlite::params_from_iter(&sid_params), |row| {
721 let message_id: i64 = row.get(0)?;
722 let content: String = row.get(1)?;
723 let session_id: String = row.get(2)?;
724 let role: String = row.get(3)?;
725 let blob: Vec<u8> = row.get(4)?;
726 let updated_at: Option<String> = row.get(5)?;
727 Ok(VectorRow {
728 id: format!("msg:{message_id}"),
729 content,
730 blob,
731 updated_at,
732 source_type: SearchSourceType::Messages,
733 filter_namespace: None,
734 filter_session_id: Some(session_id.clone()),
735 source: SearchSource::Message {
736 message_id,
737 session_id,
738 role,
739 },
740 })
741 })?;
742
743 let (message_hits, message_count) =
744 scan_vector_rows(rows, query_embedding, min_similarity, "message")?;
745 hits.extend(message_hits);
746
747 if vector_scan_warn_exceeded(message_count) {
748 tracing::warn!(
749 count = message_count,
750 "messages table exceeds vector scan threshold ({} rows)",
751 message_count
752 );
753 }
754 }
755
756 if search_episodes {
757 let (ns_clause, ns_params) = build_filter_clause("d.namespace", namespaces, 1);
758 let sql = format!(
759 "SELECT e.episode_id, e.document_id, e.search_text, e.effect_type, e.outcome, e.embedding, e.updated_at, d.namespace
760 FROM episodes e
761 JOIN documents d ON d.id = e.document_id
762 WHERE e.embedding IS NOT NULL {}",
763 ns_clause
764 );
765
766 let mut stmt = conn.prepare(&sql)?;
767 let rows = stmt.query_map(rusqlite::params_from_iter(&ns_params), |row| {
768 let episode_id: String = row.get(0)?;
769 let document_id: String = row.get(1)?;
770 let content: String = row.get(2)?;
771 let effect_type: String = row.get(3)?;
772 let outcome: String = row.get(4)?;
773 let blob: Vec<u8> = row.get(5)?;
774 let updated_at: Option<String> = row.get(6)?;
775 let namespace: String = row.get(7)?;
776 Ok(VectorRow {
777 id: episodes::episode_item_key(&episode_id),
778 content,
779 blob,
780 updated_at,
781 source_type: SearchSourceType::Episodes,
782 filter_namespace: Some(namespace),
783 filter_session_id: None,
784 source: SearchSource::Episode {
785 episode_id,
786 document_id,
787 effect_type,
788 outcome,
789 },
790 })
791 })?;
792
793 let (episode_hits, episode_count) =
794 scan_vector_rows(rows, query_embedding, min_similarity, "episode")?;
795 hits.extend(episode_hits);
796
797 if vector_scan_warn_exceeded(episode_count) {
798 tracing::warn!(
799 count = episode_count,
800 "episodes table exceeds vector scan threshold ({} rows)",
801 episode_count
802 );
803 }
804 }
805
806 Ok(rank_vector_hits(hits, pool_size))
807}
808
809fn brute_force_vector_outcome(
810 conn: &Connection,
811 query_embedding: &[f32],
812 pool_size: usize,
813 min_similarity: f64,
814 namespaces: Option<&[&str]>,
815 source_types: Option<&[SearchSourceType]>,
816 session_ids: Option<&[&str]>,
817) -> Result<VectorSearchOutcome, MemoryError> {
818 let hits = vector_search(
819 conn,
820 query_embedding,
821 pool_size,
822 min_similarity,
823 namespaces,
824 source_types,
825 session_ids,
826 )?;
827 Ok(VectorSearchOutcome {
828 requested_candidates: pool_size,
829 returned_candidates: hits.len(),
830 post_filter_candidates: hits.len(),
831 hits,
832 candidate_backend: "brute_force_f32".to_string(),
833 fallback: None,
834 exact_rerank: true,
835 degradations: Vec::new(),
836 receipt_metadata: VectorReceiptMetadata::default(),
837 })
838}
839
840#[allow(clippy::too_many_arguments)]
841fn vector_search_with_backend(
842 conn: &Connection,
843 query_embedding: &[f32],
844 pool_size: usize,
845 min_similarity: f64,
846 config: &SearchConfig,
847 context: &SearchContext,
848 namespaces: Option<&[&str]>,
849 source_types: Option<&[SearchSourceType]>,
850 session_ids: Option<&[&str]>,
851) -> Result<VectorSearchOutcome, MemoryError> {
852 if context.exactness_profile == crate::types::ExactnessProfile::PreferExact {
853 return brute_force_vector_outcome(
854 conn,
855 query_embedding,
856 pool_size,
857 min_similarity,
858 namespaces,
859 source_types,
860 session_ids,
861 );
862 }
863
864 match config.derived_vector_backend {
865 DerivedVectorBackendPolicy::Disabled => brute_force_vector_outcome(
866 conn,
867 query_embedding,
868 pool_size,
869 min_similarity,
870 namespaces,
871 source_types,
872 session_ids,
873 ),
874 DerivedVectorBackendPolicy::TurboQuantCandidateOnly => turbo_quant_vector_outcome(
875 conn,
876 query_embedding,
877 pool_size,
878 min_similarity,
879 config,
880 namespaces,
881 source_types,
882 session_ids,
883 ),
884 DerivedVectorBackendPolicy::ProveKvPoolCandidateOnly => provekv_pool_vector_outcome(
885 conn,
886 query_embedding,
887 pool_size,
888 min_similarity,
889 config,
890 namespaces,
891 source_types,
892 session_ids,
893 ),
894 }
895}
896
897#[allow(clippy::too_many_arguments)]
898fn provekv_pool_vector_outcome(
899 conn: &Connection,
900 query_embedding: &[f32],
901 pool_size: usize,
902 min_similarity: f64,
903 config: &SearchConfig,
904 namespaces: Option<&[&str]>,
905 source_types: Option<&[SearchSourceType]>,
906 session_ids: Option<&[&str]>,
907) -> Result<VectorSearchOutcome, MemoryError> {
908 if !config.turbo_quant_require_exact_rerank {
909 return Err(MemoryError::InvalidConfig {
910 field: "search.turbo_quant_require_exact_rerank",
911 reason: "proveKV pool candidate backend requires exact f32 rerank".to_string(),
912 });
913 }
914
915 let mut outcome = brute_force_vector_outcome(
916 conn,
917 query_embedding,
918 pool_size,
919 min_similarity,
920 namespaces,
921 source_types,
922 session_ids,
923 )?;
924 outcome.candidate_backend = "provekv_pool_candidate_then_exact_f32".to_string();
925 outcome.receipt_metadata.codec_family = Some("provekv_pool".to_string());
926 match crate::db::latest_ready_provekv_pool_generation(conn)? {
927 Some(row) => {
928 let item_map =
929 crate::db::load_provekv_pool_item_map(conn, &row.generation.generation_id)?;
930 let _payload =
931 crate::db::load_provekv_pool_payload(conn, &row.generation.generation_id)?;
932 outcome.receipt_metadata.artifact_generation_id = Some(row.generation.generation_id);
933 outcome.receipt_metadata.vector_artifact_manifest_digest =
934 Some(row.generation.pool_manifest_digest);
935 outcome.receipt_metadata.vector_artifact_count = Some(item_map.len());
936 outcome.degradations.push(
937 "proveKV pool generation materialized for candidate provenance; authoritative f32 exact rerank remains final"
938 .to_string(),
939 );
940 }
941 None => {
942 outcome.fallback = Some("provekv_pool_generation_not_materialized".to_string());
943 outcome.degradations.push(
944 "proveKV pool backend requested; using authoritative f32 exact path until a pool generation is materialized"
945 .to_string(),
946 );
947 }
948 }
949 Ok(outcome)
950}
951
952#[cfg(not(feature = "turbo-quant-codec"))]
953#[allow(clippy::too_many_arguments)]
954fn turbo_quant_vector_outcome(
955 conn: &Connection,
956 query_embedding: &[f32],
957 pool_size: usize,
958 min_similarity: f64,
959 _config: &SearchConfig,
960 namespaces: Option<&[&str]>,
961 source_types: Option<&[SearchSourceType]>,
962 session_ids: Option<&[&str]>,
963) -> Result<VectorSearchOutcome, MemoryError> {
964 let mut outcome = brute_force_vector_outcome(
965 conn,
966 query_embedding,
967 pool_size,
968 min_similarity,
969 namespaces,
970 source_types,
971 session_ids,
972 )?;
973 outcome.candidate_backend = "turbo_quant_candidate_then_exact_f32".to_string();
974 outcome.fallback = Some("turbo_quant_feature_disabled".to_string());
975 outcome
976 .degradations
977 .push("TurboQuant backend requested without turbo-quant-codec feature".to_string());
978 Ok(outcome)
979}
980
981#[cfg(feature = "turbo-quant-codec")]
982#[allow(clippy::too_many_arguments)]
983fn turbo_quant_vector_outcome(
984 conn: &Connection,
985 query_embedding: &[f32],
986 pool_size: usize,
987 min_similarity: f64,
988 config: &SearchConfig,
989 namespaces: Option<&[&str]>,
990 source_types: Option<&[SearchSourceType]>,
991 session_ids: Option<&[&str]>,
992) -> Result<VectorSearchOutcome, MemoryError> {
993 use crate::vector_codec::{TurboQuantCodec, VectorArtifactV1, VectorCodec};
994
995 if !config.turbo_quant_require_exact_rerank {
996 return Err(MemoryError::InvalidConfig {
997 field: "search.turbo_quant_require_exact_rerank",
998 reason: "TurboQuant candidate backend requires exact f32 rerank".to_string(),
999 });
1000 }
1001
1002 let dim = query_embedding.len();
1003 let codec = TurboQuantCodec::new(
1004 dim,
1005 config.turbo_quant_bits,
1006 config.turbo_quant_projections,
1007 config.turbo_quant_seed,
1008 )?;
1009 let profile = codec.profile().clone();
1010 let profile_digest = profile.digest();
1011 let mut metadata = VectorReceiptMetadata {
1012 codec_family: Some("turbo_quant".to_string()),
1013 codec_profile_digest: Some(profile_digest.clone()),
1014 ..VectorReceiptMetadata::default()
1015 };
1016
1017 let filtered = namespaces.is_some_and(|values| !values.is_empty())
1018 || source_types.is_some_and(|values| !values.is_empty())
1019 || session_ids.is_some_and(|values| !values.is_empty());
1020 metadata.filter_strategy = Some(if filtered {
1021 "adaptive_oversampling_after_approximate_scoring".to_string()
1022 } else {
1023 "unfiltered_top_k_heap".to_string()
1024 });
1025
1026 let raw_count = authoritative_vector_row_count(conn)?;
1027 let (current_source_snapshot_digest, current_source_row_count) =
1028 crate::db::current_source_snapshot_digest(conn, dim)?;
1029 let Some(generation) =
1030 crate::db::current_derived_vector_generation(conn, "turbo_quant", &profile_digest)?
1031 else {
1032 metadata.artifact_missing_count = Some(raw_count);
1033 metadata.vector_artifact_missing_count = Some(raw_count);
1034 let mut outcome = brute_force_vector_outcome(
1035 conn,
1036 query_embedding,
1037 pool_size,
1038 min_similarity,
1039 namespaces,
1040 source_types,
1041 session_ids,
1042 )?;
1043 outcome.candidate_backend = "turbo_quant_candidate_then_exact_f32".to_string();
1044 outcome.fallback = Some("turbo_quant_generation_missing_or_invalidated".to_string());
1045 outcome.degradations.push("No active TurboQuant artifact generation is available; authoritative raw f32 search was used".to_string());
1046 outcome.receipt_metadata = metadata;
1047 return Ok(outcome);
1048 };
1049
1050 metadata.artifact_generation_id = Some(generation.generation_id.clone());
1051 metadata.vector_artifact_manifest_digest = Some(generation.artifact_manifest_digest.clone());
1052 metadata.artifact_count = Some(generation.artifact_count);
1053
1054 let artifacts =
1055 crate::db::load_derived_vector_artifacts_by_generation(conn, &generation.generation_id)?;
1056 metadata.vector_artifact_count = Some(artifacts.len());
1057
1058 if generation.dim != dim
1059 || generation.encoding != "turbo_code_wire_v1"
1060 || generation.status != "active"
1061 || generation.source_row_count != raw_count
1062 || generation.source_row_count != current_source_row_count
1063 || generation.source_snapshot_digest != current_source_snapshot_digest
1064 || generation.artifact_count != artifacts.len()
1065 {
1066 let missing = raw_count.saturating_sub(artifacts.len());
1067 metadata.artifact_missing_count = Some(missing);
1068 metadata.vector_artifact_missing_count = Some(missing);
1069 let mut outcome = brute_force_vector_outcome(
1070 conn,
1071 query_embedding,
1072 pool_size,
1073 min_similarity,
1074 namespaces,
1075 source_types,
1076 session_ids,
1077 )?;
1078 outcome.candidate_backend = "turbo_quant_candidate_then_exact_f32".to_string();
1079 outcome.fallback = Some("turbo_quant_generation_incomplete_or_stale".to_string());
1080 outcome.degradations.push(format!(
1081 "TurboQuant generation validation failed: generation={}, status={}, dim={}, source_rows={}, artifacts={}, authoritative_rows={}, snapshot_current={}",
1082 generation.generation_id,
1083 generation.status,
1084 generation.dim,
1085 generation.source_row_count,
1086 artifacts.len(),
1087 raw_count,
1088 generation.source_snapshot_digest == current_source_snapshot_digest
1089 ));
1090 outcome.receipt_metadata = metadata;
1091 return Ok(outcome);
1092 }
1093
1094 let prepared = codec.prepare_query(query_embedding)?;
1095 let candidate_cap = if filtered {
1096 artifacts
1097 .len()
1098 .min(pool_size.saturating_mul(16).max(pool_size))
1099 } else {
1100 pool_size.min(artifacts.len())
1101 };
1102 let mut scored = BinaryHeap::with_capacity(candidate_cap.saturating_add(1));
1103 let mut corrupt_count = 0usize;
1104 let mut scanned_count = 0usize;
1105 for (seq, artifact_row) in artifacts.into_iter().enumerate() {
1106 scanned_count += 1;
1107 if artifact_row.encoding != "turbo_code_wire_v1"
1108 || artifact_row.dim != dim
1109 || artifact_row.status != "active"
1110 {
1111 corrupt_count += 1;
1112 continue;
1113 }
1114 let artifact = VectorArtifactV1::new(profile.clone(), artifact_row.encoded);
1115 if artifact.profile_digest != artifact_row.codec_profile_digest
1116 || artifact.artifact_digest != artifact_row.encoded_digest
1117 {
1118 corrupt_count += 1;
1119 continue;
1120 }
1121 let approx = match codec.score_inner_product_prepared(&artifact, &prepared) {
1122 Ok(score) if score.is_finite() => score as f64,
1123 Ok(_) => {
1124 corrupt_count += 1;
1125 continue;
1126 }
1127 Err(err) => {
1128 tracing::warn!(
1129 error = %err,
1130 item = %artifact_row.item_key,
1131 "corrupt TurboQuant artifact encountered; falling back to raw f32"
1132 );
1133 corrupt_count += 1;
1134 continue;
1135 }
1136 };
1137 if candidate_cap == 0 {
1138 continue;
1139 }
1140 let candidate = ApproxCandidate {
1141 score: approx,
1142 seq,
1143 item_key: artifact_row.item_key,
1144 };
1145 if scored.len() < candidate_cap {
1146 scored.push(candidate);
1147 } else if scored
1148 .peek()
1149 .is_some_and(|worst: &ApproxCandidate| candidate.score > worst.score)
1150 {
1151 scored.pop();
1152 scored.push(candidate);
1153 }
1154 }
1155
1156 metadata.artifact_corruption_count = Some(corrupt_count);
1157 metadata.approximate_scanned_count = Some(scanned_count);
1158 if corrupt_count > 0 {
1159 let mut outcome = brute_force_vector_outcome(
1160 conn,
1161 query_embedding,
1162 pool_size,
1163 min_similarity,
1164 namespaces,
1165 source_types,
1166 session_ids,
1167 )?;
1168 outcome.candidate_backend = "turbo_quant_candidate_then_exact_f32".to_string();
1169 outcome.fallback = Some("turbo_quant_artifact_validation_failed".to_string());
1170 outcome.degradations.push(format!(
1171 "TurboQuant artifact validation failed: {corrupt_count} corrupt artifacts in generation {}",
1172 generation.generation_id
1173 ));
1174 outcome.receipt_metadata = metadata;
1175 return Ok(outcome);
1176 }
1177
1178 let mut scored = scored.into_vec();
1179 scored.sort_by(|a, b| {
1180 b.score
1181 .partial_cmp(&a.score)
1182 .unwrap_or(std::cmp::Ordering::Equal)
1183 .then_with(|| a.seq.cmp(&b.seq))
1184 });
1185 let approximate_returned = scored.len();
1186 metadata.approximate_candidate_count = Some(approximate_returned);
1187 metadata.approximate_returned_count = Some(approximate_returned);
1188 let mut exact_hits = Vec::new();
1189 let mut raw_rows_loaded_count = 0usize;
1190 let mut missing_count = 0usize;
1191 for (approx_rank_0, candidate) in scored.into_iter().enumerate() {
1192 let Some(row) = load_vector_row_by_item_key(conn, &candidate.item_key)? else {
1193 missing_count += 1;
1194 continue;
1195 };
1196 raw_rows_loaded_count += 1;
1197 if !vector_row_matches_filters(&row, namespaces, source_types, session_ids) {
1198 continue;
1199 }
1200 let stored_embedding = crate::db::decode_f32_le(&row.blob, dim)?;
1201 let similarity = cosine_similarity(query_embedding, &stored_embedding)? as f64;
1202 if similarity >= min_similarity {
1203 exact_hits.push(VectorHit {
1204 id: row.id,
1205 content: row.content,
1206 source: row.source,
1207 similarity,
1208 updated_at: row.updated_at,
1209 source_rank: Some(approx_rank_0 + 1),
1210 source_similarity: Some(candidate.score),
1211 reranked_from_f32: true,
1212 });
1213 }
1214 }
1215 let post_filter_candidates = exact_hits.len();
1216 metadata.artifact_missing_count = Some(missing_count);
1217 metadata.vector_artifact_missing_count = Some(missing_count);
1218 metadata.vector_artifact_stale_count = Some(0);
1219 metadata.raw_rows_loaded_count = Some(raw_rows_loaded_count);
1220 metadata.exact_rerank_count = Some(raw_rows_loaded_count);
1221 let mut degradations = Vec::new();
1222 if filtered && post_filter_candidates < pool_size && candidate_cap < scanned_count {
1223 degradations.push(format!(
1224 "TurboQuant filter-aware candidate generation under-returned {post_filter_candidates} candidates for requested pool {pool_size} after scanning {scanned_count} artifacts with candidate budget {candidate_cap}"
1225 ));
1226 }
1227 if missing_count > 0 {
1228 degradations.push(format!(
1229 "TurboQuant exact rerank skipped {missing_count} candidates whose authoritative rows were missing"
1230 ));
1231 }
1232 let hits = rank_vector_hits(exact_hits, pool_size);
1233 Ok(VectorSearchOutcome {
1234 hits,
1235 candidate_backend: "turbo_quant_candidate_then_exact_f32".to_string(),
1236 requested_candidates: pool_size,
1237 returned_candidates: approximate_returned,
1238 post_filter_candidates,
1239 fallback: None,
1240 exact_rerank: true,
1241 degradations,
1242 receipt_metadata: metadata,
1243 })
1244}
1245
1246#[cfg(feature = "turbo-quant-codec")]
1247#[derive(Debug, Clone)]
1248struct ApproxCandidate {
1249 score: f64,
1250 seq: usize,
1251 item_key: String,
1252}
1253
1254#[cfg(feature = "turbo-quant-codec")]
1255impl PartialEq for ApproxCandidate {
1256 fn eq(&self, other: &Self) -> bool {
1257 self.score == other.score && self.seq == other.seq
1258 }
1259}
1260
1261#[cfg(feature = "turbo-quant-codec")]
1262impl Eq for ApproxCandidate {}
1263
1264#[cfg(feature = "turbo-quant-codec")]
1265impl PartialOrd for ApproxCandidate {
1266 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1267 Some(self.cmp(other))
1268 }
1269}
1270
1271#[cfg(feature = "turbo-quant-codec")]
1272impl Ord for ApproxCandidate {
1273 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1274 other
1275 .score
1276 .partial_cmp(&self.score)
1277 .unwrap_or(std::cmp::Ordering::Equal)
1278 .then_with(|| other.seq.cmp(&self.seq))
1279 }
1280}
1281
1282#[cfg(feature = "turbo-quant-codec")]
1283fn vector_row_matches_filters(
1284 row: &VectorRow,
1285 namespaces: Option<&[&str]>,
1286 source_types: Option<&[SearchSourceType]>,
1287 session_ids: Option<&[&str]>,
1288) -> bool {
1289 if source_types.is_some_and(|values| !values.contains(&row.source_type)) {
1290 return false;
1291 }
1292 if let Some(namespaces) = namespaces.filter(|values| !values.is_empty()) {
1293 let Some(namespace) = row.filter_namespace.as_deref() else {
1294 return false;
1295 };
1296 if !namespaces.contains(&namespace) {
1297 return false;
1298 }
1299 }
1300 if let Some(session_ids) = session_ids.filter(|values| !values.is_empty()) {
1301 let Some(session_id) = row.filter_session_id.as_deref() else {
1302 return false;
1303 };
1304 if !session_ids.contains(&session_id) {
1305 return false;
1306 }
1307 }
1308 true
1309}
1310
1311#[cfg(feature = "turbo-quant-codec")]
1312fn authoritative_vector_row_count(conn: &Connection) -> Result<usize, MemoryError> {
1313 let count: i64 = conn.query_row(
1314 "SELECT
1315 (SELECT COUNT(*) FROM facts WHERE embedding IS NOT NULL) +
1316 (SELECT COUNT(*) FROM chunks WHERE embedding IS NOT NULL) +
1317 (SELECT COUNT(*) FROM messages WHERE embedding IS NOT NULL) +
1318 (SELECT COUNT(*) FROM episodes WHERE embedding IS NOT NULL)",
1319 [],
1320 |row| row.get(0),
1321 )?;
1322 usize::try_from(count)
1323 .map_err(|err| MemoryError::Other(format!("authoritative vector count overflow: {err}")))
1324}
1325
1326#[cfg(feature = "turbo-quant-codec")]
1327fn load_vector_row_by_item_key(
1328 conn: &Connection,
1329 item_key: &str,
1330) -> Result<Option<VectorRow>, MemoryError> {
1331 let Some((domain, id)) = item_key.split_once(':') else {
1332 return Ok(None);
1333 };
1334 match domain {
1335 "fact" => conn
1336 .query_row(
1337 "SELECT id, content, namespace, embedding, updated_at
1338 FROM facts WHERE id = ?1 AND embedding IS NOT NULL",
1339 [id],
1340 |row| {
1341 let fact_id: String = row.get(0)?;
1342 let content: String = row.get(1)?;
1343 let namespace: String = row.get(2)?;
1344 let blob: Vec<u8> = row.get(3)?;
1345 let updated_at: Option<String> = row.get(4)?;
1346 Ok(VectorRow {
1347 id: format!("fact:{fact_id}"),
1348 content,
1349 blob,
1350 updated_at,
1351 source_type: SearchSourceType::Facts,
1352 filter_namespace: Some(namespace.clone()),
1353 filter_session_id: None,
1354 source: SearchSource::Fact { fact_id, namespace },
1355 })
1356 },
1357 )
1358 .optional()
1359 .map_err(MemoryError::from),
1360 "chunk" => conn
1361 .query_row(
1362 "SELECT c.id, c.content, c.document_id, d.title, c.chunk_index, c.embedding, c.created_at, d.namespace
1363 FROM chunks c
1364 JOIN documents d ON d.id = c.document_id
1365 WHERE c.id = ?1 AND c.embedding IS NOT NULL",
1366 [id],
1367 |row| {
1368 let chunk_id: String = row.get(0)?;
1369 let content: String = row.get(1)?;
1370 let document_id: String = row.get(2)?;
1371 let document_title: String = row.get(3)?;
1372 let chunk_index: i64 = row.get(4)?;
1373 let blob: Vec<u8> = row.get(5)?;
1374 let updated_at: Option<String> = row.get(6)?;
1375 let namespace: String = row.get(7)?;
1376 Ok(VectorRow {
1377 id: format!("chunk:{chunk_id}"),
1378 content,
1379 blob,
1380 updated_at,
1381 source_type: SearchSourceType::Chunks,
1382 filter_namespace: Some(namespace),
1383 filter_session_id: None,
1384 source: SearchSource::Chunk {
1385 chunk_id,
1386 document_id,
1387 document_title,
1388 chunk_index: chunk_index as usize,
1389 },
1390 })
1391 },
1392 )
1393 .optional()
1394 .map_err(MemoryError::from),
1395 "msg" => {
1396 let Ok(message_id) = id.parse::<i64>() else {
1397 return Ok(None);
1398 };
1399 conn.query_row(
1400 "SELECT id, content, session_id, role, embedding, created_at
1401 FROM messages WHERE id = ?1 AND embedding IS NOT NULL",
1402 [message_id],
1403 |row| {
1404 let message_id: i64 = row.get(0)?;
1405 let content: String = row.get(1)?;
1406 let session_id: String = row.get(2)?;
1407 let role: String = row.get(3)?;
1408 let blob: Vec<u8> = row.get(4)?;
1409 let updated_at: Option<String> = row.get(5)?;
1410 Ok(VectorRow {
1411 id: format!("msg:{message_id}"),
1412 content,
1413 blob,
1414 updated_at,
1415 source_type: SearchSourceType::Messages,
1416 filter_namespace: None,
1417 filter_session_id: Some(session_id.clone()),
1418 source: SearchSource::Message {
1419 message_id,
1420 session_id,
1421 role,
1422 },
1423 })
1424 },
1425 )
1426 .optional()
1427 .map_err(MemoryError::from)
1428 }
1429 "episode" => conn
1430 .query_row(
1431 "SELECT e.episode_id, e.document_id, e.search_text, e.effect_type, e.outcome, e.embedding, e.updated_at, d.namespace
1432 FROM episodes e
1433 JOIN documents d ON d.id = e.document_id
1434 WHERE e.episode_id = ?1 AND e.embedding IS NOT NULL",
1435 [id],
1436 |row| {
1437 let episode_id: String = row.get(0)?;
1438 let document_id: String = row.get(1)?;
1439 let content: String = row.get(2)?;
1440 let effect_type: String = row.get(3)?;
1441 let outcome: String = row.get(4)?;
1442 let blob: Vec<u8> = row.get(5)?;
1443 let updated_at: Option<String> = row.get(6)?;
1444 let namespace: String = row.get(7)?;
1445 Ok(VectorRow {
1446 id: episodes::episode_item_key(&episode_id),
1447 content,
1448 blob,
1449 updated_at,
1450 source_type: SearchSourceType::Episodes,
1451 filter_namespace: Some(namespace),
1452 filter_session_id: None,
1453 source: SearchSource::Episode {
1454 episode_id,
1455 document_id,
1456 effect_type,
1457 outcome,
1458 },
1459 })
1460 },
1461 )
1462 .optional()
1463 .map_err(MemoryError::from),
1464 _ => Ok(None),
1465 }
1466}
1467
1468fn vector_scan_warn_exceeded(count: usize) -> bool {
1469 let limit = VECTOR_SCAN_WARN_LIMIT.load(Ordering::Relaxed);
1470 limit > 0 && count > limit
1471}
1472
1473#[derive(Debug, Clone)]
1474pub(crate) struct SearchExecution {
1475 pub results: Vec<ExplainedResult>,
1476 pub receipt: Option<VectorSearchReceiptV1>,
1477}
1478
1479#[derive(Debug, Clone, Default)]
1480struct VectorReceiptMetadata {
1481 codec_family: Option<String>,
1482 codec_profile_digest: Option<String>,
1483 artifact_count: Option<usize>,
1484 artifact_corruption_count: Option<usize>,
1485 artifact_missing_count: Option<usize>,
1486 vector_artifact_manifest_digest: Option<String>,
1487 artifact_generation_id: Option<String>,
1488 approximate_scanned_count: Option<usize>,
1489 approximate_returned_count: Option<usize>,
1490 raw_rows_loaded_count: Option<usize>,
1491 filter_strategy: Option<String>,
1492 vector_artifact_count: Option<usize>,
1493 vector_artifact_missing_count: Option<usize>,
1494 vector_artifact_stale_count: Option<usize>,
1495 exact_rerank_count: Option<usize>,
1496 approximate_candidate_count: Option<usize>,
1497}
1498
1499#[derive(Debug, Clone)]
1500struct VectorSearchOutcome {
1501 hits: Vec<VectorHit>,
1502 candidate_backend: String,
1503 requested_candidates: usize,
1504 returned_candidates: usize,
1505 post_filter_candidates: usize,
1506 fallback: Option<String>,
1507 exact_rerank: bool,
1508 degradations: Vec<String>,
1509 receipt_metadata: VectorReceiptMetadata,
1510}
1511
1512fn rrf_fuse_detailed_with_context(
1513 bm25_hits: &[Bm25Hit],
1514 vector_hits: &[VectorHit],
1515 config: &SearchConfig,
1516 context: &SearchContext,
1517 top_k: usize,
1518) -> Vec<ExplainedResult> {
1519 let mut candidates: HashMap<(u8, String), RrfCandidate> = HashMap::new();
1521
1522 for (rank_0, hit) in bm25_hits.iter().enumerate() {
1523 let key = source_dedup_key(&hit.source);
1524 let rank = rank_0 + 1;
1525 candidates
1526 .entry(key)
1527 .and_modify(|candidate| {
1528 candidate.bm25_rank = Some(rank);
1529 candidate.bm25_score = Some(hit.raw_score);
1530 if candidate.updated_at.is_none() {
1531 candidate.updated_at = hit.updated_at.clone();
1532 }
1533 })
1534 .or_insert_with(|| RrfCandidate {
1535 content: hit.content.clone(),
1536 source: hit.source.clone(),
1537 updated_at: hit.updated_at.clone(),
1538 bm25_score: Some(hit.raw_score),
1539 bm25_rank: Some(rank),
1540 vector_score: None,
1541 vector_rank: None,
1542 vector_source_rank: None,
1543 vector_source_score: None,
1544 vector_reranked_from_f32: false,
1545 late_interaction_rank: None,
1546 late_interaction_score: None,
1547 });
1548 }
1549
1550 for (rank_0, hit) in vector_hits.iter().enumerate() {
1551 let key = source_dedup_key(&hit.source);
1552 let rank = rank_0 + 1;
1553 candidates
1554 .entry(key)
1555 .and_modify(|candidate| {
1556 candidate.vector_rank = Some(rank);
1557 candidate.vector_score = Some(hit.similarity);
1558 candidate.vector_source_rank = hit.source_rank.or(Some(rank));
1559 candidate.vector_source_score = hit.source_similarity.or(Some(hit.similarity));
1560 candidate.vector_reranked_from_f32 = hit.reranked_from_f32;
1561 if candidate.updated_at.is_none() {
1562 candidate.updated_at = hit.updated_at.clone();
1563 }
1564 })
1565 .or_insert_with(|| RrfCandidate {
1566 content: hit.content.clone(),
1567 source: hit.source.clone(),
1568 updated_at: hit.updated_at.clone(),
1569 bm25_score: None,
1570 bm25_rank: None,
1571 vector_score: Some(hit.similarity),
1572 vector_rank: Some(rank),
1573 vector_source_rank: hit.source_rank.or(Some(rank)),
1574 vector_source_score: hit.source_similarity.or(Some(hit.similarity)),
1575 vector_reranked_from_f32: hit.reranked_from_f32,
1576 late_interaction_rank: None,
1577 late_interaction_score: None,
1578 });
1579 }
1580
1581 let mut explained: Vec<ExplainedResult> = candidates
1582 .into_values()
1583 .map(|candidate| candidate.explained(config, context))
1584 .collect();
1585
1586 explained.sort_by(|a, b| {
1587 b.result
1588 .score
1589 .partial_cmp(&a.result.score)
1590 .unwrap_or(std::cmp::Ordering::Equal)
1591 .then_with(|| {
1592 source_dedup_key(&a.result.source).cmp(&source_dedup_key(&b.result.source))
1593 })
1594 });
1595 explained.truncate(top_k);
1596 explained
1597}
1598
1599fn rrf_fuse_detailed(
1600 bm25_hits: &[Bm25Hit],
1601 vector_hits: &[VectorHit],
1602 config: &SearchConfig,
1603 top_k: usize,
1604) -> Vec<ExplainedResult> {
1605 let context = SearchContext::default_now();
1606 rrf_fuse_detailed_with_context(bm25_hits, vector_hits, config, &context, top_k)
1607}
1608
1609pub fn rrf_fuse_with_context(
1610 bm25_hits: &[Bm25Hit],
1611 vector_hits: &[VectorHit],
1612 config: &SearchConfig,
1613 context: &SearchContext,
1614 top_k: usize,
1615) -> Vec<SearchResult> {
1616 rrf_fuse_detailed_with_context(bm25_hits, vector_hits, config, context, top_k)
1617 .into_iter()
1618 .map(|result| result.result)
1619 .collect()
1620}
1621
1622pub fn rrf_fuse(
1624 bm25_hits: &[Bm25Hit],
1625 vector_hits: &[VectorHit],
1626 config: &SearchConfig,
1627 top_k: usize,
1628) -> Vec<SearchResult> {
1629 rrf_fuse_detailed(bm25_hits, vector_hits, config, top_k)
1630 .into_iter()
1631 .map(|result| result.result)
1632 .collect()
1633}
1634
1635#[cfg(feature = "late-interaction")]
1642pub fn rrf_fuse_with_late_interaction(
1643 bm25_hits: &[Bm25Hit],
1644 vector_hits: &[VectorHit],
1645 late_interaction_scores: &[(String, f64)],
1646 config: &SearchConfig,
1647 context: &SearchContext,
1648 top_k: usize,
1649) -> Vec<ExplainedResult> {
1650 let mut candidates: HashMap<(u8, String), RrfCandidate> = HashMap::new();
1651
1652 for (rank_0, hit) in bm25_hits.iter().enumerate() {
1654 let key = source_dedup_key(&hit.source);
1655 let rank = rank_0 + 1;
1656 candidates
1657 .entry(key)
1658 .and_modify(|c| {
1659 c.bm25_rank = Some(rank);
1660 c.bm25_score = Some(hit.raw_score);
1661 if c.updated_at.is_none() {
1662 c.updated_at = hit.updated_at.clone();
1663 }
1664 })
1665 .or_insert_with(|| RrfCandidate {
1666 content: hit.content.clone(),
1667 source: hit.source.clone(),
1668 updated_at: hit.updated_at.clone(),
1669 bm25_score: Some(hit.raw_score),
1670 bm25_rank: Some(rank),
1671 vector_score: None,
1672 vector_rank: None,
1673 vector_source_rank: None,
1674 vector_source_score: None,
1675 vector_reranked_from_f32: false,
1676 late_interaction_rank: None,
1677 late_interaction_score: None,
1678 });
1679 }
1680
1681 for (rank_0, hit) in vector_hits.iter().enumerate() {
1683 let key = source_dedup_key(&hit.source);
1684 let rank = rank_0 + 1;
1685 candidates
1686 .entry(key)
1687 .and_modify(|c| {
1688 c.vector_rank = Some(rank);
1689 c.vector_score = Some(hit.similarity);
1690 c.vector_source_rank = hit.source_rank.or(Some(rank));
1691 c.vector_source_score = hit.source_similarity.or(Some(hit.similarity));
1692 c.vector_reranked_from_f32 = hit.reranked_from_f32;
1693 if c.updated_at.is_none() {
1694 c.updated_at = hit.updated_at.clone();
1695 }
1696 })
1697 .or_insert_with(|| RrfCandidate {
1698 content: hit.content.clone(),
1699 source: hit.source.clone(),
1700 updated_at: hit.updated_at.clone(),
1701 bm25_score: None,
1702 bm25_rank: None,
1703 vector_score: Some(hit.similarity),
1704 vector_rank: Some(rank),
1705 vector_source_rank: hit.source_rank.or(Some(rank)),
1706 vector_source_score: hit.source_similarity.or(Some(hit.similarity)),
1707 vector_reranked_from_f32: hit.reranked_from_f32,
1708 late_interaction_rank: None,
1709 late_interaction_score: None,
1710 });
1711 }
1712
1713 let mut li_sorted: Vec<&(String, f64)> = late_interaction_scores.iter().collect();
1716 li_sorted.sort_by(|a, b| {
1717 b.1.partial_cmp(&a.1)
1718 .unwrap_or(std::cmp::Ordering::Equal)
1719 });
1720 for (rank_0, (item_key, score)) in li_sorted.iter().enumerate() {
1721 let rank = rank_0 + 1;
1722 let matched = candidates.iter_mut().find(|(_, c)| {
1726 c.content.contains(item_key.as_str())
1727 || format!("{:?}", c.source).contains(item_key.as_str())
1728 });
1729 if let Some((_, c)) = matched {
1730 c.late_interaction_rank = Some(rank);
1731 c.late_interaction_score = Some(*score);
1732 }
1733 }
1737
1738 let mut explained: Vec<ExplainedResult> = candidates
1739 .into_values()
1740 .map(|c| c.explained(config, context))
1741 .collect();
1742
1743 explained.sort_by(|a, b| {
1744 b.result
1745 .score
1746 .partial_cmp(&a.result.score)
1747 .unwrap_or(std::cmp::Ordering::Equal)
1748 .then_with(|| {
1749 source_dedup_key(&a.result.source).cmp(&source_dedup_key(&b.result.source))
1750 })
1751 });
1752 explained.truncate(top_k);
1753 explained
1754}
1755
1756fn compute_proxy_late_interaction_scores(
1766 query_embedding: &[f32],
1767 vector_hits: &[VectorHit],
1768) -> Vec<(String, f64)> {
1769 let segment_size = 64;
1770 let query_segments: Vec<&[f32]> = query_embedding.chunks(segment_size).collect();
1771
1772 vector_hits
1773 .iter()
1774 .map(|hit| {
1775 let segment_factor = if !query_segments.is_empty() {
1776 1.0 + (query_segments.len() as f64 - 1.0) * 0.01
1777 } else {
1778 1.0
1779 };
1780 let proxy_score = hit.similarity * segment_factor;
1781 let key = format!("{:?}", hit.source);
1782 (key, proxy_score)
1783 })
1784 .collect()
1785}
1786
1787pub(crate) fn query_embedding_digest(query_embedding: &[f32]) -> String {
1788 let mut builder = DigestBuilder::new();
1789 builder
1790 .update_str("semantic-memory.query_embedding.v1")
1791 .separator()
1792 .update(&(query_embedding.len() as u64).to_le_bytes())
1793 .separator();
1794 for value in query_embedding {
1795 builder.update(&value.to_le_bytes());
1796 }
1797 format!("blake3:{}", builder.finalize().hex())
1798}
1799
1800#[cfg_attr(not(feature = "hnsw"), allow(dead_code))]
1801#[allow(clippy::too_many_arguments)]
1802fn build_receipt(
1803 context: &SearchContext,
1804 query_embedding: &[f32],
1805 search_profile: &str,
1806 candidate_backend: &str,
1807 requested_candidates: usize,
1808 returned_candidates: usize,
1809 post_filter_candidates: usize,
1810 fallback: Option<String>,
1811 exact_rerank: bool,
1812 results: &[ExplainedResult],
1813 degradations: Vec<String>,
1814) -> Option<VectorSearchReceiptV1> {
1815 build_receipt_with_metadata(
1816 context,
1817 query_embedding,
1818 search_profile,
1819 candidate_backend,
1820 requested_candidates,
1821 returned_candidates,
1822 post_filter_candidates,
1823 fallback,
1824 exact_rerank,
1825 results,
1826 degradations,
1827 VectorReceiptMetadata::default(),
1828 )
1829}
1830
1831#[allow(clippy::too_many_arguments)]
1832fn build_receipt_with_metadata(
1833 context: &SearchContext,
1834 query_embedding: &[f32],
1835 search_profile: &str,
1836 candidate_backend: &str,
1837 requested_candidates: usize,
1838 returned_candidates: usize,
1839 post_filter_candidates: usize,
1840 fallback: Option<String>,
1841 exact_rerank: bool,
1842 results: &[ExplainedResult],
1843 degradations: Vec<String>,
1844 metadata: VectorReceiptMetadata,
1845) -> Option<VectorSearchReceiptV1> {
1846 if !context.receipts_enabled() {
1847 return None;
1848 }
1849 Some(VectorSearchReceiptV1 {
1850 schema_version: "vector_search_receipt_v1".to_string(),
1851 receipt_digest: None,
1852 receipt_id: context
1853 .request_id
1854 .clone()
1855 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
1856 evaluation_time: context.evaluation_time,
1857 trace_id: context.trace_id.clone(),
1858 attempt_family_id: context.attempt_family_id.clone(),
1859 attempt_id: context.attempt_id.clone(),
1860 replay_of: context.replay_of.clone(),
1861 query_embedding_digest: Some(query_embedding_digest(query_embedding)),
1862 query_text_digest: context.query_text_digest.clone(),
1863 query_input_digest: context.query_input_digest.clone(),
1864 filter_digest: context.filter_digest.clone(),
1865 redaction_state: context.redaction_state.clone(),
1866 budget_id: context.budget_id.clone(),
1867 deadline_at: context.deadline_at,
1868 search_profile: search_profile.to_string(),
1869 candidate_backend: candidate_backend.to_string(),
1870 codec_family: metadata.codec_family.clone(),
1871 codec_profile_digest: metadata.codec_profile_digest.clone(),
1872 artifact_profile_digest: metadata.codec_profile_digest.clone(),
1873 artifact_count: metadata.artifact_count,
1874 artifact_corruption_count: metadata.artifact_corruption_count,
1875 artifact_missing_count: metadata.artifact_missing_count,
1876 vector_artifact_manifest_digest: metadata.vector_artifact_manifest_digest.clone(),
1877 artifact_generation_id: metadata.artifact_generation_id.clone(),
1878 approximate_scanned_count: metadata.approximate_scanned_count,
1879 approximate_returned_count: metadata.approximate_returned_count,
1880 raw_rows_loaded_count: metadata.raw_rows_loaded_count,
1881 filter_strategy: metadata.filter_strategy,
1882 vector_artifact_count: metadata.vector_artifact_count.or(metadata.artifact_count),
1883 vector_artifact_missing_count: metadata
1884 .vector_artifact_missing_count
1885 .or(metadata.artifact_missing_count),
1886 vector_artifact_stale_count: metadata.vector_artifact_stale_count,
1887 exact_rerank_count: metadata.exact_rerank_count.or(if exact_rerank {
1888 Some(post_filter_candidates)
1889 } else {
1890 None
1891 }),
1892 approximate_candidate_count: metadata.approximate_candidate_count,
1893 approximate: candidate_backend.contains("hnsw")
1894 || candidate_backend.contains("turbo_quant"),
1895 requested_candidates,
1896 returned_candidates,
1897 post_filter_candidates,
1898 fallback_reason: fallback.clone(),
1899 derived_candidate: if candidate_backend == "provekv_pool_candidate_then_exact_f32" {
1900 Some(crate::types::DerivedCandidateReceiptV1 {
1901 candidate_backend: candidate_backend.to_string(),
1902 codec_family: metadata.codec_family.clone(),
1903 generation_id: metadata.artifact_generation_id.clone(),
1904 embedding_snapshot_digest: None,
1905 pool_manifest_digest: metadata.vector_artifact_manifest_digest.clone(),
1906 exact_rerank,
1907 approximate: false,
1908 fallback: fallback.clone(),
1909 raw_candidate_count: returned_candidates,
1910 post_filter_count: post_filter_candidates,
1911 final_result_count: results.len(),
1912 })
1913 } else {
1914 None
1915 },
1916 fallback,
1917 exact_rerank,
1918 result_ids: results
1919 .iter()
1920 .map(|result| search_result_id(&result.result.source))
1921 .collect(),
1922 degradations,
1923 })
1924}
1925
1926#[cfg(feature = "hnsw")]
1927fn filters_are_active(
1928 namespaces: Option<&[&str]>,
1929 source_types: Option<&[SearchSourceType]>,
1930 session_ids: Option<&[&str]>,
1931) -> bool {
1932 namespaces.is_some_and(|values| !values.is_empty())
1933 || source_types.is_some_and(|values| !values.is_empty())
1934 || session_ids.is_some_and(|values| !values.is_empty())
1935}
1936
1937#[allow(clippy::too_many_arguments)]
1938pub(crate) fn hybrid_search_detailed_with_context(
1939 conn: &Connection,
1940 query: &str,
1941 query_embedding: &[f32],
1942 config: &SearchConfig,
1943 context: &SearchContext,
1944 top_k: usize,
1945 namespaces: Option<&[&str]>,
1946 source_types: Option<&[SearchSourceType]>,
1947 session_ids: Option<&[&str]>,
1948) -> Result<SearchExecution, MemoryError> {
1949 let bm25_hits = match sanitize_fts_query(query) {
1950 Some(sanitized) => bm25_search(
1951 conn,
1952 &sanitized,
1953 config.candidate_pool_size,
1954 namespaces,
1955 source_types,
1956 session_ids,
1957 )?,
1958 None => Vec::new(),
1959 };
1960
1961 let vector_outcome = vector_search_with_backend(
1962 conn,
1963 query_embedding,
1964 config.candidate_pool_size,
1965 config.min_similarity,
1966 config,
1967 context,
1968 namespaces,
1969 source_types,
1970 session_ids,
1971 )?;
1972
1973 let results = if config.late_interaction_weight > 0.0 {
1974 let li_scores = compute_proxy_late_interaction_scores(
1979 query_embedding,
1980 &vector_outcome.hits,
1981 );
1982 #[cfg(feature = "late-interaction")]
1983 {
1984 rrf_fuse_with_late_interaction(
1985 &bm25_hits,
1986 &vector_outcome.hits,
1987 &li_scores,
1988 config,
1989 context,
1990 top_k,
1991 )
1992 }
1993 #[cfg(not(feature = "late-interaction"))]
1994 {
1995 let _ = li_scores;
1996 rrf_fuse_detailed_with_context(&bm25_hits, &vector_outcome.hits, config, context, top_k)
1997 }
1998 } else {
1999 rrf_fuse_detailed_with_context(&bm25_hits, &vector_outcome.hits, config, context, top_k)
2000 };
2001 let receipt = build_receipt_with_metadata(
2002 context,
2003 query_embedding,
2004 "hybrid",
2005 &vector_outcome.candidate_backend,
2006 vector_outcome.requested_candidates,
2007 vector_outcome.returned_candidates,
2008 vector_outcome.post_filter_candidates,
2009 vector_outcome.fallback,
2010 vector_outcome.exact_rerank,
2011 &results,
2012 vector_outcome.degradations,
2013 vector_outcome.receipt_metadata,
2014 );
2015 Ok(SearchExecution { results, receipt })
2016}
2017
2018#[allow(clippy::too_many_arguments)]
2019pub(crate) fn hybrid_search_detailed(
2020 conn: &Connection,
2021 query: &str,
2022 query_embedding: &[f32],
2023 config: &SearchConfig,
2024 top_k: usize,
2025 namespaces: Option<&[&str]>,
2026 source_types: Option<&[SearchSourceType]>,
2027 session_ids: Option<&[&str]>,
2028) -> Result<Vec<ExplainedResult>, MemoryError> {
2029 let context = SearchContext::default_now();
2030 Ok(hybrid_search_detailed_with_context(
2031 conn,
2032 query,
2033 query_embedding,
2034 config,
2035 &context,
2036 top_k,
2037 namespaces,
2038 source_types,
2039 session_ids,
2040 )?
2041 .results)
2042}
2043
2044#[allow(clippy::too_many_arguments)]
2046pub fn hybrid_search_explained(
2047 conn: &Connection,
2048 query: &str,
2049 query_embedding: &[f32],
2050 config: &SearchConfig,
2051 top_k: usize,
2052 namespaces: Option<&[&str]>,
2053 source_types: Option<&[SearchSourceType]>,
2054 session_ids: Option<&[&str]>,
2055) -> Result<Vec<ExplainedResult>, MemoryError> {
2056 hybrid_search_detailed(
2057 conn,
2058 query,
2059 query_embedding,
2060 config,
2061 top_k,
2062 namespaces,
2063 source_types,
2064 session_ids,
2065 )
2066}
2067
2068#[allow(clippy::too_many_arguments)]
2070pub fn hybrid_search(
2071 conn: &Connection,
2072 query: &str,
2073 query_embedding: &[f32],
2074 config: &SearchConfig,
2075 top_k: usize,
2076 namespaces: Option<&[&str]>,
2077 source_types: Option<&[SearchSourceType]>,
2078 session_ids: Option<&[&str]>,
2079) -> Result<Vec<SearchResult>, MemoryError> {
2080 Ok(hybrid_search_detailed(
2081 conn,
2082 query,
2083 query_embedding,
2084 config,
2085 top_k,
2086 namespaces,
2087 source_types,
2088 session_ids,
2089 )?
2090 .into_iter()
2091 .map(|result| result.result)
2092 .collect())
2093}
2094
2095#[cfg(feature = "hnsw")]
2096#[derive(Clone)]
2097struct HnswCandidateSeed {
2098 source_rank: usize,
2099 source_similarity: f64,
2100}
2101
2102#[cfg(feature = "hnsw")]
2103#[allow(clippy::type_complexity)]
2104fn resolve_hnsw_hits_batched(
2105 conn: &Connection,
2106 query_embedding: &[f32],
2107 config: &SearchConfig,
2108 namespaces: Option<&[&str]>,
2109 source_types: Option<&[SearchSourceType]>,
2110 session_ids: Option<&[&str]>,
2111 hnsw_hits: &[crate::hnsw::HnswHit],
2112) -> Result<Vec<VectorHit>, MemoryError> {
2113 let search_facts = source_types
2114 .map(|st| st.contains(&SearchSourceType::Facts))
2115 .unwrap_or(true);
2116 let search_chunks = source_types
2117 .map(|st| st.contains(&SearchSourceType::Chunks))
2118 .unwrap_or(true);
2119 let search_messages = source_types
2120 .map(|st| st.contains(&SearchSourceType::Messages))
2121 .unwrap_or(false);
2122 let search_episodes = source_types
2123 .map(|st| st.contains(&SearchSourceType::Episodes))
2124 .unwrap_or(true);
2125
2126 let mut fact_entries: HashMap<String, HnswCandidateSeed> = HashMap::new();
2128 let mut chunk_entries: HashMap<String, HnswCandidateSeed> = HashMap::new();
2130 let mut message_entries: HashMap<i64, HnswCandidateSeed> = HashMap::new();
2132 let mut episode_entries: HashMap<String, HnswCandidateSeed> = HashMap::new();
2134
2135 for (rank_0, hit) in hnsw_hits.iter().enumerate() {
2136 let similarity = hit.similarity() as f64;
2137 if similarity < config.min_similarity {
2138 continue;
2139 }
2140
2141 let (domain, raw_id) = hit.parse_key()?;
2142 let seed = HnswCandidateSeed {
2143 source_rank: rank_0 + 1,
2144 source_similarity: similarity,
2145 };
2146
2147 match domain {
2148 "fact" if search_facts => {
2149 fact_entries.entry(raw_id.to_string()).or_insert(seed);
2150 }
2151 "chunk" if search_chunks => {
2152 chunk_entries.entry(raw_id.to_string()).or_insert(seed);
2153 }
2154 "msg" if search_messages => {
2155 if let Ok(message_id) = raw_id.parse::<i64>() {
2156 message_entries.entry(message_id).or_insert(seed);
2157 }
2158 }
2159 "episode" if search_episodes => {
2160 episode_entries.entry(raw_id.to_string()).or_insert(seed);
2161 }
2162 _ => {}
2163 }
2164 }
2165
2166 let mut hits = Vec::new();
2167 batch_load_fact_hits(
2168 conn,
2169 query_embedding,
2170 config,
2171 namespaces,
2172 &fact_entries,
2173 &mut hits,
2174 )?;
2175 batch_load_chunk_hits(
2176 conn,
2177 query_embedding,
2178 config,
2179 namespaces,
2180 &chunk_entries,
2181 &mut hits,
2182 )?;
2183 batch_load_message_hits(
2184 conn,
2185 query_embedding,
2186 config,
2187 session_ids,
2188 &message_entries,
2189 &mut hits,
2190 )?;
2191 batch_load_episode_hits(
2192 conn,
2193 query_embedding,
2194 config,
2195 namespaces,
2196 &episode_entries,
2197 &mut hits,
2198 )?;
2199
2200 hits.sort_by(|a, b| {
2201 b.similarity
2202 .partial_cmp(&a.similarity)
2203 .unwrap_or(std::cmp::Ordering::Equal)
2204 .then_with(|| {
2205 a.source_rank
2206 .unwrap_or(usize::MAX)
2207 .cmp(&b.source_rank.unwrap_or(usize::MAX))
2208 })
2209 });
2210 hits.truncate(config.candidate_pool_size);
2211 Ok(hits)
2212}
2213
2214#[cfg(feature = "hnsw")]
2215fn exact_similarity_from_blob(
2216 query_embedding: &[f32],
2217 blob: &[u8],
2218) -> Result<Option<f64>, MemoryError> {
2219 if blob.is_empty() {
2220 return Ok(None);
2221 }
2222 let stored = crate::db::bytes_to_embedding(blob)?;
2223 if stored.len() != query_embedding.len() {
2224 return Ok(None);
2225 }
2226 Ok(Some(cosine_similarity(query_embedding, &stored)? as f64))
2227}
2228
2229#[cfg(feature = "hnsw")]
2230#[allow(clippy::too_many_arguments)]
2231fn build_ranked_vector_hit(
2232 id: String,
2233 content: String,
2234 source: SearchSource,
2235 updated_at: Option<String>,
2236 embedding_blob: Option<Vec<u8>>,
2237 seed: &HnswCandidateSeed,
2238 query_embedding: &[f32],
2239 config: &SearchConfig,
2240) -> Result<Option<VectorHit>, MemoryError> {
2241 let similarity = if config.rerank_from_f32 {
2242 match embedding_blob {
2243 Some(blob) => exact_similarity_from_blob(query_embedding, &blob)?,
2244 None => None,
2245 }
2246 .unwrap_or(seed.source_similarity)
2247 } else {
2248 seed.source_similarity
2249 };
2250
2251 if similarity < config.min_similarity {
2252 return Ok(None);
2253 }
2254
2255 Ok(Some(VectorHit {
2256 id,
2257 content,
2258 source,
2259 similarity,
2260 updated_at,
2261 source_rank: Some(seed.source_rank),
2262 source_similarity: Some(seed.source_similarity),
2263 reranked_from_f32: config.rerank_from_f32,
2264 }))
2265}
2266
2267#[cfg(feature = "hnsw")]
2268fn batch_load_fact_hits(
2269 conn: &Connection,
2270 query_embedding: &[f32],
2271 config: &SearchConfig,
2272 namespaces: Option<&[&str]>,
2273 entries: &HashMap<String, HnswCandidateSeed>,
2275 output: &mut Vec<VectorHit>,
2276) -> Result<(), MemoryError> {
2277 if entries.is_empty() {
2278 return Ok(());
2279 }
2280
2281 let placeholders = (1..=entries.len())
2282 .map(|idx| format!("?{idx}"))
2283 .collect::<Vec<_>>()
2284 .join(", ");
2285 let sql = format!(
2286 "SELECT id, content, namespace, updated_at, embedding
2287 FROM facts
2288 WHERE id IN ({placeholders})"
2289 );
2290 let params: Vec<SqlValue> = entries
2291 .keys()
2292 .map(|id| SqlValue::Text(id.clone()))
2293 .collect();
2294 let mut stmt = conn.prepare(&sql)?;
2295 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
2296 Ok((
2297 row.get::<_, String>(0)?,
2298 row.get::<_, String>(1)?,
2299 row.get::<_, String>(2)?,
2300 row.get::<_, Option<String>>(3)?,
2301 row.get::<_, Option<Vec<u8>>>(4)?,
2302 ))
2303 })?;
2304
2305 for row in rows {
2306 let (fact_id, content, namespace, updated_at, embedding_blob) = row?;
2307 if let Some(filter) = namespaces {
2308 if !filter.contains(&namespace.as_str()) {
2309 continue;
2310 }
2311 }
2312 if let Some(seed) = entries.get(&fact_id) {
2313 if let Some(hit) = build_ranked_vector_hit(
2314 format!("fact:{fact_id}"),
2315 content,
2316 SearchSource::Fact { fact_id, namespace },
2317 updated_at,
2318 embedding_blob,
2319 seed,
2320 query_embedding,
2321 config,
2322 )? {
2323 output.push(hit);
2324 }
2325 }
2326 }
2327
2328 Ok(())
2329}
2330
2331#[cfg(feature = "hnsw")]
2332fn batch_load_chunk_hits(
2333 conn: &Connection,
2334 query_embedding: &[f32],
2335 config: &SearchConfig,
2336 namespaces: Option<&[&str]>,
2337 entries: &HashMap<String, HnswCandidateSeed>,
2339 output: &mut Vec<VectorHit>,
2340) -> Result<(), MemoryError> {
2341 if entries.is_empty() {
2342 return Ok(());
2343 }
2344
2345 let placeholders = (1..=entries.len())
2346 .map(|idx| format!("?{idx}"))
2347 .collect::<Vec<_>>()
2348 .join(", ");
2349 let sql = format!(
2350 "SELECT c.id, c.content, c.document_id, d.title, c.chunk_index, c.created_at, d.namespace, c.embedding
2351 FROM chunks c
2352 JOIN documents d ON d.id = c.document_id
2353 WHERE c.id IN ({placeholders})"
2354 );
2355 let params: Vec<SqlValue> = entries
2356 .keys()
2357 .map(|id| SqlValue::Text(id.clone()))
2358 .collect();
2359 let mut stmt = conn.prepare(&sql)?;
2360 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
2361 Ok((
2362 row.get::<_, String>(0)?,
2363 row.get::<_, String>(1)?,
2364 row.get::<_, String>(2)?,
2365 row.get::<_, String>(3)?,
2366 row.get::<_, i64>(4)?,
2367 row.get::<_, Option<String>>(5)?,
2368 row.get::<_, String>(6)?,
2369 row.get::<_, Option<Vec<u8>>>(7)?,
2370 ))
2371 })?;
2372
2373 for row in rows {
2374 let (
2375 chunk_id,
2376 content,
2377 document_id,
2378 document_title,
2379 chunk_index,
2380 updated_at,
2381 namespace,
2382 embedding_blob,
2383 ) = row?;
2384 if let Some(filter) = namespaces {
2385 if !filter.contains(&namespace.as_str()) {
2386 continue;
2387 }
2388 }
2389 if let Some(seed) = entries.get(&chunk_id) {
2390 if let Some(hit) = build_ranked_vector_hit(
2391 format!("chunk:{chunk_id}"),
2392 content,
2393 SearchSource::Chunk {
2394 chunk_id,
2395 document_id,
2396 document_title,
2397 chunk_index: chunk_index as usize,
2398 },
2399 updated_at,
2400 embedding_blob,
2401 seed,
2402 query_embedding,
2403 config,
2404 )? {
2405 output.push(hit);
2406 }
2407 }
2408 }
2409
2410 Ok(())
2411}
2412
2413#[cfg(feature = "hnsw")]
2414fn batch_load_message_hits(
2415 conn: &Connection,
2416 query_embedding: &[f32],
2417 config: &SearchConfig,
2418 session_ids: Option<&[&str]>,
2419 entries: &HashMap<i64, HnswCandidateSeed>,
2421 output: &mut Vec<VectorHit>,
2422) -> Result<(), MemoryError> {
2423 if entries.is_empty() {
2424 return Ok(());
2425 }
2426
2427 let placeholders = (1..=entries.len())
2428 .map(|idx| format!("?{idx}"))
2429 .collect::<Vec<_>>()
2430 .join(", ");
2431 let sql = format!(
2432 "SELECT id, content, session_id, role, created_at, embedding
2433 FROM messages
2434 WHERE id IN ({placeholders})"
2435 );
2436 let params: Vec<SqlValue> = entries.keys().map(|id| SqlValue::Integer(*id)).collect();
2437 let mut stmt = conn.prepare(&sql)?;
2438 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
2439 Ok((
2440 row.get::<_, i64>(0)?,
2441 row.get::<_, String>(1)?,
2442 row.get::<_, String>(2)?,
2443 row.get::<_, String>(3)?,
2444 row.get::<_, Option<String>>(4)?,
2445 row.get::<_, Option<Vec<u8>>>(5)?,
2446 ))
2447 })?;
2448
2449 for row in rows {
2450 let (message_id, content, session_id, role, updated_at, embedding_blob) = row?;
2451 if let Some(filter) = session_ids {
2452 if !filter.contains(&session_id.as_str()) {
2453 continue;
2454 }
2455 }
2456 if let Some(seed) = entries.get(&message_id) {
2457 if let Some(hit) = build_ranked_vector_hit(
2458 format!("msg:{message_id}"),
2459 content,
2460 SearchSource::Message {
2461 message_id,
2462 session_id,
2463 role,
2464 },
2465 updated_at,
2466 embedding_blob,
2467 seed,
2468 query_embedding,
2469 config,
2470 )? {
2471 output.push(hit);
2472 }
2473 }
2474 }
2475
2476 Ok(())
2477}
2478
2479#[cfg(feature = "hnsw")]
2480fn batch_load_episode_hits(
2481 conn: &Connection,
2482 query_embedding: &[f32],
2483 config: &SearchConfig,
2484 namespaces: Option<&[&str]>,
2485 entries: &HashMap<String, HnswCandidateSeed>,
2487 output: &mut Vec<VectorHit>,
2488) -> Result<(), MemoryError> {
2489 if entries.is_empty() {
2490 return Ok(());
2491 }
2492
2493 let placeholders = (1..=entries.len())
2494 .map(|idx| format!("?{idx}"))
2495 .collect::<Vec<_>>()
2496 .join(", ");
2497 let sql = format!(
2498 "SELECT e.episode_id, e.document_id, e.search_text, e.effect_type, e.outcome, e.updated_at, d.namespace, e.embedding
2499 FROM episodes e
2500 JOIN documents d ON d.id = e.document_id
2501 WHERE e.episode_id IN ({placeholders})"
2502 );
2503 let params: Vec<SqlValue> = entries
2504 .keys()
2505 .map(|id| SqlValue::Text(id.clone()))
2506 .collect();
2507 let mut stmt = conn.prepare(&sql)?;
2508 let rows = stmt.query_map(rusqlite::params_from_iter(¶ms), |row| {
2509 Ok((
2510 row.get::<_, String>(0)?,
2511 row.get::<_, String>(1)?,
2512 row.get::<_, String>(2)?,
2513 row.get::<_, String>(3)?,
2514 row.get::<_, String>(4)?,
2515 row.get::<_, Option<String>>(5)?,
2516 row.get::<_, String>(6)?,
2517 row.get::<_, Option<Vec<u8>>>(7)?,
2518 ))
2519 })?;
2520
2521 for row in rows {
2522 let (
2523 episode_id,
2524 document_id,
2525 content,
2526 effect_type,
2527 outcome,
2528 updated_at,
2529 namespace,
2530 embedding_blob,
2531 ) = row?;
2532 if let Some(filter) = namespaces {
2533 if !filter.contains(&namespace.as_str()) {
2534 continue;
2535 }
2536 }
2537 if let Some(seed) = entries.get(&episode_id) {
2538 if let Some(hit) = build_ranked_vector_hit(
2539 episodes::episode_item_key(&episode_id),
2540 content,
2541 SearchSource::Episode {
2542 episode_id,
2543 document_id,
2544 effect_type,
2545 outcome,
2546 },
2547 updated_at,
2548 embedding_blob,
2549 seed,
2550 query_embedding,
2551 config,
2552 )? {
2553 output.push(hit);
2554 }
2555 }
2556 }
2557
2558 Ok(())
2559}
2560
2561#[cfg(feature = "hnsw")]
2563#[allow(clippy::too_many_arguments)]
2564pub fn hybrid_search_with_hnsw(
2565 conn: &Connection,
2566 query: &str,
2567 query_embedding: &[f32],
2568 config: &SearchConfig,
2569 top_k: usize,
2570 namespaces: Option<&[&str]>,
2571 source_types: Option<&[SearchSourceType]>,
2572 session_ids: Option<&[&str]>,
2573 hnsw_hits: &[crate::hnsw::HnswHit],
2574) -> Result<Vec<SearchResult>, MemoryError> {
2575 Ok(hybrid_search_with_hnsw_detailed(
2576 conn,
2577 query,
2578 query_embedding,
2579 config,
2580 top_k,
2581 namespaces,
2582 source_types,
2583 session_ids,
2584 hnsw_hits,
2585 )?
2586 .into_iter()
2587 .map(|result| result.result)
2588 .collect())
2589}
2590
2591#[cfg(feature = "hnsw")]
2592#[allow(clippy::too_many_arguments)]
2593pub(crate) fn hybrid_search_with_hnsw_detailed_with_context(
2594 conn: &Connection,
2595 query: &str,
2596 query_embedding: &[f32],
2597 config: &SearchConfig,
2598 context: &SearchContext,
2599 top_k: usize,
2600 namespaces: Option<&[&str]>,
2601 source_types: Option<&[SearchSourceType]>,
2602 session_ids: Option<&[&str]>,
2603 hnsw_hits: &[crate::hnsw::HnswHit],
2604) -> Result<SearchExecution, MemoryError> {
2605 let bm25_hits = match sanitize_fts_query(query) {
2606 Some(sanitized) => bm25_search(
2607 conn,
2608 &sanitized,
2609 config.candidate_pool_size,
2610 namespaces,
2611 source_types,
2612 session_ids,
2613 )?,
2614 None => Vec::new(),
2615 };
2616
2617 let mut vector_hits = resolve_hnsw_hits_batched(
2618 conn,
2619 query_embedding,
2620 config,
2621 namespaces,
2622 source_types,
2623 session_ids,
2624 hnsw_hits,
2625 )?;
2626 let mut fallback = None;
2627 let mut degradations = Vec::new();
2628 let mut backend = "hnsw";
2629 let mut exact_rerank = config.rerank_from_f32;
2630
2631 if !hnsw_hits.is_empty()
2632 && vector_hits.len() < top_k
2633 && filters_are_active(namespaces, source_types, session_ids)
2634 {
2635 fallback = Some("hnsw_filtered_underreturn_fallback".to_string());
2636 degradations.push(format!(
2637 "HNSW returned {} post-filter vector candidates for requested top_k {}; exact filtered fallback was used",
2638 vector_hits.len(),
2639 top_k
2640 ));
2641 vector_hits = vector_search(
2642 conn,
2643 query_embedding,
2644 config.candidate_pool_size,
2645 config.min_similarity,
2646 namespaces,
2647 source_types,
2648 session_ids,
2649 )?;
2650 backend = "hnsw_then_brute_force_f32";
2651 exact_rerank = true;
2652 }
2653
2654 let results = rrf_fuse_detailed_with_context(&bm25_hits, &vector_hits, config, context, top_k);
2655 let receipt = build_receipt(
2656 context,
2657 query_embedding,
2658 "hybrid",
2659 backend,
2660 config.candidate_pool_size,
2661 hnsw_hits.len(),
2662 vector_hits.len(),
2663 fallback,
2664 exact_rerank,
2665 &results,
2666 degradations,
2667 );
2668
2669 Ok(SearchExecution { results, receipt })
2670}
2671
2672#[cfg(feature = "hnsw")]
2673#[allow(clippy::too_many_arguments)]
2674pub(crate) fn hybrid_search_with_hnsw_detailed(
2675 conn: &Connection,
2676 query: &str,
2677 query_embedding: &[f32],
2678 config: &SearchConfig,
2679 top_k: usize,
2680 namespaces: Option<&[&str]>,
2681 source_types: Option<&[SearchSourceType]>,
2682 session_ids: Option<&[&str]>,
2683 hnsw_hits: &[crate::hnsw::HnswHit],
2684) -> Result<Vec<ExplainedResult>, MemoryError> {
2685 let context = SearchContext::default_now();
2686 Ok(hybrid_search_with_hnsw_detailed_with_context(
2687 conn,
2688 query,
2689 query_embedding,
2690 config,
2691 &context,
2692 top_k,
2693 namespaces,
2694 source_types,
2695 session_ids,
2696 hnsw_hits,
2697 )?
2698 .results)
2699}
2700
2701#[cfg(feature = "hnsw")]
2703#[allow(clippy::too_many_arguments)]
2704pub fn hybrid_search_explained_with_hnsw(
2705 conn: &Connection,
2706 query: &str,
2707 query_embedding: &[f32],
2708 config: &SearchConfig,
2709 top_k: usize,
2710 namespaces: Option<&[&str]>,
2711 source_types: Option<&[SearchSourceType]>,
2712 session_ids: Option<&[&str]>,
2713 hnsw_hits: &[crate::hnsw::HnswHit],
2714) -> Result<Vec<ExplainedResult>, MemoryError> {
2715 hybrid_search_with_hnsw_detailed(
2716 conn,
2717 query,
2718 query_embedding,
2719 config,
2720 top_k,
2721 namespaces,
2722 source_types,
2723 session_ids,
2724 hnsw_hits,
2725 )
2726}
2727
2728pub(crate) fn fts_only_search_detailed(
2729 conn: &Connection,
2730 query: &str,
2731 config: &SearchConfig,
2732 top_k: usize,
2733 namespaces: Option<&[&str]>,
2734 source_types: Option<&[SearchSourceType]>,
2735 session_ids: Option<&[&str]>,
2736) -> Result<Vec<ExplainedResult>, MemoryError> {
2737 let sanitized = match sanitize_fts_query(query) {
2738 Some(value) => value,
2739 None => return Ok(Vec::new()),
2740 };
2741 let bm25_hits = bm25_search(
2742 conn,
2743 &sanitized,
2744 top_k,
2745 namespaces,
2746 source_types,
2747 session_ids,
2748 )?;
2749 Ok(rrf_fuse_detailed(&bm25_hits, &[], config, top_k))
2750}
2751
2752pub fn fts_only_search(
2754 conn: &Connection,
2755 query: &str,
2756 config: &SearchConfig,
2757 top_k: usize,
2758 namespaces: Option<&[&str]>,
2759 source_types: Option<&[SearchSourceType]>,
2760 session_ids: Option<&[&str]>,
2761) -> Result<Vec<SearchResult>, MemoryError> {
2762 Ok(fts_only_search_detailed(
2763 conn,
2764 query,
2765 config,
2766 top_k,
2767 namespaces,
2768 source_types,
2769 session_ids,
2770 )?
2771 .into_iter()
2772 .map(|result| result.result)
2773 .collect())
2774}
2775
2776#[allow(clippy::too_many_arguments)]
2777pub(crate) fn vector_only_search_detailed_with_context(
2778 conn: &Connection,
2779 query_embedding: &[f32],
2780 config: &SearchConfig,
2781 context: &SearchContext,
2782 top_k: usize,
2783 namespaces: Option<&[&str]>,
2784 source_types: Option<&[SearchSourceType]>,
2785 session_ids: Option<&[&str]>,
2786) -> Result<SearchExecution, MemoryError> {
2787 let vector_outcome = vector_search_with_backend(
2788 conn,
2789 query_embedding,
2790 top_k,
2791 config.min_similarity,
2792 config,
2793 context,
2794 namespaces,
2795 source_types,
2796 session_ids,
2797 )?;
2798 let results = rrf_fuse_detailed_with_context(&[], &vector_outcome.hits, config, context, top_k);
2799 let receipt = build_receipt_with_metadata(
2800 context,
2801 query_embedding,
2802 "vector_only",
2803 &vector_outcome.candidate_backend,
2804 vector_outcome.requested_candidates,
2805 vector_outcome.returned_candidates,
2806 vector_outcome.post_filter_candidates,
2807 vector_outcome.fallback,
2808 vector_outcome.exact_rerank,
2809 &results,
2810 vector_outcome.degradations,
2811 vector_outcome.receipt_metadata,
2812 );
2813 Ok(SearchExecution { results, receipt })
2814}
2815
2816pub(crate) fn vector_only_search_detailed(
2817 conn: &Connection,
2818 query_embedding: &[f32],
2819 config: &SearchConfig,
2820 top_k: usize,
2821 namespaces: Option<&[&str]>,
2822 source_types: Option<&[SearchSourceType]>,
2823 session_ids: Option<&[&str]>,
2824) -> Result<Vec<ExplainedResult>, MemoryError> {
2825 let context = SearchContext::default_now();
2826 Ok(vector_only_search_detailed_with_context(
2827 conn,
2828 query_embedding,
2829 config,
2830 &context,
2831 top_k,
2832 namespaces,
2833 source_types,
2834 session_ids,
2835 )?
2836 .results)
2837}
2838
2839pub fn vector_only_search(
2841 conn: &Connection,
2842 query_embedding: &[f32],
2843 config: &SearchConfig,
2844 top_k: usize,
2845 namespaces: Option<&[&str]>,
2846 source_types: Option<&[SearchSourceType]>,
2847 session_ids: Option<&[&str]>,
2848) -> Result<Vec<SearchResult>, MemoryError> {
2849 Ok(vector_only_search_detailed(
2850 conn,
2851 query_embedding,
2852 config,
2853 top_k,
2854 namespaces,
2855 source_types,
2856 session_ids,
2857 )?
2858 .into_iter()
2859 .map(|result| result.result)
2860 .collect())
2861}
2862
2863#[cfg(test)]
2864mod digest_tests {
2865 use super::query_embedding_digest;
2866
2867 #[test]
2868 fn query_embedding_digest_includes_dimension_and_bytes() {
2869 let two_dims = query_embedding_digest(&[1.0, 2.0]);
2870 let three_dims = query_embedding_digest(&[1.0, 2.0, 0.0]);
2871 let changed_byte = query_embedding_digest(&[1.0, 2.000_001]);
2872
2873 assert!(two_dims.starts_with("blake3:"));
2874 assert_eq!(two_dims.len(), 71);
2875 assert_ne!(two_dims, three_dims);
2876 assert_ne!(two_dims, changed_byte);
2877 assert_eq!(two_dims, query_embedding_digest(&[1.0, 2.0]));
2878 }
2879}
2880
2881#[cfg(feature = "hnsw")]
2883#[allow(clippy::too_many_arguments)]
2884pub fn vector_only_search_with_hnsw(
2885 conn: &Connection,
2886 query_embedding: &[f32],
2887 config: &SearchConfig,
2888 top_k: usize,
2889 namespaces: Option<&[&str]>,
2890 source_types: Option<&[SearchSourceType]>,
2891 session_ids: Option<&[&str]>,
2892 hnsw_hits: &[crate::hnsw::HnswHit],
2893) -> Result<Vec<SearchResult>, MemoryError> {
2894 Ok(vector_only_search_with_hnsw_detailed(
2895 conn,
2896 query_embedding,
2897 config,
2898 top_k,
2899 namespaces,
2900 source_types,
2901 session_ids,
2902 hnsw_hits,
2903 )?
2904 .into_iter()
2905 .map(|result| result.result)
2906 .collect())
2907}
2908
2909#[cfg(feature = "hnsw")]
2910#[allow(clippy::too_many_arguments)]
2911pub(crate) fn vector_only_search_with_hnsw_detailed_with_context(
2912 conn: &Connection,
2913 query_embedding: &[f32],
2914 config: &SearchConfig,
2915 context: &SearchContext,
2916 top_k: usize,
2917 namespaces: Option<&[&str]>,
2918 source_types: Option<&[SearchSourceType]>,
2919 session_ids: Option<&[&str]>,
2920 hnsw_hits: &[crate::hnsw::HnswHit],
2921) -> Result<SearchExecution, MemoryError> {
2922 let mut vector_hits = resolve_hnsw_hits_batched(
2923 conn,
2924 query_embedding,
2925 config,
2926 namespaces,
2927 source_types,
2928 session_ids,
2929 hnsw_hits,
2930 )?;
2931 let mut fallback = None;
2932 let mut degradations = Vec::new();
2933 let mut backend = "hnsw";
2934 let mut exact_rerank = config.rerank_from_f32;
2935
2936 if !hnsw_hits.is_empty()
2937 && vector_hits.len() < top_k
2938 && filters_are_active(namespaces, source_types, session_ids)
2939 {
2940 fallback = Some("hnsw_filtered_underreturn_fallback".to_string());
2941 degradations.push(format!(
2942 "HNSW returned {} post-filter vector candidates for requested top_k {}; exact filtered fallback was used",
2943 vector_hits.len(),
2944 top_k
2945 ));
2946 vector_hits = vector_search(
2947 conn,
2948 query_embedding,
2949 top_k,
2950 config.min_similarity,
2951 namespaces,
2952 source_types,
2953 session_ids,
2954 )?;
2955 backend = "hnsw_then_brute_force_f32";
2956 exact_rerank = true;
2957 }
2958
2959 let results = rrf_fuse_detailed_with_context(&[], &vector_hits, config, context, top_k);
2960 let receipt = build_receipt(
2961 context,
2962 query_embedding,
2963 "vector_only",
2964 backend,
2965 top_k,
2966 hnsw_hits.len(),
2967 vector_hits.len(),
2968 fallback,
2969 exact_rerank,
2970 &results,
2971 degradations,
2972 );
2973 Ok(SearchExecution { results, receipt })
2974}
2975
2976#[cfg(feature = "hnsw")]
2977#[allow(clippy::too_many_arguments)]
2978pub(crate) fn vector_only_search_with_hnsw_detailed(
2979 conn: &Connection,
2980 query_embedding: &[f32],
2981 config: &SearchConfig,
2982 top_k: usize,
2983 namespaces: Option<&[&str]>,
2984 source_types: Option<&[SearchSourceType]>,
2985 session_ids: Option<&[&str]>,
2986 hnsw_hits: &[crate::hnsw::HnswHit],
2987) -> Result<Vec<ExplainedResult>, MemoryError> {
2988 let context = SearchContext::default_now();
2989 Ok(vector_only_search_with_hnsw_detailed_with_context(
2990 conn,
2991 query_embedding,
2992 config,
2993 &context,
2994 top_k,
2995 namespaces,
2996 source_types,
2997 session_ids,
2998 hnsw_hits,
2999 )?
3000 .results)
3001}
3002
3003fn build_filter_clause(
3004 column: &str,
3005 values: Option<&[&str]>,
3006 param_offset: usize,
3007) -> (String, Vec<SqlValue>) {
3008 match values {
3009 Some(values) if !values.is_empty() => {
3010 let placeholders = (0..values.len())
3011 .map(|idx| format!("?{}", param_offset + idx))
3012 .collect::<Vec<_>>();
3013 let clause = format!(" AND {} IN ({})", column, placeholders.join(", "));
3014 let params = values
3015 .iter()
3016 .map(|value| SqlValue::Text((*value).to_string()))
3017 .collect();
3018 (clause, params)
3019 }
3020 _ => (String::new(), Vec::new()),
3021 }
3022}
3023
3024pub fn deduplicate_results(results: Vec<SearchResult>) -> Vec<SearchResult> {
3026 let mut seen = HashSet::new();
3027 results
3028 .into_iter()
3029 .filter(|result| seen.insert(source_dedup_key(&result.source)))
3030 .collect()
3031}
3032
3033#[cfg(test)]
3034mod tests {
3035 use super::*;
3036
3037 fn vector_row(id: &str) -> VectorRow {
3038 VectorRow {
3039 id: id.to_string(),
3040 content: format!("content {id}"),
3041 blob: bytemuck::cast_slice(&[1.0_f32, 0.0]).to_vec(),
3042 updated_at: None,
3043 source_type: SearchSourceType::Facts,
3044 filter_namespace: Some("default".to_string()),
3045 filter_session_id: None,
3046 source: SearchSource::Fact {
3047 fact_id: id.to_string(),
3048 namespace: "default".to_string(),
3049 },
3050 }
3051 }
3052
3053 #[test]
3054 fn timestamp_parser_accepts_sql_fractional_and_rfc3339_and_warns_by_returning_none() {
3055 assert!(parse_search_timestamp("2026-05-07 12:34:56").is_some());
3056 assert!(parse_search_timestamp("2026-05-07 12:34:56.123").is_some());
3057 assert!(parse_search_timestamp("2026-05-07T12:34:56Z").is_some());
3058 assert!(parse_search_timestamp("not-a-timestamp").is_none());
3059 }
3060
3061 #[test]
3062 fn vector_scan_hard_limit_blocks_before_unbounded_scan() {
3063 let old_warn = VECTOR_SCAN_WARN_LIMIT.swap(1, Ordering::SeqCst);
3064 let old_hard = VECTOR_SCAN_BLOCK_LIMIT.swap(2, Ordering::SeqCst);
3065 let rows = ["a", "b", "c"].into_iter().map(|id| Ok(vector_row(id)));
3066 let result = scan_vector_rows(rows, &[1.0, 0.0], -1.0, "fact");
3067 VECTOR_SCAN_WARN_LIMIT.store(old_warn, Ordering::SeqCst);
3068 VECTOR_SCAN_BLOCK_LIMIT.store(old_hard, Ordering::SeqCst);
3069
3070 match result {
3071 Err(MemoryError::VectorScanLimitExceeded {
3072 table,
3073 scanned,
3074 limit,
3075 }) => {
3076 assert_eq!(table, "fact");
3077 assert_eq!(scanned, 3);
3078 assert_eq!(limit, 2);
3079 }
3080 other => panic!("expected vector scan limit error, got {other:?}"),
3081 }
3082 }
3083}