mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-21 00:22:45 +00:00
up
This commit is contained in:
parent
48137f7f3f
commit
5792221d3f
21 changed files with 824 additions and 821 deletions
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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."""
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
from .base_file_watcher import BaseFileWatcher
|
||||
from ..component_registry import R
|
||||
|
||||
|
||||
@R.register("light")
|
||||
class LightFileWatcher(BaseFileWatcher):
|
||||
...
|
||||
133
reme2/component/file_watcher/lite_file_watcher.py
Normal file
133
reme2/component/file_watcher/lite_file_watcher.py
Normal 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])
|
||||
4
reme2/component/keyword_index/__init__.py
Normal file
4
reme2/component/keyword_index/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from .base_keyword_index import BaseKeywordIndex
|
||||
from .bm25_index import BM25Index
|
||||
|
||||
__all__ = ["BaseKeywordIndex", "BM25Index"]
|
||||
106
reme2/component/keyword_index/base_keyword_index.py
Normal file
106
reme2/component/keyword_index/base_keyword_index.py
Normal 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."""
|
||||
|
|
@ -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 = {}
|
||||
|
|
@ -30,6 +30,8 @@ class ComponentEnum(str, Enum):
|
|||
|
||||
FILE_WATCHER = "file_watcher"
|
||||
|
||||
KEYWORD_INDEX = "keyword_index"
|
||||
|
||||
SERVICE = "service"
|
||||
|
||||
CLIENT = "client"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
186
test/reme2/test_bm25_index_perf.py
Normal file
186
test/reme2/test_bm25_index_perf.py
Normal 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())
|
||||
|
|
@ -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所有测试通过!")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue