diff --git a/src/semble/cache.py b/src/semble/cache.py index 894f3e7fc..f28eca511 100644 --- a/src/semble/cache.py +++ b/src/semble/cache.py @@ -2,7 +2,6 @@ import json import logging import os -import shutil import sys from collections.abc import Sequence from pathlib import Path @@ -10,6 +9,7 @@ import orjson +from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS from semble.index.bm25 import BM25 from semble.index.dense import SelectableBasicBackend from semble.index.file_walker import walk_files @@ -89,11 +89,6 @@ def resolve_cache_folder() -> Path: return cache_dir -def clear_cache(path: str) -> None: - """Clear all exact content indexes for the given path.""" - shutil.rmtree(find_index_from_cache_folder(path).parent, ignore_errors=True) - - def save_index_to_cache(index: "SembleIndex", path: str) -> None: """Save an index to the cache folder if it was freshly built.""" if not index.loaded_from_disk: @@ -102,8 +97,6 @@ def save_index_to_cache(index: "SembleIndex", path: str) -> None: def _metadata_matches(metadata: dict, model_path: str, content: Sequence[ContentType]) -> bool: """Return True if the stored metadata is compatible with the requested parameters.""" - from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS # avoid circular import at module level - try: content_type = tuple(ContentType(s) for s in metadata["content_type"]) # chunk_size and cache_version are absent in indexes built before those fields were added; @@ -117,22 +110,28 @@ def _metadata_matches(metadata: dict, model_path: str, content: Sequence[Content return False -def get_validated_cache(path: str, model_path: str | None, content: Sequence[ContentType]) -> Path | None: - """Validates the cache folder and returns the index path.""" - index_path = find_index_from_cache_folder(path, content) - if not index_path.exists(): - return None - - persistence_path = PersistencePath.from_path(index_path) +def _load_matching_metadata( + path: str, model_path: str | None, content: Sequence[ContentType] +) -> tuple[PersistencePath, dict] | None: + """Return the cached index files and metadata for path, or None if absent or built with other settings.""" + persistence_path = PersistencePath.from_path(find_index_from_cache_folder(path, content)) if persistence_path.non_existing(): return None - + metadata = json.loads(persistence_path.metadata.read_text(encoding="utf-8")) if model_path is None: model_path = resolve_model_name() - with open(persistence_path.metadata, encoding="utf-8") as f: - metadata = json.load(f) if not _metadata_matches(metadata, model_path, content): return None + return persistence_path, metadata + + +def get_validated_cache(path: str, model_path: str | None, content: Sequence[ContentType]) -> Path | None: + """Validates the cache folder and returns the index path.""" + loaded = _load_matching_metadata(path, model_path, content) + if loaded is None: + return None + persistence_path, metadata = loaded + index_path = persistence_path.metadata.parent if is_git_url(str(path)): return index_path @@ -168,25 +167,17 @@ def load_previous_for_incremental( :return: Previous index state, or None if the cache is unavailable or invalid. """ try: - index_path = find_index_from_cache_folder(path, content) - persistence_path = PersistencePath.from_path(index_path) - if persistence_path.non_existing(): - return None - - if model_path is None: - model_path = resolve_model_name() - with open(persistence_path.metadata, encoding="utf-8") as f: - metadata = json.load(f) - if not _metadata_matches(metadata, model_path, content): + loaded = _load_matching_metadata(path, model_path, content) + if loaded is None: return None + persistence_path, metadata = loaded raw_manifest = metadata.get("files") if not raw_manifest: return None manifest = {indexed_path: FileManifestEntry(**entry) for indexed_path, entry in raw_manifest.items()} - with open(persistence_path.chunks, "rb") as f: - chunks = [Chunk.from_dict(item) for item in orjson.loads(f.read())] + chunks = [Chunk.from_dict(item) for item in orjson.loads(persistence_path.chunks.read_bytes())] vectors = SelectableBasicBackend.load(persistence_path.semantic_index).vectors bm25_index = BM25.load(persistence_path.bm25_index) diff --git a/src/semble/cli.py b/src/semble/cli.py index 5061e9647..6d4119d90 100644 --- a/src/semble/cli.py +++ b/src/semble/cli.py @@ -6,6 +6,7 @@ import re import sys import warnings +from collections.abc import Iterator from importlib.util import find_spec from pathlib import Path from shutil import rmtree @@ -49,6 +50,26 @@ def _maybe_save_index(parts: list[tuple[str, SembleIndex]]) -> None: print(f"Error saving index: {e}", file=sys.stderr) +def _add_query_args(p: argparse.ArgumentParser) -> None: + """Add the path, result-shaping, and content arguments shared by search and find-related.""" + p.add_argument( + "path", + nargs="*", + default=["."], + help="Local paths or git URLs to search together (default: current directory).", + ) + p.add_argument("-k", "--top-k", type=int, default=5, help="Number of results (default: 5).") + p.add_argument( + "--max-snippet-lines", + type=int, + default=None, + metavar="N", + help="Lines of source per result (default: full chunk). 10 = signature + body, 0 = no code.", + ) + p.add_argument("--format", choices=["json", "text"], default="json", help="Output format (default: json).") + _add_content_args(p) + + def _add_content_args(p: argparse.ArgumentParser) -> None: """Add --content and deprecated --include-text-files to a subparser.""" p.add_argument( @@ -178,15 +199,16 @@ def _run_find_related( _maybe_save_index(parts) +def _cached_index_paths(cache_folder: Path) -> Iterator[Path]: + """Yield index folders in the cache that sit under a sha256 cache key.""" + return (path for path in cache_folder.glob("*/index*") if _SHA_256_REGEX.match(path.parent.name)) + + def _clear_indexes(cache_folder: Path) -> None: """Remove all valid index entries from the cache folder.""" - indexes: set[Path] = set() - for path in cache_folder.glob("*/index*"): - if not _SHA_256_REGEX.match(path.parent.name): - continue - if PersistencePath.from_path(path).non_existing(): - continue - indexes.add(path.parent) + indexes = { + path.parent for path in _cached_index_paths(cache_folder) if not PersistencePath.from_path(path).non_existing() + } if not indexes: print(f"No indexes found to clear in `{cache_folder}`") @@ -209,9 +231,7 @@ def _clear_savings(cache_folder: Path) -> None: def _clear_orphans(cache_folder: Path) -> None: """Remove index entries whose local root_path no longer exists.""" orphans: dict[Path, str] = {} - for path in cache_folder.glob("*/index*"): - if not _SHA_256_REGEX.match(path.parent.name): - continue + for path in _cached_index_paths(cache_folder): try: with open(path / "metadata.json", encoding="utf-8") as f: metadata = json.load(f) @@ -267,22 +287,7 @@ def _cli_main() -> None: search_p = sub.add_parser("search", help="Search a codebase.") search_p.add_argument("query", help="Natural language or code query.") - search_p.add_argument( - "path", - nargs="*", - default=["."], - help="Local paths or git URLs to search together (default: current directory).", - ) - search_p.add_argument("-k", "--top-k", type=int, default=5, help="Number of results (default: 5).") - search_p.add_argument( - "--max-snippet-lines", - type=int, - default=None, - metavar="N", - help="Lines of source per result (default: full chunk). 10 = signature + body, 0 = no code.", - ) - search_p.add_argument("--format", choices=["json", "text"], default="json", help="Output format (default: json).") - _add_content_args(search_p) + _add_query_args(search_p) clear_p = sub.add_parser("clear", help="Clear the index cache.") clear_p.add_argument( @@ -294,22 +299,7 @@ def _cli_main() -> None: related_p = sub.add_parser("find-related", help="Find code similar to a specific location.") related_p.add_argument("file_path", help="File path as shown in search results.") related_p.add_argument("line", type=int, help="Line number (1-indexed).") - related_p.add_argument( - "path", - nargs="*", - default=["."], - help="Local paths or git URLs to search together (default: current directory).", - ) - related_p.add_argument("-k", "--top-k", type=int, default=5, help="Number of results (default: 5).") - related_p.add_argument( - "--max-snippet-lines", - type=int, - default=None, - metavar="N", - help="Lines of source per result (default: full chunk). 10 = signature + body, 0 = no code.", - ) - related_p.add_argument("--format", choices=["json", "text"], default="json", help="Output format (default: json).") - _add_content_args(related_p) + _add_query_args(related_p) sub.add_parser("savings", help="Show token savings and usage stats.") diff --git a/src/semble/index/bm25.py b/src/semble/index/bm25.py index 2df8d49c3..379f9315c 100644 --- a/src/semble/index/bm25.py +++ b/src/semble/index/bm25.py @@ -114,11 +114,7 @@ def load(cls, path: Path) -> "BM25": documents = data["documents"] if len(doc_order) != len(set(doc_order)) or set(documents) != set(doc_order): raise ValueError("Persisted BM25 document state is inconsistent") - index._documents = {chunk_id: Counter(counts) for chunk_id, counts in documents.items()} - for chunk_id, counts in index._documents.items(): - for term, count in counts.items(): - index.postings.setdefault(term, {})[chunk_id] = count - index._doc_lengths = {chunk_id: sum(counts.values()) for chunk_id, counts in index._documents.items()} - index._total_doc_length = sum(index._doc_lengths.values()) + for chunk_id, counts in documents.items(): + index._add_counts(chunk_id, Counter(counts), sum(counts.values())) index.set_doc_order(doc_order) return index diff --git a/src/semble/index/dense.py b/src/semble/index/dense.py index ada1f875a..aa3459f57 100644 --- a/src/semble/index/dense.py +++ b/src/semble/index/dense.py @@ -28,14 +28,9 @@ def _load_cached(model_path: str) -> StaticModel: disable_progress_bars() logging.getLogger("huggingface_hub.utils._http").addFilter(_drop_unauthenticated_warning) try: - try: - model = StaticModel.from_pretrained(model_path, force_download=False) - except ValueError: - model = StaticModel.from_pretrained(model_path, force_download=True) - finally: - disable_progress_bars() - - return model + return StaticModel.from_pretrained(model_path, force_download=False) + except ValueError: + return StaticModel.from_pretrained(model_path, force_download=True) def load_model(model_path: str | None = None) -> tuple[StaticModel, str]: diff --git a/src/semble/index/files.py b/src/semble/index/files.py index e3f4f9419..f7e46f07c 100644 --- a/src/semble/index/files.py +++ b/src/semble/index/files.py @@ -444,22 +444,16 @@ } -def _inv_mapping(mapping: dict[str, str]) -> dict[str, list[str]]: - """Invert a mapping, taking into account duplicate values.""" - inv: defaultdict[str, list[str]] = defaultdict(list) - for key, value in mapping.items(): - inv[value].append(key) - return dict(inv) - - ALL_LANGUAGES = frozenset(_EXTENSION_TO_LANGUAGE.values()) _CODE_LANGUAGES = ALL_LANGUAGES - _DOC_LANGUAGES - _CONFIG_LANGUAGES - _DATA_LANGUAGES -_LANGUAGE_TO_EXTENSION = _inv_mapping(_EXTENSION_TO_LANGUAGE) +_LANGUAGE_TO_EXTENSIONS: defaultdict[str, list[str]] = defaultdict(list) +for _extension, _language in _EXTENSION_TO_LANGUAGE.items(): + _LANGUAGE_TO_EXTENSIONS[_language].append(_extension) -_CONTENT_TYPE_LANGUAGES: dict[ContentType, frozenset[str]] = { - ContentType.CODE: frozenset(_CODE_LANGUAGES), - ContentType.DOCS: frozenset(_DOC_LANGUAGES), - ContentType.CONFIG: frozenset(_CONFIG_LANGUAGES), +_CONTENT_TYPE_LANGUAGES = { + ContentType.CODE: _CODE_LANGUAGES, + ContentType.DOCS: _DOC_LANGUAGES, + ContentType.CONFIG: _CONFIG_LANGUAGES, } @@ -470,14 +464,14 @@ def detect_language(file_name: Path) -> str | None: def get_extensions(types: Sequence[ContentType]) -> list[str]: """Returns a list of supported file extensions for the given content types.""" - languages: set[str] = set() - for content_type in types: - languages.update(_CONTENT_TYPE_LANGUAGES[content_type]) - all_extensions: set[str] = set() - for language in languages: - all_extensions.update(_LANGUAGE_TO_EXTENSION.get(language, set())) - - return sorted(all_extensions) + return sorted( + { + ext + for content_type in types + for lang in _CONTENT_TYPE_LANGUAGES[content_type] + for ext in _LANGUAGE_TO_EXTENSIONS.get(lang, []) + } + ) class FileStatus(str, Enum): @@ -488,7 +482,7 @@ class FileStatus(str, Enum): def read_file_text(file_path: Path) -> str: - """Read a file's text content, replacing invalid characters and silencing read errors.""" + """Read a file's text content, replacing invalid UTF-8 characters.""" return file_path.read_text(encoding="utf-8", errors="replace") diff --git a/src/semble/index/index.py b/src/semble/index/index.py index d24d394a3..9fb204833 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -17,11 +17,12 @@ from vicinity.backends.basic import BasicArgs from semble.cache import get_validated_cache, load_previous_for_incremental +from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS from semble.index.bm25 import BM25 from semble.index.create import create_index_from_path from semble.index.dense import SelectableBasicBackend, load_model from semble.index.files import read_file_text -from semble.index.types import CACHE_FORMAT_VERSION, FileManifestEntry, PersistencePath +from semble.index.types import CACHE_FORMAT_VERSION, FileManifestEntry, PersistencePath, PreviousIndex from semble.search import _search_semantic, search from semble.stats import save_search_stats from semble.types import CallType, Chunk, ContentType, IndexStats, SearchResult @@ -117,15 +118,10 @@ def _compute_file_sizes(self, root: Path) -> dict[str, int]: @property def stats(self) -> IndexStats: """Stats of an index.""" - language_counts: dict[str, int] = defaultdict(int) - for chunk in self.chunks: - if chunk.language: - language_counts[chunk.language] += 1 - return IndexStats( indexed_files=len(self._file_mapping), total_chunks=len(self.chunks), - languages=dict(language_counts), + languages={language: len(ids) for language, ids in self._language_mapping.items()}, ) @property @@ -167,18 +163,7 @@ def from_path( path = path.resolve() previous = load_previous_for_incremental(str(path), model_path, normalized) - bm25_index, semantic_index, chunks, manifest = create_index_from_path( - path, - model=model, - content=normalized, - display_root=path, - previous=previous, - show_progress_bar=show_progress_bar, - ) - - return SembleIndex( - model, bm25_index, semantic_index, chunks, model_path, root=path, content=normalized, manifest=manifest - ) + return _build(path, model, model_path, normalized, show_progress_bar, previous) @classmethod def merge(cls, indexes: Sequence[tuple[str, SembleIndex]]) -> SembleIndex: @@ -274,25 +259,7 @@ def from_git( raise RuntimeError(f"git clone failed for {url!r}:\n{result.stderr.strip()}") model, model_path = load_model(model_path) - resolved_path = Path(tmp_dir).resolve() - bm25_index, semantic_index, chunks, manifest = create_index_from_path( - resolved_path, - model=model, - content=normalized, - display_root=resolved_path, - show_progress_bar=show_progress_bar, - ) - - return SembleIndex( - model, - bm25_index, - semantic_index, - chunks, - model_path, - root=resolved_path, - content=normalized, - manifest=manifest, - ) + return _build(Path(tmp_dir).resolve(), model, model_path, normalized, show_progress_bar) def find_related( self, source: Chunk | SearchResult, *, top_k: int = 5, max_snippet_lines: int | None = None @@ -380,8 +347,7 @@ def load_from_disk(cls: type[SembleIndex], path: Path | str) -> SembleIndex: missing = ", ".join(str(p) for p in non_existent) raise FileNotFoundError(f"Index not found at {path}. Missing: {missing}") - with open(persistence_paths.metadata, "rb") as f: - metadata = orjson.loads(f.read()) + metadata = orjson.loads(persistence_paths.metadata.read_bytes()) found_version = metadata.get("cache_version") if found_version != CACHE_FORMAT_VERSION: raise ValueError( @@ -391,12 +357,7 @@ def load_from_disk(cls: type[SembleIndex], path: Path | str) -> SembleIndex: bm25_index = BM25.load(persistence_paths.bm25_index) semantic_index = SelectableBasicBackend.load(persistence_paths.semantic_index) - with open(persistence_paths.chunks, "rb") as f: - chunk_data = orjson.loads(f.read()) - - chunks = [] - for chunk_item in chunk_data: - chunks.append(Chunk.from_dict(chunk_item)) + chunks = [Chunk.from_dict(item) for item in orjson.loads(persistence_paths.chunks.read_bytes())] if not (len(chunks) == len(bm25_index.doc_order) == semantic_index.vectors.shape[0]): raise ValueError("Persisted index components have inconsistent document counts") root_path = metadata["root_path"] @@ -431,10 +392,7 @@ def save(self, path: Path | str) -> None: self._bm25_index.save(persistence_paths.bm25_index) self._semantic_index.save(persistence_paths.semantic_index) - with open(persistence_paths.chunks, "wb") as f: - data = orjson.dumps(self.chunks) - f.write(data) - from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS # avoid circular import at module level + persistence_paths.chunks.write_bytes(orjson.dumps(self.chunks)) root_str = None if self._root is None else str(self._root) metadata = { @@ -446,6 +404,21 @@ def save(self, path: Path | str) -> None: "cache_version": CACHE_FORMAT_VERSION, "files": self._manifest, } - with open(persistence_paths.metadata, "wb") as f: - data = orjson.dumps(metadata) - f.write(data) + persistence_paths.metadata.write_bytes(orjson.dumps(metadata)) + + +def _build( + path: Path, + model: StaticModel, + model_path: str, + content: tuple[ContentType, ...], + show_progress_bar: bool, + previous: PreviousIndex | None = None, +) -> SembleIndex: + """Index a resolved directory, storing chunk paths relative to it.""" + bm25_index, semantic_index, chunks, manifest = create_index_from_path( + path, model=model, content=content, display_root=path, previous=previous, show_progress_bar=show_progress_bar + ) + return SembleIndex( + model, bm25_index, semantic_index, chunks, model_path, root=path, content=content, manifest=manifest + ) diff --git a/src/semble/installer/agents.py b/src/semble/installer/agents.py index 345cff53f..de8fe2798 100644 --- a/src/semble/installer/agents.py +++ b/src/semble/installer/agents.py @@ -21,11 +21,7 @@ def _exists_or_denied(path: Path) -> bool: """Distinguish between existence and permission issues.""" try: path.stat() - except FileNotFoundError: - return False - # PermissionError is a subclass of OSError - # which is why this looks the way it does. - except PermissionError: + except PermissionError: # checked before OSError, its parent class return True except OSError: return False @@ -156,10 +152,6 @@ class AgentTarget: instructions_path: Path | None # None = not supported for this agent subagent_path: Path | None = None # global (user-level) sub-agent file; None = unsupported - def resolved_mcp_path(self) -> Path | None: - """Return the resolved MCP config path, or None if MCP is unsupported.""" - return self.mcp.path if self.mcp else None - def _opencode_mcp_path() -> Path: """Return the opencode config path, preferring .jsonc over .json.""" diff --git a/src/semble/installer/config.py b/src/semble/installer/config.py index 7acb9c594..b89d59e07 100644 --- a/src/semble/installer/config.py +++ b/src/semble/installer/config.py @@ -3,7 +3,7 @@ import json from functools import cache from pathlib import Path -from typing import Literal, TypeVar +from typing import Literal from semble_grammars import get_parser from tree_sitter import Node, Parser @@ -11,7 +11,6 @@ from semble.installer.agents import SEMBLE_END, SEMBLE_PIN, SEMBLE_START, Action JsonObjectResult = tuple[Node, bytes] | Literal["skipped", "error"] -_T = TypeVar("_T") _CODEX_MCP_HEADER = "[mcp_servers.semble]" _CODEX_MCP_BLOCK = f'[mcp_servers.semble]\ncommand = "uvx"\nargs = ["--from", "{SEMBLE_PIN}", "semble"]\n' @@ -64,6 +63,13 @@ def _insert_first_member(src: bytes, obj: Node, member_text: str) -> bytes: return src[: brace + 1] + b"\n" + indent + member_text.encode("utf-8") + comma + src[brace + 1 :] +def _skip_blanks_back(src: bytes, i: int) -> int: + """Move i left past any spaces and tabs immediately before it.""" + while i > 0 and src[i - 1 : i] in (b" ", b"\t"): + i -= 1 + return i + + def _delete_member(src: bytes, member: Node) -> bytes: """Remove `member` plus one adjacent comma and its leading line indentation.""" start, end = member.start_byte, member.end_byte @@ -73,17 +79,12 @@ def _delete_member(src: bytes, member: Node) -> bytes: if after < len(src) and src[after : after + 1] == b",": # prefer a trailing comma end = after + 1 else: - before = start - while before > 0 and src[before - 1 : before] in (b" ", b"\t"): - before -= 1 + before = _skip_blanks_back(src, start) if before > 0 and src[before - 1 : before] == b"\n": - before -= 1 # step over newline to find comma on preceding line - while before > 0 and src[before - 1 : before] in (b" ", b"\t"): - before -= 1 + before = _skip_blanks_back(src, before - 1) # step over newline to find comma on preceding line if before > 0 and src[before - 1 : before] == b",": start = before - 1 - while start > 0 and src[start - 1 : start] in (b" ", b"\t"): - start -= 1 + start = _skip_blanks_back(src, start) if start > 0 and src[start - 1 : start] == b"\n": start -= 1 # drop the now-empty line return src[:start] + src[end:] diff --git a/src/semble/installer/installer.py b/src/semble/installer/installer.py index 54b7b96e0..7276c963b 100644 --- a/src/semble/installer/installer.py +++ b/src/semble/installer/installer.py @@ -109,7 +109,7 @@ def _apply_subagent(agent: AgentTarget, mode: Mode) -> WriteResult | None: "MCP server", "lets the agent call semble directly as a tool", _apply_mcp, - AgentTarget.resolved_mcp_path, + lambda a: a.mcp.path if a.mcp else None, ), _Integration( IntegrationType.INSTRUCTIONS, diff --git a/src/semble/mcp.py b/src/semble/mcp.py index f03f74352..704c69164 100644 --- a/src/semble/mcp.py +++ b/src/semble/mcp.py @@ -220,23 +220,17 @@ async def _await_model(self) -> str: assert self._model_path is not None return self._model_path - def _compute_cache_key( - self, - source: str, - ref: str | None = None, - content: Sequence[ContentType] = (ContentType.CODE,), - ) -> _CacheKey: + def _compute_cache_key(self, source: str, content: Sequence[ContentType] = (ContentType.CODE,)) -> _CacheKey: """Compute the canonical key for an exact index variant.""" - is_git = is_git_url(source) - source_key = (f"{source}@{ref}" if ref else source) if is_git else str(Path(source).resolve()) + source_key = source if is_git_url(source) else str(Path(source).resolve()) normalized = tuple(content_type for content_type in ContentType if content_type in content) return source_key, normalized - def _build_index(self, source: str, ref: str | None, model_path: str, cache_key: _CacheKey) -> SembleIndex: + def _build_index(self, source: str, model_path: str, cache_key: _CacheKey) -> SembleIndex: """Build an index for the given source and cache it.""" source_key, content = cache_key index = ( - SembleIndex.from_git(source, ref=ref, model_path=model_path, content=content) + SembleIndex.from_git(source, model_path=model_path, content=content) if is_git_url(source) else SembleIndex.from_path(source_key, model_path=model_path, content=content) ) @@ -246,14 +240,14 @@ def _build_index(self, source: str, ref: str | None, model_path: str, cache_key: logger.warning("Failed to save index cache for %r", source_key, exc_info=True) return index - async def _build_tracked(self, source: str, ref: str | None, model_path: str, cache_key: _CacheKey) -> SembleIndex: + async def _build_tracked(self, source: str, model_path: str, cache_key: _CacheKey) -> SembleIndex: """Build an index and, for local paths, record when its staleness cooldown ends. The cooldown write happens after the await, i.e. back on the event loop thread, regardless of which thread `_build_index` itself ran on. """ start = time.monotonic() - index = await asyncio.to_thread(self._build_index, source, ref, model_path, cache_key) + index = await asyncio.to_thread(self._build_index, source, model_path, cache_key) if not is_git_url(source): finished = time.monotonic() self._revalidate_after[cache_key] = finished + (finished - start) * _MIN_REVALIDATE_FACTOR @@ -286,18 +280,13 @@ async def _evict_if_stale(self, cache_key: _CacheKey) -> None: if validated is None and self._tasks.get(cache_key) is cached: self.evict(cache_key) - async def get( - self, - source: str, - ref: str | None = None, - content: Sequence[ContentType] = (ContentType.CODE,), - ) -> SembleIndex: + async def get(self, source: str, content: Sequence[ContentType] = (ContentType.CODE,)) -> SembleIndex: """Return an index for the requested source, building and caching it on first access. Local paths are revalidated against the on-disk cache on every call (subject to a cooldown scaled by build time), so an entry is rebuilt once its files change. """ - cache_key = self._compute_cache_key(source, ref, content) + cache_key = self._compute_cache_key(source, content) await self._evict_if_stale(cache_key) if cache_key not in self._tasks: @@ -307,7 +296,7 @@ async def get( if len(self._tasks) >= _CACHE_MAX_SIZE: evicted_key, _ = self._tasks.popitem(last=False) self._revalidate_after.pop(evicted_key, None) - self._tasks[cache_key] = asyncio.create_task(self._build_tracked(source, ref, model_path, cache_key)) + self._tasks[cache_key] = asyncio.create_task(self._build_tracked(source, model_path, cache_key)) self._tasks.move_to_end(cache_key) task = self._tasks[cache_key] try: diff --git a/src/semble/ranking/__init__.py b/src/semble/ranking/__init__.py index 1542c085b..c77d51ed4 100644 --- a/src/semble/ranking/__init__.py +++ b/src/semble/ranking/__init__.py @@ -1,5 +1,4 @@ -from semble.ranking.boosting import apply_query_boost, boost_multi_chunk_files +from semble.ranking.boosting import apply_query_boost, boost_multi_chunk_files, resolve_alpha from semble.ranking.penalties import rerank_topk -from semble.ranking.weighting import resolve_alpha __all__ = ["apply_query_boost", "boost_multi_chunk_files", "rerank_topk", "resolve_alpha"] diff --git a/src/semble/ranking/boosting.py b/src/semble/ranking/boosting.py index fc00d058b..4edb79c3a 100644 --- a/src/semble/ranking/boosting.py +++ b/src/semble/ranking/boosting.py @@ -78,6 +78,9 @@ # Fraction of max_score added to each file's top chunk, scaled by its aggregate candidate score. _FILE_COHERENCE_BOOST_FRAC = 0.2 +_ALPHA_SYMBOL = 0.3 # lean BM25 for exact keyword matching +_ALPHA_NL = 0.5 # balanced semantic + BM25 + # Common English stopwords excluded from file-stem matching for NL queries. _STOPWORDS = frozenset( "a an and are as at be by do does for from has have how if in is it not of on or the to was" @@ -134,6 +137,13 @@ def is_symbol_query(query: str) -> bool: return _SYMBOL_QUERY_RE.match(query.strip()) is not None +def resolve_alpha(query: str, alpha: float | None) -> float: + """Return the blending weight for semantic scores, auto-detecting from query type.""" + if alpha is not None: + return alpha + return _ALPHA_SYMBOL if is_symbol_query(query) else _ALPHA_NL + + def _extract_symbol_name(query: str) -> str: """Extract the final identifier from a possibly namespace-qualified query. @@ -180,14 +190,17 @@ def _definition_tier(chunk: Chunk, names: set[str], boost_unit: float) -> float: return boost_unit * (1.5 if any(_stem_matches(stem, name.lower()) for name in names) else 1.0) -def _scan_non_candidates( +def _boost_definitions( boosted: dict[Chunk, float], names: set[str], boost_unit: float, all_chunks: list[Chunk], stem_ok: Callable[[str], bool], ) -> None: - """Boost non-candidate chunks whose lowercased file stem satisfies stem_ok (in-place).""" + """Boost candidates defining one of names, then add non-candidates whose lowercased file stem satisfies stem_ok.""" + for chunk in list(boosted): + if tier := _definition_tier(chunk, names, boost_unit): + boosted[chunk] += tier for chunk in all_chunks: if chunk in boosted: continue @@ -210,12 +223,7 @@ def _boost_symbol_definitions( names.add(query.strip()) boost_unit = max_score * _DEFINITION_BOOST_MULTIPLIER - - for chunk in list(boosted): - if tier := _definition_tier(chunk, names, boost_unit): - boosted[chunk] += tier - - _scan_non_candidates( + _boost_definitions( boosted, names, boost_unit, @@ -240,34 +248,24 @@ def _boost_embedded_symbols( return boost_unit = max_score * _DEFINITION_BOOST_MULTIPLIER * _EMBEDDED_SYMBOL_BOOST_SCALE - - for chunk in list(boosted): - if tier := _definition_tier(chunk, names, boost_unit): - boosted[chunk] += tier - symbols_lower = frozenset(s.lower() for s in names) - for chunk in all_chunks: - if chunk in boosted: - continue - stem = Path(chunk.file_path).stem.lower() + + def stem_ok(stem: str) -> bool: stem_norm = stem.replace("_", "") - if not any( + return any( stem == symbol_lower or stem_norm == symbol_lower or (len(stem) >= _EMBEDDED_STEM_MIN_LEN and symbol_lower.startswith(stem)) or (len(stem_norm) >= _EMBEDDED_STEM_MIN_LEN and symbol_lower.startswith(stem_norm)) for symbol_lower in symbols_lower - ): - continue - if tier := _definition_tier(chunk, names, boost_unit): - boosted[chunk] = tier + ) + + _boost_definitions(boosted, names, boost_unit, all_chunks, stem_ok) def _count_keyword_matches(keywords: set[str], parts: set[str]) -> int: """Count query keywords that match path parts, allowing prefix overlap (min 3 chars).""" exact = keywords & parts - if len(exact) == len(keywords): - return len(exact) n_matches = len(exact) for keyword in keywords - exact: for part in parts: diff --git a/src/semble/ranking/penalties.py b/src/semble/ranking/penalties.py index dbec0d246..c26c63fe2 100644 --- a/src/semble/ranking/penalties.py +++ b/src/semble/ranking/penalties.py @@ -1,4 +1,5 @@ import re +from functools import lru_cache from pathlib import Path from semble.types import Chunk @@ -100,16 +101,11 @@ def rerank_topk( if not scores: return [] - # Apply file-path penalties. - penalty_cache: dict[str, float] = {} - penalised: dict[Chunk, float] = {} - for chunk, score in scores.items(): - if penalise_paths: - if chunk.file_path not in penalty_cache: - penalty_cache[chunk.file_path] = _file_path_penalty(chunk.file_path) - penalised[chunk] = score * penalty_cache[chunk.file_path] - else: - penalised[chunk] = score + penalised = ( + {chunk: score * _file_path_penalty(chunk.file_path) for chunk, score in scores.items()} + if penalise_paths + else dict(scores) + ) # Sort by penalised score (highest first) — single sort. ranked = sorted(penalised, key=lambda c: -penalised[c]) @@ -140,6 +136,7 @@ def rerank_topk( return [(chunk, score) for score, chunk in selected[:top_k]] +@lru_cache(maxsize=4096) def _file_path_penalty(file_path: str) -> float: """Return a combined multiplicative penalty for all applicable path patterns.""" normalised = file_path.replace("\\", "/") diff --git a/src/semble/ranking/weighting.py b/src/semble/ranking/weighting.py deleted file mode 100644 index e023d8a19..000000000 --- a/src/semble/ranking/weighting.py +++ /dev/null @@ -1,11 +0,0 @@ -from semble.ranking.boosting import is_symbol_query - -_ALPHA_SYMBOL = 0.3 # lean BM25 for exact keyword matching -_ALPHA_NL = 0.5 # balanced semantic + BM25 - - -def resolve_alpha(query: str, alpha: float | None) -> float: - """Return the blending weight for semantic scores, auto-detecting from query type.""" - if alpha is not None: - return alpha - return _ALPHA_SYMBOL if is_symbol_query(query) else _ALPHA_NL diff --git a/src/semble/search.py b/src/semble/search.py index 8d901d699..d1c18b788 100644 --- a/src/semble/search.py +++ b/src/semble/search.py @@ -97,10 +97,9 @@ def search( semantic = _search_semantic(query, model, semantic_index, chunks, candidate_count, selector) semantic_scores: dict[Chunk, float] = {result.chunk: result.score for result in semantic} - bm25_scores = {} - for result in _search_bm25(query, bm25_index, chunks, candidate_count, selector): - if result.score: - bm25_scores[result.chunk] = result.score + bm25_scores = { + result.chunk: result.score for result in _search_bm25(query, bm25_index, chunks, candidate_count, selector) + } normalized_semantic = _rrf_scores(semantic_scores) normalized_bm25 = _rrf_scores(bm25_scores) diff --git a/src/semble/stats.py b/src/semble/stats.py index 3c6288c67..49b27bab5 100644 --- a/src/semble/stats.py +++ b/src/semble/stats.py @@ -5,19 +5,21 @@ from collections import defaultdict from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from functools import cache -from importlib import import_module from pathlib import Path -from types import ModuleType from semble.cache import resolve_cache_folder from semble.types import CallType, SearchResult +try: + import fcntl +except ImportError: # pragma: no cover + fcntl = None # type: ignore[assignment] + logger = logging.getLogger(__name__) def _get_stats_file() -> Path: - """Safely create a stats file.""" + """Return the path of the savings stats file.""" return resolve_cache_folder() / "savings.jsonl" @@ -52,15 +54,6 @@ class SavingsSummary: call_type_counts: dict[str, int] -@cache -def _import_fcntl() -> ModuleType | None: - """Return fcntl when available, otherwise None.""" - try: - return import_module("fcntl") - except ImportError: # pragma: no cover - return None - - def save_search_stats( results: list[SearchResult], call_type: CallType, @@ -91,14 +84,11 @@ def save_search_stats( stats_file = _get_stats_file() stats_file.parent.mkdir(parents=True, exist_ok=True) with stats_file.open("a") as f: - fcntl = _import_fcntl() try: if fcntl is not None: fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB) - except BlockingIOError: # pragma: no cover - return # another process holds the lock; skip this record except OSError: # pragma: no cover - return # lock contention or unsupported filesystem; skip + return # another process holds the lock, or the filesystem doesn't support locking; skip f.write(json.dumps(record) + "\n") except OSError: pass @@ -156,6 +146,11 @@ def _format_calls(calls: int) -> str: return f"{calls / 1_000:.1f}k" if calls >= 1_000 else str(calls) +def _bar(filled: int, width: int, enabled: bool) -> str: + """Render a green filled bar followed by its grey remainder.""" + return _color("32", "█" * filled, enabled) + _color("38;5;244", "░" * (width - filled), enabled) + + def _color_ratio(pct: int, enabled: bool) -> str: """Color a savings percentage according to its value.""" code = "32" if pct >= 80 else "33" if pct >= 50 else "31" @@ -180,8 +175,7 @@ def format_savings_report(path: Path | None = None) -> str: total_saved_tokens = all_time.saved_chars // 4 overall_pct = round(all_time.saved_chars / all_time.file_chars * 100) if all_time.file_chars else 0 efficiency_filled = round(overall_pct / 100 * bar_width) - efficiency_bar = _color("32", "█" * efficiency_filled, color) - efficiency_bar += _color("38;5;244", "░" * (bar_width - efficiency_filled), color) + efficiency_bar = _bar(efficiency_filled, bar_width, color) lines = [ "", @@ -206,7 +200,7 @@ def format_savings_report(path: Path | None = None) -> str: if bucket.file_chars > 0: ratio = bucket.saved_chars / bucket.file_chars filled = round(ratio * bar_width) - row_bar = _color("32", "█" * filled, color) + _color("38;5;244", "░" * (bar_width - filled), color) + row_bar = _bar(filled, bar_width, color) ratio_str = _color_ratio(round(ratio * 100), color) else: row_bar = _color("38;5;244", "░" * bar_width, color) @@ -229,7 +223,7 @@ def format_savings_report(path: Path | None = None) -> str: for i, (call_type, count) in enumerate(top, start=1): share = count / total filled = max(1, round(share * 16)) - bar = _color("32", "█" * filled, color) + _color("38;5;244", "░" * (16 - filled), color) + bar = _bar(filled, 16, color) rank = f"{i}." lines.append( f" {_color('38;5;244', f'{rank:<4}', color)} {call_type:<16} " diff --git a/tests/test_cache.py b/tests/test_cache.py index e9c197d36..076b382a4 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -12,7 +12,6 @@ _get_valid_user_cache_dir, _linux_cache_dir, _windows_cache_dir, - clear_cache, find_index_from_cache_folder, get_validated_cache, resolve_cache_folder, @@ -121,21 +120,6 @@ def test_resolve_cache_folder_semble_cache_location(tmp_path: Path) -> None: assert custom.exists() -def test_clear_cache(tmp_path: Path) -> None: - """clear_cache removes every content variant and is a no-op when none exist.""" - cache_dir = tmp_path / "repo" - index_path = cache_dir / "index" - with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): - clear_cache("/some/path") # no-op: path doesn't exist yet - index_path.mkdir(parents=True) - docs_path = cache_dir / "index-docs" - docs_path.mkdir() - with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): - clear_cache("/some/path") - assert not index_path.exists() - assert not docs_path.exists() - - def _write_metadata( path: Path, model_path: str, diff --git a/tests/test_ranking.py b/tests/test_ranking.py index 6a3339778..8407a3551 100644 --- a/tests/test_ranking.py +++ b/tests/test_ranking.py @@ -1,8 +1,7 @@ import pytest -from semble.ranking.boosting import apply_query_boost, boost_multi_chunk_files +from semble.ranking.boosting import apply_query_boost, boost_multi_chunk_files, resolve_alpha from semble.ranking.penalties import rerank_topk -from semble.ranking.weighting import resolve_alpha from tests.conftest import make_chunk