agent-cli 0.70.5__py3-none-any.whl

agent_cli/rag/_indexing.py

@@ -0,0 +1,226 @@
+"""RAG Indexing Logic."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import datetime
+import logging
+from typing import TYPE_CHECKING
+
+from agent_cli.rag._store import delete_by_file_path, get_all_metadata, upsert_docs
+from agent_cli.rag._utils import chunk_text, get_file_hash, load_document_text, should_ignore_path
+from agent_cli.rag.models import DocMetadata
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from chromadb import Collection
+
+LOGGER = logging.getLogger(__name__)
+
+
+def load_hashes_from_metadata(collection: Collection) -> tuple[dict[str, str], dict[str, float]]:
+    """Rebuild hash and mtime caches from existing DB.
+
+    Returns:
+        Tuple of (file_hashes, file_mtimes) dictionaries.
+
+    """
+    metadatas = get_all_metadata(collection)
+    file_hashes = {}
+    file_mtimes = {}
+    for meta in metadatas:
+        if meta:
+            fp = meta["file_path"]
+            file_hashes[fp] = meta["file_hash"]
+            file_mtimes[fp] = meta["file_mtime"]
+    return file_hashes, file_mtimes
+
+
+def index_file(
+    collection: Collection,
+    docs_folder: Path,
+    file_path: Path,
+    file_hashes: dict[str, str],
+    file_mtimes: dict[str, float],
+) -> bool:
+    """Index or reindex a single file.
+
+    Uses mtime-first checking for performance: only computes hash if mtime changed.
+
+    Returns:
+        True if the file was indexed (changed or new), False otherwise.
+
+    """
+    if not file_path.exists():
+        return False
+    LOGGER.info("  📄 Processing: %s", file_path.name)
+
+    try:
+        relative_path = str(file_path.relative_to(docs_folder))
+        current_mtime = file_path.stat().st_mtime
+
+        # Fast path: mtime unchanged → skip (no hash computation needed)
+        if relative_path in file_mtimes and file_mtimes[relative_path] == current_mtime:
+            return False
+
+        # mtime changed or new file: verify with hash
+        current_hash = get_file_hash(file_path)
+
+        # Hash unchanged (file was touched but not modified) → update mtime, skip
+        if relative_path in file_hashes and file_hashes[relative_path] == current_hash:
+            file_mtimes[relative_path] = current_mtime
+            return False
+
+        # Remove old chunks first (atomic-ish update)
+        remove_file(collection, docs_folder, file_path, file_hashes, file_mtimes)
+
+        # Load and chunk document
+        text = load_document_text(file_path)
+        chunks = chunk_text(text) if text and text.strip() else []
+        if not chunks:
+            return False  # Unsupported, empty, or no chunks
+
+        # Index chunks
+        ids = []
+        documents = []
+        metadatas = []
+
+        timestamp = datetime.datetime.now(datetime.UTC).isoformat()
+
+        for i, chunk in enumerate(chunks):
+            doc_id = f"{relative_path}:chunk:{i}"
+            ids.append(doc_id)
+            documents.append(chunk)
+            metadatas.append(
+                DocMetadata(
+                    source=file_path.name,
+                    file_path=relative_path,
+                    file_type=file_path.suffix,
+                    chunk_id=i,
+                    total_chunks=len(chunks),
+                    indexed_at=timestamp,
+                    file_hash=current_hash,
+                    file_mtime=current_mtime,
+                ),
+            )
+
+        # Upsert to ChromaDB in batches to avoid 502s from large payloads
+        # Use small batch size (10) to avoid overwhelming embedding servers
+        batch_size = 10
+        for i in range(0, len(ids), batch_size):
+            batch_ids = ids[i : i + batch_size]
+            batch_docs = documents[i : i + batch_size]
+            batch_meta = metadatas[i : i + batch_size]
+            upsert_docs(collection, batch_ids, batch_docs, batch_meta)
+
+        # Update tracking
+        file_hashes[relative_path] = current_hash
+        file_mtimes[relative_path] = current_mtime
+
+        LOGGER.info("  ✓ Indexed %s: %d chunks", file_path.name, len(chunks))
+        return True
+
+    except Exception:
+        LOGGER.exception("Failed to index file %s", file_path)
+        return False
+
+
+def remove_file(
+    collection: Collection,
+    docs_folder: Path,
+    file_path: Path,
+    file_hashes: dict[str, str],
+    file_mtimes: dict[str, float],
+) -> bool:
+    """Remove all chunks of a file from index.
+
+    Returns:
+        True if documents were removed (or at least untracked), False otherwise.
+
+    """
+    try:
+        relative_path = str(file_path.relative_to(docs_folder))
+        delete_by_file_path(collection, relative_path)
+
+        # If it was tracked, we consider it "removed"
+        if relative_path in file_hashes:
+            LOGGER.info("  ✓ Removed %s from index", file_path.name)
+            file_hashes.pop(relative_path, None)
+            file_mtimes.pop(relative_path, None)
+            return True
+
+        return False
+    except Exception:
+        LOGGER.exception("Error removing file %s", file_path)
+        return False
+
+
+def initial_index(
+    collection: Collection,
+    docs_folder: Path,
+    file_hashes: dict[str, str],
+    file_mtimes: dict[str, float],
+) -> None:
+    """Index all existing files on startup and remove deleted ones."""
+    LOGGER.info("🔍 Scanning existing files...")
+
+    # Snapshot of what's in the DB currently
+    paths_in_db = set(file_hashes.keys())
+    paths_found_on_disk = set()
+
+    processed_files = []
+    removed_files = []
+
+    # Gather all files first, excluding hidden and common development directories
+    all_files = [
+        p for p in docs_folder.rglob("*") if p.is_file() and not should_ignore_path(p, docs_folder)
+    ]
+
+    # 1. Index Existing Files in Parallel
+    # Use max_workers=4 to match typical local backend parallelism (e.g. llama-server -np 4)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        # Map futures to file paths
+        future_to_file = {
+            executor.submit(index_file, collection, docs_folder, f, file_hashes, file_mtimes): f
+            for f in all_files
+        }
+
+        for future in concurrent.futures.as_completed(future_to_file):
+            file_path = future_to_file[future]
+            try:
+                # Track that we found this file (regardless of index result)
+                rel_path = str(file_path.relative_to(docs_folder))
+                paths_found_on_disk.add(rel_path)
+
+                indexed = future.result()
+                if indexed:
+                    processed_files.append(file_path.name)
+            except Exception:
+                LOGGER.exception("Error processing %s", file_path.name)
+
+    # 2. Clean up Deleted Files
+    # If it's in DB but not found on disk, it was deleted offline.
+    paths_to_remove = paths_in_db - paths_found_on_disk
+
+    if paths_to_remove:
+        LOGGER.info("🧹 Cleaning up %d deleted files found in index...", len(paths_to_remove))
+        for rel_path in paths_to_remove:
+            full_path = docs_folder / rel_path
+            try:
+                if remove_file(collection, docs_folder, full_path, file_hashes, file_mtimes):
+                    removed_files.append(rel_path)
+            except Exception:
+                LOGGER.exception("Error removing stale file %s", rel_path)
+
+    if processed_files:
+        LOGGER.info("🆕 Added/Updated: %s", ", ".join(processed_files))
+
+    if removed_files:
+        LOGGER.info("🗑️ Removed: %s", ", ".join(removed_files))
+
+    LOGGER.info(
+        "✅ Initial scan complete. Indexed/Checked %d files, Removed %d stale files.",
+        len(paths_found_on_disk),
+        len(removed_files),
+    )

agent_cli/rag/_prompt.py

@@ -0,0 +1,30 @@
+"""Centralized prompts for RAG LLM calls."""
+
+RAG_PROMPT_WITH_TOOLS = """
+## Retrieved Documentation
+The following was automatically retrieved based on the user's query:
+
+<retrieved_documents>
+{context}
+</retrieved_documents>
+
+## RAG Instructions
+- Use the retrieved context ONLY if it's relevant to the question
+- If the context is irrelevant, ignore it and answer based on your knowledge
+- When using context, cite sources: [Source: filename]
+- If snippets are insufficient, call read_full_document(file_path) to get full content
+""".strip()
+
+RAG_PROMPT_NO_TOOLS = """
+## Retrieved Documentation
+The following was automatically retrieved based on the user's query:
+
+<retrieved_documents>
+{context}
+</retrieved_documents>
+
+## RAG Instructions
+- Use the retrieved context ONLY if it's relevant to the question
+- If the context is irrelevant, ignore it and answer based on your knowledge
+- When using context, cite sources: [Source: filename]
+""".strip()

agent_cli/rag/_retriever.py

@@ -0,0 +1,156 @@
+"""RAG Retrieval Logic (Functional)."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from agent_cli.core.reranker import OnnxCrossEncoder, predict_relevance
+from agent_cli.rag._store import query_docs
+from agent_cli.rag.models import RagSource, RetrievalResult
+
+if TYPE_CHECKING:
+    from chromadb import Collection
+
+LOGGER = logging.getLogger(__name__)
+
+
+def format_context(
+    ranked: list[tuple[str, dict, float]],
+    source_key: str = "source",
+    path_key: str = "file_path",
+    chunk_key: str = "chunk_id",
+) -> str:
+    """Format ranked documents as XML for context injection.
+
+    Args:
+        ranked: List of (doc, meta, score) tuples from rerank_and_filter().
+        source_key: Metadata key for source name.
+        path_key: Metadata key for file path.
+        chunk_key: Metadata key for chunk ID.
+
+    Returns:
+        XML-formatted context string.
+
+    """
+    if not ranked:
+        return ""
+
+    context_parts = []
+    for i, (doc, meta, score) in enumerate(ranked):
+        source = meta.get(source_key, "unknown")
+        path = meta.get(path_key, meta.get("doc_id", "unknown"))
+        chunk = meta.get(chunk_key, 0)
+        context_parts.append(
+            f'<document index="{i + 1}" source="{source}" '
+            f'path="{path}" chunk="{chunk}" score="{score:.3f}">\n{doc}\n</document>',
+        )
+
+    return "\n".join(context_parts)
+
+
+def rerank_and_filter(
+    reranker: OnnxCrossEncoder,
+    query: str,
+    docs: list[str],
+    metas: list[dict],
+    top_k: int,
+    min_score: float = 0.2,
+) -> list[tuple[str, dict, float]]:
+    """Rerank documents and filter by minimum score.
+
+    Args:
+        reranker: Cross-encoder model for reranking.
+        query: Search query string.
+        docs: List of document texts.
+        metas: List of metadata dicts corresponding to docs.
+        top_k: Maximum number of results to return.
+        min_score: Minimum relevance score threshold.
+
+    Returns:
+        List of (doc, meta, score) tuples, sorted by score descending.
+
+    """
+    if not docs:
+        return []
+
+    # Rerank
+    pairs = [(query, doc) for doc in docs]
+    scores = predict_relevance(reranker, pairs)
+
+    # Sort by score descending
+    ranked_all = sorted(
+        zip(docs, metas, scores, strict=False),
+        key=lambda x: x[2],
+        reverse=True,
+    )
+
+    # Filter by min_score and take top_k
+    ranked = [(d, m, s) for d, m, s in ranked_all if s >= min_score][:top_k]
+
+    # Log retrieval quality
+    filtered_count = len(ranked_all) - len([x for x in ranked_all if x[2] >= min_score])
+    top_score = ranked_all[0][2] if ranked_all else 0.0
+    LOGGER.info(
+        "Retrieval: query_len=%d, candidates=%d, returned=%d, "
+        "top_score=%.3f, min_score=%.3f, filtered=%d",
+        len(query),
+        len(docs),
+        len(ranked),
+        top_score,
+        min_score,
+        filtered_count,
+    )
+
+    return ranked
+
+
+def search_context(
+    collection: Collection,
+    reranker_model: OnnxCrossEncoder,
+    query: str,
+    top_k: int = 3,
+    min_score: float = 0.2,
+) -> RetrievalResult:
+    """Retrieve relevant context for a query using hybrid search.
+
+    Args:
+        collection: ChromaDB collection to search.
+        reranker_model: Cross-encoder model for reranking.
+        query: Search query string.
+        top_k: Maximum number of results to return.
+        min_score: Minimum relevance score threshold. Results below this are filtered out.
+
+    Returns:
+        RetrievalResult with context and sources. Empty if no results meet min_score.
+
+    """
+    # Initial retrieval - fetch more candidates for reranking
+    n_candidates = top_k * 3
+    results = query_docs(collection, query, n_results=n_candidates)
+
+    if not results["documents"] or not results["documents"][0]:
+        return RetrievalResult(context="", sources=[])
+
+    docs = results["documents"][0]
+    metas = results["metadatas"][0]  # type: ignore[index]
+
+    # Rerank and filter
+    ranked = rerank_and_filter(reranker_model, query, docs, metas, top_k, min_score)
+
+    if not ranked:
+        return RetrievalResult(context="", sources=[])
+
+    # Build context and sources
+    context = format_context(ranked)
+    sources = [
+        RagSource(
+            source=meta["source"],
+            path=meta["file_path"],
+            chunk_id=meta["chunk_id"],
+            score=float(score),
+        )
+        for _, meta, score in ranked
+    ]
+
+    return RetrievalResult(context=context, sources=sources)

