From fec6fc12848544d7311092f192e0a76d1b1376fe Mon Sep 17 00:00:00 2001
From: Teagan Glenn <that@teagantotally.rocks>
Date: Tue, 7 Oct 2025 05:35:52 -0600
Subject: [PATCH] feat(server): support similarity seeds in query media

---
 docker/pyproject.deps.toml             |  2 +-
 mcp_plex/server/tools/media_library.py | 69 +++++++++++++++++++-------
 pyproject.toml                         |  2 +-
 tests/test_server.py                   | 23 +++++++++
 uv.lock                                |  2 +-
 5 files changed, 78 insertions(+), 20 deletions(-)

diff --git a/docker/pyproject.deps.toml b/docker/pyproject.deps.toml
index 7b069ca..5f47263 100644
--- a/docker/pyproject.deps.toml
+++ b/docker/pyproject.deps.toml
@@ -1,6 +1,6 @@
 [project]
 name = "mcp-plex"
-version = "2.0.9"
+version = "2.0.10"
 requires-python = ">=3.11,<3.13"
 dependencies = [
   "fastmcp>=2.11.2",
diff --git a/mcp_plex/server/tools/media_library.py b/mcp_plex/server/tools/media_library.py
index 5354cfc..59214ba 100644
--- a/mcp_plex/server/tools/media_library.py
+++ b/mcp_plex/server/tools/media_library.py
@@ -317,6 +317,13 @@ async def query_media(
             int | None,
             Field(description="Match a TMDb identifier", examples=[568467]),
         ] = None,
+        similar_to: Annotated[
+            str | Sequence[str] | None,
+            Field(
+                description="Recommend candidates similar to these identifiers",
+                examples=[["49915"], "tt8367814"],
+            ),
+        ] = None,
         limit: Annotated[
             int,
             Field(
@@ -337,24 +344,37 @@ def _listify(value: Sequence[str] | str | None) -> list[str]:
             return [v for v in value if isinstance(v, str) and v]
 
         vector_queries: list[tuple[str, models.Document]] = []
-        if dense_query:
-            vector_queries.append(
-                (
-                    "dense",
-                    models.Document(
-                        text=dense_query, model=server.settings.dense_model
-                    ),
+        positive_point_ids: list[Any] = []
+        similar_identifiers = _listify(similar_to)
+        if similar_identifiers:
+            for identifier in similar_identifiers:
+                records = await media_helpers._find_records(
+                    server, identifier, limit=1
                 )
-            )
-        if sparse_query:
-            vector_queries.append(
-                (
-                    "sparse",
-                    models.Document(
-                        text=sparse_query, model=server.settings.sparse_model
-                    ),
+                for record in records:
+                    if record.id is not None:
+                        positive_point_ids.append(record.id)
+            if not positive_point_ids:
+                return []
+        if not positive_point_ids:
+            if dense_query:
+                vector_queries.append(
+                    (
+                        "dense",
+                        models.Document(
+                            text=dense_query, model=server.settings.dense_model
+                        ),
+                    )
+                )
+            if sparse_query:
+                vector_queries.append(
+                    (
+                        "sparse",
+                        models.Document(
+                            text=sparse_query, model=server.settings.sparse_model
+                        ),
+                    )
                 )
-            )
 
         must: list[models.FieldCondition] = []
         keyword_prefetch_conditions: list[models.FieldCondition] = []
@@ -503,7 +523,20 @@ def _listify(value: Sequence[str] | str | None) -> list[str]:
         query_obj: models.Query | None = None
         using_param: str | None = None
         prefetch_param: Sequence[models.Prefetch] | None = None
-        if vector_queries:
+        prefetch_entries: list[models.Prefetch] = []
+        if positive_point_ids:
+            recommend_query = models.RecommendQuery(
+                recommend=models.RecommendInput(positive=positive_point_ids)
+            )
+            prefetch_entries.append(
+                models.Prefetch(
+                    query=recommend_query,
+                    using="dense",
+                    limit=limit,
+                    filter=prefetch_filter,
+                )
+            )
+        if not positive_point_ids and vector_queries:
             candidate_limit = limit * 3 if len(vector_queries) > 1 else limit
             prefetch_entries = [
                 models.Prefetch(
@@ -514,6 +547,8 @@ def _listify(value: Sequence[str] | str | None) -> list[str]:
                 )
                 for name, doc in vector_queries
             ]
+
+        if prefetch_entries:
             if len(prefetch_entries) > 1:
                 query_obj = models.FusionQuery(fusion=models.Fusion.RRF)
                 using_param = None
diff --git a/pyproject.toml b/pyproject.toml
index 491006f..ff15aa5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mcp-plex"
-version = "2.0.9"
+version = "2.0.10"
 
 description = "Plex-Oriented Model Context Protocol Server"
 requires-python = ">=3.11,<3.13"
diff --git a/tests/test_server.py b/tests/test_server.py
index c269fd8..a3ce82f 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -150,6 +150,29 @@ def test_server_tools(monkeypatch):
         )
         assert episode_structured[0]["show_title"] == "Alien: Earth"
 
+        similar_structured = asyncio.run(
+            server.query_media.fn(
+                similar_to=["49915"],
+                type="episode",
+                limit=3,
+            )
+        )
+        assert similar_structured
+        assert {
+            item["plex"]["rating_key"]
+            for item in similar_structured
+            if isinstance(item.get("plex"), dict)
+        } >= {"61960"}
+
+        assert (
+            asyncio.run(
+                server.query_media.fn(
+                    similar_to="does-not-exist", type="movie", limit=1
+                )
+            )
+            == []
+        )
+
         rec = asyncio.run(server.recommend_media.fn(identifier=movie_id, limit=1))
         assert rec and rec[0]["plex"]["rating_key"] == "61960"
 
diff --git a/uv.lock b/uv.lock
index ee42ebe..37b7343 100644
--- a/uv.lock
+++ b/uv.lock
@@ -730,7 +730,7 @@ wheels = [
 
 [[package]]
 name = "mcp-plex"
-version = "2.0.9"
+version = "2.0.10"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi" },