Merge pull request #6 from Eilodon/ci-dual-layer-embedding-2l6da4

Add Layer-2 code-body chunk embeddings for semantic search

Author: Eilodon

README.md

@@ -26,9 +26,16 @@ tính graph metrics (coreness/hubs), và phục vụ qua SQLite FTS5 + semantic
 4. **Graph metrics** — `coreness` (k-core, O(V+E)) và `is_hub` để AI biết đâu là lõi hệ thống.
 5. **Incremental watcher** — hash-diff chỉ re-parse file đổi; call graph rebuild từ `call_sites` đã lưu trong DB. Quá trình parse được song song hoá (`rayon`). Khi `ci serve` khởi động và đã có index cũ, tự động chạy **incremental reindex** thay vì full index — giảm thời gian warm-up. Debounce 500ms, lọc bỏ noise (`.codeindex/`, `target/`, v.v.).
 6. **FTS5 search** — full-text search native qua SQLite triggers, BM25 dual-column.
-7. **Semantic search (Bật theo mặc định trong config)** — static code embeddings (`model2vec-rs` + `sqlite-vec`),
-   fuse với FTS bằng Reciprocal Rank Fusion (tỉ lệ 1.5x FTS / 1.0x Vector). Được kiểm soát bởi
-   `semantic_search.enabled` trong `config.json` (mặc định `true`). Tự động cấu hình trên bản build native.
+7. **Semantic search — 2 tầng (Bật theo mặc định trong config)** — static code embeddings
+   (`model2vec-rs` + `sqlite-vec`), fuse với FTS bằng Reciprocal Rank Fusion (tỉ lệ 1.5x FTS / 1.0x
+   mỗi tầng semantic). Tầng 1 embed *symbol identity* (tên + signature + docstring); Tầng 2
+   (`indexer::chunker`) embed *code body* thực tế — toàn bộ thân hàm nếu ≤30 dòng, sliding window
+   30 dòng/stride 20 nếu dài hơn, cộng với các đoạn code nằm giữa các symbol (module scaffolding,
+   field declarations) — nên một query chỉ khớp từ vựng *bên trong* thân hàm (một tên thư viện, một
+   biến, một idiom) vẫn tìm ra kết quả dù tên/docstring của symbol không chứa từ đó. `kind=semantic`
+   tự fuse 2 tầng nội bộ; `kind=hybrid` fuse cả 3 (FTS + Tầng 1 + Tầng 2) trong một lượt RRF phẳng.
+   Được kiểm soát bởi `semantic_search.enabled` trong `config.json` (mặc định `true`). Tự động cấu
+   hình trên bản build native.
 8. **`edges_ready` gating** — tool báo trung thực trạng thái index (`scanning → parsing →
    building_edges → ready`); agent không tin nhầm graph khi chưa build xong.
 
