This commit is contained in:
jinli.yl 2026-05-14 02:00:35 +08:00
parent 48137f7f3f
commit 5792221d3f
21 changed files with 824 additions and 821 deletions

View file

@ -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",
]

View file

@ -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

View file

@ -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,
)

View file

@ -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()

View file

@ -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:

View file

@ -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

View file

@ -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."""

View file

@ -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

View file

@ -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",
]

View file

@ -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()
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."""

View file

@ -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)

View file

@ -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."""

View file

@ -1,7 +0,0 @@
from .base_file_watcher import BaseFileWatcher
from ..component_registry import R
@R.register("light")
class LightFileWatcher(BaseFileWatcher):
...

View file

@ -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])

View file

@ -0,0 +1,4 @@
from .base_keyword_index import BaseKeywordIndex
from .bm25_index import BM25Index
__all__ = ["BaseKeywordIndex", "BM25Index"]

View file

@ -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."""

View file

@ -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 = {}

View file

@ -30,6 +30,8 @@ class ComponentEnum(str, Enum):
FILE_WATCHER = "file_watcher"
KEYWORD_INDEX = "keyword_index"
SERVICE = "service"
CLIENT = "client"

View file

@ -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

View file

@ -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())

View file

@ -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所有测试通过!")
print("\n所有测试通过!")