Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
24 changes: 23 additions & 1 deletion mcp_plex/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -62,6 +73,17 @@ 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(
f"Unknown dense embedding model '{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."""

Expand Down Expand Up @@ -636,6 +658,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(
Expand All @@ -647,7 +670,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):
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
18 changes: 18 additions & 0 deletions tests/test_loader_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from qdrant_client.async_qdrant_client import AsyncQdrantClient
from qdrant_client import models
import pytest

from mcp_plex import loader

Expand Down Expand Up @@ -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",
)
)
13 changes: 13 additions & 0 deletions tests/test_loader_unit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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")
2 changes: 1 addition & 1 deletion uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.