diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml new file mode 100644 index 00000000..e4100b82 --- /dev/null +++ b/.github/workflows/unittest.yml @@ -0,0 +1,55 @@ +name: Tests (reme) + +on: + push: + branches: [main, master, dev, develop, 'dev/**'] + paths: + - 'reme4/**' + - 'reme/reme2/**' + - 'tests4/**' + - 'pyproject.toml' + - '.github/workflows/tests4.yml' + pull_request: + branches: [main, master, dev, develop] + paths: + - 'reme4/**' + - 'reme/reme2/**' + - 'tests4/**' + - 'pyproject.toml' + - '.github/workflows/tests4.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + name: Unit Tests - py${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip setuptools wheel + pip install -e ".[dev]" + + - name: Run tests4 unit tests + run: | + pytest tests4/unittest \ + -v \ + --tb=long \ + -s \ + --log-cli-level=WARNING diff --git a/reme4/__init__.py b/reme4/__init__.py index 76b34dcf..8dba02ec 100644 --- a/reme4/__init__.py +++ b/reme4/__init__.py @@ -22,3 +22,5 @@ __all__ = [ "steps", "utils", ] + +__version__ = "0.4.0.0" diff --git a/reme4/components/service/http_service.py b/reme4/components/service/http_service.py index 6c7b9272..64806114 100644 --- a/reme4/components/service/http_service.py +++ b/reme4/components/service/http_service.py @@ -3,6 +3,7 @@ import asyncio import json import os +import warnings from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from typing import TYPE_CHECKING @@ -84,4 +85,16 @@ class HttpService(BaseService): self.service.post("/health")(lambda: {"status": "healthy"}) def start_service(self, app: "Application") -> None: + # uvicorn 0.41 still imports websockets.legacy / WebSocketServerProtocol + # on startup; silence those specific lines since we don't use WebSocket. + warnings.filterwarnings( + "ignore", + category=DeprecationWarning, + message=r".*websockets\.legacy is deprecated.*", + ) + warnings.filterwarnings( + "ignore", + category=DeprecationWarning, + message=r".*WebSocketServerProtocol is deprecated.*", + ) uvicorn.run(self.service, host=self.host, port=self.port, **self.kwargs) diff --git a/reme4/config/default.yaml b/reme4/config/default.yaml new file mode 100644 index 00000000..5eb447cf --- /dev/null +++ b/reme4/config/default.yaml @@ -0,0 +1,12 @@ +service: + backend: http + host: 127.0.0.1 + port: 2333 + +jobs: + - name: demo + backend: base + description: "Echo back the incoming query (smoke test)." + steps: + - backend: demo_echo + diff --git a/reme4/schema/application_config.py b/reme4/schema/application_config.py index b71e6d21..a6cdf018 100644 --- a/reme4/schema/application_config.py +++ b/reme4/schema/application_config.py @@ -29,7 +29,7 @@ class ApplicationConfig(BaseModel): app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"), description="Application display name") working_dir: str = Field(default=".reme", description="Working directory for runtime files") - enable_logo: bool = Field(default=False, description="Show ASCII logo on startup") + enable_logo: bool = Field(default=True, description="Show ASCII logo on startup") language: str = Field(default="", description="Default language for LLM interactions") log_to_console: bool = Field(default=True, description="Log to console") log_to_file: bool = Field(default=True, description="Log to file") diff --git a/reme4/steps/__init__.py b/reme4/steps/__init__.py index 2fde5b11..76f8c3a7 100644 --- a/reme4/steps/__init__.py +++ b/reme4/steps/__init__.py @@ -1,7 +1,9 @@ """steps""" +from . import demo from .base_step import BaseStep __all__ = [ "BaseStep", + "demo", ] diff --git a/reme4/steps/base_step.py b/reme4/steps/base_step.py index c79d5f5e..7cd45b55 100644 --- a/reme4/steps/base_step.py +++ b/reme4/steps/base_step.py @@ -25,6 +25,8 @@ T = TypeVar("T") class BaseStep(ABC): """Composable unit of an LLM workflow.""" + component_type = ComponentEnum.STEP + def __new__(cls, *args, **kwargs): # Snapshot init args so copy() can rebuild an equivalent instance later. instance = object.__new__(cls) diff --git a/reme4/steps/demo/__init__.py b/reme4/steps/demo/__init__.py new file mode 100644 index 00000000..057484a3 --- /dev/null +++ b/reme4/steps/demo/__init__.py @@ -0,0 +1,19 @@ +"""Demo steps for smoke-testing the application stack.""" + +from ..base_step import BaseStep +from ...components.component_registry import R + + +@R.register("demo_echo") +class DemoEchoStep(BaseStep): + """Echo back the incoming `query` field into `response.answer`.""" + + async def execute(self): + assert self.context is not None + query = self.context.get("query", "") + self.context.response.answer = f"echo: {query}" + self.context.response.metadata["step"] = self.name + return self.context.response + + +__all__ = ["DemoEchoStep"] diff --git a/tests4/unittest/test_bm25_index_perf.py b/tests4/unittest/test_bm25_index_perf.py new file mode 100644 index 00000000..b6ba0f34 --- /dev/null +++ b/tests4/unittest/test_bm25_index_perf.py @@ -0,0 +1,337 @@ +"""BM25Index performance tests for add_docs and retrieve.""" + +import asyncio +import os +import random +import tempfile +import time + +from reme4.components.keyword_index import BM25Index +from reme4.components.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", +] + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +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() -> BM25Index: + """Create and start a BM25Index using cwd as working dir, with non-filtering tokenizer.""" + index = BM25Index() + tokenizer = RegexTokenizer(filter_stopwords=False) + index.tokenizer = tokenizer + index._owned.append(tokenizer) # pylint: disable=protected-access + await index.start() + return index + + +async def _setup_index_for_retrieve(n_docs: int = 100, doc_tokens: int = 1000) -> BM25Index: + """Build an index with n_docs medium-sized docs in cwd.""" + index = await _make_index() + docs = {f"doc_{i}": _gen_random_text(doc_tokens) for i in range(n_docs)} + await index.add_docs(docs) + return index + + +def test_add_docs_small(): + """Add 100 small docs (~100 tokens each).""" + + async def run(): + docs = {f"doc_{i}": _gen_random_text(100) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = await _make_index() + 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() + + asyncio.run(run()) + + +def test_add_docs_medium(): + """Add 100 medium docs (~1000 tokens each).""" + + async def run(): + docs = {f"doc_{i}": _gen_random_text(1000) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = await _make_index() + 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() + + asyncio.run(run()) + + +def test_add_docs_large(): + """Add 100 large docs (~10000 tokens each).""" + + async def run(): + docs = {f"doc_{i}": _gen_random_text(10000) for i in range(100)} + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = await _make_index() + 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() + + asyncio.run(run()) + + +def test_retrieve_short_query(): + """Retrieve with 1-word query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = 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() + + asyncio.run(run()) + + +def test_retrieve_medium_query(): + """Retrieve with 5-word query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = 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() + + asyncio.run(run()) + + +def test_retrieve_long_query(): + """Retrieve with 20-word query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = 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() + + asyncio.run(run()) + + +def test_retrieve_very_long_query(): + """Retrieve with 100-word query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + index = 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() + + asyncio.run(run()) + + +if __name__ == "__main__": + random.seed(42) + print("=== BM25Index Performance Tests ===\n") + + print("[add_docs]") + test_add_docs_small() + test_add_docs_medium() + test_add_docs_large() + + print("\n[retrieve]") + test_retrieve_short_query() + test_retrieve_medium_query() + test_retrieve_long_query() + test_retrieve_very_long_query() + + print("\nDone.") diff --git a/tests4/unittest/test_bm25_lite.py b/tests4/unittest/test_bm25_lite.py new file mode 100644 index 00000000..75865797 --- /dev/null +++ b/tests4/unittest/test_bm25_lite.py @@ -0,0 +1,586 @@ +"""Tests for BM25Index search engine.""" + +# pylint: disable=protected-access + +import asyncio +import os +import tempfile +import warnings + +from reme4.components.keyword_index import BM25Index +from reme4.components.tokenizer import RegexTokenizer + +# Filter jieba/pkg_resources deprecation warnings +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +async def create_bm25(k1: float = 1.5, b: float = 0.75) -> BM25Index: + """Create and start a BM25Index in cwd with a non-filtering tokenizer. + + The non-filtering tokenizer keeps short test texts (e.g. "hello world") visible, + since several common test words ("hello", "我", "的") are in the default stopwords. + """ + bm25 = BM25Index(k1=k1, b=b) + # Replace the unresolved Dependency placeholder with a real tokenizer instance. + tokenizer = RegexTokenizer(filter_stopwords=False) + bm25.tokenizer = tokenizer + bm25._owned.append(tokenizer) + await bm25.start() + return bm25 + + +def test_basic_init(): + """Test BM25Index initialization.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = BM25Index() + assert bm25.k1 == 1.5 + assert bm25.b == 0.75 + assert bm25.vocab == {} + assert bm25.inverted_index == {} + assert bm25.doc_meta == {} + assert bm25.n_docs == 0 + assert bm25.avg_len == 0.0 + print("✓ test_basic_init passed") + + asyncio.run(run()) + + +def test_start_with_tokenizer(): + """Test BM25Index starts and initializes tokenizer.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + assert bm25.tokenizer is not None + assert bm25.is_started + + await bm25.close() + assert not bm25.is_started + print("✓ test_start_with_tokenizer passed") + + asyncio.run(run()) + + +def test_add_single_doc(): + """Test adding a single document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs({"doc1": "hello world"}) + + assert bm25.n_docs == 1 + assert bm25.total_len > 0 + assert "doc1" in bm25.doc_meta + + await bm25.close() + print("✓ test_add_single_doc passed") + + asyncio.run(run()) + + +def test_add_multiple_docs(): + """Test adding multiple documents.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "hello world", + "doc2": "hello python", + "doc3": "world python", + } + await bm25.add_docs(docs) + + assert bm25.n_docs == 3 + assert len(bm25.vocab) > 0 + assert len(bm25.inverted_index) > 0 + + await bm25.close() + print("✓ test_add_multiple_docs passed") + + asyncio.run(run()) + + +def test_retrieve_basic(): + """Test basic retrieval functionality.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "python programming language", + "doc2": "java programming language", + "doc3": "python data analysis", + } + await bm25.add_docs(docs) + + results = await bm25.retrieve("python", limit=3) + assert len(results) <= 3 + assert "doc1" in results or "doc3" in results + + await bm25.close() + print("✓ test_retrieve_basic passed") + + asyncio.run(run()) + + +def test_retrieve_with_limit(): + """Test retrieval with result limit.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = {f"doc{i}": f"python programming {i}" for i in range(10)} + await bm25.add_docs(docs) + + results = await bm25.retrieve("python", limit=3) + assert len(results) == 3 + + results = await bm25.retrieve("python", limit=5) + assert len(results) == 5 + + await bm25.close() + print("✓ test_retrieve_with_limit passed") + + asyncio.run(run()) + + +def test_retrieve_empty_query(): + """Test retrieval with empty or unknown query.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = {"doc1": "hello world"} + await bm25.add_docs(docs) + + results = await bm25.retrieve("", limit=3) + assert results == {} + + results = await bm25.retrieve("unknownxyz", limit=3) + assert results == {} + + await bm25.close() + print("✓ test_retrieve_empty_query passed") + + asyncio.run(run()) + + +def test_retrieve_empty_index(): + """Test retrieval from empty index.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + results = await bm25.retrieve("python", limit=3) + assert results == {} + + await bm25.close() + print("✓ test_retrieve_empty_index passed") + + asyncio.run(run()) + + +def test_update_doc(): + """Test updating an existing document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs({"doc1": "hello world python"}) + old_len = bm25.total_len + + await bm25.add_docs({"doc1": "java"}) + assert bm25.n_docs == 1 + assert bm25.total_len != old_len + + results = await bm25.retrieve("java", limit=1) + assert "doc1" in results + + await bm25.close() + print("✓ test_update_doc passed") + + asyncio.run(run()) + + +def test_remove_doc(): + """Test removing a document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "hello world", + "doc2": "hello python", + } + 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 = await bm25.retrieve("hello", limit=2) + assert "doc1" not in results + assert "doc2" in results + + await bm25.close() + print("✓ test_remove_doc passed") + + asyncio.run(run()) + + +def test_remove_nonexistent_doc(): + """Test removing a nonexistent document (should be no-op).""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs({"doc1": "hello world"}) + bm25._remove_doc("nonexistent") + assert bm25.n_docs == 1 + + await bm25.close() + print("✓ test_remove_nonexistent_doc passed") + + asyncio.run(run()) + + +def test_clear(): + """Test clearing the index.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs( + { + "doc1": "hello world", + "doc2": "hello python", + }, + ) + assert bm25.n_docs == 2 + + await bm25.clear() + assert bm25.n_docs == 0 + assert bm25.vocab == {} + assert bm25.inverted_index == {} + assert bm25.doc_meta == {} + assert bm25.total_len == 0 + assert bm25._idf_cache == {} + + await bm25.close() + print("✓ test_clear passed") + + asyncio.run(run()) + + +def test_optimize_index(): + """Test optimize_index functionality to compact vocab.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs({"doc1": "hello world"}) + bm25._remove_doc("doc1") + + assert bm25.n_docs == 0 + assert len(bm25.vocab) > 0 + + await bm25.optimize_index() + assert bm25.vocab == {} + assert bm25.inverted_index == {} + + await bm25.close() + print("✓ test_optimize_index passed") + + asyncio.run(run()) + + +def test_optimize_index_with_docs(): + """Test optimize_index with remaining documents.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs( + { + "doc1": "hello world", + "doc2": "hello python", + }, + ) + + old_vocab = bm25.vocab.copy() + bm25._remove_doc("doc1") + + await bm25.optimize_index() + + assert bm25.n_docs == 1 + assert "doc2" in bm25.doc_meta + assert len(bm25.vocab) < len(old_vocab) + + results = await bm25.retrieve("hello", limit=1) + assert "doc2" in results + + await bm25.close() + print("✓ test_optimize_index_with_docs passed") + + asyncio.run(run()) + + +def test_persistence(): + """Test dump and load persistence.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + docs = { + "doc1": "hello world", + "doc2": "hello python", + "doc3": "programming language", + } + await bm25.add_docs(docs) + + old_vocab = bm25.vocab.copy() + old_doc_meta = {k: dict(v) for k, v in bm25.doc_meta.items()} + + await bm25.dump() + await bm25.close() + + bm25_new = await create_bm25() + + assert bm25_new.vocab == old_vocab + assert bm25_new.n_docs == 3 + for doc_id in old_doc_meta: + assert doc_id in bm25_new.doc_meta + + results = await bm25_new.retrieve("hello", limit=2) + assert "doc1" in results or "doc2" in results + + await bm25_new.close() + print("✓ test_persistence passed") + + asyncio.run(run()) + + +def test_custom_params(): + """Test custom k1 and b parameters.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25(k1=2.0, b=0.5) + + assert bm25.k1 == 2.0 + assert bm25.b == 0.5 + + await bm25.add_docs({"doc1": "test document"}) + results = await bm25.retrieve("test", limit=1) + assert "doc1" in results + + await bm25.close() + print("✓ test_custom_params passed") + + asyncio.run(run()) + + +def test_chinese_text(): + """Test with Chinese text.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "我爱北京天安门", + "doc2": "北京是中国的首都", + "doc3": "上海的天气很好", + } + await bm25.add_docs(docs) + + results = await bm25.retrieve("北", limit=2) + assert len(results) <= 2 + assert "doc1" in results or "doc2" in results + + await bm25.close() + print("✓ test_chinese_text passed") + + asyncio.run(run()) + + +def test_mixed_chinese_english(): + """Test with mixed Chinese and English text.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "Python 是一种编程语言", + "doc2": "Java 编程语言", + "doc3": "Python 数据分析", + } + await bm25.add_docs(docs) + + results = await bm25.retrieve("Python", limit=3) + assert len(results) > 0 + + results = await bm25.retrieve("编", limit=2) + assert len(results) > 0 + + await bm25.close() + print("✓ test_mixed_chinese_english passed") + + asyncio.run(run()) + + +def test_idf_cache(): + """Test IDF cache functionality.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs( + { + "doc1": "hello world", + "doc2": "hello python", + }, + ) + + token = "hello" + if token in bm25.vocab: + tid = bm25.vocab[token] + idf1 = bm25._get_idf(tid) + assert tid in bm25._idf_cache + idf2 = bm25._get_idf(tid) + assert idf1 == idf2 + + await bm25.close() + print("✓ test_idf_cache passed") + + asyncio.run(run()) + + +def test_avg_len(): + """Test average document length calculation.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + assert bm25.avg_len == 0.0 + + await bm25.add_docs({"doc1": "hello world python"}) + assert bm25.avg_len > 0 + + await bm25.add_docs({"doc2": "test"}) + new_avg = bm25.avg_len + assert new_avg > 0 + + await bm25.close() + print("✓ test_avg_len passed") + + asyncio.run(run()) + + +def test_score_ordering(): + """Test that results are ordered by score descending.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + docs = { + "doc1": "python python python", + "doc2": "python python", + "doc3": "python", + } + await bm25.add_docs(docs) + + results = await bm25.retrieve("python", limit=3) + scores = list(results.values()) + + for i in range(len(scores) - 1): + assert scores[i] >= scores[i + 1] + + await bm25.close() + print("✓ test_score_ordering passed") + + asyncio.run(run()) + + +def test_empty_doc(): + """Test adding empty document.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + bm25 = await create_bm25() + + await bm25.add_docs({"doc1": ""}) + assert bm25.n_docs == 0 + + await bm25.add_docs({"doc2": " "}) + assert bm25.n_docs == 0 + + await bm25.close() + print("✓ test_empty_doc passed") + + asyncio.run(run()) + + +if __name__ == "__main__": + print("\n=== BM25Index Tests ===") + test_basic_init() + test_start_with_tokenizer() + test_add_single_doc() + test_add_multiple_docs() + test_retrieve_basic() + test_retrieve_with_limit() + test_retrieve_empty_query() + test_retrieve_empty_index() + test_update_doc() + test_remove_doc() + test_remove_nonexistent_doc() + test_clear() + test_optimize_index() + test_optimize_index_with_docs() + test_persistence() + test_custom_params() + test_chinese_text() + test_mixed_chinese_english() + test_idf_cache() + test_avg_len() + test_score_ordering() + test_empty_doc() + print("\n所有测试通过!") diff --git a/tests4/unittest/test_default_file_parser.py b/tests4/unittest/test_default_file_parser.py new file mode 100644 index 00000000..75f68ba5 --- /dev/null +++ b/tests4/unittest/test_default_file_parser.py @@ -0,0 +1,205 @@ +"""Tests for DefaultFileParser.""" + +import asyncio +import os +import tempfile + +from reme4.components.file_parser import DefaultFileParser + + +# Add parent path for import + + +def test_parse_empty_file(): + """Test parsing an empty file.""" + + async def run(): + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + temp_path = f.name + + try: + parser = DefaultFileParser() + file_node, chunks = await parser.parse(temp_path) + assert file_node.path == temp_path + assert len(chunks) == 0 + print("✓ test_parse_empty_file passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_parse_small_file(): + """Test parsing a file smaller than chunk size.""" + + async def run(): + content = "Hello World\nThis is a test\nLine 3" + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser(chunk_byte_size=10000) + _, chunks = await parser.parse(temp_path) + assert len(chunks) == 1 + assert chunks[0].start_line == 1 + assert chunks[0].end_line == 3 + assert chunks[0].text == content + print("✓ test_parse_small_file passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_parse_multiline_file(): + """Test parsing a file with multiple lines.""" + + async def run(): + lines = ["Line 1", "Line 2", "Line 3", "Line 4", "Line 5"] + content = "\n".join(lines) + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser(chunk_byte_size=10000) + _, chunks = await parser.parse(temp_path) + assert len(chunks) == 1 + assert chunks[0].start_line == 1 + assert chunks[0].end_line == 5 + print("✓ test_parse_multiline_file passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_parse_chunked_file(): + """Test parsing a file that requires multiple chunks.""" + + async def run(): + # Create content larger than chunk size + lines = ["A" * 100 for _ in range(200)] # ~20200 bytes + content = "\n".join(lines) + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser(chunk_byte_size=5000, overlap_byte_size=100) + _, chunks = await parser.parse(temp_path) + assert len(chunks) > 1, f"Expected multiple chunks, got {len(chunks)}" + # Verify overlap by checking that consecutive chunks share some content + print(f" Created {len(chunks)} chunks") + print("✓ test_parse_chunked_file passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_parse_with_custom_encoding(): + """Test parsing a file with different encodings.""" + + async def run(): + content = "你好世界\n测试内容" + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt", encoding="utf-8") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser(encoding="utf-8") + _, chunks = await parser.parse(temp_path) + assert len(chunks) >= 1 + assert "你好世界" in chunks[0].text + print("✓ test_parse_with_custom_encoding passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_file_node_properties(): + """Test FileNode has correct properties.""" + + async def run(): + content = "test content" + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser() + file_node, _ = await parser.parse(temp_path) + assert hasattr(file_node, "path") + assert hasattr(file_node, "st_mtime") + assert file_node.st_mtime > 0 + print("✓ test_file_node_properties passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_file_chunk_properties(): + """Test FileChunk has correct properties.""" + + async def run(): + content = "test content for chunk" + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + parser = DefaultFileParser() + _, chunks = await parser.parse(temp_path) + chunk = chunks[0] + assert hasattr(chunk, "path") + assert hasattr(chunk, "start_line") + assert hasattr(chunk, "end_line") + assert hasattr(chunk, "text") + assert hasattr(chunk, "id") + assert chunk.start_line >= 1 + assert chunk.end_line >= chunk.start_line + print("✓ test_file_chunk_properties passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +def test_min_chunk_and_overlap_size(): + """Test that minimum chunk and overlap sizes are enforced.""" + + async def run(): + # These values should be clamped to minimums + parser = DefaultFileParser(chunk_byte_size=1, overlap_byte_size=0) + assert parser.chunk_byte_size == 100 # minimum + assert parser.overlap_byte_size == 4 # minimum + + content = "test" + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f: + f.write(content) + temp_path = f.name + + try: + _, chunks = await parser.parse(temp_path) + assert len(chunks) == 1 + print("✓ test_min_chunk_and_overlap_size passed") + finally: + os.unlink(temp_path) + + asyncio.run(run()) + + +if __name__ == "__main__": + test_parse_empty_file() + test_parse_small_file() + test_parse_multiline_file() + test_parse_chunked_file() + test_parse_with_custom_encoding() + test_file_node_properties() + test_file_chunk_properties() + test_min_chunk_and_overlap_size() + print("\n所有测试通过!") diff --git a/tests4/unittest/test_tokenizer.py b/tests4/unittest/test_tokenizer.py new file mode 100644 index 00000000..5dc99308 --- /dev/null +++ b/tests4/unittest/test_tokenizer.py @@ -0,0 +1,169 @@ +"""Tests for Tokenizers.""" + +import asyncio + +from reme4.components.tokenizer import JiebaTokenizer, RegexTokenizer + + +async def compare_tokenizers(texts: list[str], filter_stopwords: bool = False, name: str = ""): + """Compare both tokenizers on same input.""" + jieba = JiebaTokenizer(filter_stopwords=filter_stopwords) + regex = RegexTokenizer(filter_stopwords=filter_stopwords) + + await jieba.start() + await regex.start() + + jieba_result = jieba.tokenize(texts) + regex_result = regex.tokenize(texts) + + print(f"\n--- {name} ---") + print(f"输入: {texts}") + print(f"Jieba: {jieba_result}") + print(f"Regex: {regex_result}") + + await jieba.close() + await regex.close() + + return jieba_result, regex_result + + +def test_basic_chinese(): + """Test basic Chinese text.""" + + async def run(): + jieba_result, regex_result = await compare_tokenizers( + ["我爱北京天安门", "今天天气很好"], + name="纯中文", + ) + + assert "北京" in jieba_result[0] or "天安门" in jieba_result[0] + assert "我" in regex_result[0] + print("✓ test_basic_chinese passed") + + asyncio.run(run()) + + +def test_basic_english(): + """Test basic English text.""" + + async def run(): + _, regex_result = await compare_tokenizers( + ["I love Beijing very much"], + name="英文", + ) + + assert "love" in regex_result[0] + assert "beijing" in regex_result[0] + print("✓ test_basic_english passed") + + asyncio.run(run()) + + +def test_mixed_chinese_english(): + """Test mixed Chinese-English text.""" + + async def run(): + jieba_result, regex_result = await compare_tokenizers( + ["我用 Python 学习 machine learning 和 iPhone15 Pro。"], + name="中英混合", + ) + + assert "python" in jieba_result[0] + assert "python" in regex_result[0] + print("✓ test_mixed_chinese_english passed") + + asyncio.run(run()) + + +def test_open_example(): + """Test the 'open' example.""" + + async def run(): + jieba_result, regex_result = await compare_tokenizers( + ["我觉得open很好呀,能分好次吗?"], + name="'open' 案例", + ) + + # open 保持完整 + assert "open" in jieba_result[0] + assert "open" in regex_result[0] + + # Regex 中文按字拆分 + assert "我" in regex_result[0] + assert "很" in regex_result[0] + + print("✓ test_open_example passed") + + asyncio.run(run()) + + +def test_with_stopwords(): + """Test with stopwords filtering.""" + + async def run(): + jieba_result, regex_result = await compare_tokenizers( + ["我觉得open很好呀,能分好次吗?"], + filter_stopwords=True, + name="停用词过滤", + ) + + # 停用词被过滤 + assert "吗" not in jieba_result[0] + assert "吗" not in regex_result[0] + assert "的" not in jieba_result[0] + + print("✓ test_with_stopwords passed") + + asyncio.run(run()) + + +def test_multiple_texts(): + """Test multiple texts at once.""" + + async def run(): + texts = [ + "我爱北京天安门", + "I love Python programming", + "今天学习 machine learning", + ] + jieba_result, regex_result = await compare_tokenizers(texts, name="多个文本") + + assert len(jieba_result) == 3 + assert len(regex_result) == 3 + print("✓ test_multiple_texts passed") + + asyncio.run(run()) + + +def test_tokenizer_lifecycle(): + """Test tokenizer start/close lifecycle.""" + + async def run(): + tokenizer = JiebaTokenizer(filter_stopwords=True) + + assert not tokenizer.is_started + assert len(tokenizer.stopwords) == 0 + + await tokenizer.start() + assert tokenizer.is_started + assert len(tokenizer.stopwords) > 0 + + await tokenizer.close() + assert not tokenizer.is_started + assert len(tokenizer.stopwords) == 0 + + print("✓ test_tokenizer_lifecycle passed") + + asyncio.run(run()) + + +if __name__ == "__main__": + print("\n=== Tokenizer Tests ===") + test_basic_chinese() + test_basic_english() + test_mixed_chinese_english() + test_open_example() + test_with_stopwords() + test_multiple_texts() + test_tokenizer_lifecycle() + print("\n所有测试通过!")