@@ -71,8 +78,10 @@ cargo build -p ci-cli # Đã bao gồm embeddings
 ```
 
 Model mặc định `minishlab/potion-code-16M` (256-dim, static code embeddings, pure-Rust, không
-ONNX). `search(kind="semantic")` và `kind="hybrid"` (RRF: FTS + vector) sẽ hoạt động; khi tắt,
-chúng degrade về FTS.
+ONNX) — dùng chung cho cả 2 tầng. Tầng 1 (`embedding_vecs`) index tên/signature/docstring; Tầng 2
+(`code_chunk_vecs`) index các đoạn code body thực tế (bảng quan hệ `code_chunks` lưu text/dòng, luôn
+được tạo; chỉ được embed khi feature `embeddings` bật). `search(kind="semantic")` và `kind="hybrid"`
+(RRF: FTS + vector) sẽ hoạt động; khi tắt, chúng degrade về FTS.
 
 > Lưu ý: feature `embeddings` kéo thêm dependency (tokenizers/TLS). Binary musl tĩnh phân phối
 > ở Phase IV build **không** bật feature này để giữ kích thước tối thiểu.

crates/ci-cli/src/main.rs

@@ -127,7 +127,12 @@ async fn main() -> Result<()> {
                     Ok(embedder) => {
                         ci_core::embedding::create_embedding_table(&conn, semantic.dimensions)?;
                         let n = ci_core::embedding::embed_pending(&conn, &embedder)?;
-                        println!(" {n} symbols embedded.");
+                        ci_core::embedding::create_chunk_embedding_table(
+                            &conn,
+                            semantic.dimensions,
+                        )?;
+                        let nc = ci_core::embedding::embed_pending_chunks(&conn, &embedder)?;
+                        println!(" {n} symbols, {nc} code chunks embedded.");
                     }
                     Err(e) => eprintln!("\nEmbeddings skipped: {e}"),
                 }

crates/ci-core/src/db/schema.rs

@@ -88,6 +88,26 @@ CREATE TABLE IF NOT EXISTS call_sites (
 );
 CREATE INDEX IF NOT EXISTS idx_call_sites_from   ON call_sites(from_path);
 CREATE INDEX IF NOT EXISTS idx_call_sites_callee ON call_sites(callee_name);
+
+-- Semantic search Layer 2: raw code-body slices (whole short bodies, or a
+-- sliding window over longer ones — see indexer::chunker), embedded alongside
+-- Layer 1's symbol-identity (name+signature+docstring) vectors so a query
+-- matching only implementation vocabulary (e.g. a library name used inside a
+-- function body) still has something to match against. Always created —
+-- populated only when the `embeddings` feature is enabled at build time; the
+-- companion `code_chunk_vecs` vec0 table lives in embedding.rs (needs the
+-- sqlite-vec extension registered and a runtime-configured dimension, so it
+-- can't be part of this static schema).
+CREATE TABLE IF NOT EXISTS code_chunks (
+    id         INTEGER PRIMARY KEY AUTOINCREMENT,
+    path       TEXT NOT NULL,
+    line_start INTEGER NOT NULL,
+    line_end   INTEGER NOT NULL,
+    chunk_text TEXT NOT NULL,
+    symbol_qn  TEXT,
+    file_hash  TEXT NOT NULL DEFAULT ''
+);
+CREATE INDEX IF NOT EXISTS idx_code_chunks_path ON code_chunks(path);
 ";
 
 const FTS5_SQL: &str = "
@@ -206,6 +226,46 @@ mod tests {
         assert_eq!(count, 0);
     }
 
+    #[test]
+    fn test_code_chunks_table() {
+        let conn = Connection::open_in_memory().unwrap();
+        init_db(&conn).unwrap();
+
+        let table_count: i64 = conn
+            .query_row(
+                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='code_chunks'",
+                [],
+                |r| r.get(0),
+            )
+            .unwrap();
+        assert_eq!(table_count, 1);
+
+        conn.execute(
+            "INSERT INTO code_chunks (path, line_start, line_end, chunk_text, symbol_qn, file_hash) \
+             VALUES ('a.py', 1, 3, 'def f():\n    pass', 'a.py::f', 'deadbeef')",
+            [],
+        )
+        .unwrap();
+
+        let (path, symbol_qn): (String, Option<String>) = conn
+            .query_row(
+                "SELECT path, symbol_qn FROM code_chunks WHERE line_start = 1",
+                [],
+                |r| Ok((r.get(0)?, r.get(1)?)),
+            )
+            .unwrap();
+        assert_eq!(path, "a.py");
+        assert_eq!(symbol_qn.as_deref(), Some("a.py::f"));
+
+        // symbol_qn is nullable — gap chunks have no enclosing symbol.
+        conn.execute(
+            "INSERT INTO code_chunks (path, line_start, line_end, chunk_text, file_hash) \
+             VALUES ('a.py', 4, 4, '', 'deadbeef')",
+            [],
+        )
+        .unwrap();
+    }
+
     #[test]
     fn test_symbol_metrics_history_table() {
         let conn = Connection::open_in_memory().unwrap();

crates/ci-core/src/embedding.rs

@@ -11,7 +11,11 @@ use rusqlite::Connection;
 /// True when the crate was built with the `embeddings` feature.
 pub const ENABLED: bool = cfg!(feature = "embeddings");
 
-/// The text embedded for a symbol: name + signature + docstring.
+/// The text embedded for a symbol: name + signature + docstring. This is
+/// Layer 1 of semantic search — *symbol identity*. Layer 2 (`code_chunks` /
+/// `code_chunk_vecs`, populated by `indexer::chunker`) embeds the raw code
+/// body instead, so a query that only matches implementation vocabulary (not
+/// reflected in the name/signature/docstring) still has something to match.
 pub fn symbol_doc(name: &str, signature: &str, docstring: &str) -> String {
     let mut s = String::with_capacity(name.len() + signature.len() + docstring.len() + 2);
     s.push_str(name);
@@ -55,6 +59,18 @@ mod imp {
         ))
     }
 
+    /// Create the Layer-2 KNN table for `dim`-dimensional code-chunk vectors
+    /// (idempotent). Separate from `embedding_vecs` — chunk ids and symbol ids
+    /// are unrelated key spaces.
+    pub fn create_chunk_embedding_table(conn: &Connection, dim: usize) -> rusqlite::Result<()> {
+        conn.execute_batch(&format!(
+            "CREATE VIRTUAL TABLE IF NOT EXISTS code_chunk_vecs USING vec0(
+                chunk_id INTEGER PRIMARY KEY,
+                embedding FLOAT[{dim}] distance_metric=cosine
+            );"
+        ))
+    }
+
     /// A loaded static embedding model.
     pub struct Embedder {
         model: StaticModel,
@@ -99,6 +115,18 @@ mod imp {
         Ok(())
     }
 
+    pub fn store_chunk_embedding(
+        conn: &Connection,
+        chunk_id: i64,
+        vec: &[f32],
+    ) -> rusqlite::Result<()> {
+        conn.execute(
+            "INSERT OR REPLACE INTO code_chunk_vecs(chunk_id, embedding) VALUES (?1, ?2)",
+            rusqlite::params![chunk_id, vec_to_blob(vec)],
+        )?;
+        Ok(())
+    }
+
     /// Embed every symbol that has no embedding yet; returns how many were added.
     pub fn embed_pending(conn: &Connection, embedder: &Embedder) -> rusqlite::Result<usize> {
         let rows: Vec<(i64, String, String, String)> = {
@@ -136,6 +164,65 @@ mod imp {
             .collect::<rusqlite::Result<Vec<_>>>()?;
         Ok(rows)
     }
+
+    /// Remove `code_chunk_vecs` rows whose chunk no longer exists in
+    /// `code_chunks` (file changed or was deleted since the vector was
+    /// written). Returns how many rows were pruned. Unlike `symbols` —
+    /// `code_chunks` rows are always deleted-and-reinserted as a unit per file
+    /// (see `indexer::pipeline::remove_file_rows`), so their ids never survive
+    /// a reindex of that file; without this, stale orphans would accumulate
+    /// forever and could crowd out real matches in `knn_chunks` (a KNN query
+    /// has no way to know a returned id is dangling before doing this exact
+    /// lookup).
+    pub fn prune_orphaned_chunk_vecs(conn: &Connection) -> rusqlite::Result<usize> {
+        conn.execute(
+            "DELETE FROM code_chunk_vecs WHERE chunk_id NOT IN (SELECT id FROM code_chunks)",
+            [],
+        )
+    }
+
+    /// Embed every Layer-2 code chunk that has no embedding yet; returns how
+    /// many were added. Prunes orphaned vectors first — see
+    /// `prune_orphaned_chunk_vecs`.
+    pub fn embed_pending_chunks(conn: &Connection, embedder: &Embedder) -> rusqlite::Result<usize> {
+        prune_orphaned_chunk_vecs(conn)?;
+
+        let rows: Vec<(i64, String)> = {
+            let mut stmt = conn.prepare(
+                "SELECT id, chunk_text FROM code_chunks \
+                 WHERE id NOT IN (SELECT chunk_id FROM code_chunk_vecs)",
+            )?;
+            stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
+                .collect::<rusqlite::Result<Vec<_>>>()?
+        };
+        if rows.is_empty() {
+            return Ok(0);
+        }
+        let texts: Vec<String> = rows.iter().map(|(_, text)| text.clone()).collect();
+        let vecs = embedder.embed_batch(&texts);
+        for ((id, _), v) in rows.iter().zip(vecs.iter()) {
+            store_chunk_embedding(conn, *id, v)?;
+        }
+        Ok(rows.len())
+    }
+
+    /// Nearest `k` chunk ids to `query` by cosine distance (ascending).
+    pub fn knn_chunks(
+        conn: &Connection,
+        query: &[f32],
+        k: usize,
+    ) -> rusqlite::Result<Vec<(i64, f64)>> {
+        let mut stmt = conn.prepare(
+            "SELECT chunk_id, distance FROM code_chunk_vecs \
+             WHERE embedding MATCH ?1 AND k = ?2 ORDER BY distance",
+        )?;
+        let rows = stmt
+            .query_map(rusqlite::params![vec_to_blob(query), k as i64], |r| {
+                Ok((r.get::<_, i64>(0)?, r.get::<_, f64>(1)?))
+            })?
+            .collect::<rusqlite::Result<Vec<_>>>()?;
+        Ok(rows)
+    }
 }
 
 // ---------------------------------------------------------------------------
@@ -151,6 +238,10 @@ mod imp {
         Ok(())
     }
 
+    pub fn create_chunk_embedding_table(_conn: &Connection, _dim: usize) -> rusqlite::Result<()> {
+        Ok(())
+    }
+
     /// Stub embedder — `load` always fails, so callers keep `None` and degrade.
     pub struct Embedder;
 
@@ -178,10 +269,24 @@ mod imp {
     pub fn knn(_c: &Connection, _q: &[f32], _k: usize) -> rusqlite::Result<Vec<(i64, f64)>> {
         Ok(Vec::new())
     }
+    pub fn store_chunk_embedding(_c: &Connection, _id: i64, _v: &[f32]) -> rusqlite::Result<()> {
+        Ok(())
+    }
+    pub fn prune_orphaned_chunk_vecs(_c: &Connection) -> rusqlite::Result<usize> {
+        Ok(0)
+    }
+    pub fn embed_pending_chunks(_c: &Connection, _e: &Embedder) -> rusqlite::Result<usize> {
+        Ok(0)
+    }
+    pub fn knn_chunks(_c: &Connection, _q: &[f32], _k: usize) -> rusqlite::Result<Vec<(i64, f64)>> {
+        Ok(Vec::new())
+    }
 }
 
 pub use imp::{
-    Embedder, create_embedding_table, embed_pending, knn, register_extension, store_embedding,
+    Embedder, create_chunk_embedding_table, create_embedding_table, embed_pending,
+    embed_pending_chunks, knn, knn_chunks, prune_orphaned_chunk_vecs, register_extension,
+    store_chunk_embedding, store_embedding,
 };
 
 #[cfg(test)]
@@ -216,6 +321,53 @@ mod tests {
         assert_eq!(hits[0].0, 2, "nearest should be id 2");
     }
 
+    #[cfg(feature = "embeddings")]
+    #[test]
+    fn vec0_knn_chunks_with_synthetic_vectors() {
+        use rusqlite::Connection;
+        register_extension();
+        let conn = Connection::open_in_memory().unwrap();
+        crate::db::schema::init_db(&conn).unwrap();
+        create_chunk_embedding_table(&conn, 3).unwrap();
+
+        // Layer-2 chunk vectors live in their own table/key-space from symbols.
+        store_chunk_embedding(&conn, 10, &[1.0, 0.0, 0.0]).unwrap();
+        store_chunk_embedding(&conn, 20, &[0.0, 1.0, 0.0]).unwrap();
+        store_chunk_embedding(&conn, 30, &[0.0, 0.0, 1.0]).unwrap();
+
+        let hits = knn_chunks(&conn, &[0.0, 0.0, 0.9], 2).unwrap();
+        assert_eq!(hits.len(), 2);
+        assert_eq!(hits[0].0, 30, "nearest should be chunk id 30");
+    }
+
+    #[cfg(feature = "embeddings")]
+    #[test]
+    fn prune_orphaned_chunk_vecs_removes_only_dangling_rows() {
+        use rusqlite::Connection;
+        register_extension();
+        let conn = Connection::open_in_memory().unwrap();
+        crate::db::schema::init_db(&conn).unwrap();
+        create_chunk_embedding_table(&conn, 3).unwrap();
+
+        // id 2 has a matching code_chunks row; id 1 is an orphan (e.g. left
+        // over from a file that was since reindexed with new chunk ids).
+        conn.execute(
+            "INSERT INTO code_chunks (id, path, line_start, line_end, chunk_text, file_hash) \
+             VALUES (2, 'a.py', 1, 1, 'pass', '')",
+            [],
+        )
+        .unwrap();
+        store_chunk_embedding(&conn, 1, &[1.0, 0.0, 0.0]).unwrap();
+        store_chunk_embedding(&conn, 2, &[0.0, 1.0, 0.0]).unwrap();
+
+        let pruned = prune_orphaned_chunk_vecs(&conn).unwrap();
+        assert_eq!(pruned, 1, "exactly the dangling id-1 row must be pruned");
+
+        let hits = knn_chunks(&conn, &[0.0, 1.0, 0.0], 10).unwrap();
+        assert_eq!(hits.len(), 1);
+        assert_eq!(hits[0].0, 2);
+    }
+
     /// KNN latency benchmark: 100k synthetic 256-dim vectors, topK=10.
     /// Run with: cargo test -p ci-core --features embeddings -- --ignored --nocapture bench_knn_latency
     #[cfg(feature = "embeddings")]