feat(workspace): add DefaultIndexer class
This commit is contained in:
alt-glitch
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user