source-kb 0.2.2__py3-none-any.whl

core/rag/indexer.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+"""Index builder — writes chunks to ChromaDB vector database.
+
+Migrated from engine/rag/indexer.py. Removes print() calls and global config;
+accepts Config object explicitly. Supports checkpoint resume, chunk change
+detection, batch embedding, and incremental rebuild.
+"""
+
+import gc
+import hashlib
+import json
+import logging
+import random
+import time
+import traceback
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import chromadb
+from filelock import FileLock
+
+from core.rag.chunker import Chunk
+from core.rag.embedder import get_embeddings
+
+if TYPE_CHECKING:
+    from core.config import Config
+
+logger = logging.getLogger(__name__)
+
+BATCH_SIZE = 100
+_PROGRESS_VERSION = 3
+_DELETE_PAGE_SIZE = 5000
+
+
+def _get_client(config: "Config", kb_name: str | None = None):
+    """Get ChromaDB persistent client."""
+    chroma_path = config.chroma_dir(kb_name)
+    chroma_path.mkdir(parents=True, exist_ok=True)
+    return chromadb.PersistentClient(path=str(chroma_path))
+
+
+def _progress_file(config: "Config", kb_name: str) -> Path:
+    """Progress file path for checkpoint resume."""
+    return config.chroma_dir(kb_name) / "index_progress.json"
+
+
+def embed_with_retry(
+    texts: list[str], config: "Config", max_retries: int = 3
+) -> list[list[float]]:
+    """Batch embedding with retry and jitter for transient failures."""
+    for attempt in range(max_retries):
+        try:
+            return get_embeddings(texts, config)
+        except RuntimeError:
+            raise
+        except Exception as e:
+            if attempt < max_retries - 1:
+                wait = 2 ** attempt + random.uniform(0, 1)
+                logger.warning("Embedding retry %d/%d: %s", attempt + 1, max_retries, e)
+                time.sleep(wait)
+            else:
+                raise
+
+
+def _embed_config_key(config: "Config") -> str:
+    """Return a string identifying the current embedding configuration."""
+    return f"{config.embed_backend}:{config.embed_model}"
+
+
+def _chunk_id(source: str, index: int) -> str:
+    """Unified chunk ID: {source}_{index}."""
+    return f"{source}_{index}"
+
+
+def _compute_chunks_hash(chunks: list[Chunk], config: "Config") -> str:
+    """Compute content fingerprint for chunks."""
+    h = hashlib.sha256()
+    h.update(_embed_config_key(config).encode())
+    h.update(str(len(chunks)).encode())
+    for c in chunks:
+        h.update(c.text.encode("utf-8", errors="replace"))
+    return h.hexdigest()
+
+
+def _save_progress(config: "Config", kb_name: str, count: int, chunks_hash: str):
+    """Save indexing progress checkpoint."""
+    pf = _progress_file(config, kb_name)
+    lock = FileLock(str(pf) + ".lock")
+    with lock:
+        pf.write_text(json.dumps({
+            "version": _PROGRESS_VERSION,
+            "kb": kb_name, "count": count,
+            "chunks_hash": chunks_hash, "ts": time.time(),
+        }))
+
+
+def _load_progress(config: "Config", kb_name: str, chunks_hash: str) -> int:
+    """Load progress; returns checkpoint count if hash and version match."""
+    try:
+        pf = _progress_file(config, kb_name)
+        lock = FileLock(str(pf) + ".lock")
+        with lock:
+            data = json.loads(pf.read_text())
+        if (data.get("version") == _PROGRESS_VERSION
+                and data.get("kb") == kb_name
+                and data.get("chunks_hash") == chunks_hash):
+            return data.get("count", 0)
+    except Exception as e:
+        logger.debug("Failed to load progress for %s: %s", kb_name, e)
+    return 0
+
+
+def _delete_chunks_by_source(collection, source: str) -> int:
+    """Delete all chunks matching a specific source value (paginated)."""
+    deleted = 0
+    while True:
+        try:
+            results = collection.get(where={"source": source}, limit=_DELETE_PAGE_SIZE)
+        except Exception as e:
+            logger.debug("delete_chunks_by_source query failed for %s: %s", source, e)
+            break
+        if not results["ids"]:
+            break
+        collection.delete(ids=results["ids"])
+        deleted += len(results["ids"])
+    return deleted
+
+
+def _delete_chunks_by_category(collection, category: str) -> int:
+    """Delete all chunks matching a category (paginated), with prefix fallback."""
+    deleted = 0
+    try:
+        while True:
+            results = collection.get(where={"category": category}, limit=_DELETE_PAGE_SIZE)
+            if not results["ids"]:
+                break
+            collection.delete(ids=results["ids"])
+            deleted += len(results["ids"])
+    except Exception as e:
+        logger.debug("delete_chunks_by_category where-query failed for %s, falling back to prefix scan: %s", category, e)
+        while True:
+            try:
+                all_results = collection.get(limit=_DELETE_PAGE_SIZE, include=["metadatas"])
+            except Exception as e2:
+                logger.debug("prefix scan query failed: %s", e2)
+                break
+            if not all_results["ids"]:
+                break
+            ids_to_delete = [
+                all_results["ids"][i]
+                for i, meta in enumerate(all_results["metadatas"])
+                if meta.get("source", "").startswith(f"{category}/")
+                or meta.get("source", "").startswith(f"{category}\\")
+            ]
+            if not ids_to_delete:
+                break
+            collection.delete(ids=ids_to_delete)
+            deleted += len(ids_to_delete)
+            if len(all_results["ids"]) < _DELETE_PAGE_SIZE:
+                break
+    return deleted
+
+
+def _insert_chunks(
+    collection, chunks: list[Chunk], config: "Config",
+    kb_name: str, start: int, chunks_hash: str,
+) -> bool:
+    """Embed and insert chunks into collection with checkpoint resume."""
+    total = len(chunks)
+    for i in range(start, total, BATCH_SIZE):
+        batch = chunks[i:i + BATCH_SIZE]
+        batch_end = min(i + BATCH_SIZE, total)
+
+        try:
+            texts = [c.text for c in batch]
+            embeddings = embed_with_retry(texts, config)
+            collection.add(
+                ids=[_chunk_id(c.metadata.get("source", "unknown"), j)
+                     for j, c in enumerate(batch, start=i)],
+                embeddings=embeddings,
+                documents=texts,
+                metadatas=[c.metadata for c in batch],
+            )
+        except Exception as e:
+            logger.error("[%s] batch %d-%d failed: %s", kb_name, i, batch_end, e)
+            logger.debug("Traceback:\n%s", traceback.format_exc())
+            _save_progress(config, kb_name, i, chunks_hash)
+            return False
+
+        _save_progress(config, kb_name, batch_end, chunks_hash)
+        gc.collect()
+        logger.info("[%s] %d/%d", kb_name, batch_end, total)
+
+    return True
+
+
+def build_index(
+    chunks: list[Chunk],
+    collection_name: str,
+    config: "Config",
+    kb_name: str = "unknown",
+    module: str | None = None,
+    files: list[str] | None = None,
+):
+    """Build ChromaDB index with unified chunk IDs.
+
+    Modes:
+      - module=None, files=None: full rebuild (delete entire collection)
+      - module="x": module-level rebuild (delete by category, then insert)
+      - module="x", files=["a.md"]: file-level rebuild (delete by source)
+    """
+    total = len(chunks)
+    if total == 0:
+        logger.info("[%s] no chunks to index", kb_name)
+        return
+
+    client = _get_client(config, kb_name)
+
+    if module:
+        _build_incremental(client, collection_name, chunks, config, kb_name, module, files)
+    else:
+        _build_full(client, collection_name, chunks, config, kb_name)
+
+
+def _build_incremental(client, collection_name, chunks, config, kb_name, module, files):
+    """Incremental mode: get or create collection, delete targeted chunks."""
+    total = len(chunks)
+    try:
+        collection = client.get_collection(collection_name)
+    except Exception:
+        collection = client.get_or_create_collection(
+            name=collection_name, metadata={"hnsw:space": "cosine"})
+
+    if files:
+        deleted = 0
+        for fname in files:
+            source_posix = f"{module}/{fname}"
+            deleted += _delete_chunks_by_source(collection, source_posix)
+        logger.info("[%s] deleted %d old chunks for %d file(s)", kb_name, deleted, len(files))
+    else:
+        deleted = _delete_chunks_by_category(collection, module)
+        logger.info("[%s] deleted %d old chunks for module %s", kb_name, deleted, module)
+
+    chunks_hash = _compute_chunks_hash(chunks, config)
+    _insert_chunks(collection, chunks, config, kb_name, 0, chunks_hash)
+    logger.info("[%s] done! inserted %d chunks for %s", kb_name, total, module)
+
+
+def _build_full(client, collection_name, chunks, config, kb_name):
+    """Full rebuild mode with checkpoint resume."""
+    total = len(chunks)
+    chunks_hash = _compute_chunks_hash(chunks, config)
+    saved = _load_progress(config, kb_name, chunks_hash)
+
+    if saved > 0:
+        logger.info("[%s] resuming from %d/%d", kb_name, saved, total)
+        collection = client.get_collection(collection_name)
+        start = saved
+    else:
+        logger.info("[%s] full rebuild...", kb_name)
+        try:
+            client.delete_collection(collection_name)
+        except Exception as e:
+            logger.debug("delete_collection(%s) skipped: %s", collection_name, e)
+        collection = client.create_collection(
+            name=collection_name, metadata={"hnsw:space": "cosine"})
+        start = 0
+
+    if start >= total:
+        logger.info("[%s] already complete: %d/%d", kb_name, start, total)
+        return
+
+    ok = _insert_chunks(collection, chunks, config, kb_name, start, chunks_hash)
+    if ok:
+        logger.info("[%s] done! %d/%d", kb_name, collection.count(), total)
+
+
+def get_collection(collection_name: str, config: "Config", kb_name: str | None = None):
+    """Get an existing collection.
+
+    Raises:
+        RuntimeError: If the collection does not exist.
+    """
+    client = _get_client(config, kb_name)
+    try:
+        return client.get_collection(collection_name)
+    except Exception:
+        raise RuntimeError(
+            f"Collection '{collection_name}' not found. "
+            f"Run index build first."
+        )

