diff --git a/tests/workspace/test_default_indexer.py b/tests/workspace/test_default_indexer.py new file mode 100644 index 0000000000..6a2bb02b8f --- /dev/null +++ b/tests/workspace/test_default_indexer.py @@ -0,0 +1,85 @@ +"""Tests for DefaultIndexer — verifies it satisfies BaseIndexer contract.""" + +from __future__ import annotations + +from workspace.base import BaseIndexer +from workspace.default import DefaultIndexer +from workspace.types import IndexSummary, SearchResult + + +def test_default_indexer_is_base_indexer_subclass(): + assert issubclass(DefaultIndexer, BaseIndexer) + + +def test_default_indexer_indexes_and_searches(make_workspace_config, write_file): + cfg = make_workspace_config() + write_file( + cfg.workspace_root / "docs" / "hello.md", "# Hello\n\nWorld of workspace.\n" + ) + + indexer = DefaultIndexer(cfg) + summary = indexer.index() + + assert isinstance(summary, IndexSummary) + assert summary.files_indexed == 1 + assert summary.files_errored == 0 + + results = indexer.search("workspace") + assert isinstance(results, list) + assert len(results) > 0 + assert all(isinstance(r, SearchResult) for r in results) + + +def test_default_indexer_search_respects_limit(make_workspace_config, write_file): + cfg = make_workspace_config() + for i in range(5): + write_file( + cfg.workspace_root / "docs" / f"doc{i}.md", + f"# Doc {i}\n\nThis document talks about testing limit param.\n", + ) + + indexer = DefaultIndexer(cfg) + indexer.index() + + results = indexer.search("document", limit=2) + assert len(results) <= 2 + + +def test_default_indexer_status_returns_dict(make_workspace_config, write_file): + cfg = make_workspace_config() + write_file(cfg.workspace_root / "docs" / "a.md", "# A\n\nContent.\n") + + indexer = DefaultIndexer(cfg) + indexer.index() + + status = indexer.status() + assert isinstance(status, dict) + assert "file_count" in status + assert "chunk_count" in status + assert "db_path" in status + + +def test_default_indexer_index_is_idempotent(make_workspace_config, write_file): + cfg = make_workspace_config() + write_file(cfg.workspace_root / "docs" / "a.md", "# A\n\nContent A.\n") + + indexer = DefaultIndexer(cfg) + first = indexer.index() + assert first.files_indexed == 1 + + second = indexer.index() + assert second.files_indexed == 0 + assert second.files_skipped >= 1 + + +def test_default_indexer_progress_callback(make_workspace_config, write_file): + cfg = make_workspace_config() + write_file(cfg.workspace_root / "docs" / "a.md", "# A\n\nContent.\n") + + calls = [] + indexer = DefaultIndexer(cfg) + indexer.index(progress=lambda cur, total, path: calls.append((cur, total, path))) + + assert len(calls) > 0 + assert calls[0][0] == 1 + assert calls[0][1] >= 1 diff --git a/workspace/default.py b/workspace/default.py new file mode 100644 index 0000000000..e01e35dc3e --- /dev/null +++ b/workspace/default.py @@ -0,0 +1,546 @@ +"""DefaultIndexer — built-in Chonkie + SQLite FTS5 workspace backend.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import logging +import re +import time +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Literal + +from workspace.base import BaseIndexer, ProgressCallback +from workspace.config import ChunkingConfig, WorkspaceConfig +from workspace.constants import ( + CHUNKING_PLAN_VERSION, + CODE_SUFFIXES, + MARKDOWN_SUFFIXES, + WORKSPACE_SUBDIRS, + get_index_dir, + resolve_path_prefix, +) +from workspace.files import discover_workspace_files, seed_hermesignore +from workspace.store import SQLiteFTS5Store +from workspace.types import ( + ChunkRecord, + FileRecord, + IndexingError, + IndexSummary, + SearchResult, +) + +PipelineKind = Literal["markdown", "code", "plain"] + +log = logging.getLogger(__name__) + +_replace = dataclasses.replace + +_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE) +_NEWLINE_RE = re.compile(r"\n") +_MAX_ERRORS = 50 + + +class DefaultIndexer(BaseIndexer): + def __init__(self, config: WorkspaceConfig) -> None: + self._config = config + + def index(self, *, progress: ProgressCallback | None = None) -> IndexSummary: + self._require_chonkie() + + start = time.monotonic() + self._ensure_workspace_dirs() + config_sig = self._config_signature() + + files_indexed = 0 + files_skipped = 0 + files_errored = 0 + chunks_created = 0 + errors: list[IndexingError] = [] + + discovery = discover_workspace_files(self._config) + files_skipped += discovery.filtered_count + all_files = discovery.files + total = len(all_files) + disk_paths: set[str] = set() + + pipelines = self._build_pipelines(self._config.knowledgebase.chunking) + + with SQLiteFTS5Store(self._config.workspace_root) as store: + for i, (root_path, file_path) in enumerate(all_files): + abs_path = str(file_path.resolve()) + disk_paths.add(abs_path) + write_started = False + + if progress: + progress(i + 1, total, abs_path) + + try: + content_hash = _file_hash(file_path) + existing = store.get_file_record(abs_path) + if ( + existing + and existing.content_hash == content_hash + and existing.config_signature == config_sig + ): + files_skipped += 1 + continue + + text = _read_file_text(file_path) + if text is None: + files_errored += 1 + _append_error( + errors, + IndexingError( + path=abs_path, + stage="read", + error_type="EncodingError", + message="Could not decode file with sufficient confidence", + ), + ) + continue + + if not text.strip(): + files_skipped += 1 + continue + + suffix = file_path.suffix.lower() + chunk_records = self._process_file( + abs_path, text, suffix, pipelines + ) + + stat = file_path.stat() + record = FileRecord( + abs_path=abs_path, + root_path=root_path, + content_hash=content_hash, + config_signature=config_sig, + size_bytes=stat.st_size, + modified_at=datetime.fromtimestamp( + stat.st_mtime, tz=timezone.utc + ).isoformat(), + indexed_at=datetime.now(tz=timezone.utc).isoformat(), + chunk_count=len(chunk_records), + ) + + store.conn.execute("SAVEPOINT workspace_file_update") + write_started = True + store.delete_chunks_for_file(abs_path) + store.upsert_file(record) + if chunk_records: + store.insert_chunks(chunk_records) + store.conn.execute("RELEASE SAVEPOINT workspace_file_update") + store.commit() + write_started = False + + files_indexed += 1 + chunks_created += len(chunk_records) + + except Exception as exc: + if write_started: + try: + store.conn.execute( + "ROLLBACK TO SAVEPOINT workspace_file_update" + ) + store.conn.execute( + "RELEASE SAVEPOINT workspace_file_update" + ) + except Exception: + log.warning( + "Failed to roll back workspace update for %s", + abs_path, + exc_info=True, + ) + files_errored += 1 + stage = "read" if isinstance(exc, FileNotFoundError) else "store" + _append_error( + errors, + IndexingError( + path=abs_path, + stage=stage, + error_type=type(exc).__name__, + message=str(exc), + ), + ) + log.warning("Failed to index %s: %s", abs_path, exc, exc_info=True) + continue + + if discovery.complete: + pruned = _prune_stale(store, disk_paths) + else: + pruned = 0 + log.warning( + "Workspace discovery was incomplete; skipping stale prune for this run" + ) + store.commit() + + elapsed = time.monotonic() - start + return IndexSummary( + files_indexed=files_indexed, + files_skipped=files_skipped, + files_pruned=pruned, + files_errored=files_errored, + chunks_created=chunks_created, + duration_seconds=elapsed, + errors=errors, + errors_truncated=files_errored > _MAX_ERRORS, + ) + + def search( + self, + query: str, + *, + limit: int = 20, + path_prefix: str | None = None, + file_glob: str | None = None, + ) -> list[SearchResult]: + if limit is None: + limit = self._config.knowledgebase.search.default_limit + + resolved_prefix = resolve_path_prefix(path_prefix) + + with SQLiteFTS5Store(self._config.workspace_root) as store: + return store.search( + query, + limit=limit, + path_prefix=resolved_prefix, + file_glob=file_glob, + ) + + def status(self) -> dict[str, Any]: + with SQLiteFTS5Store(self._config.workspace_root) as store: + return store.status() + + def _build_pipelines(self, ch: ChunkingConfig) -> dict[PipelineKind, Any]: + from chonkie import Pipeline + + overlap_kwargs = dict( + tokenizer="word", + context_size=ch.overlap, + mode="token", + method="suffix", + merge=False, + ) + return { + "markdown": ( + Pipeline() + .process_with("markdown", tokenizer="word") + .chunk_with("recursive", tokenizer="word", chunk_size=ch.chunk_size) + .refine_with("overlap", **overlap_kwargs) + ), + "code": ( + Pipeline() + .chunk_with( + "code", + tokenizer="word", + chunk_size=ch.chunk_size, + language="auto", + ) + .refine_with("overlap", **overlap_kwargs) + ), + "plain": ( + Pipeline() + .chunk_with("recursive", tokenizer="word", chunk_size=ch.chunk_size) + .refine_with("overlap", **overlap_kwargs) + ), + } + + def _process_file( + self, + abs_path: str, + text: str, + suffix: str, + pipelines: dict[PipelineKind, Any], + ) -> list[ChunkRecord]: + if suffix in MARKDOWN_SUFFIXES: + return self._process_markdown(abs_path, text, pipelines) + elif suffix in CODE_SUFFIXES: + return self._process_code(abs_path, text, pipelines) + else: + return self._process_plain(abs_path, text, pipelines) + + def _process_markdown( + self, + abs_path: str, + text: str, + pipelines: dict[PipelineKind, Any], + ) -> list[ChunkRecord]: + from chonkie.types import MarkdownDocument + + result = pipelines["markdown"].run(texts=text) + assert isinstance(result, MarkdownDocument), ( + f"markdown pipeline returned {type(result).__name__}" + ) + doc = result + + headings = _scan_headings(text) + line_offsets = _build_line_offsets(text) + candidates: list[ChunkRecord] = [] + + for chunk in doc.chunks: + if not chunk.text.strip(): + continue + sc, ec = chunk.start_index, chunk.end_index + candidates.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=0, + content=chunk.text, + token_count=chunk.token_count, + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=_nearest_heading(headings, sc), + kind="markdown_text", + context=chunk.context, + ) + ) + + for code in doc.code: + if not code.content.strip(): + continue + sc, ec = code.start_index, code.end_index + metadata = ( + json.dumps({"language": code.language}) if code.language else None + ) + candidates.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=0, + content=code.content, + token_count=len(code.content.split()), + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=_nearest_heading(headings, sc), + kind="markdown_code", + chunk_metadata=metadata, + ) + ) + + for table in doc.tables: + if not table.content.strip(): + continue + sc, ec = table.start_index, table.end_index + candidates.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=0, + content=table.content, + token_count=len(table.content.split()), + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=_nearest_heading(headings, sc), + kind="markdown_table", + ) + ) + + for image in doc.images: + if not image.alias: + continue + sc, ec = image.start_index, image.end_index + candidates.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=0, + content=image.alias, + token_count=len(image.alias.split()), + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=_nearest_heading(headings, sc), + kind="markdown_image", + ) + ) + + candidates.sort(key=lambda c: c.start_char) + return [_replace(c, chunk_index=i) for i, c in enumerate(candidates)] + + def _process_code( + self, + abs_path: str, + text: str, + pipelines: dict[PipelineKind, Any], + ) -> list[ChunkRecord]: + from chonkie.types import Document + + result = pipelines["code"].run(texts=text) + assert isinstance(result, Document), ( + f"code pipeline returned {type(result).__name__}" + ) + doc = result + line_offsets = _build_line_offsets(text) + records: list[ChunkRecord] = [] + for i, chunk in enumerate(doc.chunks): + sc, ec = chunk.start_index, chunk.end_index + records.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=i, + content=chunk.text, + token_count=chunk.token_count, + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=None, + kind="code", + chunk_metadata=None, + context=chunk.context, + ) + ) + return records + + def _process_plain( + self, + abs_path: str, + text: str, + pipelines: dict[PipelineKind, Any], + ) -> list[ChunkRecord]: + from chonkie.types import Document + + result = pipelines["plain"].run(texts=text) + assert isinstance(result, Document), ( + f"plain pipeline returned {type(result).__name__}" + ) + doc = result + line_offsets = _build_line_offsets(text) + records: list[ChunkRecord] = [] + for i, chunk in enumerate(doc.chunks): + sc, ec = chunk.start_index, chunk.end_index + records.append( + ChunkRecord( + chunk_id=_make_id(), + abs_path=abs_path, + chunk_index=i, + content=chunk.text, + token_count=chunk.token_count, + start_line=_offset_to_line(line_offsets, sc), + end_line=_offset_to_line(line_offsets, max(0, ec - 1)), + start_char=sc, + end_char=ec, + section=None, + kind="text", + context=chunk.context, + ) + ) + return records + + def _ensure_workspace_dirs(self) -> None: + root = self._config.workspace_root + root.mkdir(parents=True, exist_ok=True) + for sub in WORKSPACE_SUBDIRS: + (root / sub).mkdir(exist_ok=True) + get_index_dir(root).mkdir(parents=True, exist_ok=True) + seed_hermesignore(root) + + def _require_chonkie(self) -> None: + try: + import chonkie # noqa: F401 + except ImportError: + raise RuntimeError( + "Chonkie is required for workspace indexing. " + "Install it with: pip install hermes-agent[workspace]" + ) + + def _config_signature(self) -> str: + ch = self._config.knowledgebase.chunking + blob = json.dumps( + { + "chunk_size": ch.chunk_size, + "overlap": ch.overlap, + "overlap_mode": "token", + "overlap_method": "suffix", + "code_chunker": "production_v1", + "chunking_plan_version": CHUNKING_PLAN_VERSION, + }, + sort_keys=True, + ) + return hashlib.sha256(blob.encode()).hexdigest()[:16] + + +def _append_error(errors: list[IndexingError], error: IndexingError) -> None: + if len(errors) < _MAX_ERRORS: + errors.append(error) + + +def _read_file_text(path: Path) -> str | None: + raw = path.read_bytes() + try: + return raw.decode("utf-8") + except UnicodeDecodeError: + pass + try: + from charset_normalizer import from_bytes + + result = from_bytes(raw).best() + if result is None or result.encoding is None: + return None + if result.coherence < 0.5: + return None + return str(result) + except ImportError: + log.debug("charset-normalizer not installed, skipping non-UTF8 file: %s", path) + return None + + +def _scan_headings(text: str) -> list[tuple[int, str]]: + return [(m.start(), m.group(0).strip()) for m in _HEADING_RE.finditer(text)] + + +def _nearest_heading(headings: list[tuple[int, str]], char_offset: int) -> str | None: + best = None + for offset, heading in headings: + if offset <= char_offset: + best = heading + else: + break + return best + + +def _build_line_offsets(text: str) -> list[int]: + return [0] + [m.end() for m in _NEWLINE_RE.finditer(text)] + + +def _offset_to_line(offsets: list[int], char_offset: int) -> int: + lo, hi = 0, len(offsets) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + if offsets[mid] <= char_offset: + lo = mid + else: + hi = mid - 1 + return lo + 1 + + +def _file_hash(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for block in iter(lambda: f.read(65536), b""): + h.update(block) + return h.hexdigest() + + +def _make_id() -> str: + return f"chnk_{uuid.uuid4().hex[:12]}" + + +def _prune_stale(store: SQLiteFTS5Store, disk_paths: set[str]) -> int: + indexed = store.all_indexed_paths() + stale = indexed - disk_paths + for path in stale: + store.delete_file(path) + return len(stale)