1#[allow(unused_imports)]
23use zeph_db::sql;
24use zeph_memory::{FieldCondition, FieldValue, QdrantOps, VectorFilter, VectorPoint, VectorStore};
25
26use zeph_common::{EmbeddingVector, Normalized};
27
28use crate::error::Result;
29
30const CODE_COLLECTION: &str = "zeph_code_chunks";
31
32#[derive(Clone)]
46pub struct CodeStore {
47 ops: QdrantOps,
48 collection: String,
49 pool: zeph_db::DbPool,
50}
51
52pub struct ChunkInsert<'a> {
58 pub file_path: &'a str,
60 pub language: &'a str,
62 pub node_type: &'a str,
64 pub entity_name: Option<&'a str>,
66 pub line_start: usize,
68 pub line_end: usize,
70 pub code: &'a str,
72 pub scope_chain: &'a str,
74 pub content_hash: &'a str,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct NodeKind(pub String);
94
95impl std::fmt::Display for NodeKind {
96 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 f.write_str(&self.0)
98 }
99}
100
101impl AsRef<str> for NodeKind {
102 fn as_ref(&self) -> &str {
103 &self.0
104 }
105}
106
107impl From<String> for NodeKind {
108 fn from(s: String) -> Self {
109 Self(s)
110 }
111}
112
113impl From<&str> for NodeKind {
114 fn from(s: &str) -> Self {
115 Self(s.to_owned())
116 }
117}
118
119#[derive(Debug)]
124pub struct SearchHit {
125 pub code: String,
127 pub file_path: String,
129 pub line_range: (usize, usize),
131 pub score: f32,
133 pub node_type: NodeKind,
135 pub language: crate::languages::Lang,
137 pub entity_name: Option<String>,
139 pub scope_chain: String,
141}
142
143impl CodeStore {
144 #[must_use]
164 pub fn with_ops(ops: QdrantOps, pool: zeph_db::DbPool) -> Self {
165 Self {
166 ops,
167 collection: CODE_COLLECTION.into(),
168 pool,
169 }
170 }
171
172 #[tracing::instrument(name = "index.store.ensure_collection", skip_all)]
178 pub async fn ensure_collection(&self, vector_size: u64) -> Result<()> {
179 self.ops
180 .ensure_collection_with_quantization(
181 &self.collection,
182 vector_size,
183 &["language", "file_path", "node_type"],
184 )
185 .await?;
186 Ok(())
187 }
188
189 #[tracing::instrument(name = "index.store.upsert_chunk", skip_all)]
195 pub async fn upsert_chunk(&self, chunk: &ChunkInsert<'_>, vector: Vec<f32>) -> Result<String> {
196 tracing::Span::current().record("file_path", chunk.file_path);
197 let point_id = uuid::Uuid::new_v4().to_string();
198
199 let payload = serde_json::json!({
200 "file_path": chunk.file_path,
201 "language": chunk.language,
202 "node_type": chunk.node_type,
203 "entity_name": chunk.entity_name,
204 "line_start": chunk.line_start,
205 "line_end": chunk.line_end,
206 "code": chunk.code,
207 "scope_chain": chunk.scope_chain,
208 "content_hash": chunk.content_hash,
209 });
210
211 let payload_map = match payload {
212 serde_json::Value::Object(m) => m.into_iter().collect(),
213 _ => std::collections::HashMap::new(),
214 };
215
216 VectorStore::upsert(
217 &self.ops,
218 &self.collection,
219 vec![VectorPoint {
220 id: point_id.clone(),
221 vector,
222 payload: payload_map,
223 }],
224 )
225 .await?;
226
227 let line_start = i64::try_from(chunk.line_start)?;
228 let line_end = i64::try_from(chunk.line_end)?;
229
230 zeph_db::query(
231 sql!("INSERT INTO chunk_metadata \
232 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type, entity_name) \
233 VALUES (?, ?, ?, ?, ?, ?, ?, ?) \
234 ON CONFLICT(file_path, content_hash) DO UPDATE SET \
235 qdrant_id = excluded.qdrant_id, \
236 line_start = excluded.line_start, line_end = excluded.line_end, \
237 language = excluded.language, node_type = excluded.node_type, \
238 entity_name = excluded.entity_name"),
239 )
240 .bind(&point_id)
241 .bind(chunk.file_path)
242 .bind(chunk.content_hash)
243 .bind(line_start)
244 .bind(line_end)
245 .bind(chunk.language)
246 .bind(chunk.node_type)
247 .bind(chunk.entity_name)
248 .execute(&self.pool)
249 .await?;
250
251 Ok(point_id)
252 }
253
254 #[tracing::instrument(name = "index.store.upsert_chunks_batch", skip_all)]
263 pub async fn upsert_chunks_batch(
264 &self,
265 chunks: Vec<(ChunkInsert<'_>, Vec<f32>)>,
266 ) -> Result<Vec<String>> {
267 tracing::Span::current().record("chunk_count", chunks.len());
268 if chunks.is_empty() {
269 return Ok(Vec::new());
270 }
271
272 let mut point_ids: Vec<String> = Vec::with_capacity(chunks.len());
273 let mut points: Vec<VectorPoint> = Vec::with_capacity(chunks.len());
274
275 for (chunk, vector) in &chunks {
276 let point_id = uuid::Uuid::new_v4().to_string();
277
278 let payload = serde_json::json!({
279 "file_path": chunk.file_path,
280 "language": chunk.language,
281 "node_type": chunk.node_type,
282 "entity_name": chunk.entity_name,
283 "line_start": chunk.line_start,
284 "line_end": chunk.line_end,
285 "code": chunk.code,
286 "scope_chain": chunk.scope_chain,
287 "content_hash": chunk.content_hash,
288 });
289
290 let payload_map = match payload {
291 serde_json::Value::Object(m) => m.into_iter().collect(),
292 _ => std::collections::HashMap::new(),
293 };
294
295 points.push(VectorPoint {
296 id: point_id.clone(),
297 vector: vector.clone(),
298 payload: payload_map,
299 });
300 point_ids.push(point_id);
301 }
302
303 VectorStore::upsert(&self.ops, &self.collection, points).await?;
304
305 let mut tx = self.pool.begin().await?;
306 for (idx, (chunk, _)) in chunks.iter().enumerate() {
307 let point_id = &point_ids[idx];
308 let line_start = i64::try_from(chunk.line_start)?;
309 let line_end = i64::try_from(chunk.line_end)?;
310
311 zeph_db::query(
312 sql!("INSERT INTO chunk_metadata \
313 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type, entity_name) \
314 VALUES (?, ?, ?, ?, ?, ?, ?, ?) \
315 ON CONFLICT(file_path, content_hash) DO UPDATE SET \
316 qdrant_id = excluded.qdrant_id, \
317 line_start = excluded.line_start, line_end = excluded.line_end, \
318 language = excluded.language, node_type = excluded.node_type, \
319 entity_name = excluded.entity_name"),
320 )
321 .bind(point_id)
322 .bind(chunk.file_path)
323 .bind(chunk.content_hash)
324 .bind(line_start)
325 .bind(line_end)
326 .bind(chunk.language)
327 .bind(chunk.node_type)
328 .bind(chunk.entity_name)
329 .execute(&mut *tx)
330 .await?;
331 }
332 tx.commit().await?;
333
334 Ok(point_ids)
335 }
336
337 #[tracing::instrument(name = "index.store.chunk_exists", skip_all, fields(%content_hash))]
343 pub async fn chunk_exists(&self, content_hash: &str) -> Result<bool> {
344 let row: (i64,) = zeph_db::query_as(sql!(
345 "SELECT COUNT(*) FROM chunk_metadata WHERE content_hash = ?"
346 ))
347 .bind(content_hash)
348 .fetch_one(&self.pool)
349 .await?;
350 Ok(row.0 > 0)
351 }
352
353 #[tracing::instrument(name = "index.store.existing_hashes", skip_all)]
362 pub async fn existing_hashes(
363 &self,
364 hashes: &[&str],
365 ) -> Result<std::collections::HashSet<String>> {
366 tracing::Span::current().record("hash_count", hashes.len());
367 if hashes.is_empty() {
368 return Ok(std::collections::HashSet::new());
369 }
370
371 let mut result = std::collections::HashSet::new();
372
373 for chunk in hashes.chunks(900) {
374 let placeholders = zeph_db::placeholder_list(1, chunk.len());
375 let sql = format!(
376 "SELECT content_hash FROM chunk_metadata WHERE content_hash IN ({placeholders})"
377 );
378 let mut query = zeph_db::query_scalar::<_, String>(zeph_db::sqlx::AssertSqlSafe(sql));
379 for hash in chunk {
380 query = query.bind(*hash);
381 }
382 let rows: Vec<String> = query.fetch_all(&self.pool).await?;
383 result.extend(rows);
384 }
385
386 Ok(result)
387 }
388
389 #[tracing::instrument(name = "index.store.remove_file_chunks", skip_all)]
395 pub async fn remove_file_chunks(&self, file_path: &str) -> Result<usize> {
396 tracing::Span::current().record("file_path", file_path);
397 let ids: Vec<(String,)> = zeph_db::query_as(sql!(
398 "SELECT qdrant_id FROM chunk_metadata WHERE file_path = ?"
399 ))
400 .bind(file_path)
401 .fetch_all(&self.pool)
402 .await?;
403
404 if ids.is_empty() {
405 return Ok(0);
406 }
407
408 let point_ids: Vec<String> = ids.iter().map(|(id,)| id.clone()).collect();
409
410 VectorStore::delete_by_ids(&self.ops, &self.collection, point_ids).await?;
411
412 let count = ids.len();
413 zeph_db::query(sql!("DELETE FROM chunk_metadata WHERE file_path = ?"))
414 .bind(file_path)
415 .execute(&self.pool)
416 .await?;
417
418 Ok(count)
419 }
420
421 #[tracing::instrument(name = "index.store.search", skip_all)]
435 pub async fn search(
436 &self,
437 query_vector: EmbeddingVector<Normalized>,
438 limit: usize,
439 language_filter: Option<String>,
440 ) -> Result<Vec<SearchHit>> {
441 let limit_u64 = u64::try_from(limit)?;
442 let filter = language_filter.map(|lang| VectorFilter {
443 must: vec![FieldCondition {
444 field: "language".into(),
445 value: FieldValue::Text(lang),
446 }],
447 must_not: vec![],
448 });
449
450 let results = VectorStore::search(
451 &self.ops,
452 &self.collection,
453 query_vector.into_inner(),
454 limit_u64,
455 filter,
456 )
457 .await?;
458
459 Ok(results
460 .into_iter()
461 .filter_map(|p| SearchHit::from_payload(&p))
462 .collect())
463 }
464
465 #[tracing::instrument(name = "index.store.indexed_files", skip_all)]
471 pub async fn indexed_files(&self) -> Result<Vec<String>> {
472 let rows: Vec<(String,)> =
473 zeph_db::query_as(sql!("SELECT DISTINCT file_path FROM chunk_metadata"))
474 .fetch_all(&self.pool)
475 .await?;
476 Ok(rows.into_iter().map(|(p,)| p).collect())
477 }
478}
479
480impl SearchHit {
481 fn from_payload(point: &zeph_memory::ScoredVectorPoint) -> Option<Self> {
482 let get_str = |key: &str| -> Option<String> {
483 point
484 .payload
485 .get(key)
486 .and_then(serde_json::Value::as_str)
487 .map(ToOwned::to_owned)
488 };
489 let get_usize = |key: &str| -> Option<usize> {
490 point
491 .payload
492 .get(key)
493 .and_then(serde_json::Value::as_i64)
494 .and_then(|v| usize::try_from(v).ok())
495 };
496
497 let language_str = get_str("language")?;
498 let language = crate::languages::Lang::from_id(&language_str)?;
499 Some(Self {
500 code: get_str("code")?,
501 file_path: get_str("file_path")?,
502 line_range: (get_usize("line_start")?, get_usize("line_end")?),
503 score: point.score,
504 node_type: NodeKind::from(get_str("node_type")?),
505 language,
506 entity_name: get_str("entity_name"),
507 scope_chain: get_str("scope_chain").unwrap_or_default(),
508 })
509 }
510}
511
512#[cfg(test)]
513mod tests {
514 use super::*;
515 use zeph_memory::ScoredVectorPoint;
516
517 fn make_scored_point(payload: serde_json::Value, score: f32) -> ScoredVectorPoint {
518 let map = match payload {
519 serde_json::Value::Object(m) => m.into_iter().collect(),
520 _ => std::collections::HashMap::new(),
521 };
522 ScoredVectorPoint {
523 id: "test-id".to_string(),
524 score,
525 payload: map,
526 }
527 }
528
529 #[test]
530 fn search_hit_from_payload_full() {
531 let point = make_scored_point(
532 serde_json::json!({
533 "code": "fn foo() {}",
534 "file_path": "src/lib.rs",
535 "line_start": 10,
536 "line_end": 12,
537 "language": "rust",
538 "node_type": "function_item",
539 "entity_name": "foo",
540 "scope_chain": "mod::foo"
541 }),
542 0.9,
543 );
544 let hit = SearchHit::from_payload(&point).unwrap();
545 assert_eq!(hit.code, "fn foo() {}");
546 assert_eq!(hit.file_path, "src/lib.rs");
547 assert_eq!(hit.line_range, (10, 12));
548 assert!((hit.score - 0.9).abs() < f32::EPSILON);
549 assert_eq!(hit.node_type.as_ref(), "function_item");
550 assert_eq!(hit.language, crate::languages::Lang::Rust);
551 assert_eq!(hit.entity_name, Some("foo".to_string()));
552 assert_eq!(hit.scope_chain, "mod::foo");
553 }
554
555 #[test]
556 fn search_hit_from_payload_no_entity_name() {
557 let point = make_scored_point(
558 serde_json::json!({
559 "code": "struct Bar {}",
560 "file_path": "src/bar.rs",
561 "line_start": 1,
562 "line_end": 3,
563 "language": "rust",
564 "node_type": "struct_item",
565 "scope_chain": ""
566 }),
567 0.7,
568 );
569 let hit = SearchHit::from_payload(&point).unwrap();
570 assert!(hit.entity_name.is_none());
571 assert_eq!(hit.node_type.as_ref(), "struct_item");
572 }
573
574 #[test]
575 fn search_hit_from_payload_missing_required_field_returns_none() {
576 let point = make_scored_point(
578 serde_json::json!({
579 "file_path": "src/lib.rs",
580 "line_start": 1,
581 "line_end": 2,
582 "language": "rust",
583 "node_type": "function_item"
584 }),
585 0.5,
586 );
587 assert!(SearchHit::from_payload(&point).is_none());
588 }
589
590 async fn setup_pool() -> zeph_db::DbPool {
591 zeph_db::DbConfig {
592 url: ":memory:".to_string(),
593 ..Default::default()
594 }
595 .connect()
596 .await
597 .unwrap()
598 }
599
600 #[tokio::test]
601 async fn chunk_exists_returns_false_then_true() {
602 let pool = setup_pool().await;
603
604 let exists = zeph_db::query_as::<_, (i64,)>(sql!(
605 "SELECT COUNT(*) FROM chunk_metadata WHERE content_hash = ?"
606 ))
607 .bind("abc123")
608 .fetch_one(&pool)
609 .await
610 .unwrap();
611 assert_eq!(exists.0, 0);
612
613 zeph_db::query(sql!(
614 "INSERT INTO chunk_metadata \
615 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type) \
616 VALUES (?, ?, ?, ?, ?, ?, ?)"
617 ))
618 .bind("q1")
619 .bind("src/main.rs")
620 .bind("abc123")
621 .bind(1_i64)
622 .bind(10_i64)
623 .bind("rust")
624 .bind("function_item")
625 .execute(&pool)
626 .await
627 .unwrap();
628
629 let exists = zeph_db::query_as::<_, (i64,)>(sql!(
630 "SELECT COUNT(*) FROM chunk_metadata WHERE content_hash = ?"
631 ))
632 .bind("abc123")
633 .fetch_one(&pool)
634 .await
635 .unwrap();
636 assert!(exists.0 > 0);
637 }
638
639 #[tokio::test]
640 async fn remove_file_chunks_cleans_sqlite() {
641 let pool = setup_pool().await;
642
643 for i in 0..3 {
644 zeph_db::query(sql!(
645 "INSERT INTO chunk_metadata \
646 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type) \
647 VALUES (?, ?, ?, ?, ?, ?, ?)"
648 ))
649 .bind(format!("q{i}"))
650 .bind("src/lib.rs")
651 .bind(format!("hash{i}"))
652 .bind(1_i64)
653 .bind(10_i64)
654 .bind("rust")
655 .bind("function_item")
656 .execute(&pool)
657 .await
658 .unwrap();
659 }
660
661 let ids: Vec<(String,)> = zeph_db::query_as(sql!(
662 "SELECT qdrant_id FROM chunk_metadata WHERE file_path = ?"
663 ))
664 .bind("src/lib.rs")
665 .fetch_all(&pool)
666 .await
667 .unwrap();
668 assert_eq!(ids.len(), 3);
669
670 zeph_db::query(sql!("DELETE FROM chunk_metadata WHERE file_path = ?"))
671 .bind("src/lib.rs")
672 .execute(&pool)
673 .await
674 .unwrap();
675
676 let remaining: (i64,) = zeph_db::query_as(sql!(
677 "SELECT COUNT(*) FROM chunk_metadata WHERE file_path = ?"
678 ))
679 .bind("src/lib.rs")
680 .fetch_one(&pool)
681 .await
682 .unwrap();
683 assert_eq!(remaining.0, 0);
684 }
685
686 #[tokio::test]
687 async fn indexed_files_distinct() {
688 let pool = setup_pool().await;
689
690 for (i, path) in ["src/a.rs", "src/b.rs", "src/a.rs"].iter().enumerate() {
691 zeph_db::query(sql!(
692 "INSERT INTO chunk_metadata \
693 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type) \
694 VALUES (?, ?, ?, ?, ?, ?, ?) \
695 ON CONFLICT(qdrant_id) DO UPDATE SET \
696 file_path = excluded.file_path, content_hash = excluded.content_hash, \
697 line_start = excluded.line_start, line_end = excluded.line_end, \
698 language = excluded.language, node_type = excluded.node_type"
699 ))
700 .bind(format!("q{i}"))
701 .bind(path)
702 .bind(format!("hash{i}"))
703 .bind(1_i64)
704 .bind(10_i64)
705 .bind("rust")
706 .bind("function_item")
707 .execute(&pool)
708 .await
709 .unwrap();
710 }
711
712 let rows: Vec<(String,)> =
713 zeph_db::query_as(sql!("SELECT DISTINCT file_path FROM chunk_metadata"))
714 .fetch_all(&pool)
715 .await
716 .unwrap();
717 let files: Vec<String> = rows.into_iter().map(|(p,)| p).collect();
718 assert_eq!(files.len(), 2);
719 assert!(files.contains(&"src/a.rs".to_string()));
720 assert!(files.contains(&"src/b.rs".to_string()));
721 }
722
723 #[tokio::test]
727 async fn upsert_same_file_path_and_hash_is_idempotent() {
728 let pool = setup_pool().await;
729
730 for i in 0..2_u32 {
731 zeph_db::query(sql!(
732 "INSERT INTO chunk_metadata \
733 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type) \
734 VALUES (?, ?, ?, ?, ?, ?, ?) \
735 ON CONFLICT(file_path, content_hash) DO UPDATE SET \
736 qdrant_id = excluded.qdrant_id, \
737 line_start = excluded.line_start, line_end = excluded.line_end, \
738 language = excluded.language, node_type = excluded.node_type, \
739 entity_name = excluded.entity_name"
740 ))
741 .bind(format!("q{i}"))
742 .bind("src/lib.rs")
743 .bind("dedup_hash")
744 .bind(1_i64)
745 .bind(5_i64)
746 .bind("rust")
747 .bind("function_item")
748 .execute(&pool)
749 .await
750 .unwrap();
751 }
752
753 let count: (i64,) = zeph_db::query_as(sql!(
754 "SELECT COUNT(*) FROM chunk_metadata \
755 WHERE file_path = 'src/lib.rs' AND content_hash = 'dedup_hash'"
756 ))
757 .fetch_one(&pool)
758 .await
759 .unwrap();
760
761 assert_eq!(count.0, 1, "duplicate upsert must not produce a second row");
762
763 let qdrant_id: (String,) = zeph_db::query_as(sql!(
765 "SELECT qdrant_id FROM chunk_metadata \
766 WHERE file_path = 'src/lib.rs' AND content_hash = 'dedup_hash'"
767 ))
768 .fetch_one(&pool)
769 .await
770 .unwrap();
771 assert_eq!(
772 qdrant_id.0, "q1",
773 "qdrant_id must reflect the latest upsert"
774 );
775 }
776
777 #[tokio::test]
778 async fn existing_hashes_empty_input_returns_empty_set() {
779 let pool = setup_pool().await;
780 let ops = zeph_memory::QdrantOps::new("http://127.0.0.1:1", None).unwrap();
781 let store = CodeStore::with_ops(ops, pool);
782 let result = store.existing_hashes(&[]).await.unwrap();
783 assert!(result.is_empty());
784 }
785
786 #[tokio::test]
787 async fn existing_hashes_chunking_above_900() {
788 let pool = setup_pool().await;
789
790 for i in 0..901_usize {
792 zeph_db::query(sql!(
793 "INSERT INTO chunk_metadata \
794 (qdrant_id, file_path, content_hash, line_start, line_end, language, node_type) \
795 VALUES (?, ?, ?, ?, ?, ?, ?)"
796 ))
797 .bind(format!("q{i}"))
798 .bind("src/lib.rs")
799 .bind(format!("hash{i:04}"))
800 .bind(1_i64)
801 .bind(2_i64)
802 .bind("rust")
803 .bind("function_item")
804 .execute(&pool)
805 .await
806 .unwrap();
807 }
808
809 let all_hashes: Vec<String> = (0..901).map(|i| format!("hash{i:04}")).collect();
810 let refs: Vec<&str> = all_hashes.iter().map(String::as_str).collect();
811
812 let ops = zeph_memory::QdrantOps::new("http://127.0.0.1:1", None).unwrap();
813 let store = CodeStore::with_ops(ops, pool);
814 let result = store.existing_hashes(&refs).await.unwrap();
815
816 assert_eq!(result.len(), 901);
817 assert!(result.contains("hash0000"));
819 assert!(result.contains("hash0900"));
820 }
821}