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