diff --git a/mcp_plex/loader.py b/mcp_plex/loader.py index 3d52d7b..6a7275b 100644 --- a/mcp_plex/loader.py +++ b/mcp_plex/loader.py @@ -10,25 +10,25 @@ import click import httpx from fastembed import SparseTextEmbedding, TextEmbedding -from qdrant_client.async_qdrant_client import AsyncQdrantClient from qdrant_client import models +from qdrant_client.async_qdrant_client import AsyncQdrantClient from .types import ( AggregatedItem, ExternalIDs, IMDbTitle, + PlexGuid, PlexItem, + PlexPerson, TMDBEpisode, TMDBItem, TMDBMovie, TMDBShow, - PlexGuid, - PlexPerson, ) try: # Only import plexapi when available; the sample data mode does not require it. - from plexapi.server import PlexServer from plexapi.base import PlexPartialObject + from plexapi.server import PlexServer except Exception: PlexServer = None # type: ignore[assignment] PlexPartialObject = object # type: ignore[assignment] diff --git a/mcp_plex/server.py b/mcp_plex/server.py index 526a421..f6df13f 100644 --- a/mcp_plex/server.py +++ b/mcp_plex/server.py @@ -2,16 +2,21 @@ from __future__ import annotations import asyncio -import os import json +import os from collections import OrderedDict -from typing import Any, Annotated +from typing import Annotated, Any +from fastembed import SparseTextEmbedding, TextEmbedding from fastmcp.server import FastMCP -from qdrant_client.async_qdrant_client import AsyncQdrantClient -from qdrant_client import models -from fastembed import TextEmbedding, SparseTextEmbedding from pydantic import Field +from qdrant_client import models +from qdrant_client.async_qdrant_client import AsyncQdrantClient + +try: + from sentence_transformers import CrossEncoder +except Exception: + CrossEncoder = None try: from sentence_transformers import CrossEncoder @@ -151,7 +156,7 @@ async def search_media( ) -> list[dict[str, Any]]: """Hybrid similarity search across media items using dense and sparse vectors.""" dense_task = asyncio.to_thread(lambda: list(_dense_model.embed([query]))[0]) - sparse_task = asyncio.to_thread(lambda: _sparse_model.query_embed(query)) + sparse_task = asyncio.to_thread(lambda: next(_sparse_model.query_embed(query))) dense_vec, sparse_vec = await asyncio.gather(dense_task, sparse_task) named_dense = models.NamedVector(name="dense", vector=dense_vec) sv = models.SparseVector( diff --git a/mcp_plex/types.py b/mcp_plex/types.py index 28040f5..532f05d 100644 --- a/mcp_plex/types.py +++ b/mcp_plex/types.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import List, Optional, Literal +from typing import List, Literal, Optional from pydantic import BaseModel, Field diff --git a/tests/test_load_from_plex.py b/tests/test_load_from_plex.py index 44c4c81..d7b757b 100644 --- a/tests/test_load_from_plex.py +++ b/tests/test_load_from_plex.py @@ -1,5 +1,6 @@ import asyncio import types + import httpx from mcp_plex import loader diff --git a/tests/test_loader_cli.py b/tests/test_loader_cli.py index 80d6249..da46024 100644 --- a/tests/test_loader_cli.py +++ b/tests/test_loader_cli.py @@ -1,6 +1,7 @@ import asyncio -from click.testing import CliRunner + import pytest +from click.testing import CliRunner from mcp_plex import loader diff --git a/tests/test_loader_integration.py b/tests/test_loader_integration.py index d75b34c..c2b97ee 100644 --- a/tests/test_loader_integration.py +++ b/tests/test_loader_integration.py @@ -2,9 +2,10 @@ from pathlib import Path from types import SimpleNamespace -from mcp_plex import loader from qdrant_client import models +from mcp_plex import loader + class DummyTextEmbedding: def __init__(self, name: str): diff --git a/tests/test_loader_unit.py b/tests/test_loader_unit.py index eeb3a03..c6fe106 100644 --- a/tests/test_loader_unit.py +++ b/tests/test_loader_unit.py @@ -1,17 +1,17 @@ -import types import asyncio -import httpx - +import types from pathlib import Path +import httpx + from mcp_plex.loader import ( - _extract_external_ids, - _load_from_sample, _build_plex_item, + _extract_external_ids, _fetch_imdb, + _fetch_tmdb_episode, _fetch_tmdb_movie, _fetch_tmdb_show, - _fetch_tmdb_episode, + _load_from_sample, ) diff --git a/tests/test_server.py b/tests/test_server.py index 70f53c4..f93e717 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1,15 +1,16 @@ -from typing import Any import asyncio -from pathlib import Path import importlib -import types import json -import time import sys +import time +import types +from pathlib import Path +from typing import Any + import pytest +from qdrant_client import models from mcp_plex import loader -from qdrant_client import models class DummyTextEmbedding: @@ -51,7 +52,7 @@ def passage_embed(self, texts): def query_embed(self, text): time.sleep(0.1) - return DummySparseVector([0], [1.0]) + yield DummySparseVector([0], [1.0]) class DummyReranker: