1use std::time::Duration;
42
43use crate::error::{KernelError, Result};
44use reqwest::header::CONTENT_TYPE;
45use serde::Deserialize;
46
47use super::{AsyncVectorIndex, SearchHit};
48
49pub struct ElasticsearchVectorIndex {
57 client: reqwest::Client,
58 base_url: String,
61 index: String,
62 dim: usize,
63}
64
65impl ElasticsearchVectorIndex {
66 pub async fn new(url: &str, index: &str, dim: usize) -> Result<Self> {
70 validate_index_name(index)?;
71 let client = reqwest::Client::builder()
72 .connect_timeout(Duration::from_secs(5))
77 .timeout(Duration::from_secs(30))
78 .build()
79 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))?;
80 let idx = Self {
81 client,
82 base_url: url.trim_end_matches('/').to_string(),
83 index: index.to_string(),
84 dim,
85 };
86 idx.ensure_index().await?;
87 Ok(idx)
88 }
89
90 pub async fn delete_index(&self) -> Result<()> {
92 let resp = self.delete(&format!("/{}", &self.index)).await?;
93 if !resp.status().is_success() && resp.status().as_u16() != 404 {
95 return Err(self.status_err(resp).await);
96 }
97 Ok(())
98 }
99
100 async fn ensure_index(&self) -> Result<()> {
102 let head = self
104 .client
105 .head(format!("{}/{}", &self.base_url, &self.index))
106 .send()
107 .await
108 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))?;
109 if head.status().as_u16() == 200 {
110 return Ok(());
111 }
112 if head.status().as_u16() != 404 {
114 return Err(self.status_err(head).await);
115 }
116 let body = serde_json::json!({
117 "mappings": {
118 "properties": {
119 "vector": {
120 "type": "dense_vector",
121 "dims": self.dim,
122 "index": true,
123 "similarity": "cosine"
124 },
125 "ext_id": { "type": "long" }
126 }
127 }
128 });
129 let resp = self.put(&format!("/{}", &self.index), body).await?;
130 if !resp.status().is_success() {
131 return Err(self.status_err(resp).await);
132 }
133 Ok(())
134 }
135
136 fn parse_id(_id: &str) -> Option<u64> {
139 _id.parse::<u64>().ok()
140 }
141
142 async fn put(&self, path: &str, body: serde_json::Value) -> Result<reqwest::Response> {
145 self.client
146 .put(format!("{}{}", &self.base_url, path))
147 .json(&body)
148 .send()
149 .await
150 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))
151 }
152
153 async fn delete(&self, path: &str) -> Result<reqwest::Response> {
154 self.client
155 .delete(format!("{}{}", &self.base_url, path))
156 .send()
157 .await
158 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))
159 }
160
161 async fn ndjson(&self, path: &str, body: String) -> Result<reqwest::Response> {
162 self.client
163 .post(format!("{}{}", &self.base_url, path))
164 .header(CONTENT_TYPE, "application/x-ndjson")
165 .body(body)
166 .send()
167 .await
168 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))
169 }
170
171 async fn status_err(&self, resp: reqwest::Response) -> KernelError {
172 let status = resp.status();
173 let body = resp.text().await.unwrap_or_default();
174 let body = truncate_error_body(&redact_credentials(&body));
178 KernelError::Embedding(format!(
179 "elasticsearch returned status {status} for index `{}` [url redacted]: {}",
180 &self.index, body
181 ))
182 }
183
184 async fn submit_bulk(&self, body: String, op: &str) -> Result<()> {
187 let resp = self.ndjson("/_bulk?refresh=wait_for", body).await?;
188 if !resp.status().is_success() {
189 return Err(self.status_err(resp).await);
190 }
191 let parsed: BulkResponse = decode(resp).await?;
192 if parsed.errors {
193 return Err(KernelError::Embedding(format!(
194 "elasticsearch bulk {op} reported per-item errors [url redacted]: {}",
195 first_failing_bulk_item(&parsed.items)
196 )));
197 }
198 Ok(())
199 }
200}
201
202const BULK_CHUNK_SIZE: usize = 500;
207
208#[async_trait::async_trait]
209impl AsyncVectorIndex for ElasticsearchVectorIndex {
210 async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
211 if vectors.len() != ids.len() {
212 return Err(KernelError::Embedding(format!(
213 "vectors.len() ({}) must equal ids.len() ({})",
214 vectors.len(),
215 ids.len()
216 )));
217 }
218 if vectors.is_empty() {
219 return Ok(());
220 }
221 for (vchunk, idchunk) in vectors
227 .chunks(BULK_CHUNK_SIZE)
228 .zip(ids.chunks(BULK_CHUNK_SIZE))
229 {
230 let mut body = String::new();
231 for (v, &id) in vchunk.iter().zip(idchunk.iter()) {
232 body.push_str(
233 &serde_json::to_string(&serde_json::json!({
234 "index": { "_index": &self.index, "_id": id.to_string() }
235 }))
236 .map_err(|e| KernelError::Embedding(format!("bulk encode: {e}")))?,
237 );
238 body.push('\n');
239 body.push_str(
240 &serde_json::to_string(&serde_json::json!({
241 "ext_id": id,
242 "vector": v
243 }))
244 .map_err(|e| KernelError::Embedding(format!("bulk encode: {e}")))?,
245 );
246 body.push('\n');
247 }
248 self.submit_bulk(body, "upsert").await?;
249 }
250 Ok(())
251 }
252
253 async fn remove(&self, ids: &[u64]) -> Result<()> {
254 if ids.is_empty() {
255 return Ok(());
256 }
257 for idchunk in ids.chunks(BULK_CHUNK_SIZE) {
258 let mut body = String::new();
259 for &id in idchunk {
260 body.push_str(
261 &serde_json::to_string(&serde_json::json!({
262 "delete": { "_index": &self.index, "_id": id.to_string() }
263 }))
264 .map_err(|e| KernelError::Embedding(format!("bulk encode: {e}")))?,
265 );
266 body.push('\n');
267 }
268 self.submit_bulk(body, "delete").await?;
271 }
272 Ok(())
273 }
274
275 async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
279 let num_candidates = knn_num_candidates(k);
280 let body = serde_json::json!({
281 "knn": {
282 "field": "vector",
283 "query_vector": query,
284 "k": k,
285 "num_candidates": num_candidates
286 },
287 "_source": false,
288 "size": k
289 });
290 let resp = self
291 .client
292 .post(format!("{}/{}/_search", &self.base_url, &self.index))
293 .json(&body)
294 .send()
295 .await
296 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))?;
297 if !resp.status().is_success() {
298 return Err(self.status_err(resp).await);
299 }
300 let parsed: SearchResponse = decode(resp).await?;
301 Ok(parsed
302 .hits
303 .hits
304 .into_iter()
305 .filter_map(|h| {
306 Self::parse_id(&h._id).map(|id| SearchHit {
307 id,
308 score: h._score,
309 })
310 })
311 .collect())
312 }
313
314 async fn search_filtered(
315 &self,
316 query: &[f32],
317 k: usize,
318 allowlist: &[u64],
319 ) -> Result<Vec<SearchHit>> {
320 if allowlist.is_empty() {
324 return Ok(vec![]);
325 }
326 let num_candidates = knn_num_candidates(k);
327 let allowlist: Vec<u64> = allowlist.to_vec();
328 let body = serde_json::json!({
329 "knn": {
330 "field": "vector",
331 "query_vector": query,
332 "k": k,
333 "num_candidates": num_candidates,
334 "filter": [{ "terms": { "ext_id": allowlist } }]
335 },
336 "_source": false,
337 "size": k
338 });
339 let resp = self
340 .client
341 .post(format!("{}/{}/_search", &self.base_url, &self.index))
342 .json(&body)
343 .send()
344 .await
345 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))?;
346 if !resp.status().is_success() {
347 return Err(self.status_err(resp).await);
348 }
349 let parsed: SearchResponse = decode(resp).await?;
350 Ok(parsed
351 .hits
352 .hits
353 .into_iter()
354 .filter_map(|h| {
355 Self::parse_id(&h._id).map(|id| SearchHit {
356 id,
357 score: h._score,
358 })
359 })
360 .collect())
361 }
362
363 async fn len(&self) -> Result<usize> {
364 let resp = self
365 .client
366 .post(format!("{}/{}/_count", &self.base_url, &self.index))
367 .json(&serde_json::json!({}))
368 .send()
369 .await
370 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))?;
371 if !resp.status().is_success() {
372 return Err(self.status_err(resp).await);
373 }
374 let parsed: CountResponse = decode(resp).await?;
375 Ok(parsed.count as usize)
376 }
377
378 fn dim(&self) -> usize {
379 self.dim
380 }
381}
382
383pub fn redact_credentials(s: &str) -> String {
395 let mut out = String::with_capacity(s.len());
396 let mut rest = s;
397 loop {
398 match rest.find("://") {
399 None => {
400 out.push_str(rest);
401 break;
402 }
403 Some(idx) => {
404 out.push_str(&rest[..idx + 3]);
406 let after = &rest[idx + 3..];
407 let auth_end = after.find(['/', '?', '#']).unwrap_or(after.len());
411 let auth = &after[..auth_end];
412 if let Some(at) = auth.rfind('@') {
413 out.push_str("<redacted>@");
414 out.push_str(&auth[at + 1..]);
415 } else {
416 out.push_str(auth);
417 }
418 rest = &after[auth_end..];
419 }
420 }
421 }
422 out
423}
424
425const MAX_KNN_CANDIDATES: usize = 1_000;
433
434fn knn_num_candidates(k: usize) -> usize {
442 let base = k.max(1).saturating_mul(10);
443 base.min(MAX_KNN_CANDIDATES).max(k)
444}
445
446const ERROR_BODY_MAX_CHARS: usize = 1024;
452
453fn truncate_error_body(s: &str) -> String {
461 if s.chars().count() <= ERROR_BODY_MAX_CHARS {
462 return s.to_string();
463 }
464 let cut = s
467 .char_indices()
468 .nth(ERROR_BODY_MAX_CHARS)
469 .map(|(i, _)| i)
470 .unwrap_or(s.len());
471 format!("{}... [truncated]", &s[..cut])
472}
473
474fn validate_index_name(index: &str) -> Result<()> {
483 if index.is_empty() {
484 return Err(KernelError::Embedding(
485 "elasticsearch index name must not be empty".into(),
486 ));
487 }
488 if index == "." || index == ".." {
491 return Err(KernelError::Embedding(format!(
492 "elasticsearch index name must not be `.` or `..` (reserved): `{}`",
493 index
494 )));
495 }
496 if index.len() > 255 {
497 return Err(KernelError::Embedding(format!(
498 "elasticsearch index name exceeds 255 bytes ({} bytes)",
499 index.len()
500 )));
501 }
502 match index.as_bytes()[0] {
503 b'_' | b'-' | b'+' => {
504 return Err(KernelError::Embedding(format!(
505 "elasticsearch index name must not start with `_`, `-`, or `+`: `{}`",
506 index
507 )));
508 }
509 _ => {}
510 }
511 if let Some(bad) = index.bytes().find(|&c| {
512 !(c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, b'_' | b'-' | b'.'))
513 }) {
514 return Err(KernelError::Embedding(format!(
515 "elasticsearch index name contains an illegal byte 0x{bad:02x} (`{}`): \
516 only lowercase a-z, 0-9, `_`, `-`, `.` are allowed",
517 index
518 )));
519 }
520 Ok(())
521}
522
523fn first_failing_bulk_item(items: &[serde_json::Value]) -> String {
530 for item in items {
531 if let Some(detail) = item.as_object().and_then(|o| o.values().next()) {
532 let status = detail.get("status").and_then(|v| v.as_i64()).unwrap_or(0);
533 let has_error = detail.get("error").is_some();
534 if status >= 400 || has_error {
535 return redact_credentials(&item.to_string());
536 }
537 }
538 }
539 "(no failing item found)".into()
540}
541
542async fn decode<T: serde::de::DeserializeOwned>(resp: reqwest::Response) -> Result<T> {
544 resp.json::<T>()
545 .await
546 .map_err(|e| KernelError::Embedding(redact_credentials(&e.to_string())))
547}
548
549#[derive(Deserialize)]
550struct SearchResponse {
551 hits: SearchHits,
552}
553
554#[derive(Deserialize)]
555struct SearchHits {
556 hits: Vec<SearchInnerHit>,
557}
558
559#[derive(Deserialize)]
560struct SearchInnerHit {
561 _id: String,
562 _score: f32,
563}
564
565#[derive(Deserialize)]
566struct CountResponse {
567 count: u64,
568}
569
570#[derive(Deserialize)]
571struct BulkResponse {
572 errors: bool,
573 #[serde(default)]
574 items: Vec<serde_json::Value>,
575}
576
577#[cfg(test)]
578mod tests {
579 use super::*;
580 use crate::embedding::AsyncVectorIndex;
581
582 const DIM: usize = 4;
583
584 fn unique_index() -> String {
585 format!("llm_kernel_test_{}", std::process::id())
586 }
587
588 fn offline_index(base_url: &str, dim: usize) -> ElasticsearchVectorIndex {
592 ElasticsearchVectorIndex {
593 client: reqwest::Client::new(),
594 base_url: base_url.trim_end_matches('/').to_string(),
595 index: "llm_kernel_test_offline".to_string(),
596 dim,
597 }
598 }
599
600 #[test]
601 fn parse_id_accepts_numeric_and_drops_rest() {
602 assert_eq!(ElasticsearchVectorIndex::parse_id("42"), Some(42));
603 assert_eq!(ElasticsearchVectorIndex::parse_id("0"), Some(0));
604 assert_eq!(
605 ElasticsearchVectorIndex::parse_id("18446744073709551615"),
606 Some(u64::MAX)
607 );
608 assert_eq!(ElasticsearchVectorIndex::parse_id("abc"), None);
610 assert_eq!(ElasticsearchVectorIndex::parse_id(""), None);
611 assert_eq!(ElasticsearchVectorIndex::parse_id("1.5"), None);
612 }
613
614 #[test]
615 fn redact_credentials_strips_userinfo() {
616 let cases = [
617 ("http://u:pw@host:9200", "http://<redacted>@host:9200"),
618 (
619 "https://elastic:secret@es.local/x",
620 "https://<redacted>@es.local/x",
621 ),
622 ("http://localhost:9200", "http://localhost:9200"),
623 ("http://user@host", "http://<redacted>@host"),
624 ("https://u:p@ss@host:9200", "https://<redacted>@host:9200"),
627 ("no url here", "no url here"),
628 (
631 "index 中文 — http://u:pw@h:9200",
632 "index 中文 — http://<redacted>@h:9200",
633 ),
634 ];
635 for (input, expected) in cases {
636 assert_eq!(redact_credentials(input), expected, "input = {input:?}");
637 }
638 assert!(!redact_credentials("https://u:secret@host").contains("secret"));
640 let leaked = redact_credentials("https://u:p@ss@host:9200");
642 assert!(!leaked.contains("p@ss"), "password tail leaked: {leaked}");
643 assert!(!leaked.contains("ss@"), "password tail leaked: {leaked}");
644 }
645
646 #[test]
647 fn validate_index_name_accepts_and_rejects() {
648 for ok in ["docs", "docs_v2", "my-index", "idx.2026", "a", "a.b-c_d"] {
650 assert!(
651 validate_index_name(ok).is_ok(),
652 "{ok:?} should be a valid index name"
653 );
654 }
655 for name in [
657 "", "Docs", "with space", "comma,idx", "_underscore", "-dash", "+plus", "bad/slash", "한글", ".", "..", ] {
669 assert!(
670 validate_index_name(name).is_err(),
671 "{name:?} should be rejected"
672 );
673 }
674 assert!(validate_index_name(&"a".repeat(255)).is_ok());
676 assert!(validate_index_name(&"a".repeat(256)).is_err());
677 }
678
679 #[test]
682 fn knn_num_candidates_scales_caps_and_floors() {
683 assert_eq!(knn_num_candidates(1), 10);
685 assert_eq!(knn_num_candidates(5), 50);
686 assert_eq!(knn_num_candidates(50), 500);
687 assert_eq!(knn_num_candidates(100), MAX_KNN_CANDIDATES);
689 assert_eq!(knn_num_candidates(200), MAX_KNN_CANDIDATES);
691 assert!(knn_num_candidates(200) >= 200);
692 assert_eq!(knn_num_candidates(0), 10);
694 }
695
696 #[test]
698 fn truncate_error_body_leaves_short_body_unchanged() {
699 assert_eq!(truncate_error_body(""), "");
700 assert_eq!(truncate_error_body("short error"), "short error");
701 let at_cap: String = "a".repeat(ERROR_BODY_MAX_CHARS);
703 let out = truncate_error_body(&at_cap);
704 assert_eq!(out.chars().count(), ERROR_BODY_MAX_CHARS);
705 assert!(!out.contains("[truncated]"));
706 }
707
708 #[test]
712 fn truncate_error_body_caps_huge_body_with_marker() {
713 let huge: String = "a".repeat(ERROR_BODY_MAX_CHARS + 500);
715 let out = truncate_error_body(&huge);
716 assert!(out.ends_with("... [truncated]"));
717 let kept = out.strip_suffix("... [truncated]").unwrap();
718 assert_eq!(kept.chars().count(), ERROR_BODY_MAX_CHARS);
719
720 let cjk: String = "중".repeat(ERROR_BODY_MAX_CHARS + 10);
724 let out_cjk = truncate_error_body(&cjk);
725 assert!(out_cjk.contains("[truncated]"));
728 }
731
732 #[test]
736 fn truncate_error_body_keeps_credentials_redacted() {
737 let with_cred = "error: see https://u:super-secret@host/idx for details";
740 let out = truncate_error_body(&redact_credentials(with_cred));
741 assert!(!out.contains("super-secret"), "credential leaked: {out}");
742 assert!(out.contains("<redacted>"));
743
744 let padding: String = "x".repeat(ERROR_BODY_MAX_CHARS + 50);
750 let long_cred = format!("{padding} then https://u:p@ss@host:9200");
751 let redacted = redact_credentials(&long_cred);
752 let out2 = truncate_error_body(&redacted);
753 assert!(
754 !out2.contains("p@ss") && !out2.contains("super-secret"),
755 "credential tail leaked: {out2}"
756 );
757 }
758
759 #[test]
763 fn first_failing_bulk_item_picks_failing_and_redacts() {
764 let items = vec![
766 serde_json::json!({ "index": { "_id": "1", "status": 200 } }),
767 serde_json::json!({
768 "index": { "_id": "2", "status": 400, "error": { "type": "mapper", "reason": "bad" } }
769 }),
770 ];
771 let s = first_failing_bulk_item(&items);
772 assert!(
773 s.contains("\"_id\":\"2\""),
774 "should name the failing item: {s}"
775 );
776 assert!(s.contains("400"));
777 let err_only = vec![serde_json::json!({
779 "delete": { "_id": "9", "error": { "type": "x", "reason": "y" } }
780 })];
781 assert!(first_failing_bulk_item(&err_only).contains("\"_id\":\"9\""));
782 let with_url = vec![serde_json::json!({
784 "index": { "_id": "3", "status": 500, "error": { "reason": "see https://u:secret@host" } }
785 })];
786 let leaked = first_failing_bulk_item(&with_url);
787 assert!(!leaked.contains("secret"), "credential leaked: {leaked}");
788 assert!(leaked.contains("<redacted>"));
789 let none = vec![serde_json::json!({ "index": { "_id": "1", "status": 200 } })];
791 assert_eq!(first_failing_bulk_item(&none), "(no failing item found)");
792 }
793
794 #[test]
798 fn credentialed_url_error_redacts_password() {
799 let credentialed = "https://elastic:super-secret-pw@es.internal:9200/idx";
800 let raw = format!("error sending request for url ({credentialed}): connection refused");
802 let redacted = redact_credentials(&raw);
803 assert!(
804 !redacted.contains("super-secret-pw"),
805 "password leaked in redacted error: {redacted}"
806 );
807 assert!(redacted.contains("<redacted>"));
808 }
809
810 #[tokio::test]
813 async fn empty_allowlist_returns_empty_without_network() {
814 let idx = offline_index("http://0.0.0.0:1", DIM);
815 let res = idx.search_filtered(&[1.0, 0.0, 0.0, 0.0], 5, &[]).await;
818 assert!(res.is_ok(), "empty allowlist must not error: {res:?}");
819 assert!(res.unwrap().is_empty());
820 }
821
822 async fn run_live_conformance(idx: &ElasticsearchVectorIndex) -> Result<()> {
825 if idx.dim() != DIM {
826 return Err(KernelError::Embedding("dim mismatch".into()));
827 }
828 if !idx.is_empty().await? {
829 return Err(KernelError::Embedding("not empty at start".into()));
830 }
831 idx.add(
832 &[vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
833 &[1, 2],
834 )
835 .await?;
836 if idx.len().await? != 2 {
837 return Err(KernelError::Embedding("len != 2 after add".into()));
838 }
839
840 let hits = idx.search(&[1.0, 0.0, 0.0, 0.0], 1).await?;
841 if hits.len() != 1 || hits[0].id != 1 {
842 return Err(KernelError::Embedding("nearest neighbor != id 1".into()));
843 }
844
845 let filtered = idx.search_filtered(&[1.0, 0.0, 0.0, 0.0], 2, &[2]).await?;
846 if filtered.len() != 1 || filtered[0].id != 2 {
847 return Err(KernelError::Embedding("filtered search != id 2".into()));
848 }
849
850 idx.add(&[vec![0.9, 0.1, 0.0, 0.0]], &[1]).await?;
852 if idx.len().await? != 2 {
853 return Err(KernelError::Embedding("len != 2 after re-add".into()));
854 }
855
856 idx.remove(&[1]).await?;
857 if idx.len().await? != 1 {
858 return Err(KernelError::Embedding("len != 1 after remove".into()));
859 }
860 let after = idx.search(&[1.0, 0.0, 0.0, 0.0], 5).await?;
861 if after.iter().any(|h| h.id == 1) {
862 return Err(KernelError::Embedding(
863 "id 1 still present after remove".into(),
864 ));
865 }
866 Ok(())
867 }
868
869 #[tokio::test]
873 async fn live_elastic_conformance() {
874 let url = match std::env::var("LLMKERNEL_ELASTIC_URL") {
875 Ok(u) => u,
876 Err(_) => {
877 eprintln!("skipped: LLMKERNEL_ELASTIC_URL unset (no live Elasticsearch)");
878 return;
879 }
880 };
881
882 let index = unique_index();
883 let idx = match ElasticsearchVectorIndex::new(&url, &index, DIM).await {
884 Ok(i) => i,
885 Err(e) => panic!("connect + create index: {e:?}"),
886 };
887 let result = run_live_conformance(&idx).await;
890 let _ = idx.delete_index().await;
891 result.expect("elasticsearch conformance failed");
892 }
893}