core/rag/loader.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+"""Document loader — loads markdown and source files from knowledge base directory.
+
+Migrated from engine/rag/loader.py. Removes print() calls (no CLI I/O in core/).
+"""
+
+import logging
+from pathlib import Path
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Document:
+    """A loaded document with content and metadata."""
+    content: str
+    metadata: dict
+
+
+def load_documents(
+    knowledge_dir: str | Path,
+    module: str | None = None,
+    files: list[str] | None = None,
+) -> list[Document]:
+    """Recursively load .md files from a directory.
+
+    Args:
+        knowledge_dir: Root knowledge directory.
+        module: If set, only load from knowledge_dir/module/.
+        files: If set (requires module), only load these filenames.
+
+    Returns:
+        List of Document objects with content and metadata.
+    """
+    knowledge_dir = Path(knowledge_dir)
+    if not knowledge_dir.exists():
+        logger.warning("Directory does not exist: %s", knowledge_dir)
+        return []
+
+    if module:
+        scan_dir = knowledge_dir / module
+        if not scan_dir.exists():
+            logger.warning("Module directory does not exist: %s", scan_dir)
+            return []
+    else:
+        scan_dir = knowledge_dir
+
+    docs: list[Document] = []
+    for path in sorted(scan_dir.rglob("*.md")):
+        if path.is_symlink():
+            continue
+        if ".meta" in path.parts:
+            continue
+        if files and path.name not in files:
+            continue
+
+        content = _read_file(path)
+        if not content:
+            continue
+
+        rel = path.relative_to(knowledge_dir)
+        category = rel.parts[0] if len(rel.parts) > 1 else "general"
+
+        docs.append(Document(
+            content=content,
+            metadata={
+                "source": rel.as_posix(),
+                "category": category,
+            },
+        ))
+
+    logger.info("Loaded %d documents from %s", len(docs), scan_dir)
+    return docs
+
+
+def _read_file(path: Path) -> str | None:
+    """Read a file, handling encoding automatically."""
+    for enc in ("utf-8", "gbk"):
+        try:
+            text = path.read_text(encoding=enc).strip()
+            if enc != "utf-8":
+                logger.warning("Used %s fallback encoding for: %s", enc, path.name)
+            return text if text else None
+        except (UnicodeDecodeError, UnicodeError):
+            continue
+    logger.error("Skipping undecodable file: %s", path.name)
+    return None

