From 0316eaa15b3fc4a59a634ba6870e76795253f032 Mon Sep 17 00:00:00 2001 From: Hendrixx-RE Date: Mon, 20 Jul 2026 02:52:35 +0530 Subject: [PATCH 1/3] feat(python): add local cross-encoder reranker support --- sdks/python/sdk/pyproject.toml | 3 + sdks/python/sdk/src/moss/__init__.py | 3 +- sdks/python/sdk/src/moss/__init__.pyi | 6 + .../python/sdk/src/moss/client/moss_client.py | 104 ++++++++++++++---- 4 files changed, 95 insertions(+), 21 deletions(-) diff --git a/sdks/python/sdk/pyproject.toml b/sdks/python/sdk/pyproject.toml index a656fcc7..2c6e5af5 100644 --- a/sdks/python/sdk/pyproject.toml +++ b/sdks/python/sdk/pyproject.toml @@ -32,6 +32,9 @@ dependencies = [ ] [project.optional-dependencies] +rerank = [ + "sentence-transformers>=3.0.0", +] dev = [ "pytest>=8.4.2", "pytest-asyncio>=1.2.0", diff --git a/sdks/python/sdk/src/moss/__init__.py b/sdks/python/sdk/src/moss/__init__.py index 13e3ef2d..e8673ed7 100644 --- a/sdks/python/sdk/src/moss/__init__.py +++ b/sdks/python/sdk/src/moss/__init__.py @@ -31,12 +31,11 @@ JobPhase, JobProgress, JobStatusResponse, - QueryOptions, QueryResultDocumentInfo, SearchResult, ) -from .client.moss_client import MossClient +from .client.moss_client import MossClient, QueryOptions __version__ = "1.0.0b19" diff --git a/sdks/python/sdk/src/moss/__init__.pyi b/sdks/python/sdk/src/moss/__init__.pyi index 6aa41699..fcc0c844 100644 --- a/sdks/python/sdk/src/moss/__init__.pyi +++ b/sdks/python/sdk/src/moss/__init__.pyi @@ -171,12 +171,18 @@ class QueryOptions: top_k: Optional[int] alpha: Optional[float] filter: Optional[dict] + rerank: bool + rerank_top_k: Optional[int] + rerank_model: Optional[str] def __init__( self, embedding: Optional[Sequence[float]] = ..., top_k: Optional[int] = ..., alpha: Optional[float] = ..., filter: Optional[dict] = ..., + rerank: bool = ..., + rerank_top_k: Optional[int] = ..., + rerank_model: Optional[str] = ..., ) -> None: ... diff --git a/sdks/python/sdk/src/moss/client/moss_client.py b/sdks/python/sdk/src/moss/client/moss_client.py index 51bd2833..1e76dd7c 100644 --- a/sdks/python/sdk/src/moss/client/moss_client.py +++ b/sdks/python/sdk/src/moss/client/moss_client.py @@ -17,13 +17,35 @@ MutationOptions, MutationResult, JobStatusResponse, - QueryOptions, QueryResultDocumentInfo, SearchResult, ) logger = logging.getLogger(__name__) +from typing import Sequence + +class QueryOptions: + """Options for search queries.""" + def __init__( + self, + embedding: Optional[Sequence[float]] = None, + top_k: Optional[int] = None, + alpha: Optional[float] = None, + filter: Optional[dict] = None, + rerank: bool = False, + rerank_top_k: Optional[int] = None, + rerank_model: Optional[str] = None, + ): + self.embedding = embedding + self.top_k = top_k + self.alpha = alpha + self.filter = filter + self.rerank = rerank + self.rerank_top_k = rerank_top_k + self.rerank_model = rerank_model + + def _get_manage_url() -> str: """Manage URL, overridable via env for local development.""" @@ -189,24 +211,28 @@ async def query( Otherwise, falls back to the cloud query API. Args: - options: Query options (top_k, alpha, embedding, filter). Example filter: - QueryOptions(filter={"$and": [ - {"field": "city", "condition": {"$eq": "NYC"}}, - {"field": "price", "condition": {"$lt": "50"}}, - ]}) + options: Query options (top_k, alpha, embedding, filter, rerank, etc). """ is_loaded = await asyncio.to_thread(self._manager.has_index, name) - if is_loaded: - return await self._query_local(name, query, options) + rerank = getattr(options, "rerank", False) + override_top_k = getattr(options, "rerank_top_k", 50) if rerank else None - if getattr(options, "filter", None) is not None: - logger.warning( - "Metadata filter ignored: filtering is only supported for locally loaded indexes. " - "Call load_index('%s') first.", - name, - ) - return await self._query_cloud(name, query, options) + if is_loaded: + result = await self._query_local(name, query, options, override_top_k) + else: + if getattr(options, "filter", None) is not None: + logger.warning( + "Metadata filter ignored: filtering is only supported for locally loaded indexes. " + "Call load_index('%s') first.", + name, + ) + result = await self._query_cloud(name, query, options, override_top_k) + + if rerank: + result = await self._rerank_results(query, result, options) + + return result # -- Internal --------------------------------------------------- @@ -215,10 +241,9 @@ async def _query_local( name: str, query: str, options: Optional[QueryOptions], + override_top_k: Optional[int] = None, ) -> SearchResult: - top_k = getattr(options, "top_k", None) - if top_k is None: - top_k = 5 + top_k = override_top_k if override_top_k is not None else (getattr(options, "top_k", None) or 5) alpha = getattr(options, "alpha", None) if alpha is None: alpha = 0.8 @@ -258,9 +283,10 @@ async def _query_cloud( name: str, query: str, options: Optional[QueryOptions], + override_top_k: Optional[int] = None, ) -> SearchResult: """Fallback: query via the cloud API when the index is not loaded locally.""" - top_k = getattr(options, "top_k", None) or 10 + top_k = override_top_k if override_top_k is not None else (getattr(options, "top_k", None) or 10) query_embedding = getattr(options, "embedding", None) request_body: Dict[str, Any] = { @@ -306,6 +332,46 @@ def _dict_to_search_result(data: dict) -> SearchResult: time_taken_ms=data.get("timeTakenMs"), ) + async def _rerank_results( + self, query: str, search_result: SearchResult, options: Optional[QueryOptions] + ) -> SearchResult: + if not search_result.docs: + return search_result + + try: + from sentence_transformers import CrossEncoder + except ImportError: + raise ImportError( + "The 'sentence-transformers' package is required for reranking. " + "Install it with: pip install 'moss[rerank]'" + ) + + model_name = getattr(options, "rerank_model", None) or "cross-encoder/ms-marco-MiniLM-L-6-v2" + + def do_rerank(): + if not hasattr(self.__class__, "_cross_encoder_cache"): + self.__class__._cross_encoder_cache = {} + + if model_name not in self.__class__._cross_encoder_cache: + self.__class__._cross_encoder_cache[model_name] = CrossEncoder(model_name) + + model = self.__class__._cross_encoder_cache[model_name] + + pairs = [[query, doc.text] for doc in search_result.docs] + scores = model.predict(pairs) + + for doc, score in zip(search_result.docs, scores): + doc.score = float(score) + + search_result.docs.sort(key=lambda d: d.score, reverse=True) + + original_top_k = getattr(options, "top_k", None) or 5 + search_result.docs = search_result.docs[:original_top_k] + return search_result + + return await asyncio.to_thread(do_rerank) + + def _resolve_model_id( self, docs: List[DocumentInfo], From decc3ec076b73ddff75be3cae721b476e11eb084 Mon Sep 17 00:00:00 2001 From: Hendrixx-RE Date: Mon, 20 Jul 2026 03:06:18 +0530 Subject: [PATCH 2/3] "linter fix" --- sdks/python/sdk/README.md | 25 ++ sdks/python/sdk/src/moss/__init__.py | 8 +- sdks/python/sdk/src/moss/__init__.pyi | 27 -- .../python/sdk/src/moss/client/moss_client.py | 73 +++-- sdks/python/sdk/tests/conftest.py | 5 +- sdks/python/sdk/tests/test_client_extended.py | 5 +- sdks/python/sdk/tests/test_cloud_fallback.py | 4 +- .../sdk/tests/test_create_index_versions.py | 75 +++-- sdks/python/sdk/tests/test_e2e.py | 35 ++- sdks/python/sdk/tests/test_hot_reload.py | 12 +- .../sdk/tests/test_metadata_filter_e2e.py | 168 +++++++++--- sdks/python/sdk/tests/test_rerank.py | 85 ++++++ sdks/python/sdk/tests/test_search.py | 259 ++++++++++++------ sdks/python/sdk/tests/test_types.py | 10 +- 14 files changed, 558 insertions(+), 233 deletions(-) create mode 100644 sdks/python/sdk/tests/test_rerank.py diff --git a/sdks/python/sdk/README.md b/sdks/python/sdk/README.md index 5f5ef4af..b8ebefb2 100644 --- a/sdks/python/sdk/README.md +++ b/sdks/python/sdk/README.md @@ -186,6 +186,31 @@ results = await client.query( For a complete runnable example, see [`examples/python/metadata_filtering.py`](../../examples/python/metadata_filtering.py). +### Reranking (High Precision) + +If you need higher precision at the top of the result list than first-stage retrieval gives, you can enable the optional local cross-encoder reranker. + +When `rerank=True` is provided, Moss will fetch a larger pool of candidates (`rerank_top_k`, default 50) and locally re-score them using a `sentence-transformers` cross-encoder before returning the final `top_k` results. + +First, install the optional dependencies: +```bash +pip install 'moss[rerank]' +``` + +Then, enable it in your query: +```python +results = await client.query( + "my-docs", + "How to process a refund?", + QueryOptions( + top_k=5, + rerank=True, + rerank_top_k=50, # Optional: number of candidates to fetch for reranking + # rerank_model="cross-encoder/ms-marco-MiniLM-L-6-v2" # Optional override + ), +) +``` + ## šŸ§  Providing custom embeddings Already using your own embedding model? Supply vectors directly when managing diff --git a/sdks/python/sdk/src/moss/__init__.py b/sdks/python/sdk/src/moss/__init__.py index e8673ed7..c0f0e55c 100644 --- a/sdks/python/sdk/src/moss/__init__.py +++ b/sdks/python/sdk/src/moss/__init__.py @@ -24,13 +24,13 @@ IndexInfo, IndexStatus, IndexStatusValues, - ModelRef, - MutationOptions, - MutationResult, - JobStatus, JobPhase, JobProgress, + JobStatus, JobStatusResponse, + ModelRef, + MutationOptions, + MutationResult, QueryResultDocumentInfo, SearchResult, ) diff --git a/sdks/python/sdk/src/moss/__init__.pyi b/sdks/python/sdk/src/moss/__init__.pyi index fcc0c844..2248d824 100644 --- a/sdks/python/sdk/src/moss/__init__.pyi +++ b/sdks/python/sdk/src/moss/__init__.pyi @@ -2,57 +2,45 @@ from __future__ import annotations from typing import ClassVar, Dict, List, Optional, Sequence - class MossClient: """Semantic search client for vector similarity operations.""" DEFAULT_MODEL_ID: ClassVar[str] def __init__(self, project_id: str, project_key: str) -> None: ... - async def create_index( self, name: str, docs: List[DocumentInfo], model_id: Optional[str] = ..., ) -> MutationResult: ... - async def add_docs( self, name: str, docs: List[DocumentInfo], options: Optional[MutationOptions] = None, ) -> MutationResult: ... - async def delete_docs( self, name: str, doc_ids: List[str], ) -> MutationResult: ... - async def get_job_status(self, job_id: str) -> JobStatusResponse: ... - async def get_index(self, name: str) -> IndexInfo: ... - async def list_indexes(self) -> List[IndexInfo]: ... - async def delete_index(self, name: str) -> bool: ... - async def get_docs( self, name: str, options: Optional[GetDocumentsOptions] = None, ) -> List[DocumentInfo]: ... - async def load_index( self, name: str, auto_refresh: bool = False, polling_interval_in_seconds: int = 600, ) -> str: ... - async def unload_index(self, name: str) -> None: ... - async def query( self, name: str, @@ -60,7 +48,6 @@ class MossClient: options: Optional[QueryOptions] = None, ) -> SearchResult: ... - class MutationResult: """Return value from create_index/add_docs/delete_docs.""" @@ -68,7 +55,6 @@ class MutationResult: index_name: str doc_count: int - class MutationOptions: """Options for add_docs (e.g. upsert behavior).""" @@ -76,7 +62,6 @@ class MutationOptions: def __init__(self, upsert: Optional[bool] = None) -> None: ... - class GetDocumentsOptions: """Options for get_docs (e.g. filter by document IDs).""" @@ -84,7 +69,6 @@ class GetDocumentsOptions: def __init__(self, doc_ids: Optional[List[str]] = None) -> None: ... - class JobStatus: """Enum-like class for job status values.""" @@ -96,7 +80,6 @@ class JobStatus: value: str - class JobPhase: """Enum-like class for job phase values.""" @@ -109,7 +92,6 @@ class JobPhase: value: str - class JobProgress: """Progress update for a job.""" @@ -118,7 +100,6 @@ class JobProgress: progress: float current_phase: Optional[JobPhase] - class JobStatusResponse: """Full status response from get_job_status.""" @@ -131,13 +112,11 @@ class JobStatusResponse: updated_at: str completed_at: Optional[str] - class ModelRef: id: str version: str def __init__(self, id: str, version: str) -> None: ... - class QueryResultDocumentInfo: id: str text: str @@ -151,7 +130,6 @@ class QueryResultDocumentInfo: score: float = ..., ) -> None: ... - class DocumentInfo: id: str text: str @@ -165,7 +143,6 @@ class DocumentInfo: embedding: Optional[Sequence[float]] = ..., ) -> None: ... - class QueryOptions: embedding: Optional[Sequence[float]] top_k: Optional[int] @@ -185,7 +162,6 @@ class QueryOptions: rerank_model: Optional[str] = ..., ) -> None: ... - class IndexInfo: id: str name: str @@ -207,7 +183,6 @@ class IndexInfo: model: ModelRef, ) -> None: ... - class SearchResult: docs: List[QueryResultDocumentInfo] query: str @@ -221,7 +196,6 @@ class SearchResult: time_taken_ms: Optional[int] = None, ) -> None: ... - class IndexStatus: NotStarted: ClassVar[str] Building: ClassVar[str] @@ -229,7 +203,6 @@ class IndexStatus: Failed: ClassVar[str] def __init__(self, value: str) -> None: ... - IndexStatusValues: Dict[str, str] __version__: str diff --git a/sdks/python/sdk/src/moss/client/moss_client.py b/sdks/python/sdk/src/moss/client/moss_client.py index 1e76dd7c..941acb27 100644 --- a/sdks/python/sdk/src/moss/client/moss_client.py +++ b/sdks/python/sdk/src/moss/client/moss_client.py @@ -4,29 +4,29 @@ import logging import os import uuid -from typing import Any, Dict, List, Optional +from typing import Any, ClassVar, Dict, List, Optional, Sequence import httpx from moss_core import ( CLOUD_API_MANAGE_URL, - ManageClient, DocumentInfo, GetDocumentsOptions, IndexInfo, IndexManager, + JobStatusResponse, + ManageClient, MutationOptions, MutationResult, - JobStatusResponse, QueryResultDocumentInfo, SearchResult, ) logger = logging.getLogger(__name__) -from typing import Sequence class QueryOptions: """Options for search queries.""" + def __init__( self, embedding: Optional[Sequence[float]] = None, @@ -37,16 +37,31 @@ def __init__( rerank_top_k: Optional[int] = None, rerank_model: Optional[str] = None, ): + if top_k is not None and (not isinstance(top_k, int) or top_k < 1): + raise ValueError("top_k must be an integer >= 1") + if alpha is not None and ( + not isinstance(alpha, (int, float)) or not (0.0 <= alpha <= 1.0) + ): + raise ValueError("alpha must be a float between 0.0 and 1.0") + if embedding is not None: + try: + embedding = [float(x) for x in embedding] + except (TypeError, ValueError): + raise ValueError("embedding must be a sequence of numbers") + if rerank_top_k is not None and ( + not isinstance(rerank_top_k, int) or rerank_top_k < 1 + ): + raise ValueError("rerank_top_k must be an integer >= 1") + self.embedding = embedding self.top_k = top_k - self.alpha = alpha + self.alpha = float(alpha) if alpha is not None else None self.filter = filter - self.rerank = rerank + self.rerank = bool(rerank) self.rerank_top_k = rerank_top_k self.rerank_model = rerank_model - def _get_manage_url() -> str: """Manage URL, overridable via env for local development.""" return os.getenv("MOSS_CLOUD_API_MANAGE_URL", CLOUD_API_MANAGE_URL) @@ -83,6 +98,7 @@ class MossClient: """ DEFAULT_MODEL_ID = "moss-minilm" + _cross_encoder_cache: ClassVar[Dict[str, Any]] = {} def __init__(self, project_id: str, project_key: str) -> None: self._project_id = project_id @@ -215,8 +231,10 @@ async def query( """ is_loaded = await asyncio.to_thread(self._manager.has_index, name) - rerank = getattr(options, "rerank", False) - override_top_k = getattr(options, "rerank_top_k", 50) if rerank else None + rerank = getattr(options, "rerank", False) is True + override_top_k = ( + (getattr(options, "rerank_top_k", None) or 50) if rerank else None + ) if is_loaded: result = await self._query_local(name, query, options, override_top_k) @@ -228,10 +246,10 @@ async def query( name, ) result = await self._query_cloud(name, query, options, override_top_k) - + if rerank: result = await self._rerank_results(query, result, options) - + return result # -- Internal --------------------------------------------------- @@ -243,7 +261,11 @@ async def _query_local( options: Optional[QueryOptions], override_top_k: Optional[int] = None, ) -> SearchResult: - top_k = override_top_k if override_top_k is not None else (getattr(options, "top_k", None) or 5) + top_k = ( + override_top_k + if override_top_k is not None + else (getattr(options, "top_k", None) or 5) + ) alpha = getattr(options, "alpha", None) if alpha is None: alpha = 0.8 @@ -286,7 +308,11 @@ async def _query_cloud( override_top_k: Optional[int] = None, ) -> SearchResult: """Fallback: query via the cloud API when the index is not loaded locally.""" - top_k = override_top_k if override_top_k is not None else (getattr(options, "top_k", None) or 10) + top_k = ( + override_top_k + if override_top_k is not None + else (getattr(options, "top_k", None) or 10) + ) query_embedding = getattr(options, "embedding", None) request_body: Dict[str, Any] = { @@ -346,32 +372,37 @@ async def _rerank_results( "Install it with: pip install 'moss[rerank]'" ) - model_name = getattr(options, "rerank_model", None) or "cross-encoder/ms-marco-MiniLM-L-6-v2" + model_name = ( + getattr(options, "rerank_model", None) + or "cross-encoder/ms-marco-MiniLM-L-6-v2" + ) - def do_rerank(): + def do_rerank() -> SearchResult: if not hasattr(self.__class__, "_cross_encoder_cache"): self.__class__._cross_encoder_cache = {} if model_name not in self.__class__._cross_encoder_cache: - self.__class__._cross_encoder_cache[model_name] = CrossEncoder(model_name) + self.__class__._cross_encoder_cache[model_name] = CrossEncoder( + model_name + ) model = self.__class__._cross_encoder_cache[model_name] - pairs = [[query, doc.text] for doc in search_result.docs] + local_docs = search_result.docs + pairs = [[query, doc.text] for doc in local_docs] scores = model.predict(pairs) - for doc, score in zip(search_result.docs, scores): + for doc, score in zip(local_docs, scores): doc.score = float(score) - search_result.docs.sort(key=lambda d: d.score, reverse=True) + local_docs.sort(key=lambda d: d.score, reverse=True) original_top_k = getattr(options, "top_k", None) or 5 - search_result.docs = search_result.docs[:original_top_k] + search_result.docs = local_docs[:original_top_k] return search_result return await asyncio.to_thread(do_rerank) - def _resolve_model_id( self, docs: List[DocumentInfo], diff --git a/sdks/python/sdk/tests/conftest.py b/sdks/python/sdk/tests/conftest.py index 28651233..1189811e 100644 --- a/sdks/python/sdk/tests/conftest.py +++ b/sdks/python/sdk/tests/conftest.py @@ -6,14 +6,13 @@ import os import warnings from pathlib import Path - from unittest.mock import MagicMock, patch -from moss import MossClient - import pytest from dotenv import load_dotenv +from moss import MossClient + # Load .env from project root. project_env_path = Path(__file__).parent.parent / ".env" load_dotenv(project_env_path) diff --git a/sdks/python/sdk/tests/test_client_extended.py b/sdks/python/sdk/tests/test_client_extended.py index 2b53f916..4889f582 100644 --- a/sdks/python/sdk/tests/test_client_extended.py +++ b/sdks/python/sdk/tests/test_client_extended.py @@ -2,14 +2,13 @@ from __future__ import annotations +from importlib.metadata import version from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from importlib.metadata import version - -from moss import __version__, MossClient +from moss import MossClient, __version__ from moss.client.moss_client import _get_manage_url, _get_query_url diff --git a/sdks/python/sdk/tests/test_cloud_fallback.py b/sdks/python/sdk/tests/test_cloud_fallback.py index 938465bc..e9aa0531 100644 --- a/sdks/python/sdk/tests/test_cloud_fallback.py +++ b/sdks/python/sdk/tests/test_cloud_fallback.py @@ -51,9 +51,7 @@ async def cloud_fallback_index(self, moss_client): index_name = generate_unique_index_name("test-cloud-fallback") # Create the index with documents - docs = [ - DocumentInfo(id=doc["id"], text=doc["text"]) for doc in TEST_DOCUMENTS - ] + docs = [DocumentInfo(id=doc["id"], text=doc["text"]) for doc in TEST_DOCUMENTS] await moss_client.create_index(index_name, docs, TEST_MODEL_ID) yield index_name diff --git a/sdks/python/sdk/tests/test_create_index_versions.py b/sdks/python/sdk/tests/test_create_index_versions.py index 9f81a71b..14321b47 100644 --- a/sdks/python/sdk/tests/test_create_index_versions.py +++ b/sdks/python/sdk/tests/test_create_index_versions.py @@ -8,7 +8,8 @@ import string import time from dataclasses import dataclass -from importlib.metadata import PackageNotFoundError, version as pkg_version +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as pkg_version from typing import List, Optional import pytest @@ -17,6 +18,7 @@ from .constants import TEST_MODEL_ID, TEST_PROJECT_ID, TEST_PROJECT_KEY + # Define the structure to hold benchmark results @dataclass class BenchResult: @@ -25,6 +27,7 @@ class BenchResult: success: bool error: str = "" + # Generate dummy documents with random deterministic text of fixed length 200 def _gen_docs(n: int, seed: int = 1337) -> List[DocumentInfo]: """Generate n dummy documents.""" @@ -34,22 +37,27 @@ def _gen_docs(n: int, seed: int = 1337) -> List[DocumentInfo]: for i in range(n): txt = "".join(rng.choices(alphabet, k=200)).strip() try: - docs.append(DocumentInfo(id=f"doc_{i}", text=f"Document {i}. {txt}", metadata={})) + docs.append( + DocumentInfo(id=f"doc_{i}", text=f"Document {i}. {txt}", metadata={}) + ) except TypeError: docs.append(DocumentInfo(id=f"doc_{i}", text=f"Document {i}. {txt}")) return docs + # Try to create index and measure time taken -async def _try_create_index(client: MossClient, index_name: str, doc_count: int) -> BenchResult: +async def _try_create_index( + client: MossClient, index_name: str, doc_count: int +) -> BenchResult: """Try to create index, return result with timing.""" docs = _gen_docs(doc_count) - + # Cleanup before try: await client.delete_index(index_name) except Exception: pass - + t0 = time.perf_counter() try: result = client.create_index(index_name, docs, TEST_MODEL_ID) @@ -57,67 +65,80 @@ async def _try_create_index(client: MossClient, index_name: str, doc_count: int) success = bool(await result) else: success = bool(result) - + t1 = time.perf_counter() time_ms = (t1 - t0) * 1000.0 - + # Cleanup after try: await client.delete_index(index_name) except Exception: pass - - return BenchResult(doc_count=doc_count, time_ms=time_ms, success=success, error="") - + + return BenchResult( + doc_count=doc_count, time_ms=time_ms, success=success, error="" + ) + except Exception as e: t1 = time.perf_counter() time_ms = (t1 - t0) * 1000.0 err_msg = str(e) - + # Cleanup after (best effort) try: await client.delete_index(index_name) except Exception: pass - - return BenchResult(doc_count=doc_count, time_ms=time_ms, success=False, error=err_msg) + + return BenchResult( + doc_count=doc_count, time_ms=time_ms, success=False, error=err_msg + ) + # The main benchmark test @pytest.mark.asyncio async def test_create_index_versions_bench(): """Benchmark create_index across different doc counts.""" - doc_counts = [600 , 800, 1000] - + doc_counts = [600, 800, 1000] + try: sdk_version = pkg_version("moss") except PackageNotFoundError: sdk_version = "unknown" - + client = MossClient(TEST_PROJECT_ID, TEST_PROJECT_KEY) results = [] - + for i, n in enumerate(doc_counts, 1): - print(f"Running experiment {i}/{len(doc_counts)}: create_index with {n} documents...", flush=True, end=" ") + print( + f"Running experiment {i}/{len(doc_counts)}: create_index with {n} documents...", + flush=True, + end=" ", + ) index_name = f"createindex-bench-{sdk_version.replace('.', '_')}-{n}" result = await _try_create_index(client, index_name, n) results.append(result) - + if result.success: print(f"āœ… {result.time_ms:.1f}ms", flush=True) else: print("āŒ failed", flush=True) - + # Print table print("\n" + "=" * 110) print(" CREATE_INDEX BENCHMARK") print("=" * 110) - print(f"{'sdk_version':<16} | {'docs':>6} | {'time_taken':>14} | {'status':<15} | notes") + print( + f"{'sdk_version':<16} | {'docs':>6} | {'time_taken':>14} | {'status':<15} | notes" + ) print("-" * 110) - + for r in results: if r.success: time_str = f"{r.time_ms:.1f} ms" if r.time_ms else "N/A" - print(f"{sdk_version:<16} | {r.doc_count:>6} | {time_str:>14} | {'āœ… success':<15} |") + print( + f"{sdk_version:<16} | {r.doc_count:>6} | {time_str:>14} | {'āœ… success':<15} |" + ) else: # Classify error err_lower = r.error.lower() @@ -127,8 +148,10 @@ async def test_create_index_versions_bench(): note = "not supported (payload too large)" else: note = r.error[:60] + "..." if len(r.error) > 60 else r.error - + time_str = f"{r.time_ms:.1f} ms" if r.time_ms else "N/A" - print(f"{sdk_version:<16} | {r.doc_count:>6} | {time_str:>14} | {'āŒ failed':<15} | {note}") - + print( + f"{sdk_version:<16} | {r.doc_count:>6} | {time_str:>14} | {'āŒ failed':<15} | {note}" + ) + print("=" * 110 + "\n") diff --git a/sdks/python/sdk/tests/test_e2e.py b/sdks/python/sdk/tests/test_e2e.py index edbc375e..c6a70259 100644 --- a/sdks/python/sdk/tests/test_e2e.py +++ b/sdks/python/sdk/tests/test_e2e.py @@ -19,15 +19,14 @@ from typing import List import pytest - from moss_core import DocumentInfo from moss import GetDocumentsOptions, MossClient from .constants import ( + TEST_MODEL_ID, TEST_PROJECT_ID, TEST_PROJECT_KEY, - TEST_MODEL_ID, generate_unique_index_name, ) @@ -50,15 +49,19 @@ def generate_docs(count: int, seed: int = 42) -> List[DocumentInfo]: docs = [] for i in range(count): topic = rng.choice(DOC_TOPICS) - docs.append(DocumentInfo( - id=f"doc-{i}", - text=f"{topic} (variation {i}, seed {rng.randint(0, 100000)})", - )) + docs.append( + DocumentInfo( + id=f"doc-{i}", + text=f"{topic} (variation {i}, seed {rng.randint(0, 100000)})", + ) + ) return docs def generate_docs_with_embeddings( - count: int, dimension: int, seed: int = 42, + count: int, + dimension: int, + seed: int = 42, ) -> List[DocumentInfo]: rng = random.Random(seed) docs = [] @@ -68,16 +71,19 @@ def generate_docs_with_embeddings( magnitude = sum(x * x for x in embedding) ** 0.5 if magnitude > 0: embedding = [x / magnitude for x in embedding] - docs.append(DocumentInfo( - id=f"doc-{i}", - text=f"{topic} (variation {i})", - embedding=embedding, - )) + docs.append( + DocumentInfo( + id=f"doc-{i}", + text=f"{topic} (variation {i})", + embedding=embedding, + ) + ) return docs # -- Fixtures ---------------------------------------------------------- + @pytest.fixture def client(): return MossClient(TEST_PROJECT_ID, TEST_PROJECT_KEY) @@ -85,6 +91,7 @@ def client(): # -- Tests ------------------------------------------------------------- + @pytest.mark.e2e class TestBulkCreateIndex: """Test create_index with large doc counts.""" @@ -284,7 +291,9 @@ async def test_get_docs_from_large_index(self, client): all_docs = await client.get_docs(index_name) assert len(all_docs) == 5_000 - specific = await client.get_docs(index_name, GetDocumentsOptions(doc_ids=["doc-0", "doc-100", "doc-999"])) + specific = await client.get_docs( + index_name, GetDocumentsOptions(doc_ids=["doc-0", "doc-100", "doc-999"]) + ) assert len(specific) == 3 ids = {d.id for d in specific} assert ids == {"doc-0", "doc-100", "doc-999"} diff --git a/sdks/python/sdk/tests/test_hot_reload.py b/sdks/python/sdk/tests/test_hot_reload.py index 6649d951..29c22a4b 100644 --- a/sdks/python/sdk/tests/test_hot_reload.py +++ b/sdks/python/sdk/tests/test_hot_reload.py @@ -72,9 +72,7 @@ async def test_load_index_with_auto_refresh_enabled( assert loaded_name == test_index @pytest.mark.asyncio - async def test_accept_custom_polling_interval( - self, moss_client, test_index - ): + async def test_accept_custom_polling_interval(self, moss_client, test_index): """Should accept custom polling interval.""" # Load with a custom polling interval (5 minutes) loaded_name = await moss_client.load_index( @@ -183,9 +181,7 @@ async def test_index(self, moss_client): pass @pytest.mark.asyncio - async def test_query_cloud_when_index_not_loaded( - self, moss_client, test_index - ): + async def test_query_cloud_when_index_not_loaded(self, moss_client, test_index): """Should query cloud when index is not loaded locally.""" # Query without loading - should fall back to cloud results = await moss_client.query( @@ -199,9 +195,7 @@ async def test_query_cloud_when_index_not_loaded( assert len(results.docs) > 0 @pytest.mark.asyncio - async def test_query_locally_after_loading_index( - self, moss_client, test_index - ): + async def test_query_locally_after_loading_index(self, moss_client, test_index): """Should query locally after loading index.""" # Load index locally await moss_client.load_index(test_index) diff --git a/sdks/python/sdk/tests/test_metadata_filter_e2e.py b/sdks/python/sdk/tests/test_metadata_filter_e2e.py index 9b971aa2..92e95442 100644 --- a/sdks/python/sdk/tests/test_metadata_filter_e2e.py +++ b/sdks/python/sdk/tests/test_metadata_filter_e2e.py @@ -27,12 +27,59 @@ ) FILTER_DOCS = [ - DocumentInfo(id="f1", text="Cheap coffee shop in downtown New York with great espresso", metadata={"city": "NYC", "price": "12", "category": "food", "location": "40.7580,-73.9855"}), - DocumentInfo(id="f2", text="High-end sushi restaurant in central Tokyo with omakase menu", metadata={"city": "Tokyo", "price": "45", "category": "food", "location": "35.6762,139.6503"}), - DocumentInfo(id="f3", text="Weekly tech meetup and hackathon event in New York coworking space", metadata={"city": "NYC", "price": "0", "category": "tech", "location": "40.6892,-74.0445"}), - DocumentInfo(id="f4", text="Art museum and cultural exhibition in the heart of Paris", metadata={"city": "Paris", "price": "20", "category": "culture", "location": "48.8566,2.3522"}), - DocumentInfo(id="f5", text="Popular street food market in Tokyo with yakitori and ramen stalls", metadata={"city": "Tokyo", "price": "8", "category": "food", "location": "35.6595,139.7004"}), - DocumentInfo(id="f6", text="General document with no metadata attached for edge case testing"), + DocumentInfo( + id="f1", + text="Cheap coffee shop in downtown New York with great espresso", + metadata={ + "city": "NYC", + "price": "12", + "category": "food", + "location": "40.7580,-73.9855", + }, + ), + DocumentInfo( + id="f2", + text="High-end sushi restaurant in central Tokyo with omakase menu", + metadata={ + "city": "Tokyo", + "price": "45", + "category": "food", + "location": "35.6762,139.6503", + }, + ), + DocumentInfo( + id="f3", + text="Weekly tech meetup and hackathon event in New York coworking space", + metadata={ + "city": "NYC", + "price": "0", + "category": "tech", + "location": "40.6892,-74.0445", + }, + ), + DocumentInfo( + id="f4", + text="Art museum and cultural exhibition in the heart of Paris", + metadata={ + "city": "Paris", + "price": "20", + "category": "culture", + "location": "48.8566,2.3522", + }, + ), + DocumentInfo( + id="f5", + text="Popular street food market in Tokyo with yakitori and ramen stalls", + metadata={ + "city": "Tokyo", + "price": "8", + "category": "food", + "location": "35.6595,139.7004", + }, + ), + DocumentInfo( + id="f6", text="General document with no metadata attached for edge case testing" + ), ] @@ -42,9 +89,12 @@ def moss_client(): async def query_ids(client, name, filt): - result = await client.query(name, "food and restaurants", QueryOptions(top_k=10, filter=filt)) + result = await client.query( + name, "food and restaurants", QueryOptions(top_k=10, filter=filt) + ) return sorted(d.id for d in result.docs) + class TestMetadataFilterE2E: @pytest_asyncio.fixture(scope="class") @@ -60,93 +110,131 @@ async def loaded_index(self, moss_client): @pytest.mark.asyncio async def test_eq(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$eq": "NYC"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "city", "condition": {"$eq": "NYC"}} + ) assert ids == ["f1", "f3"] @pytest.mark.asyncio async def test_ne(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$ne": "NYC"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "city", "condition": {"$ne": "NYC"}} + ) assert ids == ["f2", "f4", "f5"] @pytest.mark.asyncio async def test_gt_numeric(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "price", "condition": {"$gt": "10"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "price", "condition": {"$gt": "10"}} + ) assert ids == ["f1", "f2", "f4"] @pytest.mark.asyncio async def test_lt_int_coercion(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "price", "condition": {"$lt": 15}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "price", "condition": {"$lt": 15}} + ) assert ids == ["f1", "f3", "f5"] @pytest.mark.asyncio async def test_gte(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "price", "condition": {"$gte": "20"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "price", "condition": {"$gte": "20"}} + ) assert ids == ["f2", "f4"] @pytest.mark.asyncio async def test_lte(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "price", "condition": {"$lte": "8"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "price", "condition": {"$lte": "8"}} + ) assert ids == ["f3", "f5"] @pytest.mark.asyncio async def test_in(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$in": ["NYC", "Paris"]}}) + ids = await query_ids( + moss_client, + loaded_index, + {"field": "city", "condition": {"$in": ["NYC", "Paris"]}}, + ) assert ids == ["f1", "f3", "f4"] @pytest.mark.asyncio async def test_nin(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$nin": ["NYC"]}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "city", "condition": {"$nin": ["NYC"]}} + ) assert ids == ["f2", "f4", "f5"] @pytest.mark.asyncio async def test_and(self, moss_client, loaded_index): - filt = {"$and": [ - {"field": "city", "condition": {"$eq": "NYC"}}, - {"field": "category", "condition": {"$eq": "food"}}, - ]} + filt = { + "$and": [ + {"field": "city", "condition": {"$eq": "NYC"}}, + {"field": "category", "condition": {"$eq": "food"}}, + ] + } ids = await query_ids(moss_client, loaded_index, filt) assert ids == ["f1"] @pytest.mark.asyncio async def test_or(self, moss_client, loaded_index): - filt = {"$or": [ - {"field": "city", "condition": {"$eq": "Paris"}}, - {"field": "category", "condition": {"$eq": "tech"}}, - ]} + filt = { + "$or": [ + {"field": "city", "condition": {"$eq": "Paris"}}, + {"field": "category", "condition": {"$eq": "tech"}}, + ] + } ids = await query_ids(moss_client, loaded_index, filt) assert ids == ["f3", "f4"] @pytest.mark.asyncio async def test_nested_and_or(self, moss_client, loaded_index): - filt = {"$and": [ - {"$or": [ - {"field": "city", "condition": {"$eq": "NYC"}}, - {"field": "city", "condition": {"$eq": "Tokyo"}}, - ]}, - {"field": "category", "condition": {"$eq": "food"}}, - ]} + filt = { + "$and": [ + { + "$or": [ + {"field": "city", "condition": {"$eq": "NYC"}}, + {"field": "city", "condition": {"$eq": "Tokyo"}}, + ] + }, + {"field": "category", "condition": {"$eq": "food"}}, + ] + } ids = await query_ids(moss_client, loaded_index, filt) assert ids == ["f1", "f2", "f5"] @pytest.mark.asyncio async def test_no_matches(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$eq": "Berlin"}}) + ids = await query_ids( + moss_client, loaded_index, {"field": "city", "condition": {"$eq": "Berlin"}} + ) assert ids == [] @pytest.mark.asyncio async def test_skips_docs_without_metadata(self, moss_client, loaded_index): - ids = await query_ids(moss_client, loaded_index, {"field": "city", "condition": {"$ne": "nonexistent"}}) + ids = await query_ids( + moss_client, + loaded_index, + {"field": "city", "condition": {"$ne": "nonexistent"}}, + ) assert "f6" not in ids @pytest.mark.asyncio async def test_no_filter_returns_all(self, moss_client, loaded_index): - result = await moss_client.query(loaded_index, "food and restaurants", QueryOptions(top_k=10)) + result = await moss_client.query( + loaded_index, "food and restaurants", QueryOptions(top_k=10) + ) assert len(result.docs) == 6 @pytest.mark.asyncio async def test_near_within_range(self, moss_client, loaded_index): # 10km around Times Square ā€” f1 (~0km) and f3 (~8.7km, Statue of Liberty area) - ids = await query_ids(moss_client, loaded_index, {"field": "location", "condition": {"$near": "40.7580,-73.9855,10000"}}) + ids = await query_ids( + moss_client, + loaded_index, + {"field": "location", "condition": {"$near": "40.7580,-73.9855,10000"}}, + ) assert "f1" in ids assert "f3" in ids assert "f2" not in ids # Tokyo @@ -155,12 +243,20 @@ async def test_near_within_range(self, moss_client, loaded_index): @pytest.mark.asyncio async def test_near_excludes_far(self, moss_client, loaded_index): # 5km around Times Square ā€” only f1 (~0km), not f3 (~8.7km) - ids = await query_ids(moss_client, loaded_index, {"field": "location", "condition": {"$near": "40.7580,-73.9855,5000"}}) + ids = await query_ids( + moss_client, + loaded_index, + {"field": "location", "condition": {"$near": "40.7580,-73.9855,5000"}}, + ) assert "f1" in ids assert "f3" not in ids @pytest.mark.asyncio async def test_near_skips_docs_without_location(self, moss_client, loaded_index): # f6 has no metadata at all ā€” should not appear - ids = await query_ids(moss_client, loaded_index, {"field": "location", "condition": {"$near": "40.7580,-73.9855,10000000"}}) + ids = await query_ids( + moss_client, + loaded_index, + {"field": "location", "condition": {"$near": "40.7580,-73.9855,10000000"}}, + ) assert "f6" not in ids diff --git a/sdks/python/sdk/tests/test_rerank.py b/sdks/python/sdk/tests/test_rerank.py new file mode 100644 index 00000000..0a42833d --- /dev/null +++ b/sdks/python/sdk/tests/test_rerank.py @@ -0,0 +1,85 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from moss import QueryResultDocumentInfo, SearchResult +from moss.client.moss_client import MossClient, QueryOptions + + +def test_query_options_validation(): + # Test valid options + opts = QueryOptions( + top_k=5, alpha=0.5, embedding=[0.1, 0.2], rerank=True, rerank_top_k=10 + ) + assert opts.top_k == 5 + assert opts.alpha == 0.5 + assert opts.embedding == [0.1, 0.2] + assert opts.rerank is True + assert opts.rerank_top_k == 10 + + # Test invalid top_k + with pytest.raises(ValueError, match="top_k must be an integer >= 1"): + QueryOptions(top_k=-1) + + # Test invalid alpha + with pytest.raises(ValueError, match="alpha must be a float between 0.0 and 1.0"): + QueryOptions(alpha=2.0) + + # Test invalid embedding + with pytest.raises(ValueError, match="embedding must be a sequence of numbers"): + QueryOptions(embedding=["invalid"]) + + # Test invalid rerank_top_k + with pytest.raises(ValueError, match="rerank_top_k must be an integer >= 1"): + QueryOptions(rerank_top_k=0) + + +@pytest.mark.asyncio +async def test_rerank_results(): + client = MossClient("test", "key") + + # Create fake docs + doc1 = QueryResultDocumentInfo(id="1", text="Bad match", score=0.9) + doc2 = QueryResultDocumentInfo(id="2", text="Good match", score=0.8) + doc3 = QueryResultDocumentInfo(id="3", text="Okay match", score=0.85) + + search_result = SearchResult( + docs=[doc1, doc2, doc3], query="test query", index_name="idx", time_taken_ms=10 + ) + + opts = QueryOptions(top_k=2, rerank=True, rerank_model="mock-model") + + # We need to mock sentence_transformers.CrossEncoder + mock_model = MagicMock() + # Let's say the cross encoder scores doc2 the highest, then doc3, then doc1 + mock_model.predict.return_value = [0.1, 0.99, 0.5] + + with patch.dict( + "sys.modules", + { + "sentence_transformers": MagicMock( + CrossEncoder=MagicMock(return_value=mock_model) + ) + }, + ): + result = await client._rerank_results("test query", search_result, opts) + + # It should slice to top_k=2, so doc1 should be dropped + assert len(result.docs) == 2 + + # First should be doc2 + assert result.docs[0].id == "2" + assert result.docs[0].score == pytest.approx(0.99, abs=1e-5) + + # Second should be doc3 + assert result.docs[1].id == "3" + assert result.docs[1].score == pytest.approx(0.5, abs=1e-5) + + # Ensure predict was called with pairs + mock_model.predict.assert_called_once_with( + [ + ["test query", "Bad match"], + ["test query", "Good match"], + ["test query", "Okay match"], + ] + ) diff --git a/sdks/python/sdk/tests/test_search.py b/sdks/python/sdk/tests/test_search.py index 10af17a5..879fccef 100644 --- a/sdks/python/sdk/tests/test_search.py +++ b/sdks/python/sdk/tests/test_search.py @@ -4,19 +4,22 @@ import statistics import time from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as pkg_version from pathlib import Path from typing import Any, Dict, List, Optional, Set -from importlib.metadata import PackageNotFoundError, version as pkg_version -from moss import MossClient, DocumentInfo, SearchResult, QueryOptions import pytest +from moss import DocumentInfo, MossClient, QueryOptions, SearchResult + from .constants import TEST_PROJECT_ID, TEST_PROJECT_KEY # ------------------- Configuration ------------------- EMBEDDING_MODEL = "moss-minilm" + # Interface for Experiment Configuration @dataclass class ExperimentConfig: @@ -25,6 +28,7 @@ class ExperimentConfig: alpha: Optional[float] top_k: int = 10 + # Interface for Run Statistics @dataclass class RunStats: @@ -39,38 +43,41 @@ class RunStats: p95_latency_ms: Optional[float] error_reason: Optional[str] = None + # Dataclass for Caching the dataset to avoid re-loading @dataclass class DatasetCache: """Cache for dataset data to avoid re-loading across experiments.""" + dataset_path: Optional[Path] = None qrels: Optional[Dict[str, Dict[str, float]]] = None queries: Optional[Dict[str, str]] = None + # -------------------- Data Loaders -------------------- -#Loads the corpus from the corpus.jsonl file and combines the title and text into a single string + +# Loads the corpus from the corpus.jsonl file and combines the title and text into a single string def load_corpus(path: Path) -> List[DocumentInfo]: """Loads the entire mini-corpus into memory at once.""" if not path.exists(): pytest.skip(f"File not found: {path}") - + print(f"Loading docs from {path.name}...") docs = [] with path.open("r", encoding="utf-8") as f: for line in f: data = json.loads(line) # Combine title + text for best search results - text = (str(data.get("title", "")) + " " + str(data.get("text", ""))).strip() - docs.append(DocumentInfo( - id=str(data["_id"]), - text=text, - metadata={} - )) + text = ( + str(data.get("title", "")) + " " + str(data.get("text", "")) + ).strip() + docs.append(DocumentInfo(id=str(data["_id"]), text=text, metadata={})) print(f" Loaded {len(docs)} documents.") return docs -#-------- Loads the qrels from the qrels/test.tsv file -------- + +# -------- Loads the qrels from the qrels/test.tsv file -------- # Example struct of qrels : {"query-id": {"doc-id": score}} def load_qrels(path: Path) -> Dict[str, Dict[str, float]]: """Loads relevance (Query ID -> Doc ID -> Score).""" @@ -79,7 +86,7 @@ def load_qrels(path: Path) -> Dict[str, Dict[str, float]]: # Cache file path cache_path = path.parent / f"{path.stem}_cache.json" - + # Check if cache exists and is newer than source file if cache_path.exists(): cache_mtime = cache_path.stat().st_mtime @@ -90,8 +97,10 @@ def load_qrels(path: Path) -> Dict[str, Dict[str, float]]: print(f"Loading qrels from cache: {cache_path.name}...") with cache_path.open("r", encoding="utf-8") as f: qrels = json.load(f) - qrels = {qid: {did: float(score) for did, score in docs.items()} - for qid, docs in qrels.items()} + qrels = { + qid: {did: float(score) for did, score in docs.items()} + for qid, docs in qrels.items() + } print(f" Loaded {len(qrels)} query judgments from cache.") return qrels @@ -110,16 +119,17 @@ def load_qrels(path: Path) -> Dict[str, Dict[str, float]]: if len(parts) >= 3: qid, did, score = parts[0], parts[1], float(parts[2]) qrels.setdefault(qid, {})[did] = score - + # Save to cache print(f" Caching qrels to {cache_path.name}...") with cache_path.open("w", encoding="utf-8") as f: json.dump(qrels, f, indent=2) print(f" Loaded {len(qrels)} query judgments.") - + return qrels -#-------- Loads the queries from the queries.jsonl -------- + +# -------- Loads the queries from the queries.jsonl -------- # Example struct of queries : {"_id": "text"} def load_queries(path: Path, valid_qids: Set[str]) -> Dict[str, str]: """Loads queries that exist in the Qrels.""" @@ -132,39 +142,52 @@ def load_queries(path: Path, valid_qids: Set[str]) -> Dict[str, str]: queries[qid] = str(data["text"]) return queries + # -------------------- Metrics Utils -------------------- -#-------- Calculates the NDCG metric -------- + +# -------- Calculates the NDCG metric -------- # NDCG (Normalized Discounted Cumulative Gain) measures ranking quality by considering # both relevance and position. Higher positions and higher relevance scores contribute more. -# +# # Formula: NDCG@k = DCG@k / IDCG@k # - DCG@k = Ī£(i=0 to k-1) (2^rel_i - 1) / log2(i + 2) # * Sums relevance gains (2^rel - 1) discounted by position (log2(i+2)) # - IDCG@k = DCG of the ideal ranking (documents sorted by relevance descending) # - Result: Score between 0.0 (worst) and 1.0 (perfect ranking) -def calculate_ndcg(retrieved_ids: List[str], true_rels: Dict[str, float], k: int) -> float: +def calculate_ndcg( + retrieved_ids: List[str], true_rels: Dict[str, float], k: int +) -> float: dcg = 0.0 idcg = 0.0 - + # 1. DCG for i, doc_id in enumerate(retrieved_ids[:k]): rel = true_rels.get(doc_id, 0) if rel > 0: dcg += (2**rel - 1) / math.log2(i + 2) - + # 2. IDCG (Ideal ordering) ideal_scores = sorted(true_rels.values(), reverse=True)[:k] for i, rel in enumerate(ideal_scores): idcg += (2**rel - 1) / math.log2(i + 2) - + return (dcg / idcg) if idcg > 0 else 0.0 + # -------------------- Core Execution -------------------- # Run a single experiment scenario, based on the provided config -async def run_scenario(client: MossClient, config: ExperimentConfig, queries: Dict, qrels: Dict, index_name: str) -> RunStats: - print(f"\nRunning Experiment: {config.name} (alpha={config.alpha}, dataset={config.dataset_path.name}) \n") - +async def run_scenario( + client: MossClient, + config: ExperimentConfig, + queries: Dict, + qrels: Dict, + index_name: str, +) -> RunStats: + print( + f"\nRunning Experiment: {config.name} (alpha={config.alpha}, dataset={config.dataset_path.name}) \n" + ) + latencies = [] hits = 0 mrr_sum = 0.0 @@ -174,19 +197,19 @@ async def run_scenario(client: MossClient, config: ExperimentConfig, queries: Di current_sdk_version = pkg_version("moss") except PackageNotFoundError: current_sdk_version = "unknown" - + for qid, text in queries.items(): # 1. Search options = QueryOptions(top_k=config.top_k, alpha=config.alpha) res: SearchResult = await client.query(index_name, text, options) - + # 2. Latency Measurement latencies.append(res.time_taken_ms) - + # 3. Score retrieved_ids = [d.id for d in res.docs] relevant_docs = qrels.get(qid, {}) - + # Check Hit & MRR is_hit = False for rank, doc_id in enumerate(retrieved_ids, 1): @@ -195,7 +218,7 @@ async def run_scenario(client: MossClient, config: ExperimentConfig, queries: Di hits += 1 mrr_sum += 1.0 / rank is_hit = True - + # Check NDCG ndcg_sum += calculate_ndcg(retrieved_ids, relevant_docs, config.top_k) @@ -209,63 +232,131 @@ async def run_scenario(client: MossClient, config: ExperimentConfig, queries: Di mrr=mrr_sum / count, ndcg=ndcg_sum / count, avg_latency_ms=statistics.mean(latencies), - p95_latency_ms=(sorted(latencies)[int(len(latencies) * 0.95)] if len(latencies) >= 2 else float('nan')) + p95_latency_ms=( + sorted(latencies)[int(len(latencies) * 0.95)] + if len(latencies) >= 2 + else float("nan") + ), ) + async def main(): print("\n==========================================") print(" MOSS BENCHMARK: Multi-Dataset Testing") print("==========================================\n") - - client = MossClient( - TEST_PROJECT_ID, - TEST_PROJECT_KEY - ) + + client = MossClient(TEST_PROJECT_ID, TEST_PROJECT_KEY) # Define the base dataset path dynamically BASE_DATASET_PATH = Path(__file__).resolve().parents[3] / "test-dataset" # 1. Define Experiments experiments = [ - ExperimentConfig("Keyword Search(alpha=0)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=0), - ExperimentConfig("Keyword Search(alpha=0)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=0), - ExperimentConfig("Keyword Search(alpha=0)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=0), - - ExperimentConfig("Fusion Search(alpha=0.2)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=0.2), - ExperimentConfig("Fusion Search(alpha=0.2)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=0.2), - ExperimentConfig("Fusion Search(alpha=0.2)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=0.2), - - ExperimentConfig("Fusion Search(alpha=0.4)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=0.4), - ExperimentConfig("Fusion Search(alpha=0.4)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=0.4), - ExperimentConfig("Fusion Search(alpha=0.4)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=0.4), - - ExperimentConfig("Fusion Search(alpha=0.6)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=0.6), - ExperimentConfig("Fusion Search(alpha=0.6)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=0.6), - ExperimentConfig("Fusion Search(alpha=0.6)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=0.6), - - ExperimentConfig("Fusion Search(alpha=0.8)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=0.8), - ExperimentConfig("Fusion Search(alpha=0.8)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=0.8), - ExperimentConfig("Fusion Search(alpha=0.8)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=0.8), - - ExperimentConfig("Embedding Search(alpha=1)", dataset_path=BASE_DATASET_PATH / "full_scifact", alpha=1), - ExperimentConfig("Embedding Search(alpha=1)", dataset_path=BASE_DATASET_PATH / "full_nfcorpus", alpha=1), - ExperimentConfig("Embedding Search(alpha=1)", dataset_path=BASE_DATASET_PATH / "mini_msmarco", alpha=1), - + ExperimentConfig( + "Keyword Search(alpha=0)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=0, + ), + ExperimentConfig( + "Keyword Search(alpha=0)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=0, + ), + ExperimentConfig( + "Keyword Search(alpha=0)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=0, + ), + ExperimentConfig( + "Fusion Search(alpha=0.2)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=0.2, + ), + ExperimentConfig( + "Fusion Search(alpha=0.2)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=0.2, + ), + ExperimentConfig( + "Fusion Search(alpha=0.2)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=0.2, + ), + ExperimentConfig( + "Fusion Search(alpha=0.4)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=0.4, + ), + ExperimentConfig( + "Fusion Search(alpha=0.4)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=0.4, + ), + ExperimentConfig( + "Fusion Search(alpha=0.4)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=0.4, + ), + ExperimentConfig( + "Fusion Search(alpha=0.6)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=0.6, + ), + ExperimentConfig( + "Fusion Search(alpha=0.6)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=0.6, + ), + ExperimentConfig( + "Fusion Search(alpha=0.6)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=0.6, + ), + ExperimentConfig( + "Fusion Search(alpha=0.8)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=0.8, + ), + ExperimentConfig( + "Fusion Search(alpha=0.8)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=0.8, + ), + ExperimentConfig( + "Fusion Search(alpha=0.8)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=0.8, + ), + ExperimentConfig( + "Embedding Search(alpha=1)", + dataset_path=BASE_DATASET_PATH / "full_scifact", + alpha=1, + ), + ExperimentConfig( + "Embedding Search(alpha=1)", + dataset_path=BASE_DATASET_PATH / "full_nfcorpus", + alpha=1, + ), + ExperimentConfig( + "Embedding Search(alpha=1)", + dataset_path=BASE_DATASET_PATH / "mini_msmarco", + alpha=1, + ), ] # 2. Sort experiments by dataset to minimize I/O experiments = sorted(experiments, key=lambda e: e.dataset_path) - + # --------------------------------------------------------- # PHASE 1: INDEX PREPARATION (Delete -> Create) # --------------------------------------------------------- # Identify all unique datasets required by the experiments unique_dataset_paths = sorted(list({exp.dataset_path for exp in experiments})) - + # Map dataset paths to their created index names dataset_to_index_map: Dict[Path, str] = {} failed_datasets: Dict[Path, Dict[str, Any]] = {} - + # Get current version for safe index naming try: current_sdk_version = pkg_version("moss") @@ -273,8 +364,10 @@ async def main(): current_sdk_version = "unknown" safe_version = current_sdk_version.replace(".", "_").replace("-", "_") - print(f"\nšŸš€ PRE-COMPUTING INDICES for {len(unique_dataset_paths)} unique datasets...") - + print( + f"\nšŸš€ PRE-COMPUTING INDICES for {len(unique_dataset_paths)} unique datasets..." + ) + for dpath in unique_dataset_paths: corpus_path = dpath / "corpus.jsonl" if not corpus_path.exists(): @@ -290,7 +383,7 @@ async def main(): # 3. Force Delete Existing Index print(f" šŸ—‘ļø Cleaning old index: {index_name}...") try: - # Assuming delete_index is the API method. + # Assuming delete_index is the API method. # Wrap in try/except in case it throws an error if index doesn't exist. await client.delete_index(index_name) except Exception: @@ -381,7 +474,9 @@ async def main(): continue if exp.dataset_path not in dataset_to_index_map: - print(f"Skipping {exp.name}: Index was not created during the preparation phase for {exp.dataset_path.name}") + print( + f"Skipping {exp.name}: Index was not created during the preparation phase for {exp.dataset_path.name}" + ) continue index_name = dataset_to_index_map[exp.dataset_path] @@ -391,16 +486,18 @@ async def main(): if exp.dataset_path != dataset_cache.dataset_path: qrels_path = exp.dataset_path / "qrels.tsv" queries_path = exp.dataset_path / "queries.jsonl" - + if not qrels_path.exists() or not queries_path.exists(): print(f"Skipping {exp.name}: Missing qrels/queries") continue print(f"\nšŸ“‚ Loading test data for: {exp.dataset_path.name}") dataset_cache.qrels = load_qrels(qrels_path) - dataset_cache.queries = load_queries(queries_path, set(dataset_cache.qrels.keys())) + dataset_cache.queries = load_queries( + queries_path, set(dataset_cache.qrels.keys()) + ) dataset_cache.dataset_path = exp.dataset_path - + # CRITICAL STEP: Load the index specifically for this dataset print(f"šŸ”Œ Connecting to index: {index_name}") try: @@ -432,18 +529,14 @@ async def main(): # ----- Run Experiment ----- stats = await run_scenario( - client, - exp, - dataset_cache.queries, - dataset_cache.qrels, - index_name + client, exp, dataset_cache.queries, dataset_cache.qrels, index_name ) results.append(stats) # --------------------------------------------------------- # PHASE 3: REPORTING # --------------------------------------------------------- - + # Print Table header = f"{'Experiment':<25} | {'SDK Version':<12} | {'Dataset':<30} | {'HitRate':>8} | {'NDCG':>8} | {'MRR':>8} | {'Avg(ms)':>8} | {'P95(ms)':>8} | {'Reason':<40}" separator = "=" * len(header) @@ -452,7 +545,7 @@ async def main(): print(separator) print(header) print("-" * len(header)) - + for stats in results: if stats.error_reason: print( @@ -463,20 +556,20 @@ async def main(): f"{stats.experiment:<25} | {stats.sdk_version:<12} | {stats.dataset_name:<30} | {stats.hit_rate:>8.3f} | {stats.ndcg:>8.3f} | {stats.mrr:>8.3f} | {stats.avg_latency_ms:>8.1f} | {stats.p95_latency_ms:>8.1f} | {'':<40}" ) print(separator) - + # --------------------------------------------------------- # PHASE 4: CLEANUP (Delete all created indices) # --------------------------------------------------------- print("\n" + "=" * 60) print(" CLEANING UP INDICES") print("=" * 60) - + all_indices = set(dataset_to_index_map.values()) # Also include indices from failed datasets (they might have been created before failing) for failure_info in failed_datasets.values(): if "index_name" in failure_info: all_indices.add(failure_info["index_name"]) - + if all_indices: print(f"\nšŸ—‘ļø Deleting {len(all_indices)} indices...") for index_name in sorted(all_indices): @@ -487,7 +580,7 @@ async def main(): print(f" āš ļø Failed to delete {index_name}: {e}") else: print("\n No indices to delete.") - + print("\nDone.") @@ -501,4 +594,4 @@ async def test_search_benchmark(): try: asyncio.run(main()) except KeyboardInterrupt: - pass \ No newline at end of file + pass diff --git a/sdks/python/sdk/tests/test_types.py b/sdks/python/sdk/tests/test_types.py index f5694c28..96cb4d34 100644 --- a/sdks/python/sdk/tests/test_types.py +++ b/sdks/python/sdk/tests/test_types.py @@ -58,7 +58,7 @@ async def test_get_docs_with_doc_ids(self): patch("moss.client.moss_client.ManageClient") as mock_manage, patch("moss.client.moss_client.IndexManager") as _, ): - from moss import MossClient, GetDocumentsOptions + from moss import GetDocumentsOptions, MossClient mock_manage_instance = mock_manage.return_value mock_manage_instance.get_docs = MagicMock(return_value=[]) @@ -126,7 +126,7 @@ async def test_create_index_with_document_info(self): patch("moss.client.moss_client.ManageClient") as mock_manage, patch("moss.client.moss_client.IndexManager") as _, ): - from moss import MossClient, DocumentInfo + from moss import DocumentInfo, MossClient mock_manage_instance = mock_manage.return_value mock_result = MagicMock() @@ -154,7 +154,7 @@ async def test_create_index_with_metadata(self): patch("moss.client.moss_client.ManageClient") as mock_manage, patch("moss.client.moss_client.IndexManager") as _, ): - from moss import MossClient, DocumentInfo + from moss import DocumentInfo, MossClient mock_manage_instance = mock_manage.return_value mock_result = MagicMock() @@ -188,7 +188,7 @@ async def test_create_index_with_embeddings(self): patch("moss.client.moss_client.ManageClient") as mock_manage, patch("moss.client.moss_client.IndexManager") as _, ): - from moss import MossClient, DocumentInfo + from moss import DocumentInfo, MossClient mock_manage_instance = mock_manage.return_value mock_result = MagicMock() @@ -211,7 +211,7 @@ class TestSearchResultUsage: """Tests for SearchResult handling.""" def test_search_result_structure(self): - from moss import SearchResult, QueryResultDocumentInfo + from moss import QueryResultDocumentInfo, SearchResult doc1 = QueryResultDocumentInfo(id="d1", text="First", score=0.95) doc2 = QueryResultDocumentInfo(id="d2", text="Second", score=0.85) From 6f56b8a175ecd367304fd6ca0cfacc4325ddd541 Mon Sep 17 00:00:00 2001 From: Hendrixx-RE Date: Tue, 21 Jul 2026 02:14:12 +0530 Subject: [PATCH 3/3] fix type-error --- sdks/python/sdk/tests/test_rerank.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/sdk/tests/test_rerank.py b/sdks/python/sdk/tests/test_rerank.py index 0a42833d..fdaf3560 100644 --- a/sdks/python/sdk/tests/test_rerank.py +++ b/sdks/python/sdk/tests/test_rerank.py @@ -27,7 +27,7 @@ def test_query_options_validation(): # Test invalid embedding with pytest.raises(ValueError, match="embedding must be a sequence of numbers"): - QueryOptions(embedding=["invalid"]) + QueryOptions(embedding=["invalid"]) # type: ignore # Test invalid rerank_top_k with pytest.raises(ValueError, match="rerank_top_k must be an integer >= 1"):