diff --git a/.github/workflows/build-reusable.yml b/.github/workflows/build-reusable.yml index 73e2e1b..1b51e3e 100644 --- a/.github/workflows/build-reusable.yml +++ b/.github/workflows/build-reusable.yml @@ -47,8 +47,8 @@ jobs: - name: Run ty type checker run: | - # Run ty on core packages and apps, excluding tests - ty check --exclude "tests/**" packages/leann-core/src apps + # Run ty on core packages, apps, and tests + ty check packages/leann-core/src apps tests build: needs: [lint, type-check] diff --git a/packages/leann-core/src/leann/metadata_filter.py b/packages/leann-core/src/leann/metadata_filter.py index 1bf4ac1..5a8ffbd 100644 --- a/packages/leann-core/src/leann/metadata_filter.py +++ b/packages/leann-core/src/leann/metadata_filter.py @@ -7,7 +7,7 @@ operators for different data types including numbers, strings, booleans, and lis """ import logging -from typing import Any, Union +from typing import Any, Optional, Union logger = logging.getLogger(__name__) @@ -47,7 +47,7 @@ class MetadataFilterEngine: } def apply_filters( - self, search_results: list[dict[str, Any]], metadata_filters: MetadataFilters + self, search_results: list[dict[str, Any]], metadata_filters: Optional[MetadataFilters] ) -> list[dict[str, Any]]: """ Apply metadata filters to a list of search results. diff --git a/tests/test_basic.py b/tests/test_basic.py index 651111f..0268a70 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -91,7 +91,7 @@ def test_large_index(): builder.build_index(index_path) searcher = LeannSearcher(index_path) - results = searcher.search(["word10 word20"], top_k=10) - assert len(results[0]) == 10 + results = searcher.search("word10 word20", top_k=10) + assert len(results) == 10 # Cleanup searcher.cleanup() diff --git a/tests/test_cli_prompt_template.py b/tests/test_cli_prompt_template.py index 981bb78..774e29f 100644 --- a/tests/test_cli_prompt_template.py +++ b/tests/test_cli_prompt_template.py @@ -123,7 +123,7 @@ class TestPromptTemplateStoredInEmbeddingOptions: cli = LeannCLI() # Mock load_documents to return a document so builder is created - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() @@ -175,7 +175,7 @@ class TestPromptTemplateStoredInEmbeddingOptions: cli = LeannCLI() # Mock load_documents to return a document so builder is created - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() @@ -230,7 +230,7 @@ class TestPromptTemplateStoredInEmbeddingOptions: cli = LeannCLI() # Mock load_documents to return a document so builder is created - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() @@ -307,7 +307,7 @@ class TestPromptTemplateStoredInEmbeddingOptions: cli = LeannCLI() # Mock load_documents to return a document so builder is created - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() @@ -376,7 +376,7 @@ class TestPromptTemplateStoredInEmbeddingOptions: cli = LeannCLI() # Mock load_documents to return a document so builder is created - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() @@ -432,7 +432,7 @@ class TestPromptTemplateFlowsToComputeEmbeddings: cli = LeannCLI() # Mock load_documents to return a simple document - cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) + cli.load_documents = Mock(return_value=[{"text": "test content", "metadata": {}}]) # type: ignore[assignment] parser = cli.create_parser() diff --git a/tests/test_prompt_template_e2e.py b/tests/test_prompt_template_e2e.py index 80c9cce..6cb40b1 100644 --- a/tests/test_prompt_template_e2e.py +++ b/tests/test_prompt_template_e2e.py @@ -67,7 +67,7 @@ def check_lmstudio_available() -> bool: return False -def get_lmstudio_first_model() -> str: +def get_lmstudio_first_model() -> str | None: """Get the first available model from LM Studio.""" try: response = requests.get("http://localhost:1234/v1/models", timeout=5.0) @@ -91,6 +91,7 @@ class TestPromptTemplateOpenAI: model_name = get_lmstudio_first_model() if not model_name: pytest.skip("No models loaded in LM Studio") + assert model_name is not None # Type narrowing for type checker texts = ["artificial intelligence", "machine learning"] prompt_template = "search_query: " @@ -120,6 +121,7 @@ class TestPromptTemplateOpenAI: model_name = get_lmstudio_first_model() if not model_name: pytest.skip("No models loaded in LM Studio") + assert model_name is not None # Type narrowing for type checker text = "machine learning" base_url = "http://localhost:1234/v1" @@ -271,6 +273,7 @@ class TestLMStudioSDK: model_name = get_lmstudio_first_model() if not model_name: pytest.skip("No models loaded in LM Studio") + assert model_name is not None # Type narrowing for type checker try: from leann.embedding_compute import _query_lmstudio_context_limit diff --git a/tests/test_prompt_template_persistence.py b/tests/test_prompt_template_persistence.py index eefda04..55e24e3 100644 --- a/tests/test_prompt_template_persistence.py +++ b/tests/test_prompt_template_persistence.py @@ -581,7 +581,7 @@ class TestQueryTemplateApplicationInComputeEmbedding: # Create a concrete implementation for testing class TestSearcher(BaseSearcher): - def search(self, query_vectors, top_k, complexity, beam_width=1, **kwargs): + def search(self, query, top_k, complexity=64, beam_width=1, prune_ratio=0.0, recompute_embeddings=False, pruning_strategy="global", zmq_port=None, **kwargs): return {"labels": [], "distances": []} searcher = object.__new__(TestSearcher) @@ -625,7 +625,7 @@ class TestQueryTemplateApplicationInComputeEmbedding: # Create a concrete implementation for testing class TestSearcher(BaseSearcher): - def search(self, query_vectors, top_k, complexity, beam_width=1, **kwargs): + def search(self, query, top_k, complexity=64, beam_width=1, prune_ratio=0.0, recompute_embeddings=False, pruning_strategy="global", zmq_port=None, **kwargs): return {"labels": [], "distances": []} searcher = object.__new__(TestSearcher) @@ -671,7 +671,7 @@ class TestQueryTemplateApplicationInComputeEmbedding: from leann.searcher_base import BaseSearcher class TestSearcher(BaseSearcher): - def search(self, query_vectors, top_k, complexity, beam_width=1, **kwargs): + def search(self, query, top_k, complexity=64, beam_width=1, prune_ratio=0.0, recompute_embeddings=False, pruning_strategy="global", zmq_port=None, **kwargs): return {"labels": [], "distances": []} searcher = object.__new__(TestSearcher) @@ -710,7 +710,7 @@ class TestQueryTemplateApplicationInComputeEmbedding: from leann.searcher_base import BaseSearcher class TestSearcher(BaseSearcher): - def search(self, query_vectors, top_k, complexity, beam_width=1, **kwargs): + def search(self, query, top_k, complexity=64, beam_width=1, prune_ratio=0.0, recompute_embeddings=False, pruning_strategy="global", zmq_port=None, **kwargs): return {"labels": [], "distances": []} searcher = object.__new__(TestSearcher) @@ -774,7 +774,7 @@ class TestQueryTemplateApplicationInComputeEmbedding: from leann.searcher_base import BaseSearcher class TestSearcher(BaseSearcher): - def search(self, query_vectors, top_k, complexity, beam_width=1, **kwargs): + def search(self, query, top_k, complexity=64, beam_width=1, prune_ratio=0.0, recompute_embeddings=False, pruning_strategy="global", zmq_port=None, **kwargs): return {"labels": [], "distances": []} searcher = object.__new__(TestSearcher) diff --git a/tests/test_readme_examples.py b/tests/test_readme_examples.py index a562a02..371d13c 100644 --- a/tests/test_readme_examples.py +++ b/tests/test_readme_examples.py @@ -97,17 +97,15 @@ def test_backend_options(): with tempfile.TemporaryDirectory() as temp_dir: # Use smaller model in CI to avoid memory issues - if os.environ.get("CI") == "true": - model_args = { - "embedding_model": "sentence-transformers/all-MiniLM-L6-v2", - "dimensions": 384, - } - else: - model_args = {} + is_ci = os.environ.get("CI") == "true" + embedding_model = "sentence-transformers/all-MiniLM-L6-v2" if is_ci else "facebook/contriever" + dimensions = 384 if is_ci else None # Test HNSW backend (as shown in README) hnsw_path = str(Path(temp_dir) / "test_hnsw.leann") - builder_hnsw = LeannBuilder(backend_name="hnsw", **model_args) + builder_hnsw = LeannBuilder( + backend_name="hnsw", embedding_model=embedding_model, dimensions=dimensions + ) builder_hnsw.add_text("Test document for HNSW backend") builder_hnsw.build_index(hnsw_path) assert Path(hnsw_path).parent.exists() @@ -115,7 +113,9 @@ def test_backend_options(): # Test DiskANN backend (mentioned as available option) diskann_path = str(Path(temp_dir) / "test_diskann.leann") - builder_diskann = LeannBuilder(backend_name="diskann", **model_args) + builder_diskann = LeannBuilder( + backend_name="diskann", embedding_model=embedding_model, dimensions=dimensions + ) builder_diskann.add_text("Test document for DiskANN backend") builder_diskann.build_index(diskann_path) assert Path(diskann_path).parent.exists()