agent_cli/rag/_store.py

@@ -0,0 +1,48 @@
+"""ChromaDB functional interface."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from agent_cli.core.chroma import delete_where, upsert
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from chromadb import Collection
+
+    from agent_cli.rag.models import DocMetadata
+
+LOGGER = logging.getLogger(__name__)
+
+
+def upsert_docs(
+    collection: Collection,
+    ids: list[str],
+    documents: list[str],
+    metadatas: Sequence[DocMetadata],
+) -> None:
+    """Upsert documents into the collection."""
+    upsert(collection, ids=ids, documents=documents, metadatas=metadatas)
+
+
+def delete_by_file_path(collection: Collection, file_path: str) -> None:
+    """Delete all chunks associated with a file path."""
+    delete_where(collection, {"file_path": file_path})
+
+
+def query_docs(collection: Collection, text: str, n_results: int) -> dict[str, Any]:
+    """Query the collection."""
+    return collection.query(query_texts=[text], n_results=n_results)
+
+
+def get_all_metadata(collection: Collection) -> list[dict[str, Any]]:
+    """Retrieve all metadata from the collection."""
+    result = collection.get(include=["metadatas"])
+    return result.get("metadatas", []) or []  # type: ignore[return-value]
+
+
+def count_docs(collection: Collection) -> int:
+    """Return total number of documents."""
+    return collection.count()

agent_cli/rag/_utils.py

@@ -0,0 +1,218 @@
+"""Utility functions for RAG: Document loading and chunking."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+# Configure logging
+LOGGER = logging.getLogger(__name__)
+
+# Non-hidden directories to ignore (hidden dirs already caught by startswith(".") check)
+DEFAULT_IGNORE_DIRS: frozenset[str] = frozenset(
+    {
+        "__pycache__",
+        "venv",
+        "env",
+        "htmlcov",
+        "node_modules",
+        "build",
+        "dist",
+    },
+)
+
+# Non-hidden files to ignore (hidden files already caught by startswith(".") check)
+DEFAULT_IGNORE_FILES: frozenset[str] = frozenset(
+    {
+        "Thumbs.db",
+    },
+)
+
+
+def should_ignore_path(path: Path, base_folder: Path) -> bool:
+    """Check if a path should be ignored during indexing.
+
+    Ignores:
+    - Any path component starting with '.' (hidden files/dirs)
+    - Common development directories (__pycache__, node_modules, venv, etc.)
+    - .egg-info directories
+    - OS metadata files (Thumbs.db)
+
+    Args:
+        path: The file path to check.
+        base_folder: The base folder for computing relative paths.
+
+    Returns:
+        True if the path should be ignored, False otherwise.
+
+    """
+    rel_parts = path.relative_to(base_folder).parts
+
+    for part in rel_parts:
+        # Hidden files/directories (starting with .)
+        if part.startswith("."):
+            return True
+        # Common ignore directories
+        if part in DEFAULT_IGNORE_DIRS:
+            return True
+        # .egg-info directories
+        if part.endswith(".egg-info"):
+            return True
+
+    # Check specific file patterns
+    return path.name in DEFAULT_IGNORE_FILES
+
+
+# Files to read as plain text directly (fast path)
+TEXT_EXTENSIONS = {
+    ".txt",
+    ".md",
+    ".json",
+    ".py",
+    ".js",
+    ".ts",
+    ".yaml",
+    ".yml",
+    ".rs",
+    ".go",
+    ".c",
+    ".cpp",
+    ".h",
+    ".sh",
+    ".toml",
+    ".rst",
+    ".ini",
+    ".cfg",
+}
+
+# Files to convert using MarkItDown (rich documents)
+MARKITDOWN_EXTENSIONS = {
+    ".pdf",
+    ".docx",
+    ".pptx",
+    ".xlsx",
+    ".html",
+    ".htm",
+    ".csv",
+    ".xml",
+}
+
+SUPPORTED_EXTENSIONS = TEXT_EXTENSIONS | MARKITDOWN_EXTENSIONS
+
+
+def load_document_text(file_path: Path) -> str | None:
+    """Load text from a file path."""
+    suffix = file_path.suffix.lower()
+
+    try:
+        if suffix in TEXT_EXTENSIONS:
+            return file_path.read_text(errors="ignore")
+
+        if suffix in MARKITDOWN_EXTENSIONS:
+            from markitdown import MarkItDown  # noqa: PLC0415
+
+            md = MarkItDown()
+            result = md.convert(str(file_path))
+            return result.text_content
+
+        return None  # Unsupported
+    except Exception:
+        LOGGER.exception("Failed to load %s", file_path)
+        return None
+
+
+# Separators ordered by preference (most semantic first)
+SEPARATORS = ("\n\n", "\n", ". ", ", ", " ")
+
+
+def _find_break_point(text: str, start: int, end: int, min_chunk: int) -> int:
+    """Find a good break point near end, preferring semantic boundaries.
+
+    Searches backwards from end to find the last occurrence of a separator.
+    Only accepts separators that would create a chunk of at least min_chunk size.
+    If none qualify, falls back to the best available earlier separator before
+    finally splitting at the exact end. Returns the position after the separator
+    (so the separator stays with the preceding chunk).
+    """
+    min_pos = start + min_chunk
+    fallback_point = -1
+    for sep in SEPARATORS:
+        pos = text.rfind(sep, start, end)
+        if pos <= start:
+            continue
+        candidate = pos + len(sep)
+        if pos >= min_pos:
+            return candidate
+        fallback_point = max(fallback_point, candidate)
+    if fallback_point != -1:
+        return fallback_point
+    # No separator found at acceptable position, break at end (character-level split)
+    return end
+
+
+def chunk_text(text: str, chunk_size: int = 1200, overlap: int = 200) -> list[str]:
+    r"""Split text into chunks, preferring semantic boundaries.
+
+    Strategy:
+    1. Slice the original text directly (no split/join, so no char loss)
+    2. Find break points at separators: \n\n, \n, ". ", ", ", " "
+    3. Fall back to character-level breaks when no separator found
+    4. Overlap by starting next chunk earlier in the text
+
+    Args:
+        text: The text to chunk.
+        chunk_size: Maximum chunk size in characters (default 1200, ~300 words).
+        overlap: Overlap between chunks in characters for context continuity.
+
+    Returns:
+        List of text chunks.
+
+    Raises:
+        ValueError: If chunk_size <= 0 or overlap >= chunk_size.
+
+    """
+    if chunk_size <= 0:
+        msg = f"chunk_size must be positive, got {chunk_size}"
+        raise ValueError(msg)
+    if overlap >= chunk_size:
+        msg = f"overlap ({overlap}) must be less than chunk_size ({chunk_size})"
+        raise ValueError(msg)
+
+    if not text or not text.strip():
+        return []
+
+    text = text.strip()
+    if len(text) <= chunk_size:
+        return [text]
+
+    # Only accept separators that use at least half the chunk budget
+    min_chunk = chunk_size // 2
+
+    chunks = []
+    start = 0
+
+    while start < len(text):
+        end = start + chunk_size
+
+        if end >= len(text):
+            # Last chunk - take everything remaining
+            chunks.append(text[start:])
+            break
+
+        # Find a good break point
+        break_point = _find_break_point(text, start, end, min_chunk)
+        chunks.append(text[start:break_point])
+
+        # Next chunk starts with overlap (but must make progress)
+        start = max(start + 1, break_point - overlap)
+
+    return chunks
+
+
+def get_file_hash(file_path: Path) -> str:
+    """Get hash of file content."""
+    return hashlib.md5(file_path.read_bytes(), usedforsecurity=False).hexdigest()