From e8360729761a4c5687d248b7ca2b39cca82d2275 Mon Sep 17 00:00:00 2001 From: Teagan Glenn Date: Tue, 23 Sep 2025 08:37:42 -0600 Subject: [PATCH 1/2] fix: derive qdrant model params without private api --- AGENTS.md | 3 +++ mcp_plex/loader.py | 25 ++++++++++++++++++++++++- pyproject.toml | 2 +- tests/test_loader_integration.py | 18 ++++++++++++++++++ tests/test_loader_unit.py | 13 +++++++++++++ uv.lock | 2 +- 6 files changed, 60 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 1b1d034..51dec09 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -27,6 +27,9 @@ split into batches of five. - Qdrant upserts are batched and network errors are logged so large loads can proceed even when individual batches fail. +- Qdrant model metadata is tracked locally in a supported mapping so the loader + no longer depends on the private `_get_model_params` helper removed in newer + clients. ## User Queries The project should handle natural-language searches and recommendations such as: diff --git a/mcp_plex/loader.py b/mcp_plex/loader.py index d69d4f6..ce374a4 100644 --- a/mcp_plex/loader.py +++ b/mcp_plex/loader.py @@ -47,6 +47,17 @@ _imdb_batch_limit: int = 5 _qdrant_batch_size: int = 1000 +# Known Qdrant-managed dense embedding models with their dimensionality and +# similarity metric. To support a new server-side embedding model, add an entry +# here with the appropriate vector size and `models.Distance` value. +_DENSE_MODEL_PARAMS: dict[str, tuple[int, models.Distance]] = { + "BAAI/bge-small-en-v1.5": (384, models.Distance.COSINE), + "BAAI/bge-base-en-v1.5": (768, models.Distance.COSINE), + "BAAI/bge-large-en-v1.5": (1024, models.Distance.COSINE), + "text-embedding-3-small": (1536, models.Distance.COSINE), + "text-embedding-3-large": (3072, models.Distance.COSINE), +} + async def _gather_in_batches( tasks: Sequence[Awaitable[T]], batch_size: int @@ -62,6 +73,18 @@ async def _gather_in_batches( return results +def _resolve_dense_model_params(model_name: str) -> tuple[int, models.Distance]: + """Look up Qdrant vector parameters for a known dense embedding model.""" + + try: + return _DENSE_MODEL_PARAMS[model_name] + except KeyError as exc: + raise ValueError( + "Unknown dense embedding model" + f" '{model_name}'. Update _DENSE_MODEL_PARAMS with the model's size and distance." + ) from exc + + async def _fetch_imdb(client: httpx.AsyncClient, imdb_id: str) -> Optional[IMDbTitle]: """Fetch metadata for an IMDb ID with caching and retry logic.""" @@ -636,6 +659,7 @@ async def run( ) ) + dense_size, dense_distance = _resolve_dense_model_params(dense_model_name) if qdrant_url is None and qdrant_host is None: qdrant_url = ":memory:" client = AsyncQdrantClient( @@ -647,7 +671,6 @@ async def run( https=qdrant_https, prefer_grpc=qdrant_prefer_grpc, ) - dense_size, dense_distance = client._get_model_params(dense_model_name) collection_name = "media-items" created_collection = False if not await client.collection_exists(collection_name): diff --git a/pyproject.toml b/pyproject.toml index 36512b4..9dde796 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "mcp-plex" -version = "0.26.27" +version = "0.26.28" description = "Plex-Oriented Model Context Protocol Server" requires-python = ">=3.11,<3.13" diff --git a/tests/test_loader_integration.py b/tests/test_loader_integration.py index 3359461..5a10d8d 100644 --- a/tests/test_loader_integration.py +++ b/tests/test_loader_integration.py @@ -6,6 +6,7 @@ from qdrant_client.async_qdrant_client import AsyncQdrantClient from qdrant_client import models +import pytest from mcp_plex import loader @@ -97,3 +98,20 @@ def test_run_upserts_in_batches(monkeypatch): asyncio.run(_run_loader(sample_dir)) assert CaptureClient.upsert_calls == 2 assert len(CaptureClient.captured_points) == 2 + + +def test_run_raises_for_unknown_dense_model(): + sample_dir = Path(__file__).resolve().parents[1] / "sample-data" + + with pytest.raises(ValueError, match="Unknown dense embedding model"): + asyncio.run( + loader.run( + None, + None, + None, + sample_dir, + None, + None, + dense_model_name="not-a-real/model", + ) + ) diff --git a/tests/test_loader_unit.py b/tests/test_loader_unit.py index cbaf974..9ea0589 100644 --- a/tests/test_loader_unit.py +++ b/tests/test_loader_unit.py @@ -7,6 +7,7 @@ import httpx from qdrant_client import models +import pytest from mcp_plex import loader from mcp_plex.imdb_cache import IMDbCache @@ -22,6 +23,7 @@ _load_imdb_retry_queue, _persist_imdb_retry_queue, _process_imdb_retry_queue, + _resolve_dense_model_params, resolve_tmdb_season_number, ) from mcp_plex.types import TMDBSeason, TMDBShow @@ -482,3 +484,14 @@ async def upsert(self, collection_name: str, points, **kwargs): monkeypatch.setattr(loader, "_qdrant_batch_size", 1) asyncio.run(loader._upsert_in_batches(client, "c", points)) assert client.calls == 3 + + +def test_resolve_dense_model_params_known_model(): + size, distance = _resolve_dense_model_params("BAAI/bge-small-en-v1.5") + assert size == 384 + assert distance is models.Distance.COSINE + + +def test_resolve_dense_model_params_unknown_model(): + with pytest.raises(ValueError, match="Unknown dense embedding model"): + _resolve_dense_model_params("not-a-real/model") diff --git a/uv.lock b/uv.lock index 0b1f334..015b753 100644 --- a/uv.lock +++ b/uv.lock @@ -690,7 +690,7 @@ wheels = [ [[package]] name = "mcp-plex" -version = "0.26.27" +version = "0.26.28" source = { editable = "." } dependencies = [ { name = "fastapi" }, From 725c5fd3ec31db895eef595567a4e2e0f2077828 Mon Sep 17 00:00:00 2001 From: Teagan Glenn Date: Tue, 23 Sep 2025 08:38:37 -0600 Subject: [PATCH 2/2] Update mcp_plex/loader.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- mcp_plex/loader.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mcp_plex/loader.py b/mcp_plex/loader.py index ce374a4..e4221ee 100644 --- a/mcp_plex/loader.py +++ b/mcp_plex/loader.py @@ -80,8 +80,7 @@ def _resolve_dense_model_params(model_name: str) -> tuple[int, models.Distance]: return _DENSE_MODEL_PARAMS[model_name] except KeyError as exc: raise ValueError( - "Unknown dense embedding model" - f" '{model_name}'. Update _DENSE_MODEL_PARAMS with the model's size and distance." + f"Unknown dense embedding model '{model_name}'. Update _DENSE_MODEL_PARAMS with the model's size and distance." ) from exc