1use std::collections::BTreeMap;
7use std::sync::Arc;
8
9use serde::{Deserialize, Serialize};
10
11use crate::redaction::SharedEventRedactor;
12use crate::{EventId, SessionEventKind, SessionMeta, StoredEvent};
13
14pub const DEFAULT_SEARCH_LIMIT: usize = 50;
15pub const MAX_SEARCH_LIMIT: usize = 500;
16const RRF_K: f32 = 60.0;
17
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum SearchMode {
21 Fts,
22 Semantic,
23 #[default]
24 Hybrid,
25}
26
27#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
28pub struct SearchFilter {
29 #[serde(default)]
30 pub tenant_id: Option<String>,
31 #[serde(default)]
32 pub project_scope: Option<String>,
33 #[serde(default)]
34 pub session_id: Option<String>,
35}
36
37#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
38pub struct SearchQuery {
39 pub query: String,
40 #[serde(default)]
41 pub mode: SearchMode,
42 #[serde(default)]
43 pub filter: SearchFilter,
44 #[serde(default)]
45 pub limit: Option<usize>,
46}
47
48impl SearchQuery {
49 pub fn validate(&self) -> Result<(), String> {
50 if self.query.trim().is_empty() {
51 return Err("search query must be non-empty".to_string());
52 }
53 if self.query.chars().any(|character| character == '\0') {
54 return Err("search query must not contain NUL".to_string());
55 }
56 let has_scope = [
57 self.filter.tenant_id.as_deref(),
58 self.filter.project_scope.as_deref(),
59 self.filter.session_id.as_deref(),
60 ]
61 .into_iter()
62 .flatten()
63 .any(|scope| !scope.trim().is_empty());
64 if !has_scope {
65 return Err(
66 "search requires tenant_id, project_scope, or session_id scope".to_string(),
67 );
68 }
69 Ok(())
70 }
71
72 pub fn limit(&self) -> usize {
73 self.limit
74 .unwrap_or(DEFAULT_SEARCH_LIMIT)
75 .clamp(1, MAX_SEARCH_LIMIT)
76 }
77}
78
79#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
80pub struct SearchHit {
81 pub session_id: String,
82 pub event_id: EventId,
83 pub kind: SessionEventKind,
84 pub score: f32,
85 #[serde(skip_serializing_if = "Option::is_none")]
86 pub fts_score: Option<f32>,
87 #[serde(skip_serializing_if = "Option::is_none")]
88 pub semantic_score: Option<f32>,
89 pub snippet: String,
90 pub event: StoredEvent,
91}
92
93#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
94pub struct SearchResponse {
95 pub requested_mode: SearchMode,
96 pub effective_mode: SearchMode,
97 pub embedding_backend: String,
98 pub semantic_floor: bool,
99 #[serde(skip_serializing_if = "Option::is_none")]
100 pub fallback_reason: Option<String>,
101 pub hits: Vec<SearchHit>,
102}
103
104pub trait Embedder: Send + Sync {
110 fn embed(&self, text: &str) -> Vec<f32>;
111 fn dim(&self) -> usize;
112 fn name(&self) -> &str;
113
114 fn is_semantic(&self) -> bool {
130 false
131 }
132
133 fn embed_batch(&self, texts: &[String]) -> Vec<Vec<f32>> {
134 texts.iter().map(|text| self.embed(text)).collect()
135 }
136}
137
138pub struct LexicalEmbedder {
140 dim: usize,
141}
142
143impl LexicalEmbedder {
144 pub fn new(dim: usize) -> Self {
145 Self { dim: dim.max(16) }
146 }
147
148 fn add_feature(&self, vector: &mut [f32], feature: &str, weight: f32) {
149 let hash = fnv1a(feature.as_bytes(), 0);
150 let bucket = (hash % self.dim as u64) as usize;
151 let sign = if fnv1a(feature.as_bytes(), 0x9e37_79b9_7f4a_7c15) & 1 == 0 {
152 1.0
153 } else {
154 -1.0
155 };
156 vector[bucket] += sign * weight;
157 }
158}
159
160impl Default for LexicalEmbedder {
161 fn default() -> Self {
162 Self::new(256)
163 }
164}
165
166impl Embedder for LexicalEmbedder {
167 fn embed(&self, text: &str) -> Vec<f32> {
168 let mut vector = vec![0.0; self.dim];
169 for token in word_tokens(text) {
170 self.add_feature(&mut vector, &token, 1.0);
171 }
172 for gram in char_ngrams(text, 3) {
173 self.add_feature(&mut vector, &gram, 0.35);
174 }
175 l2_normalize(&mut vector);
176 vector
177 }
178
179 fn dim(&self) -> usize {
180 self.dim
181 }
182
183 #[allow(clippy::unnecessary_literal_bound)]
184 fn name(&self) -> &str {
185 "lexical-hash"
186 }
187
188 }
193
194pub fn default_embedder() -> Arc<dyn Embedder> {
195 Arc::new(LexicalEmbedder::default())
196}
197
198pub fn cosine(left: &[f32], right: &[f32]) -> f32 {
199 if left.is_empty() || left.len() != right.len() {
200 return 0.0;
201 }
202 let mut dot = 0.0;
203 let mut left_norm = 0.0;
204 let mut right_norm = 0.0;
205 for (left, right) in left.iter().zip(right.iter()) {
206 if !left.is_finite() || !right.is_finite() {
207 return 0.0;
208 }
209 dot += left * right;
210 left_norm += left * left;
211 right_norm += right * right;
212 }
213 if left_norm <= 0.0 || right_norm <= 0.0 {
214 return 0.0;
215 }
216 (dot / (left_norm.sqrt() * right_norm.sqrt())).clamp(-1.0, 1.0)
217}
218
219pub fn l2_normalize(vector: &mut [f32]) {
220 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
221 if norm > 0.0 {
222 for value in vector {
223 *value /= norm;
224 }
225 }
226}
227
228pub fn event_search_text(event: &StoredEvent) -> String {
229 let mut parts = Vec::new();
230 parts.push(event.kind.discriminator().replace('_', " "));
231 if let Some(actor) = event.actor.as_deref() {
232 parts.push(actor.to_string());
233 }
234 collect_json_strings(&event.payload, &mut parts);
235 parts.join("\n")
236}
237
238pub(crate) fn redacted_search_document(
239 redactor: Option<&SharedEventRedactor>,
240 meta: &SessionMeta,
241 event: &StoredEvent,
242) -> String {
243 redacted_search_document_parts(
244 redactor,
245 meta.title.as_deref(),
246 meta.cwd.as_deref(),
247 meta.model.as_deref(),
248 meta.project_scope.as_deref(),
249 event,
250 )
251}
252
253pub(crate) fn redacted_search_document_parts(
254 redactor: Option<&SharedEventRedactor>,
255 title: Option<&str>,
256 cwd: Option<&str>,
257 model: Option<&str>,
258 project_scope: Option<&str>,
259 event: &StoredEvent,
260) -> String {
261 let mut metadata = serde_json::json!({
262 "title": title,
263 "cwd": cwd,
264 "model": model,
265 "project_scope": project_scope,
266 });
267 if let Some(redactor) = redactor {
268 redactor.redact_json_in_place(&mut metadata);
269 }
270 search_document_parts(
271 metadata.get("title").and_then(serde_json::Value::as_str),
272 metadata.get("cwd").and_then(serde_json::Value::as_str),
273 metadata.get("model").and_then(serde_json::Value::as_str),
274 metadata
275 .get("project_scope")
276 .and_then(serde_json::Value::as_str),
277 event,
278 )
279}
280
281pub(crate) fn search_document_parts(
282 title: Option<&str>,
283 cwd: Option<&str>,
284 model: Option<&str>,
285 project_scope: Option<&str>,
286 event: &StoredEvent,
287) -> String {
288 let event_text = event_search_text(event);
289 [title, cwd, model, project_scope, Some(event_text.as_str())]
290 .into_iter()
291 .flatten()
292 .collect::<Vec<_>>()
293 .join("\n")
294}
295
296pub fn snippet(text: &str, query: &str, max_chars: usize) -> String {
297 let text = text.trim();
298 if text.chars().count() <= max_chars {
299 return text.to_string();
300 }
301 let folded = text.to_lowercase();
302 let needle = word_tokens(query).into_iter().next().unwrap_or_default();
303 let byte_anchor = if needle.is_empty() {
304 0
305 } else {
306 folded.find(&needle).unwrap_or(0)
307 };
308 let mut original_byte_anchor = byte_anchor.min(text.len());
309 while original_byte_anchor > 0 && !text.is_char_boundary(original_byte_anchor) {
310 original_byte_anchor -= 1;
311 }
312 #[expect(
313 clippy::string_slice,
314 reason = "original_byte_anchor is walked back to a char boundary by the loop above"
315 )]
316 let char_anchor = text[..original_byte_anchor].chars().count();
317 let start = char_anchor.saturating_sub(max_chars / 3);
318 let excerpt = text.chars().skip(start).take(max_chars).collect::<String>();
319 format!(
320 "{}{}{}",
321 if start > 0 { "…" } else { "" },
322 excerpt,
323 if start + max_chars < text.chars().count() {
324 "…"
325 } else {
326 ""
327 }
328 )
329}
330
331pub(crate) fn lexical_score(query: &str, text: &str) -> f32 {
332 let query_tokens = word_tokens(query);
333 if query_tokens.is_empty() {
334 return 0.0;
335 }
336 let text_tokens = word_tokens(text);
337 let frequencies =
338 text_tokens
339 .into_iter()
340 .fold(BTreeMap::<String, usize>::new(), |mut counts, token| {
341 *counts.entry(token).or_default() += 1;
342 counts
343 });
344 if query_tokens
345 .iter()
346 .any(|token| !frequencies.contains_key(token))
347 {
348 return 0.0;
349 }
350 let matched = query_tokens
351 .iter()
352 .filter_map(|token| frequencies.get(token))
353 .map(|count| 1.0 + (*count as f32).ln())
354 .sum::<f32>();
355 let exact = text
356 .to_lowercase()
357 .contains(query.trim().to_lowercase().as_str());
358 matched / query_tokens.len() as f32 + if exact { 1.0 } else { 0.0 }
359}
360
361pub(crate) fn combined_score(
362 mode: SearchMode,
363 fts_rank: Option<usize>,
364 semantic_rank: Option<usize>,
365 fts_score: Option<f32>,
366 semantic_score: Option<f32>,
367) -> f32 {
368 match mode {
369 SearchMode::Fts => fts_score.unwrap_or_default(),
370 SearchMode::Semantic => semantic_score.unwrap_or_default(),
371 SearchMode::Hybrid => {
372 fts_rank
373 .map(|rank| 1.0 / (RRF_K + rank as f32 + 1.0))
374 .unwrap_or_default()
375 + semantic_rank
376 .map(|rank| 1.0 / (RRF_K + rank as f32 + 1.0))
377 .unwrap_or_default()
378 }
379 }
380}
381
382pub(crate) fn ranks(scores: &[f32]) -> BTreeMap<usize, usize> {
383 let mut ranked = scores
384 .iter()
385 .copied()
386 .enumerate()
387 .filter(|(_, score)| *score > 0.0)
388 .collect::<Vec<_>>();
389 ranked.sort_by(|(left_index, left), (right_index, right)| {
390 right
391 .total_cmp(left)
392 .then_with(|| left_index.cmp(right_index))
393 });
394 ranked
395 .into_iter()
396 .enumerate()
397 .map(|(rank, (index, _))| (index, rank))
398 .collect()
399}
400
401pub(crate) fn vector_blob(vector: &[f32]) -> Vec<u8> {
402 let mut bytes = Vec::with_capacity(std::mem::size_of_val(vector));
403 for value in vector {
404 bytes.extend_from_slice(&value.to_le_bytes());
405 }
406 bytes
407}
408
409pub(crate) fn vector_from_blob(bytes: &[u8], dim: usize) -> Option<Vec<f32>> {
410 if bytes.len() != dim.checked_mul(std::mem::size_of::<f32>())? {
411 return None;
412 }
413 Some(
414 bytes
415 .chunks_exact(4)
416 .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
417 .collect(),
418 )
419}
420
421pub(crate) fn word_tokens(text: &str) -> Vec<String> {
422 let mut tokens = Vec::new();
423 let mut current = String::new();
424 let mut previous_lower = false;
425 let flush = |current: &mut String, tokens: &mut Vec<String>| {
426 if !current.is_empty() {
427 tokens.push(std::mem::take(current));
428 }
429 };
430 for character in text.chars() {
431 if character.is_alphanumeric() {
432 if character.is_uppercase() && previous_lower {
433 flush(&mut current, &mut tokens);
434 }
435 current.extend(character.to_lowercase());
436 previous_lower = character.is_lowercase() || character.is_numeric();
437 } else {
438 flush(&mut current, &mut tokens);
439 previous_lower = false;
440 }
441 }
442 flush(&mut current, &mut tokens);
443 tokens
444}
445
446pub(crate) fn fts_literal_query(query: &str) -> String {
447 word_tokens(query)
448 .into_iter()
449 .map(|token| format!("\"{}\"", token.replace('"', "\"\"")))
450 .collect::<Vec<_>>()
451 .join(" AND ")
452}
453
454fn char_ngrams(text: &str, width: usize) -> Vec<String> {
455 if width == 0 {
456 return Vec::new();
457 }
458 let mut normalized = String::with_capacity(text.len() + 2);
459 normalized.push(' ');
460 let mut previous_space = true;
461 for character in text.chars() {
462 if character.is_whitespace() {
463 if !previous_space {
464 normalized.push(' ');
465 previous_space = true;
466 }
467 } else {
468 normalized.extend(character.to_lowercase());
469 previous_space = false;
470 }
471 }
472 if !previous_space {
473 normalized.push(' ');
474 }
475 let characters = normalized.chars().collect::<Vec<_>>();
476 characters
477 .windows(width)
478 .map(|window| window.iter().collect())
479 .collect()
480}
481
482fn fnv1a(bytes: &[u8], seed: u64) -> u64 {
483 const FNV_PRIME: u64 = 0x0000_0100_0000_01B3;
484 let mut hash = seed ^ 0xcbf2_9ce4_8422_2325;
485 for byte in bytes {
486 hash ^= u64::from(*byte);
487 hash = hash.wrapping_mul(FNV_PRIME);
488 }
489 hash
490}
491
492fn collect_json_strings(value: &serde_json::Value, parts: &mut Vec<String>) {
493 match value {
494 serde_json::Value::String(text) => parts.push(text.clone()),
495 serde_json::Value::Array(items) => {
496 for item in items {
497 collect_json_strings(item, parts);
498 }
499 }
500 serde_json::Value::Object(fields) => {
501 for value in fields.values() {
502 collect_json_strings(value, parts);
503 }
504 }
505 serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {}
506 }
507}
508
509pub mod conformance {
520 use super::{cosine, Embedder};
521
522 pub struct RetrievalCase {
529 pub query: &'static str,
530 pub related: &'static str,
531 pub decoy: &'static str,
532 }
533
534 pub const RETRIEVAL_CASES: &[RetrievalCase] = &[
537 RetrievalCase {
538 query: "find authentication code",
539 related: "verifyCredentials checks the caller identity before granting access",
540 decoy: "findCode looks up a numeric status code in a lookup table",
541 },
542 RetrievalCase {
543 query: "where do we log a user in",
544 related: "authenticate establishes a session for the account",
545 decoy: "logWarning writes a user-facing message to the log file",
546 },
547 RetrievalCase {
548 query: "how are invoices sent to customers",
549 related: "deliverReceipt mails the billing document to the account owner",
550 decoy: "customerNotes stores free-form text attached to invoices",
551 },
552 RetrievalCase {
553 query: "retry a failed network request",
554 related: "backoffScheduler reattempts the call after an exponential delay",
555 decoy: "networkRequestLog records every request that failed validation",
556 },
557 RetrievalCase {
558 query: "limit how many requests a client may send",
559 related: "throttle rejects traffic above the configured quota",
560 decoy: "requestClient sends a limited set of headers with each call",
561 },
562 RetrievalCase {
563 query: "clean up temporary files on shutdown",
564 related: "releaseScratchStorage removes working directories when the process exits",
565 decoy: "temporaryFileCache keeps a clean list of the files it has opened",
566 },
567 ];
568
569 pub fn bar(cases: usize) -> usize {
573 cases.saturating_sub(1)
574 }
575
576 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
579 pub struct SemanticClaimAudit {
580 pub claims_semantic: bool,
582 pub wins: usize,
584 pub cases: usize,
586 pub clears_bar: bool,
588 pub honest: bool,
590 }
591
592 pub fn audit_semantic_claim(embedder: &dyn Embedder) -> SemanticClaimAudit {
598 let mut wins = 0usize;
599 for case in RETRIEVAL_CASES {
600 let query = embedder.embed(case.query);
601 let related = cosine(&query, &embedder.embed(case.related));
602 let decoy = cosine(&query, &embedder.embed(case.decoy));
603 if related > decoy {
604 wins += 1;
605 }
606 }
607 let cases = RETRIEVAL_CASES.len();
608 let clears_bar = wins >= bar(cases);
609 let claims_semantic = embedder.is_semantic();
610 SemanticClaimAudit {
611 claims_semantic,
612 wins,
613 cases,
614 clears_bar,
615 honest: claims_semantic == clears_bar,
616 }
617 }
618
619 pub fn assert_semantic_claim_is_earned(embedder: &dyn Embedder) {
627 let audit = audit_semantic_claim(embedder);
628 assert!(
629 audit.honest,
630 "backend `{}` reports is_semantic() == {} but won {}/{} retrieval cases (bar is {}). \
631 {}",
632 embedder.name(),
633 audit.claims_semantic,
634 audit.wins,
635 audit.cases,
636 bar(audit.cases),
637 if audit.claims_semantic {
638 "The claim is not earned: either fix the backend or drop the override."
639 } else {
640 "The backend outgrew its disclaimer: override is_semantic() to true so callers \
641 stop being told they are on the lexical floor."
642 }
643 );
644 }
645}
646
647#[cfg(test)]
648mod tests {
649 use super::*;
650
651 #[test]
652 fn lexical_embedder_is_deterministic_and_related() {
653 let embedder = LexicalEmbedder::default();
654 let query = embedder.embed("rate limiting middleware");
655 assert_eq!(query, embedder.embed("rate limiting middleware"));
656 assert!(
657 cosine(&query, &embedder.embed("API rate limiter"))
658 > cosine(&query, &embedder.embed("markdown table renderer"))
659 );
660 }
661
662 #[test]
663 fn fts_queries_are_literal_and_identifier_aware() {
664 assert_eq!(
665 fts_literal_query("getUserByID OR token*"),
666 "\"get\" AND \"user\" AND \"by\" AND \"id\" AND \"or\" AND \"token\""
667 );
668 }
669
670 #[test]
671 fn vector_blob_round_trips() {
672 let vector = vec![-1.0, 0.25, 4.0];
673 assert_eq!(vector_from_blob(&vector_blob(&vector), 3), Some(vector));
674 assert_eq!(vector_from_blob(&[0, 1], 3), None);
675 }
676
677 #[test]
678 fn search_requires_an_explicit_scope() {
679 let error = SearchQuery {
680 query: "needle".to_string(),
681 mode: SearchMode::Fts,
682 filter: SearchFilter::default(),
683 limit: None,
684 }
685 .validate()
686 .expect_err("unscoped search must be rejected");
687 assert!(error.contains("requires"));
688 }
689
690 #[test]
691 fn unicode_snippet_anchor_never_slices_at_a_folded_byte_offset() {
692 let text = format!("{}needle", "İ".repeat(300));
693 let rendered = snippet(&text, "needle", 40);
694 assert!(rendered.contains("needle"));
695 }
696
697 struct OracleEmbedder;
701
702 impl OracleEmbedder {
703 fn axis(text: &str) -> usize {
706 let cases = conformance::RETRIEVAL_CASES;
707 for (index, case) in cases.iter().enumerate() {
708 if case.query == text || case.related == text {
709 return index;
710 }
711 if case.decoy == text {
712 return cases.len() + index;
713 }
714 }
715 cases.len() * 2
716 }
717 }
718
719 impl Embedder for OracleEmbedder {
720 fn embed(&self, text: &str) -> Vec<f32> {
721 let mut vector = vec![0.0; conformance::RETRIEVAL_CASES.len() * 2 + 1];
722 vector[Self::axis(text)] = 1.0;
723 vector
724 }
725
726 fn dim(&self) -> usize {
727 conformance::RETRIEVAL_CASES.len() * 2 + 1
728 }
729
730 #[allow(clippy::unnecessary_literal_bound)]
731 fn name(&self) -> &str {
732 "oracle"
733 }
734
735 fn is_semantic(&self) -> bool {
736 true
737 }
738 }
739
740 struct ModestOracle;
744
745 impl Embedder for ModestOracle {
746 fn embed(&self, text: &str) -> Vec<f32> {
747 OracleEmbedder.embed(text)
748 }
749
750 fn dim(&self) -> usize {
751 OracleEmbedder.dim()
752 }
753
754 #[allow(clippy::unnecessary_literal_bound)]
755 fn name(&self) -> &str {
756 "modest-oracle"
757 }
758 }
759
760 struct BoastfulEmbedder;
762
763 impl Embedder for BoastfulEmbedder {
764 fn embed(&self, _text: &str) -> Vec<f32> {
765 vec![1.0, 0.0]
766 }
767
768 fn dim(&self) -> usize {
769 2
770 }
771
772 #[allow(clippy::unnecessary_literal_bound)]
773 fn name(&self) -> &str {
774 "boastful"
775 }
776
777 fn is_semantic(&self) -> bool {
778 true
779 }
780 }
781
782 #[test]
783 fn lexical_floor_inherits_an_honest_default() {
784 let audit = conformance::audit_semantic_claim(&LexicalEmbedder::default());
787 assert!(!audit.claims_semantic, "the floor must not claim meaning");
788 assert!(
789 !audit.clears_bar,
790 "surface matching cleared a bar built to defeat it ({}/{}); the corpus has decayed",
791 audit.wins, audit.cases
792 );
793 assert!(
797 audit.wins <= 1,
798 "the floor won {}/{} cases; decoys are no longer out-matching the related documents \
799 on surface form, so this corpus is drifting toward vacuity",
800 audit.wins,
801 audit.cases
802 );
803 conformance::assert_semantic_claim_is_earned(&LexicalEmbedder::default());
804 }
805
806 #[test]
807 fn a_backend_that_ranks_may_claim() {
808 let audit = conformance::audit_semantic_claim(&OracleEmbedder);
809 assert_eq!(audit.wins, audit.cases);
810 conformance::assert_semantic_claim_is_earned(&OracleEmbedder);
811 }
812
813 #[test]
814 #[should_panic(expected = "The claim is not earned")]
815 fn a_backend_that_cannot_rank_may_not_claim() {
816 conformance::assert_semantic_claim_is_earned(&BoastfulEmbedder);
817 }
818
819 #[test]
820 #[should_panic(expected = "outgrew its disclaimer")]
821 fn a_backend_that_outgrows_its_disclaimer_is_caught() {
822 conformance::assert_semantic_claim_is_earned(&ModestOracle);
823 }
824}