diff --git a/reme2/component/__init__.py b/reme2/component/__init__.py index 7cf0a611..cda27c37 100644 --- a/reme2/component/__init__.py +++ b/reme2/component/__init__.py @@ -8,6 +8,7 @@ from . import client from . import embedding from . import file_parser from . import file_watcher +from . import keyword_index from . import job from . import service from .application_context import ApplicationContext @@ -34,6 +35,7 @@ __all__ = [ "embedding", "file_parser", "file_watcher", + "keyword_index", "job", "service", ] diff --git a/reme2/component/base_component.py b/reme2/component/base_component.py index 7b7c122e..c7df66db 100644 --- a/reme2/component/base_component.py +++ b/reme2/component/base_component.py @@ -2,6 +2,7 @@ import asyncio from abc import ABC +from pathlib import Path from typing import TYPE_CHECKING from ..enumeration import ComponentEnum @@ -20,11 +21,11 @@ class BaseComponent(ABC): component_type = ComponentEnum.BASE def __init__( - self, - name: str | None = None, - backend: str = "", - app_context: "ApplicationContext | None" = None, - **kwargs, + self, + name: str | None = None, + backend: str = "", + app_context: "ApplicationContext | None" = None, + **kwargs, ) -> None: self.name: str = name or self.__class__.__name__ self.backend: str = backend @@ -37,6 +38,25 @@ class BaseComponent(ABC): self._is_started: bool = False self._lock: asyncio.Lock = asyncio.Lock() + @property + def is_started(self) -> bool: + return self._is_started + + def get_component(self, component_type: ComponentEnum, name: str): + """Get a component by type and name from app_context.""" + if self.app_context is None: + raise ValueError("app_context is not set") + component_dict = self.app_context.components.get(component_type, {}) + if name not in component_dict: + raise ValueError(f"{component_type.value} '{name}' not found.") + return component_dict[name] + + @property + def working_path(self) -> Path: + if self.app_context is None: + return Path.cwd() + return Path(self.app_context.app_config.working_dir) + async def _start(self) -> None: """Start the component.""" @@ -64,10 +84,6 @@ class BaseComponent(ABC): await self.close() await self.start() - @property - def is_started(self) -> bool: - return self._is_started - async def __call__(self, **kwargs): raise NotImplementedError diff --git a/reme2/component/base_step.py b/reme2/component/base_step.py index e149fa4a..3412cff6 100644 --- a/reme2/component/base_step.py +++ b/reme2/component/base_step.py @@ -60,7 +60,7 @@ class BaseStep(BaseComponent): return result - def _get_component(self, key: ComponentEnum, name: str, attr: str | None = None): + def get_component(self, key: ComponentEnum, name: str, attr: str | None = None): assert self.app_context is not None comp = self.app_context.components[key][name] return getattr(comp, attr) if attr else comp @@ -79,7 +79,7 @@ class BaseStep(BaseComponent): @property def as_llm(self) -> ChatModelBase: name = self.kwargs.get("as_llm", "default") - return name if isinstance(name, ChatModelBase) else self._get_component(ComponentEnum.AS_LLM, name, "model") + return name if isinstance(name, ChatModelBase) else self.get_component(ComponentEnum.AS_LLM, name, "model") @property def as_llm_formatter(self) -> FormatterBase: @@ -87,7 +87,7 @@ class BaseStep(BaseComponent): return ( name if isinstance(name, FormatterBase) - else self._get_component( + else self.get_component( ComponentEnum.AS_LLM_FORMATTER, name, "formatter", @@ -100,7 +100,7 @@ class BaseStep(BaseComponent): return ( name if isinstance(name, TokenCounterBase) - else self._get_component( + else self.get_component( ComponentEnum.AS_TOKEN_COUNTER, name, "token_counter", @@ -110,7 +110,7 @@ class BaseStep(BaseComponent): @property def file_store(self) -> BaseFileStore: name = self.kwargs.get("file_store", "default") - return name if isinstance(name, BaseFileStore) else self._get_component(ComponentEnum.FILE_STORE, name) + return name if isinstance(name, BaseFileStore) else self.get_component(ComponentEnum.FILE_STORE, name) @property def embedding(self) -> BaseEmbeddingModel: @@ -118,7 +118,7 @@ class BaseStep(BaseComponent): return ( name if isinstance(name, BaseEmbeddingModel) - else self._get_component( + else self.get_component( ComponentEnum.EMBEDDING_MODEL, name, ) diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 1c4ebebe..04004a3d 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -45,8 +45,7 @@ class BaseEmbeddingModel(BaseComponent): @property def cache_path(self) -> Path: - working_dir = self.app_context.app_config.working_dir if self.app_context else "" - return Path(working_dir) / "embedding_cache" / f"{self.name}.npz" + return self.working_path / "embedding_cache" / f"{self.name}.npz" async def _start(self) -> None: self._embedding_cache.clear() diff --git a/reme2/component/file_parser/base_file_parser.py b/reme2/component/file_parser/base_file_parser.py index 8b6f65dc..a37d1795 100644 --- a/reme2/component/file_parser/base_file_parser.py +++ b/reme2/component/file_parser/base_file_parser.py @@ -1,10 +1,4 @@ -"""Abstract base class for file parsers. - -Single-pass parser: `parse(path)` returns `(FileNode, list[FileChunk])`. -Embeddings are NOT attached here — the file_store owns the embedding -pipeline (it has the cached chunks needed for hash-diff and the -embedding model handle). -""" +"""Abstract base class for file parsers.""" from abc import abstractmethod from pathlib import Path @@ -25,14 +19,7 @@ class BaseFileParser(BaseComponent): def _get_relative_path(self, path: str | Path) -> str: - """Get path relative to working_dir. - - Args: - path: Absolute or relative path to the file. - - Returns: - Path relative to working_dir. - """ + """Get path relative to working_dir.""" file_path = Path(path).absolute() working_path = Path(self.working_dir).absolute() try: diff --git a/reme2/component/file_parser/default_file_parser.py b/reme2/component/file_parser/default_file_parser.py index d644c955..b1f30bae 100644 --- a/reme2/component/file_parser/default_file_parser.py +++ b/reme2/component/file_parser/default_file_parser.py @@ -1,5 +1,3 @@ -"""Default file parser with byte-based chunking.""" - from pathlib import Path import aiofiles @@ -7,6 +5,7 @@ import aiofiles from .base_file_parser import BaseFileParser from ..component_registry import R from ...schema import FileChunk, FileNode +from ...utils.common_utils import hash_text @R.register("default") @@ -17,43 +16,37 @@ class DefaultFileParser(BaseFileParser): super().__init__(**kwargs) self.encoding = encoding self.chunk_byte_size = max(100, chunk_byte_size) - # overlap >= 4 ensures truncated multibyte UTF-8 chars decode correctly in next chunk self.overlap_byte_size = max(4, overlap_byte_size) async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: file_path = Path(path) stat = file_path.stat() - relative_path = self._get_relative_path(path) + rel_path = self._get_relative_path(path) async with aiofiles.open(file_path, "rb") as f: data = await f.read() if not data: - return FileNode(path=relative_path, st_mtime=stat.st_mtime), [] + return FileNode(path=rel_path, st_mtime=stat.st_mtime), [] - # Find all newline byte positions newline_positions = [i for i, b in enumerate(data) if b == ord(b"\n")] - chunks: list[FileChunk] = [] step = self.chunk_byte_size - self.overlap_byte_size - start_byte = 0 + start = 0 - while start_byte < len(data): - end_byte = min(start_byte + self.chunk_byte_size, len(data)) - chunk_data = data[start_byte:end_byte] - text = chunk_data.decode(self.encoding, errors="ignore") - - # Calculate line numbers from byte positions - start_line = sum(1 for p in newline_positions if p < start_byte) + 1 - end_line = sum(1 for p in newline_positions if p < end_byte) + 1 + while start < len(data): + end = min(start + self.chunk_byte_size, len(data)) + text = data[start:end].decode(self.encoding, errors="ignore") + start_line = sum(1 for p in newline_positions if p < start) + 1 + end_line = sum(1 for p in newline_positions if p < end) + 1 chunks.append(FileChunk( - path=relative_path, + path=rel_path, start_line=start_line, end_line=end_line, text=text, - )) + ).set_hash_id()) - start_byte += step if end_byte < len(data) else len(data) + start += step if end < len(data) else len(data) - return FileNode(path=relative_path, st_mtime=stat.st_mtime), chunks + return FileNode(path=rel_path, st_mtime=stat.st_mtime, chunk_ids=[c.id for c in chunks]), chunks \ No newline at end of file diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 68d98d9d..906b97e2 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -1,80 +1,73 @@ -"""Abstract base class for file stores.""" - import re from abc import abstractmethod -from pathlib import Path -from .bm25_lite import BM25Lite from ..base_component import BaseComponent from ..embedding import BaseEmbeddingModel +from ..keyword_index import BaseKeywordIndex from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode class BaseFileStore(BaseComponent): - """Abstract file-store engine.""" - component_type = ComponentEnum.FILE_STORE def __init__( self, store_name: str, embedding_model: str = "default", - tokenizer: str = "default", + keyword_index: str = "default", fts_enabled: bool = True, **kwargs, ): - """Initialize the file store.""" super().__init__(**kwargs) if not re.match(r"^[a-zA-Z0-9_]+$", store_name): raise ValueError(f"Invalid store name '{store_name}'. Only alphanumeric and underscores allowed.") self.store_name = store_name or self.name self._embedding_model_name = embedding_model - self.tokenizer = tokenizer + self._keyword_index_name = keyword_index self.fts_enabled = fts_enabled self.embedding_model: BaseEmbeddingModel | None = None + self.keyword_index: BaseKeywordIndex | None = None self.vector_enabled = bool(embedding_model) - self.working_dir = self.app_context.app_config.working_dir if self.app_context else "" - self.store_path = Path(self.working_dir) / "file_store" / store_name + self.store_path = self.working_path / self.component_type.value / store_name self.store_path.mkdir(parents=True, exist_ok=True) if not self.vector_enabled and not self.fts_enabled: raise ValueError("At least one of embedding_model or fts_enabled must be set.") - self.bm25: BM25Lite | None = None + self.file_nodes: dict[str, FileNode] = {} async def _start(self) -> None: - """Start the file store.""" - if self.vector_enabled and self.app_context is not None: - model_dict = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {}) - if self._embedding_model_name not in model_dict: - raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.") - model = model_dict[self._embedding_model_name] - if not isinstance(model, BaseEmbeddingModel): - raise TypeError(f"Expected BaseEmbeddingModel, got {type(model).__name__}") - self.embedding_model = model - + if self.vector_enabled: + self.embedding_model = self.get_component(ComponentEnum.EMBEDDING_MODEL, self._embedding_model_name) if self.fts_enabled: - self.bm25 = BM25Lite(index_dir=self.store_path, tokenizer=self.tokenizer, app_context=self.app_context) - if self.bm25 is not None: - await self.bm25.start() + self.keyword_index = self.get_component(ComponentEnum.KEYWORD_INDEX, self._keyword_index_name) + await self.load_file_nodes() async def _close(self) -> None: - """Close the file store and release resources.""" - self.embedding_model = None - if self.fts_enabled and self.bm25 is not None: - await self.bm25.close() + if self.vector_enabled: + self.embedding_model = None + if self.fts_enabled: + self.keyword_index = None + await self.dump_file_nodes() - # Composite operations + async def load_file_nodes(self): + ... - async def upsert_file(self, node: FileNode, chunks: list[FileChunk]) -> None: - """Upsert a node and its associated chunks.""" + async def dump_file_nodes(self): + ... - async def delete(self, path: str) -> None: - """Delete a node and all its associated chunks by path.""" + async def upsert_file( + self, + file: tuple[FileNode, list[FileChunk]] | list[tuple[FileNode, list[FileChunk]]], + ) -> None: + """Upsert a file and its chunks into the store.""" - async def reindex(self) -> None: - """Re-index all nodes and chunks in the store.""" + async def delete_by_path(self, path: str | list[str]) -> None: + """Delete files by their paths from the store.""" + + async def clear(self): + """Clear the store of all files and chunks.""" @abstractmethod async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: @@ -83,3 +76,6 @@ class BaseFileStore(BaseComponent): @abstractmethod async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: """Perform full-text keyword search.""" + + async def graph_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk] | None: + """Perform graph search.""" diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index 6ca09f8f..9292e850 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -20,35 +20,30 @@ class LocalFileStore(BaseFileStore): def __init__(self, encoding: str = "utf-8", **kwargs): super().__init__(**kwargs) - self._encoding = encoding - self._nodes: dict[str, FileNode] = {} - self._chunks: dict[str, FileChunk] = {} - self._nodes_file = self.store_path / "nodes.jsonl" - self._chunks_file = self.store_path / "chunks.jsonl" + self.encoding = encoding + self.file_chunks: dict[str, FileChunk] = {} + self.nodes_path = self.store_path / "file_nodes.jsonl" + self.chunks_path = self.store_path / "file_chunks.jsonl" # Lifecycle async def _start(self) -> None: await super()._start() - await self._load(self._nodes_file, self._nodes, FileNode, "path") - await self._load(self._chunks_file, self._chunks, FileChunk, "id") + await self._load_jsonl(self.chunks_path, self.file_chunks, FileChunk, "id") self.logger.info(f"LocalFileStore '{self.store_name}' ready: " - f"{len(self._nodes)} nodes, {len(self._chunks)} chunks") + f"{len(self.file_nodes)} nodes, {len(self.file_chunks)} chunks") async def _close(self) -> None: - await self._dump(self._nodes_file, list(self._nodes.values())) - await self._dump(self._chunks_file, list(self._chunks.values())) - self._nodes.clear() - self._chunks.clear() + await self._dump_jsonl(self.chunks_path, list(self.file_chunks.values())) + self.file_chunks.clear() await super()._close() - async def _load(self, file: Path, target: dict, model: type[BaseModel], key: str) -> None: + async def _load_jsonl(self, file: Path, target: dict, model: type[BaseModel], key: str) -> None: if not file.exists(): return - target.clear() try: - async with aiofiles.open(file, encoding=self._encoding) as f: + async with aiofiles.open(file, encoding=self.encoding) as f: async for line in f: line = line.strip() if line: @@ -57,65 +52,76 @@ class LocalFileStore(BaseFileStore): except Exception as e: self.logger.exception(f"Failed to load {file}: {e}") - async def _dump(self, file: Path, items: list[BaseModel]) -> None: + async def _dump_jsonl(self, file: Path, items: list[BaseModel]) -> None: try: content = "\n".join(o.model_dump_json() for o in items) tmp = file.with_suffix(".tmp") - async with aiofiles.open(tmp, "w", encoding=self._encoding) as f: + async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: await f.write(content) tmp.replace(file) except Exception as e: self.logger.exception(f"Failed to write {file}: {e}") - # Node operations + # Base class interface - async def upsert_node(self, node: FileNode) -> None: - self._nodes[node.path] = node + async def load_file_nodes(self) -> None: + await self._load_jsonl(self.nodes_path, self.file_nodes, FileNode, "path") - async def get_node_by_path(self, path: str) -> FileNode | None: - return self._nodes.get(path) + async def dump_file_nodes(self) -> None: + await self._dump_jsonl(self.nodes_path, list(self.file_nodes.values())) - async def delete_node_by_path(self, path: str) -> FileNode | None: - return self._nodes.pop(path, None) + async def upsert_file( + self, + file: tuple[FileNode, list[FileChunk]] | list[tuple[FileNode, list[FileChunk]]], + ) -> None: + if isinstance(file, tuple): + file = [file] + for node, chunks in file: + old_node = self.file_nodes.pop(node.path, None) + cached = {} + if old_node and self.vector_enabled: + for cid in old_node.chunk_ids: + old = self.file_chunks.pop(cid, None) + if old and old.embedding: + cached[cid] = old.embedding - # Chunk operations + node.chunk_ids = [] + needs_embed = [] + for c in chunks: + if self.vector_enabled and not c.embedding: + if c.id in cached: + c.embedding = cached[c.id] + elif c.text: + needs_embed.append(c) + node.chunk_ids.append(c.id) + self.file_chunks[c.id] = c + self.file_nodes[node.path] = node - async def upsert_chunks_with_path(self, path: str, chunks: list[FileChunk]) -> None: - existing = await self.get_chunks_by_path(path) - cached = {c.id: c.embedding for c in existing if c.embedding} + if needs_embed and self.embedding_model: + await self.embedding_model.get_node_embeddings(needs_embed) - await self.delete_chunks_by_path(path) - if not chunks: - return + if self.fts_enabled and self.keyword_index: + await self.keyword_index.add_docs({c.id: c.text for c in chunks if c.text}) - needs_embed: list[FileChunk] = [] - for c in chunks: - if c.embedding: - continue - if c.id in cached: - c.embedding = cached[c.id] - elif c.text: - needs_embed.append(c) + async def delete_by_path(self, path: str | list[str]) -> None: + if isinstance(path, str): + path = [path] + deleted_chunk_ids: list[str] = [] + for p in path: + node = self.file_nodes.pop(p, None) + if node: + for cid in node.chunk_ids: + self.file_chunks.pop(cid, None) + deleted_chunk_ids.append(cid) - if needs_embed: - embeddings = await self.get_embeddings([c.text for c in needs_embed]) - if embeddings: - for c, emb in zip(needs_embed, embeddings): - if emb: - c.embedding = emb + if self.fts_enabled and self.keyword_index and deleted_chunk_ids: + await self.keyword_index.delete_docs(deleted_chunk_ids) - for c in chunks: - self._chunks[c.id] = c - - async def get_chunks_by_path(self, path: str) -> list[FileChunk]: - chunks = [c for c in self._chunks.values() if c.path == path] - chunks.sort(key=lambda c: c.start_line) - return chunks - - async def delete_chunks_by_path(self, path: str) -> None: - stale = [cid for cid, c in self._chunks.items() if c.path == path] - for cid in stale: - del self._chunks[cid] + async def clear(self) -> None: + self.file_nodes.clear() + self.file_chunks.clear() + if self.fts_enabled and self.keyword_index: + await self.keyword_index.clear() # Search @@ -127,15 +133,12 @@ class LocalFileStore(BaseFileStore): if not query_embedding: return [] - candidates = [c for c in self._chunks.values() if c.embedding] - emb_missing = [c for c in self._chunks.values() if not c.embedding] - if emb_missing: - logger.warning(f"Embedding missing for {len(emb_missing)} chunks") + candidates = [c for c in self.file_chunks.values() if c.embedding] if not candidates: return [] - chunk_embs = np.array([c.embedding for c in candidates]) - similarities = batch_cosine_similarity(np.array([query_embedding]), chunk_embs)[0] + candidate_embeddings = np.array([c.embedding for c in candidates]) + similarities = batch_cosine_similarity(np.array([query_embedding]), candidate_embeddings)[0] results = [ c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}}) @@ -145,31 +148,18 @@ class LocalFileStore(BaseFileStore): return results[:limit] async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - if not self.fts_enabled or self.bm25 is None: + if not self.fts_enabled or self.keyword_index is None: return [] query = query.strip() if not query: return [] - doc_id_score_dict = self.bm25.retrieve(query, limit=limit) + doc_id_score_dict = await self.keyword_index.retrieve(query, limit=limit) results = [] for doc_id, score in doc_id_score_dict.items(): - chunk = self._chunks.get(doc_id) + chunk = self.file_chunks.get(doc_id) if chunk: results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}})) return results - - # Reindex - - async def reindex_vector(self) -> None: - if not self.vector_enabled or self.embedding_model is None: - return - - await self.embedding_model.get_node_embeddings(list(self._chunks.values())) - - async def reindex_keyword(self) -> None: - if not self.fts_enabled or self.bm25 is None: - return - diff --git a/reme2/component/file_watcher/__init__.py b/reme2/component/file_watcher/__init__.py index 69adcb01..0b11d995 100644 --- a/reme2/component/file_watcher/__init__.py +++ b/reme2/component/file_watcher/__init__.py @@ -1,11 +1,9 @@ """File watcher implementations for monitoring file system changes.""" from .base_file_watcher import BaseFileWatcher -from .full_file_watcher import FullFileWatcher -from .light_file_watcher import LightFileWatcher +from .lite_file_watcher import LiteFileWatcher __all__ = [ "BaseFileWatcher", - "FullFileWatcher", - "LightFileWatcher", + "LiteFileWatcher", ] diff --git a/reme2/component/file_watcher/base_file_watcher.py b/reme2/component/file_watcher/base_file_watcher.py index 08beb5cc..c1204958 100644 --- a/reme2/component/file_watcher/base_file_watcher.py +++ b/reme2/component/file_watcher/base_file_watcher.py @@ -1,6 +1,16 @@ +import asyncio +from abc import abstractmethod +from pathlib import Path + +from watchfiles import Change + from ..base_component import BaseComponent -from ...enumeration import ComponentEnum from ..file_parser import BaseFileParser +from ..file_store import BaseFileStore +from ...enumeration import ComponentEnum +from ...utils import get_logger + +logger = get_logger() class BaseFileWatcher(BaseComponent): @@ -8,31 +18,94 @@ class BaseFileWatcher(BaseComponent): def __init__( self, - watch_path: list[str], - recursive: bool = False, - file_parser: str = "default", - file_store: str = "default", + watch_paths: list[str] | str, + suffix_filters: list[str] | None = None, + recursive: bool = True, + force_polling: bool = True, debounce: int = 2000, poll_delay_ms: int = 2000, + file_store: str = "default", + file_parser: str = "default", **kwargs, ): super().__init__(**kwargs) - self.watch_path: list[str] = watch_path + watch_paths = [watch_paths] if isinstance(watch_paths, str) else watch_paths + self.watch_paths: list[Path] = [self.working_path / x for x in watch_paths if (self.working_path / x).exists()] + self.suffix_filters: list[str] = suffix_filters or ["md"] self.recursive: bool = recursive - self.file_parser_name: str = file_parser - self.file_store: str = file_store + self.force_polling: bool = force_polling self.debounce: int = debounce self.poll_delay_ms: int = poll_delay_ms - + self.file_store_name: str = file_store + self.file_parser_name: str = file_parser + self._stop_event: asyncio.Event = asyncio.Event() + self._background_task: asyncio.Task | None = None + self.file_store: BaseFileStore | None = None self.file_parser: BaseFileParser | None = None + self._retry_interval: float = 10 - async def _start(self) -> None: - await super()._start() - if self.app_context is not None: - file_parser_dict = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) - if self.file_parser not in file_parser_dict: - raise ValueError(f"File parser {self.file_parser} not found") - self.file_parser = file_parser_dict[self.file_parser] + async def _start(self): + self.file_store = self.get_component(ComponentEnum.FILE_STORE, self.file_store_name) + self.file_parser = self.get_component(ComponentEnum.FILE_PARSER, self.file_parser_name) - async def _close(self) -> None: - await super()._close() \ No newline at end of file + async def background_task(): + await self.update_store() + await self.watch_loop() + + self._background_task = asyncio.create_task(background_task()) + logger.info(f"Started watching: {self.watch_paths}") + + async def _close(self): + self._stop_event.set() + if self._background_task: + await self._background_task + logger.info("Stopped watching") + + def watch_filter(self, _change: Change, path: str) -> bool: + if not self.suffix_filters: + return True + return any(path.endswith("." + s.strip(".")) for s in self.suffix_filters) + + async def scan_existing_files(self) -> list[Path]: + files: list[Path] = [] + for path in self.watch_paths: + if not path.exists(): + continue + if path.is_file(): + if self.watch_filter(Change.added, str(path)): + files.append(path) + else: + items = path.rglob("*") if self.recursive else path.iterdir() + files.extend(p for p in items if p.is_file() and self.watch_filter(Change.added, str(p))) + return files + + async def clear_store(self): + if self.file_store is None: + raise ValueError("file_store is not initialized!") + await self.file_store.clear() + + async def reset_store(self): + if self.file_store is None: + raise ValueError("file_store is not initialized!") + await self.file_store.clear() + await self.on_added(await self.scan_existing_files()) + + @abstractmethod + async def watch_loop(self): + """Watch for file changes and update the store.""" + + @abstractmethod + async def update_store(self): + """Update the store with the latest file changes.""" + + @abstractmethod + async def on_added(self, path: Path | list[Path]): + """Handle file added event.""" + + @abstractmethod + async def on_modified(self, path: Path | list[Path]): + """Handle file modified event.""" + + @abstractmethod + async def on_deleted(self, path: Path | list[Path]): + """Handle file deleted event.""" diff --git a/reme2/component/file_watcher/base_file_watcher_old.py b/reme2/component/file_watcher/base_file_watcher_old.py deleted file mode 100644 index babf7320..00000000 --- a/reme2/component/file_watcher/base_file_watcher_old.py +++ /dev/null @@ -1,427 +0,0 @@ -"""Base file watcher with watchfiles integration. - -Per change → single-step pipeline: - added/modified → node, chunks = parser.parse(path) - file_store.upsert(node, chunks) - deleted → file_store.delete_chunks + delete_node - -The parser owns: chunking and edge extraction (text-only, no embedding). -The file_store owns: persistence of node + chunks AND the embedding -pipeline — `upsert` hash-diffs incoming chunks against its persisted -ones, reuses cached embeddings for unchanged blocks, and calls the -embedding API only for new-hash blocks. -The watcher owns: scheduling, cancel-on-modify, retry/timeout, and -startup recovery (re-parsing files whose mtime drifted while offline). -""" - -import asyncio -import os -import time -from pathlib import Path - -from watchfiles import Change, awatch - -from ..base_component import BaseComponent -from ..file_parser import BaseFileParser -from ..file_store import BaseFileStore -from ...enumeration import ComponentEnum - - -class BaseFileWatcher(BaseComponent): - """Watches a directory and feeds the file_store via single-step parse+upsert.""" - - component_type = ComponentEnum.FILE_WATCHER - - _META_DIR = ".reme" - - def __init__( - self, - recursive: bool = False, - debounce: int = 2000, - file_store: str = "default", - default_parser: str | None = None, - rebuild_index_on_start: bool = False, - poll_delay_ms: int = 2000, - parse_max_attempts: int = 3, - parse_retry_backoff: float = 2.0, - parse_task_timeout: float = 300.0, - **kwargs, - ): - super().__init__(**kwargs) - self._file_store_name: str = file_store - self._default_parser_name: str | None = default_parser - self.file_store: BaseFileStore | None = None - self._suffix_to_parser: dict[str, BaseFileParser] = {} - self._default_parser: BaseFileParser | None = None - self.recursive: bool = recursive - self.debounce: int = debounce - self.rebuild_index_on_start: bool = rebuild_index_on_start - self.poll_delay_ms: int = poll_delay_ms - - self._stop_event = asyncio.Event() - self._watch_task: asyncio.Task | None = None - - # Parse pipeline configuration. - self._parse_max_attempts: int = max(1, int(parse_max_attempts)) - self._parse_retry_backoff: float = float(parse_retry_backoff) - self._parse_task_timeout: float = float(parse_task_timeout) - self._tasks: dict[str, asyncio.Task] = {} - self._mgmt_lock: asyncio.Lock | None = None - self._run_lock: asyncio.Lock | None = None - self._last_failure: dict[str, float] = {} - - @property - def watch_path(self) -> str: - """Vault root the watcher is bound to (== file_store.working_dir). - - Available after _start. Reads the live value from the bound - store so config has a single source of truth. - """ - if self.file_store is None or self.file_store.working_dir is None: - raise RuntimeError( - "watch_path requires a started file_store with working_dir set", - ) - return str(self.file_store.working_dir) - - @property - def _meta_path(self) -> Path: - return (Path(self.watch_path) / self._META_DIR).resolve() - - # -- Lifecycle ---------------------------------------------------------- - - async def _start(self): - """Resolve file_store + parsers, sync any disk drift, start watch loop.""" - if self._file_store_name: - assert self.app_context is not None, "app_context must be provided" - - stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {}) - if self._file_store_name not in stores: - raise ValueError(f"File store '{self._file_store_name}' not found.") - store = stores[self._file_store_name] - if not isinstance(store, BaseFileStore): - raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}") - if store.working_dir is None: - raise ValueError( - f"File store '{self._file_store_name}' has no working_dir; " - f"set components.file_store.{self._file_store_name}.working_dir in config", - ) - self.file_store = store - - parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {}) - for parser in parsers.values(): - if isinstance(parser, BaseFileParser): - for suffix in parser.suffixes: - self._suffix_to_parser[suffix] = parser - - if self._default_parser_name and self._default_parser_name in parsers: - parser = parsers[self._default_parser_name] - if isinstance(parser, BaseFileParser): - self._default_parser = parser - - if not self._suffix_to_parser and not self._default_parser: - self.logger.warning("No file parsers registered") - - async def _initialize_and_watch(): - await self._initial_sync_and_recovery() - await self._watch_loop() - - self._mgmt_lock = asyncio.Lock() - self._run_lock = asyncio.Lock() - self._last_failure.clear() - self._stop_event.clear() - self._watch_task = asyncio.create_task(_initialize_and_watch()) - self.logger.info(f"Started watching: {self.watch_path}") - - async def _close(self): - """Stop watching and release resources.""" - self._stop_event.set() - if self._watch_task and not self._watch_task.done(): - self._watch_task.cancel() - try: - await self._watch_task - except asyncio.CancelledError: - pass - - # Cancel all in-flight parse tasks before tearing down state. - tasks = list(self._tasks.values()) - for t in tasks: - if not t.done(): - t.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - self._tasks.clear() - self._last_failure.clear() - - self._watch_task = None - self._stop_event.clear() - self._mgmt_lock = None - self._run_lock = None - self.file_store = None - self._suffix_to_parser.clear() - self._default_parser = None - self.logger.info("Stopped watching") - - # -- Startup sync ------------------------------------------------------- - - async def _initial_sync_and_recovery(self) -> None: - """At startup: file_store has its persisted state loaded already; we - diff against current disk for any changes that happened while we - were offline, then re-enqueue them through the parse pipeline.""" - await self._sync_with_disk() - - async def _sync_with_disk(self) -> bool: - """Diff cached graph (in file_store) vs disk state; reindex changes.""" - if self.file_store is None: - return False - - watch_path = Path(self.watch_path) - if not watch_path.exists(): - self.logger.warning(f"Watch path does not exist: {watch_path}") - return False - - on_disk: dict[str, float] = {} - if watch_path.is_file(): - on_disk[str(watch_path.resolve())] = watch_path.stat().st_mtime - elif watch_path.is_dir(): - iterator = watch_path.rglob("*") if self.recursive else watch_path.iterdir() - for fp in iterator: - if not fp.is_file(): - continue - abs_path = str(fp.resolve()) - if not self._watch_filter(abs_path): - continue - if not self._get_parser(abs_path): - continue - try: - on_disk[abs_path] = fp.stat().st_mtime - except OSError: - continue - - cached_paths = set(self.file_store.nodes.keys()) - disk_paths = set(on_disk.keys()) - - added = disk_paths - cached_paths - deleted = cached_paths - disk_paths - modified: set[str] = set() - for p in disk_paths & cached_paths: - cached_node = self.file_store.get_node_by_path(p) - if cached_node is None or cached_node.st_mtime != on_disk[p]: - modified.add(p) - - unchanged = len(disk_paths & cached_paths) - len(modified) - if not (added or deleted or modified): - self.logger.info(f"Index up-to-date ({unchanged} files cached, 0 changes)") - return False - - self.logger.info( - f"Incremental sync: +{len(added)} ~{len(modified)} -{len(deleted)} (unchanged {unchanged})", - ) - - changes: set[tuple[Change, str]] = set() - for p in deleted: - changes.add((Change.deleted, p)) - for p in added: - changes.add((Change.added, p)) - for p in modified: - changes.add((Change.modified, p)) - await self.on_changes(changes) - return True - - # -- Watch loop --------------------------------------------------------- - - async def _interruptible_sleep(self, seconds: float): - try: - await asyncio.wait_for(self._stop_event.wait(), timeout=seconds) - except asyncio.TimeoutError: - pass - - async def _watch_loop(self): - while not self._stop_event.is_set(): - if not Path(self.watch_path).exists(): - self.logger.warning(f"Watch path does not exist: {self.watch_path}, waiting 10 seconds...") - await self._interruptible_sleep(10) - continue - - try: - self.logger.info(f"Starting watch on: {self.watch_path}") - async for changes in awatch( - self.watch_path, - watch_filter=lambda _, p: self._watch_filter(p), - recursive=self.recursive, - debounce=self.debounce, - poll_delay_ms=self.poll_delay_ms, - stop_event=self._stop_event, - ): - if self._stop_event.is_set(): - break - await self.on_changes(changes) - - except FileNotFoundError as e: - self.logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...") - if not self._stop_event.is_set(): - await self._interruptible_sleep(10) - - except Exception as e: - self.logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True) - if not self._stop_event.is_set(): - await self._interruptible_sleep(10) - - def _watch_filter(self, path: str) -> bool: - resolved = Path(path).resolve() - if resolved == self._meta_path: - return False - return self._meta_path not in resolved.parents - - def _get_parser(self, path: str) -> BaseFileParser | None: - suffix = Path(path).suffix.lower() - return self._suffix_to_parser.get(suffix, self._default_parser) - - # -- Change dispatch ---------------------------------------------------- - - async def on_changes(self, changes: set[tuple[Change, str]]) -> None: - if self.file_store is None: - self.logger.warning("File store not initialized, skipping changes") - return - - for change_type, path in changes: - try: - if change_type in (Change.added, Change.modified): - await self._on_modified(path) - elif change_type == Change.deleted: - await self._on_deleted(path) - except FileNotFoundError: - self.logger.warning(f"File not found: {path}, skipping") - except PermissionError: - self.logger.warning(f"Permission denied: {path}, skipping") - except Exception as e: - self.logger.opt(exception=True).error("Error processing {p}: {err}", p=path, err=str(e)) - - async def _on_modified(self, path: str) -> None: - if not Path(path).is_file(): - return - parser = self._get_parser(path) - if parser is None: - self.logger.debug(f"No parser for {path}, skipping") - return - await self._submit_parse_task(path, parser) - - async def _on_deleted(self, path: str) -> None: - # Cancel any in-flight parse task FIRST so a delayed cancellation - # can't race the cleanup writes below. - assert self.file_store is not None, "_on_deleted requires file_store" - targets = [path, *self._descendant_indexed_paths(path)] - for p in targets: - await self._cancel_parse_task(p) - await self.file_store.delete_chunks_by_path(p) - await self.file_store.delete_node_by_path(p) - if len(targets) == 1: - self.logger.info(f"Deleted {path}") - else: - self.logger.info(f"Deleted directory {path} ({len(targets) - 1} indexed children)") - - def _descendant_indexed_paths(self, path: str) -> list[str]: - """Indexed file paths that live beneath `path` as a directory. - - watchfiles emits a single Change.deleted for a removed directory - (no per-file events), so we have to find the orphans ourselves. - Resolves both sides to handle symlinks and trailing-separator - drift between the watcher and the file_store keys. - """ - assert self.file_store is not None - try: - prefix = str(Path(path).resolve()) + os.sep - except OSError: - prefix = path.rstrip(os.sep) + os.sep - return [p for p in self.file_store.nodes if p.startswith(prefix)] - - # -- Parse task pipeline ------------------------------------------------ - - async def _submit_parse_task(self, path: str, parser: BaseFileParser) -> None: - """Cancel any prior task + spawn a new one. Atomic under _mgmt_lock.""" - assert self._mgmt_lock is not None, "_submit_parse_task before _start()" - async with self._mgmt_lock: - await self._cancel_locked(path) - self._tasks[path] = asyncio.create_task( - self._run_parse_task(path, parser), - ) - - async def _cancel_parse_task(self, path: str) -> None: - """Cancel + await any in-flight parse task for `path`.""" - assert self._mgmt_lock is not None, "_cancel_parse_task before _start()" - async with self._mgmt_lock: - await self._cancel_locked(path) - - async def _cancel_locked(self, path: str) -> None: - """Cancel + await the task for `path`. Caller holds `_mgmt_lock`.""" - task = self._tasks.pop(path, None) - if task is not None and not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass - - async def flush_parse_tasks(self) -> None: - """Wait until all parse tasks finish (loops because retries spawn new ones).""" - while self._tasks: - await asyncio.gather(*list(self._tasks.values()), return_exceptions=True) - - def pending_parse_count(self) -> int: - return sum(1 for t in self._tasks.values() if not t.done()) - - def failed_parse_paths(self) -> dict[str, float]: - return dict(self._last_failure) - - async def _run_parse_task(self, path: str, parser: BaseFileParser) -> None: - """One per-path parse task: parse → upsert. - - The store owns hash-diff + embedding: it inspects its own - existing chunks for `path`, reuses cached embeddings for - unchanged blocks, and only calls the embedding API for new-hash - blocks. The whole upsert is one call from this side. - """ - try: - for attempt in range(1, self._parse_max_attempts + 1): - try: - assert self.file_store is not None - node, chunks = await asyncio.wait_for( - parser.parse(path), - timeout=self._parse_task_timeout, - ) - - assert self._run_lock is not None - async with self._run_lock: - await asyncio.wait_for( - self.file_store.upsert(node, chunks), - timeout=self._parse_task_timeout, - ) - self._last_failure.pop(path, None) - - self.logger.info( - f"indexed {path}: chunks={len(chunks)} edges={len(node.edges)}", - ) - return - except asyncio.CancelledError: - raise - except FileNotFoundError: - self.logger.debug(f"parse task {path}: file gone, skipping") - self._last_failure.pop(path, None) - return - except Exception as e: - self._last_failure[path] = time.time() - if attempt < self._parse_max_attempts: - backoff = self._parse_retry_backoff * attempt - self.logger.warning( - f"parse task {path} attempt {attempt} failed " - f"({type(e).__name__}: {e}); retry in {backoff:.1f}s", - ) - await asyncio.sleep(backoff) - continue - self.logger.error( - f"parse task {path} giving up after {attempt} attempts: " - f"{type(e).__name__}: {e}", - ) - return - finally: - if self._tasks.get(path) is asyncio.current_task(): - self._tasks.pop(path, None) diff --git a/reme2/component/file_watcher/full_file_watcher.py b/reme2/component/file_watcher/full_file_watcher.py deleted file mode 100644 index c6030424..00000000 --- a/reme2/component/file_watcher/full_file_watcher.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Full file watcher for all supported file types.""" - -from .base_file_watcher import BaseFileWatcher -from ..component_registry import R - - -@R.register("full") -class FullFileWatcher(BaseFileWatcher): - """Watches all supported file types and delegates parsing to registered parsers.""" diff --git a/reme2/component/file_watcher/light_file_watcher.py b/reme2/component/file_watcher/light_file_watcher.py deleted file mode 100644 index d0115b89..00000000 --- a/reme2/component/file_watcher/light_file_watcher.py +++ /dev/null @@ -1,7 +0,0 @@ -from .base_file_watcher import BaseFileWatcher -from ..component_registry import R - - -@R.register("light") -class LightFileWatcher(BaseFileWatcher): - ... diff --git a/reme2/component/file_watcher/lite_file_watcher.py b/reme2/component/file_watcher/lite_file_watcher.py new file mode 100644 index 00000000..2a0331a9 --- /dev/null +++ b/reme2/component/file_watcher/lite_file_watcher.py @@ -0,0 +1,133 @@ +import asyncio +from pathlib import Path + +from watchfiles import Change, awatch + +from .base_file_watcher import BaseFileWatcher +from ..component_registry import R +from ...schema import FileChunk, FileNode +from ...utils import get_logger + +logger = get_logger() + + +@R.register("lite") +class LiteFileWatcher(BaseFileWatcher): + + async def _interruptible_sleep(self): + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=self._retry_interval) + except asyncio.TimeoutError: + pass + + async def watch_loop(self): + if not self.watch_paths: + logger.warning("No watch paths specified") + return + + while not self._stop_event.is_set(): + valid_paths = [p for p in self.watch_paths if p.exists()] + if not valid_paths: + logger.warning(f"No valid paths, retrying in {self._retry_interval}s...") + await self._interruptible_sleep() + continue + + invalid_paths = set(self.watch_paths) - set(valid_paths) + if invalid_paths: + logger.warning(f"Skipping invalid paths: {invalid_paths}") + + try: + logger.info(f"Watching: {valid_paths}") + async for changes in awatch( + *valid_paths, + watch_filter=self.watch_filter, + recursive=self.recursive, + force_polling=self.force_polling, + debounce=self.debounce, + poll_delay_ms=self.poll_delay_ms, + stop_event=self._stop_event, + ): + if self._stop_event.is_set(): + break + added = [Path(p) for c, p in changes if c == Change.added] + modified = [Path(p) for c, p in changes if c == Change.modified] + deleted = [Path(p) for c, p in changes if c == Change.deleted] + if added: + logger.info(f"Detected {len(added)} added file(s)") + await self.on_added(added) + if modified: + logger.info(f"Detected {len(modified)} modified file(s)") + await self.on_modified(modified) + if deleted: + logger.info(f"Detected {len(deleted)} deleted file(s)") + await self.on_deleted(deleted) + + except Exception: + logger.exception(f"Watch error, retrying in {self._retry_interval}s...") + if not self._stop_event.is_set(): + await self._interruptible_sleep() + + async def update_store(self): + if self.file_store is None: + raise ValueError("file_store is not initialized!") + + existing_paths: dict[str, float] = {str(p): p.stat().st_mtime for p in await self.scan_existing_files()} + indexed_paths: dict[str, float] = {p: n.st_mtime for p, n in self.file_store.file_nodes.items()} + + existing_keys = set(existing_paths.keys()) + indexed_keys = set(indexed_paths.keys()) + + paths_to_delete = indexed_keys - existing_keys + paths_to_add = existing_keys - indexed_keys + paths_to_modify = [p for p in existing_keys & indexed_keys if existing_paths[p] != indexed_paths[p]] + + if paths_to_modify: + logger.info(f"Updating {len(paths_to_modify)} modified file(s)") + await self.on_modified([Path(p) for p in paths_to_modify]) + + if paths_to_delete: + logger.info(f"Removing {len(paths_to_delete)} deleted file(s)") + await self.on_deleted([Path(p) for p in paths_to_delete]) + + if paths_to_add: + logger.info(f"Indexing {len(paths_to_add)} new file(s)") + await self.on_added([Path(p) for p in paths_to_add]) + + if not paths_to_modify and not paths_to_delete and not paths_to_add: + logger.info("Store is up to date") + + async def on_added(self, path: Path | list[Path]): + if self.file_parser is None or self.file_store is None: + raise RuntimeError("file_parser or file_store is not initialized!") + + paths = [path] if isinstance(path, Path) else path + parsed: list[tuple[FileNode, list[FileChunk]]] = [] + for p in paths: + if p.is_file(): + logger.info(f"Adding file: {p}") + parsed.append(await self.file_parser.parse(p)) + if parsed: + await self.file_store.delete_by_path([str(p) for p in paths if p.is_file()]) + await self.file_store.upsert_file(parsed) + + async def on_modified(self, path: Path | list[Path]): + if self.file_parser is None or self.file_store is None: + raise RuntimeError("file_parser or file_store is not initialized!") + + paths = [path] if isinstance(path, Path) else path + parsed: list[tuple[FileNode, list[FileChunk]]] = [] + for p in paths: + if p.is_file(): + logger.info(f"Updating file: {p}") + parsed.append(await self.file_parser.parse(p)) + if parsed: + await self.file_store.delete_by_path([str(p) for p in paths if p.is_file()]) + await self.file_store.upsert_file(parsed) + + async def on_deleted(self, path: Path | list[Path]): + if self.file_store is None: + raise RuntimeError("file_store is not initialized!") + + paths = [path] if isinstance(path, Path) else path + logger.info(f"Deleting {len(paths)} file(s)") + await self.file_store.delete_by_path([str(p) for p in paths]) diff --git a/reme2/component/keyword_index/__init__.py b/reme2/component/keyword_index/__init__.py new file mode 100644 index 00000000..fc08222c --- /dev/null +++ b/reme2/component/keyword_index/__init__.py @@ -0,0 +1,4 @@ +from .base_keyword_index import BaseKeywordIndex +from .bm25_index import BM25Index + +__all__ = ["BaseKeywordIndex", "BM25Index"] diff --git a/reme2/component/keyword_index/base_keyword_index.py b/reme2/component/keyword_index/base_keyword_index.py new file mode 100644 index 00000000..82b957be --- /dev/null +++ b/reme2/component/keyword_index/base_keyword_index.py @@ -0,0 +1,106 @@ +"""Abstract base class for keyword index implementations.""" + +from abc import abstractmethod +from pathlib import Path + +from ..base_component import BaseComponent +from ..tokenizer import BaseTokenizer +from ...enumeration import ComponentEnum + + +class BaseKeywordIndex(BaseComponent): + """Abstract base class for keyword index implementations.""" + + component_type = ComponentEnum.KEYWORD_INDEX + + def __init__(self, tokenizer: str = "default", **kwargs): + super().__init__(**kwargs) + self.tokenizer_name = tokenizer + self.tokenizer: BaseTokenizer | None = None + self.index_path = self.working_path / self.component_type.value + self.index_path.mkdir(parents=True, exist_ok=True) + + async def _start(self) -> None: + """Initialize tokenizer and load existing index if available.""" + if self.app_context is None: + from ..tokenizer import RegexTokenizer + self.tokenizer = RegexTokenizer(filter_stopwords=False) + else: + self.tokenizer = self.get_component(ComponentEnum.TOKENIZER, self.tokenizer_name) + + if self.tokenizer is not None: + await self.tokenizer.start() + + if self.index_file.exists(): + await self.load() + self.logger.info(f"Loaded index from {self.index_path}") + + async def _close(self) -> None: + """Save index and cleanup tokenizer on shutdown.""" + await self.dump() + self.logger.info(f"Saved index to {self.index_path}") + + if self.tokenizer is not None: + await self.tokenizer.close() + + @property + def index_file(self) -> Path: + """Path to the index pickle file based on tokenizer name.""" + if self.tokenizer is None: + raise RuntimeError("Tokenizer not initialized. Call start() first.") + name = type(self.tokenizer).__name__.replace("Tokenizer", "").lower() + return self.index_path / f"bm25_{name}.pkl" + + def _tokenize(self, text: str) -> list[str]: + if self.tokenizer is None: + raise RuntimeError("Tokenizer not initialized. Call start() first.") + return self.tokenizer.tokenize([text])[0] + + @abstractmethod + async def add_docs(self, docs_dict: dict[str, str]) -> None: + """Index or update multiple documents. + + Args: + docs_dict: Mapping of document ID to document content. + """ + + @abstractmethod + async def delete_docs(self, doc_ids: list[str]) -> None: + """Remove documents from the index. + + Args: + doc_ids: List of document IDs to remove. + """ + + @abstractmethod + async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: + """Search for documents matching the query. + + Args: + query: Search query string. + limit: Maximum number of results to return. + + Returns: + Dict mapping document IDs to scores, sorted by score descending. + """ + + @abstractmethod + async def dump(self) -> None: + """Persist index to disk.""" + + @abstractmethod + async def load(self) -> None: + """Load index from disk.""" + + @abstractmethod + async def clear(self) -> None: + """Reset the index to empty state.""" + + async def reset_index(self, docs_dict: dict[str, str]) -> None: + """Reset the index and re-add documents.""" + await self.clear() + await self.add_docs(docs_dict) + await self.dump() + + async def optimize_index(self) -> None: + """Optimize index for performance.""" diff --git a/reme2/component/file_store/bm25_lite.py b/reme2/component/keyword_index/bm25_index.py similarity index 72% rename from reme2/component/file_store/bm25_lite.py rename to reme2/component/keyword_index/bm25_index.py index 1265a6a2..17dfa72d 100644 --- a/reme2/component/file_store/bm25_lite.py +++ b/reme2/component/keyword_index/bm25_index.py @@ -1,6 +1,6 @@ -"""Lightweight BM25 search engine with persistent index support. +"""BM25 search engine with persistent index support. -BM25Lite implements the Okapi BM25 ranking algorithm for text retrieval. +BM25Index implements the Okapi BM25 ranking algorithm for text retrieval. It uses an inverted index for efficient document lookup and supports incremental updates, persistence via pickle, and automatic vocab compaction. """ @@ -8,12 +8,9 @@ incremental updates, persistence via pickle, and automatic vocab compaction. import math import pickle from collections import Counter -from pathlib import Path from typing import TypedDict -from ..base_component import BaseComponent -from ..tokenizer import BaseTokenizer -from ...enumeration import ComponentEnum +from .base_keyword_index import BaseKeywordIndex class DocMeta(TypedDict): @@ -27,8 +24,8 @@ class DocMeta(TypedDict): token_ids: set[int] -class BM25Lite(BaseComponent): - """Lightweight BM25 search engine with file-based persistence. +class BM25Index(BaseKeywordIndex): + """BM25 search engine with file-based persistence. BM25 (Best Matching 25) is a probabilistic ranking function that scores documents based on term frequency and document length normalization. @@ -47,34 +44,42 @@ class BM25Lite(BaseComponent): tokenizer: Name of tokenizer component to use. """ - def __init__( - self, - index_dir: str | Path, - k1: float = 1.5, - b: float = 0.75, - tokenizer: str = "default", - **kwargs): + def __init__(self, k1: float = 1.5, b: float = 0.75, **kwargs): super().__init__(**kwargs) self.k1 = k1 self.b = b - self.index_dir = Path(index_dir) - self.index_dir.mkdir(parents=True, exist_ok=True) self.vocab: dict[str, int] = {} self.inverted_index: dict[int, dict[str, int]] = {} self.doc_meta: dict[str, DocMeta] = {} self.total_len = 0 - self._idf_cache: dict[int, float] = {} - self.tokenizer_name = tokenizer - self._tokenizer: BaseTokenizer | None = None - def clear(self): - """Reset the index to empty state.""" - self.vocab = {} - self.inverted_index = {} - self.doc_meta = {} - self.total_len = 0 - self._idf_cache = {} + def _tokens_to_ids(self, tokens: list[str]) -> list[int]: + ids = [] + for token in tokens: + token = token.strip() + if token: + ids.append(self.vocab.setdefault(token, len(self.vocab))) + return ids + + def _remove_doc(self, doc_id: str) -> None: + if doc_id not in self.doc_meta: + return + meta = self.doc_meta[doc_id] + self.total_len -= meta["len"] + for tid in meta["token_ids"]: + if tid in self.inverted_index: + self.inverted_index[tid].pop(doc_id, None) + if not self.inverted_index[tid]: + del self.inverted_index[tid] + del self.doc_meta[doc_id] + + def _get_idf(self, token_id: int) -> float: + if token_id in self._idf_cache: + return self._idf_cache[token_id] + df = len(self.inverted_index.get(token_id, {})) + self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0 + return self._idf_cache[token_id] @property def n_docs(self) -> int: @@ -86,55 +91,7 @@ class BM25Lite(BaseComponent): """Average document length in tokens.""" return self.total_len / self.n_docs if self.n_docs > 0 else 0.0 - @property - def index_file(self) -> Path: - """Path to the index pickle file based on tokenizer name.""" - if self._tokenizer is None: - raise RuntimeError("Tokenizer not initialized. Call start() first.") - name = type(self._tokenizer).__name__.replace("Tokenizer", "").lower() - return self.index_dir / f"index_{name}.pkl" - - async def _start(self) -> None: - """Initialize tokenizer and load existing index if available.""" - if self.app_context is not None: - tokenizer_dict: dict = self.app_context.components.get(ComponentEnum.TOKENIZER, {}) - if self.tokenizer_name in tokenizer_dict: - self._tokenizer = tokenizer_dict[self.tokenizer_name] - - if self._tokenizer is None: - from ..tokenizer import RegexTokenizer - self._tokenizer = RegexTokenizer(filter_stopwords=False) - - if self._tokenizer is not None: - await self._tokenizer.start() - - if self.index_file.exists(): - self.load() - self.logger.info(f"Loaded index from {self.index_dir}") - - async def _close(self) -> None: - """Save index and cleanup tokenizer on shutdown.""" - if self.inverted_index: - self.dump() - self.logger.info(f"Saved index to {self.index_dir}") - - if self._tokenizer is not None: - await self._tokenizer.close() - - def _tokenize(self, text: str) -> list[str]: - if self._tokenizer is None: - raise RuntimeError("Tokenizer not initialized. Call start() first.") - return self._tokenizer.tokenize([text])[0] - - def _tokens_to_ids(self, tokens: list[str]) -> list[int]: - ids = [] - for token in tokens: - token = token.strip() - if token: - ids.append(self.vocab.setdefault(token, len(self.vocab))) - return ids - - def add_docs(self, docs_dict: dict[str, str]): + async def add_docs(self, docs_dict: dict[str, str]) -> None: """Index or update multiple documents. Args: @@ -159,60 +116,17 @@ class BM25Lite(BaseComponent): self._idf_cache = {} - def reindex(self): - """Rebuild vocab to remove unused tokens and compact token IDs.""" - # Collect all token IDs still in use - used_token_ids: set[int] = set() - for tid in self.inverted_index: - used_token_ids.add(tid) + async def delete_docs(self, doc_ids: list[str]) -> None: + """Remove documents from the index. - if not used_token_ids: - self.clear() - return - - # Build new vocab with compact IDs - old_to_new: dict[int, int] = {} - new_vocab: dict[str, int] = {} - for token, old_tid in self.vocab.items(): - if old_tid in used_token_ids: - new_tid = len(new_vocab) - new_vocab[token] = new_tid - old_to_new[old_tid] = new_tid - - # Rebuild inverted_index with new token IDs - new_inverted_index: dict[int, dict[str, int]] = {} - for old_tid, postings in self.inverted_index.items(): - new_tid = old_to_new[old_tid] - new_inverted_index[new_tid] = postings - - # Update doc_meta token_ids - for doc_id, meta in self.doc_meta.items(): - meta["token_ids"] = {old_to_new[old_tid] for old_tid in meta["token_ids"] if old_tid in old_to_new} - - self.vocab = new_vocab - self.inverted_index = new_inverted_index + Args: + doc_ids: List of document IDs to remove. + """ + for doc_id in doc_ids: + self._remove_doc(doc_id) self._idf_cache = {} - def _remove_doc(self, doc_id: str): - if doc_id not in self.doc_meta: - return - meta = self.doc_meta[doc_id] - self.total_len -= meta["len"] - for tid in meta["token_ids"]: - if tid in self.inverted_index: - self.inverted_index[tid].pop(doc_id, None) - if not self.inverted_index[tid]: - del self.inverted_index[tid] - del self.doc_meta[doc_id] - - def _get_idf(self, token_id: int) -> float: - if token_id in self._idf_cache: - return self._idf_cache[token_id] - df = len(self.inverted_index.get(token_id, {})) - self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0 - return self._idf_cache[token_id] - - def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: + async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]: """Search for documents matching the query. Args: @@ -240,7 +154,7 @@ class BM25Lite(BaseComponent): return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:limit]) if scores else {} - def dump(self): + async def dump(self) -> None: """Persist index to disk via pickle.""" with open(self.index_file, "wb") as f: pickle.dump( @@ -255,7 +169,7 @@ class BM25Lite(BaseComponent): f, ) - def load(self): + async def load(self) -> None: """Load index from disk. Clears index on failure.""" try: with open(self.index_file, "rb") as f: @@ -270,4 +184,46 @@ class BM25Lite(BaseComponent): except Exception as e: self.logger.exception(f"Failed to load index: {e}") self.index_file.unlink(missing_ok=True) - self.clear() + await self.clear() + + async def clear(self) -> None: + """Reset the index to empty state.""" + self.vocab = {} + self.inverted_index = {} + self.doc_meta = {} + self.total_len = 0 + self._idf_cache = {} + + async def optimize_index(self) -> None: + """Rebuild vocab to remove unused tokens and compact token IDs.""" + # Collect all token IDs still in use + used_token_ids: set[int] = set() + for tid in self.inverted_index: + used_token_ids.add(tid) + + if not used_token_ids: + await self.clear() + return + + # Build new vocab with compact IDs + old_to_new: dict[int, int] = {} + new_vocab: dict[str, int] = {} + for token, old_tid in self.vocab.items(): + if old_tid in used_token_ids: + new_tid = len(new_vocab) + new_vocab[token] = new_tid + old_to_new[old_tid] = new_tid + + # Rebuild inverted_index with new token IDs + new_inverted_index: dict[int, dict[str, int]] = {} + for old_tid, postings in self.inverted_index.items(): + new_tid = old_to_new[old_tid] + new_inverted_index[new_tid] = postings + + # Update doc_meta token_ids + for doc_id, meta in self.doc_meta.items(): + meta["token_ids"] = {old_to_new[old_tid] for old_tid in meta["token_ids"] if old_tid in old_to_new} + + self.vocab = new_vocab + self.inverted_index = new_inverted_index + self._idf_cache = {} diff --git a/reme2/enumeration/component_enum.py b/reme2/enumeration/component_enum.py index 1a40c6e9..d3878720 100644 --- a/reme2/enumeration/component_enum.py +++ b/reme2/enumeration/component_enum.py @@ -30,6 +30,8 @@ class ComponentEnum(str, Enum): FILE_WATCHER = "file_watcher" + KEYWORD_INDEX = "keyword_index" + SERVICE = "service" CLIENT = "client" diff --git a/reme2/schema/file_chunk.py b/reme2/schema/file_chunk.py index 3ccb7103..4f8077cc 100644 --- a/reme2/schema/file_chunk.py +++ b/reme2/schema/file_chunk.py @@ -12,3 +12,8 @@ class FileChunk(EmbNode): @property def score(self) -> float: return self.scores.get("score", 0.0) + + def set_hash_id(self): + from ..utils import hash_text + self.id = hash_text(" ".join([self.path, str(self.start_line), str(self.end_line), self.text])) + return self.id \ No newline at end of file diff --git a/test/reme2/test_bm25_index_perf.py b/test/reme2/test_bm25_index_perf.py new file mode 100644 index 00000000..de7ee00d --- /dev/null +++ b/test/reme2/test_bm25_index_perf.py @@ -0,0 +1,186 @@ +"""BM25Index performance tests for add_docs and retrieve.""" + +import asyncio +import random +import time +import tempfile +from pathlib import Path + +from reme2.component.keyword_index.bm25_index import BM25Index +from reme2.component.tokenizer import RegexTokenizer + +# A small vocab of realistic-looking words for generating random text +_VOCAB = [ + "algorithm", "data", "machine", "learning", "model", "network", "neural", "training", + "optimization", "gradient", "loss", "function", "parameter", "weight", "bias", "layer", + "activation", "relu", "sigmoid", "softmax", "backpropagation", "forward", "pass", + "batch", "epoch", "iteration", "convergence", "divergence", "regularization", "dropout", + "attention", "transformer", "encoder", "decoder", "embedding", "token", "vector", + "matrix", "tensor", "computation", "graph", "node", "edge", "vertex", "path", + "search", "retrieval", "index", "query", "document", "corpus", "term", "frequency", + "inverse", "score", "rank", "relevance", "precision", "recall", "f1", "metric", + "evaluation", "benchmark", "dataset", "sample", "feature", "label", "class", "predict", + "classification", "regression", "clustering", "dimension", "reduction", "pca", "tsne", + "visualization", "matplotlib", "plot", "chart", "histogram", "scatter", "line", "bar", + "database", "sql", "query", "table", "row", "column", "index", "primary", "foreign", + "key", "constraint", "schema", "migration", "version", "control", "git", "commit", + "branch", "merge", "conflict", "resolution", "review", "approve", "reject", "pull", + "request", "issue", "bug", "fix", "feature", "enhancement", "refactor", "test", + "deploy", "production", "staging", "development", "environment", "configuration", + "setting", "variable", "constant", "global", "local", "scope", "closure", "callback", + "promise", "async", "await", "synchronous", "asynchronous", "concurrent", "parallel", + "thread", "process", "memory", "cache", "buffer", "queue", "stack", "heap", "pool", +] + + +def _gen_random_text(n_tokens: int) -> str: + """Generate random text with approximately n_tokens words.""" + words = random.choices(_VOCAB, k=n_tokens) + return " ".join(words) + + +def _gen_random_query(n_words: int) -> str: + """Generate a random query with n_words words.""" + words = random.choices(_VOCAB, k=n_words) + return " ".join(words) + + +async def _make_index(tmp_dir: Path) -> BM25Index: + """Create and start a BM25Index instance, bypassing app_context.""" + index = BM25Index.__new__(BM25Index) + # Manually run BaseComponent.__init__ fields + index.name = "perf_test" + index.backend = "" + index.app_context = None + index.kwargs = {} + index.logger = __import__("logging").getLogger("perf_test") + index._is_started = False + index._lock = asyncio.Lock() + # BaseKeywordIndex fields + index.tokenizer_name = "default" + index.index_path = tmp_dir / "keyword_index" + index.index_path.mkdir(parents=True, exist_ok=True) + # BM25Index fields + index.k1 = 1.5 + index.b = 0.75 + index.vocab = {} + index.inverted_index = {} + index.doc_meta = {} + index.total_len = 0 + index._idf_cache = {} + # Init tokenizer + index.tokenizer = RegexTokenizer(filter_stopwords=False) + await index.tokenizer.start() + index._is_started = True + return index + + +async def test_add_docs_small(): + """Add 100 small docs (~100 tokens each).""" + docs = {f"doc_{i}": _gen_random_text(100) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp: + index = await _make_index(Path(tmp)) + t0 = time.perf_counter() + await index.add_docs(docs) + elapsed = time.perf_counter() - t0 + print(f" add_docs (100 docs x ~100 tokens): {elapsed:.4f}s") + await index.close() + + +async def test_add_docs_medium(): + """Add 100 medium docs (~1000 tokens each).""" + docs = {f"doc_{i}": _gen_random_text(1000) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp: + index = await _make_index(Path(tmp)) + t0 = time.perf_counter() + await index.add_docs(docs) + elapsed = time.perf_counter() - t0 + print(f" add_docs (100 docs x ~1000 tokens): {elapsed:.4f}s") + await index.close() + + +async def test_add_docs_large(): + """Add 100 large docs (~10000 tokens each).""" + docs = {f"doc_{i}": _gen_random_text(10000) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp: + index = await _make_index(Path(tmp)) + t0 = time.perf_counter() + await index.add_docs(docs) + elapsed = time.perf_counter() - t0 + print(f" add_docs (100 docs x ~10000 tokens): {elapsed:.4f}s") + await index.close() + + +async def _setup_index_for_retrieve(n_docs: int = 100, doc_tokens: int = 1000) -> tuple[BM25Index, str]: + """Build an index with n_docs medium-sized docs, return (index, tmp_dir).""" + tmp = tempfile.mkdtemp() + index = await _make_index(Path(tmp)) + docs = {f"doc_{i}": _gen_random_text(doc_tokens) for i in range(n_docs)} + await index.add_docs(docs) + return index, tmp + + +async def test_retrieve_short_query(): + """Retrieve with 1-word query.""" + index, tmp = await _setup_index_for_retrieve() + query = _gen_random_query(1) + t0 = time.perf_counter() + await index.retrieve(query, limit=10) + elapsed = time.perf_counter() - t0 + print(f" retrieve (1-word query, 100 docs): {elapsed:.6f}s") + await index.close() + + +async def test_retrieve_medium_query(): + """Retrieve with 5-word query.""" + index, tmp = await _setup_index_for_retrieve() + query = _gen_random_query(5) + t0 = time.perf_counter() + await index.retrieve(query, limit=10) + elapsed = time.perf_counter() - t0 + print(f" retrieve (5-word query, 100 docs): {elapsed:.6f}s") + await index.close() + + +async def test_retrieve_long_query(): + """Retrieve with 20-word query.""" + index, tmp = await _setup_index_for_retrieve() + query = _gen_random_query(20) + t0 = time.perf_counter() + await index.retrieve(query, limit=10) + elapsed = time.perf_counter() - t0 + print(f" retrieve (20-word query, 100 docs): {elapsed:.6f}s") + await index.close() + + +async def test_retrieve_very_long_query(): + """Retrieve with 100-word query.""" + index, tmp = await _setup_index_for_retrieve() + query = _gen_random_query(100) + t0 = time.perf_counter() + await index.retrieve(query, limit=10) + elapsed = time.perf_counter() - t0 + print(f" retrieve (100-word query, 100 docs): {elapsed:.6f}s") + await index.close() + + +async def main(): + random.seed(42) + print("=== BM25Index Performance Tests ===\n") + + print("[add_docs]") + await test_add_docs_small() + await test_add_docs_medium() + await test_add_docs_large() + + print("\n[retrieve]") + await test_retrieve_short_query() + await test_retrieve_medium_query() + await test_retrieve_long_query() + await test_retrieve_very_long_query() + + print("\nDone.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/test/reme2/test_bm25_lite.py b/test/reme2/test_bm25_lite.py index a04a950a..c0d234a5 100644 --- a/test/reme2/test_bm25_lite.py +++ b/test/reme2/test_bm25_lite.py @@ -1,4 +1,4 @@ -"""Tests for BM25Lite search engine.""" +"""Tests for BM25Index search engine.""" import asyncio import tempfile @@ -13,22 +13,22 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent)) warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") -from reme2.component.file_store.bm25_lite import BM25Lite +from reme2.component.keyword_index.bm25_index import BM25Index -async def create_bm25(index_dir: Path, k1: float = 1.5, b: float = 0.75) -> BM25Lite: - """Create and start a BM25Lite instance.""" - bm25 = BM25Lite(index_dir=index_dir, k1=k1, b=b) +async def create_bm25(index_dir: Path, k1: float = 1.5, b: float = 0.75) -> BM25Index: + """Create and start a BM25Index instance.""" + bm25 = BM25Index(index_dir=index_dir, k1=k1, b=b) await bm25.start() return bm25 def test_basic_init(): - """Test BM25Lite initialization.""" + """Test BM25Index initialization.""" async def run(): with tempfile.TemporaryDirectory() as tmpdir: - bm25 = BM25Lite(index_dir=tmpdir) + bm25 = BM25Index(index_dir=tmpdir) assert bm25.k1 == 1.5 assert bm25.b == 0.75 assert bm25.vocab == {} @@ -42,12 +42,12 @@ def test_basic_init(): def test_start_with_tokenizer(): - """Test BM25Lite starts and initializes tokenizer.""" + """Test BM25Index starts and initializes tokenizer.""" async def run(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - assert bm25._tokenizer is not None + assert bm25.tokenizer is not None assert bm25.is_started await bm25.close() @@ -64,7 +64,7 @@ def test_add_single_doc(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({"doc1": "hello world"}) + await bm25.add_docs({"doc1": "hello world"}) assert bm25.n_docs == 1 assert bm25.total_len > 0 @@ -88,7 +88,7 @@ def test_add_multiple_docs(): "doc2": "hello python", "doc3": "world python", } - bm25.add_docs(docs) + await bm25.add_docs(docs) assert bm25.n_docs == 3 assert len(bm25.vocab) > 0 @@ -112,9 +112,9 @@ def test_retrieve_basic(): "doc2": "java programming language", "doc3": "python data analysis", } - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("python", limit=3) + results = await bm25.retrieve("python", limit=3) assert len(results) <= 3 assert "doc1" in results or "doc3" in results @@ -134,12 +134,12 @@ def test_retrieve_with_limit(): docs = { f"doc{i}": f"python programming {i}" for i in range(10) } - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("python", limit=3) + results = await bm25.retrieve("python", limit=3) assert len(results) == 3 - results = bm25.retrieve("python", limit=5) + results = await bm25.retrieve("python", limit=5) assert len(results) == 5 await bm25.close() @@ -156,12 +156,12 @@ def test_retrieve_empty_query(): bm25 = await create_bm25(Path(tmpdir)) docs = {"doc1": "hello world"} - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("", limit=3) + results = await bm25.retrieve("", limit=3) assert results == {} - results = bm25.retrieve("unknownxyz", limit=3) + results = await bm25.retrieve("unknownxyz", limit=3) assert results == {} await bm25.close() @@ -177,7 +177,7 @@ def test_retrieve_empty_index(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - results = bm25.retrieve("python", limit=3) + results = await bm25.retrieve("python", limit=3) assert results == {} await bm25.close() @@ -193,14 +193,14 @@ def test_update_doc(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({"doc1": "hello world python"}) + await bm25.add_docs({"doc1": "hello world python"}) old_len = bm25.total_len - bm25.add_docs({"doc1": "java"}) + await bm25.add_docs({"doc1": "java"}) assert bm25.n_docs == 1 assert bm25.total_len != old_len - results = bm25.retrieve("java", limit=1) + results = await bm25.retrieve("java", limit=1) assert "doc1" in results await bm25.close() @@ -220,14 +220,14 @@ def test_remove_doc(): "doc1": "hello world", "doc2": "hello python", } - bm25.add_docs(docs) + await bm25.add_docs(docs) assert bm25.n_docs == 2 bm25._remove_doc("doc1") assert bm25.n_docs == 1 assert "doc1" not in bm25.doc_meta - results = bm25.retrieve("hello", limit=2) + results = await bm25.retrieve("hello", limit=2) assert "doc1" not in results assert "doc2" in results @@ -244,7 +244,7 @@ def test_remove_nonexistent_doc(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({"doc1": "hello world"}) + await bm25.add_docs({"doc1": "hello world"}) bm25._remove_doc("nonexistent") assert bm25.n_docs == 1 @@ -261,13 +261,13 @@ def test_clear(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({ + await bm25.add_docs({ "doc1": "hello world", "doc2": "hello python", }) assert bm25.n_docs == 2 - bm25.clear() + await bm25.clear() assert bm25.n_docs == 0 assert bm25.vocab == {} assert bm25.inverted_index == {} @@ -281,37 +281,37 @@ def test_clear(): asyncio.run(run()) -def test_reindex(): - """Test reindex functionality to compact vocab.""" +def test_optimize_index(): + """Test optimize_index functionality to compact vocab.""" async def run(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({"doc1": "hello world"}) + await bm25.add_docs({"doc1": "hello world"}) bm25._remove_doc("doc1") assert bm25.n_docs == 0 assert len(bm25.vocab) > 0 - bm25.reindex() + await bm25.optimize_index() assert bm25.vocab == {} assert bm25.inverted_index == {} await bm25.close() - print("✓ test_reindex passed") + print("✓ test_optimize_index passed") asyncio.run(run()) -def test_reindex_with_docs(): - """Test reindex with remaining documents.""" +def test_optimize_index_with_docs(): + """Test optimize_index with remaining documents.""" async def run(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({ + await bm25.add_docs({ "doc1": "hello world", "doc2": "hello python", }) @@ -319,17 +319,17 @@ def test_reindex_with_docs(): old_vocab = bm25.vocab.copy() bm25._remove_doc("doc1") - bm25.reindex() + await bm25.optimize_index() assert bm25.n_docs == 1 assert "doc2" in bm25.doc_meta assert len(bm25.vocab) < len(old_vocab) - results = bm25.retrieve("hello", limit=1) + results = await bm25.retrieve("hello", limit=1) assert "doc2" in results await bm25.close() - print("✓ test_reindex_with_docs passed") + print("✓ test_optimize_index_with_docs passed") asyncio.run(run()) @@ -347,7 +347,7 @@ def test_persistence(): "doc2": "hello python", "doc3": "programming language", } - bm25.add_docs(docs) + await bm25.add_docs(docs) old_vocab = bm25.vocab.copy() old_doc_meta = {k: dict(v) for k, v in bm25.doc_meta.items()} @@ -362,7 +362,7 @@ def test_persistence(): for doc_id, meta in old_doc_meta.items(): assert doc_id in bm25_new.doc_meta - results = bm25_new.retrieve("hello", limit=2) + results = await bm25_new.retrieve("hello", limit=2) assert "doc1" in results or "doc2" in results await bm25_new.close() @@ -381,8 +381,8 @@ def test_custom_params(): assert bm25.k1 == 2.0 assert bm25.b == 0.5 - bm25.add_docs({"doc1": "test document"}) - results = bm25.retrieve("test", limit=1) + await bm25.add_docs({"doc1": "test document"}) + results = await bm25.retrieve("test", limit=1) assert "doc1" in results await bm25.close() @@ -403,9 +403,9 @@ def test_chinese_text(): "doc2": "北京是中国的首都", "doc3": "上海的天气很好", } - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("北京", limit=2) + results = await bm25.retrieve("北京", limit=2) assert len(results) <= 2 assert "doc1" in results or "doc2" in results @@ -427,12 +427,12 @@ def test_mixed_chinese_english(): "doc2": "Java 编程语言", "doc3": "Python 数据分析", } - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("Python", limit=3) + results = await bm25.retrieve("Python", limit=3) assert len(results) > 0 - results = bm25.retrieve("编程", limit=2) + results = await bm25.retrieve("编程", limit=2) assert len(results) > 0 await bm25.close() @@ -448,7 +448,7 @@ def test_idf_cache(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({ + await bm25.add_docs({ "doc1": "hello world", "doc2": "hello python", }) @@ -476,10 +476,10 @@ def test_avg_len(): assert bm25.avg_len == 0.0 - bm25.add_docs({"doc1": "hello world python"}) + await bm25.add_docs({"doc1": "hello world python"}) assert bm25.avg_len > 0 - bm25.add_docs({"doc2": "test"}) + await bm25.add_docs({"doc2": "test"}) new_avg = bm25.avg_len assert new_avg > 0 @@ -501,9 +501,9 @@ def test_score_ordering(): "doc2": "python python", "doc3": "python", } - bm25.add_docs(docs) + await bm25.add_docs(docs) - results = bm25.retrieve("python", limit=3) + results = await bm25.retrieve("python", limit=3) scores = list(results.values()) for i in range(len(scores) - 1): @@ -522,10 +522,10 @@ def test_empty_doc(): with tempfile.TemporaryDirectory() as tmpdir: bm25 = await create_bm25(Path(tmpdir)) - bm25.add_docs({"doc1": ""}) + await bm25.add_docs({"doc1": ""}) assert bm25.n_docs == 0 - bm25.add_docs({"doc2": " "}) + await bm25.add_docs({"doc2": " "}) assert bm25.n_docs == 0 await bm25.close() @@ -535,7 +535,7 @@ def test_empty_doc(): if __name__ == "__main__": - print("\n=== BM25Lite Tests ===") + print("\n=== BM25Index Tests ===") test_basic_init() test_start_with_tokenizer() test_add_single_doc() @@ -548,8 +548,8 @@ if __name__ == "__main__": test_remove_doc() test_remove_nonexistent_doc() test_clear() - test_reindex() - test_reindex_with_docs() + test_optimize_index() + test_optimize_index_with_docs() test_persistence() test_custom_params() test_chinese_text() @@ -558,4 +558,4 @@ if __name__ == "__main__": test_avg_len() test_score_ordering() test_empty_doc() - print("\n所有测试通过!") \ No newline at end of file + print("\n所有测试通过!")