core/rag/retriever.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+"""Retriever — semantic search from ChromaDB vector index.
+
+Migrated from engine/rag/retriever.py. Removes global config dependency;
+accepts Config object explicitly. Returns ranked results above similarity threshold.
+"""
+
+import logging
+from typing import TYPE_CHECKING
+
+from core.rag.embedder import get_embeddings
+from core.rag.indexer import get_collection
+
+if TYPE_CHECKING:
+    from core.config import Config
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve(
+    query: str,
+    config: "Config",
+    top_k: int | None = None,
+    collection_name: str | None = None,
+    kb_name: str | None = None,
+    min_score: float | None = None,
+) -> list[dict]:
+    """Retrieve chunks most relevant to the query.
+
+    Args:
+        query: Query text for semantic search.
+        config: Config object providing retrieval settings.
+        top_k: Maximum number of results to return.
+        collection_name: ChromaDB collection name.
+        kb_name: Knowledge base name (used to locate ChromaDB).
+        min_score: Minimum similarity threshold (cosine similarity).
+
+    Returns:
+        List of dicts with keys: text, metadata, score.
+        Sorted by score descending.
+
+    Raises:
+        ValueError: If collection_name cannot be determined.
+        RuntimeError: If the collection does not exist.
+    """
+    if top_k is None:
+        top_k = config.top_k
+    if min_score is None:
+        min_score = config.similarity_threshold
+
+    if not collection_name:
+        kb_cfg = config.get_kb(kb_name) if kb_name else {}
+        collection_name = kb_cfg.get("collection") if isinstance(kb_cfg, dict) else None
+        if not collection_name:
+            raise ValueError(
+                "collection_name is required — either pass it directly or "
+                "configure 'collection' in the knowledge_base config"
+            )
+
+    collection = get_collection(collection_name, config, kb_name=kb_name)
+    query_embedding = get_embeddings([query], config)[0]
+
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=top_k,
+        include=["documents", "metadatas", "distances"],
+    )
+
+    chunks: list[dict] = []
+    for i in range(len(results["ids"][0])):
+        score = 1 - results["distances"][0][i]  # cosine similarity
+        if score < min_score:
+            continue
+        chunks.append({
+            "text": results["documents"][0][i],
+            "metadata": results["metadatas"][0][i],
+            "score": score,
+        })
+
+    logger.debug("Retrieved %d chunks (query=%r, top_k=%d)", len(chunks), query[:50], top_k)
+    return chunks

core/skeleton/__init__.py

@@ -0,0 +1,11 @@
+"""core.skeleton — skeleton extraction, file classification, split planning, and merge.
+
+Sub-modules:
+    parsers/    — Language-specific skeleton parsers (registry pattern)
+    extract     — Skeleton extraction orchestrator
+    file_list   — File list extraction + coverage check
+    classify    — Rule-based + LLM-assisted file classification
+    split       — SplitConfig + compute_splits + plan_splits
+    metadata    — Global metadata pre-generation
+    merge       — Shard merge (fast) + post-merge refinement
+"""