501 lines
19 KiB
Python
501 lines
19 KiB
Python
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import requests
|
||
|
|
import xml.etree.ElementTree as ET
|
||
|
|
from plexapi.server import PlexServer
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ToolSpec:
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
inputSchema: dict[str, Any] = field(
|
||
|
|
default_factory=lambda: {"type": "object", "properties": {}, "required": []}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_plex_datetime(dt_str: str | None) -> str | None:
|
||
|
|
if not dt_str:
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
return datetime.fromisoformat(dt_str).isoformat()
|
||
|
|
except (ValueError, TypeError):
|
||
|
|
return dt_str
|
||
|
|
|
||
|
|
|
||
|
|
class PlexService:
|
||
|
|
def __init__(
|
||
|
|
self, baseurl: str, token: str = "", username: str = "", password: str = ""
|
||
|
|
):
|
||
|
|
self.baseurl = baseurl
|
||
|
|
self._token = token
|
||
|
|
self._username = username
|
||
|
|
self._password = password
|
||
|
|
self._plex: PlexServer | None = None
|
||
|
|
self._connected = False
|
||
|
|
|
||
|
|
def _ensure_connected(self) -> PlexServer:
|
||
|
|
if self._plex is not None:
|
||
|
|
return self._plex
|
||
|
|
token = self._token
|
||
|
|
if not token and self._username and self._password:
|
||
|
|
token = self._authenticate()
|
||
|
|
self._plex = PlexServer(self.baseurl, token)
|
||
|
|
self._connected = True
|
||
|
|
return self._plex
|
||
|
|
|
||
|
|
def _authenticate(self) -> str:
|
||
|
|
url = "https://plex.tv/users/sign_in.xml"
|
||
|
|
headers = {
|
||
|
|
"X-Plex-Client-Identifier": "PLEXAGE LIMITSETTER V1",
|
||
|
|
"X-Plex-Product": "Plex Age Limit Setter",
|
||
|
|
"X-Plex-Version": "1.0",
|
||
|
|
}
|
||
|
|
resp = requests.post(
|
||
|
|
url, auth=(self._username, self._password), headers=headers
|
||
|
|
)
|
||
|
|
resp.raise_for_status()
|
||
|
|
root = ET.fromstring(resp.content)
|
||
|
|
token = root.get("authenticationToken")
|
||
|
|
if not token:
|
||
|
|
raise RuntimeError("Could not retrieve authentication token from Plex.tv")
|
||
|
|
return token
|
||
|
|
|
||
|
|
# ── library introspection ──────────────────────────────────
|
||
|
|
|
||
|
|
def list_libraries(self) -> list[dict[str, Any]]:
|
||
|
|
plex = self._ensure_connected()
|
||
|
|
return [
|
||
|
|
{"title": s.title, "type": s.type, "key": s.key, "uuid": s.uuid}
|
||
|
|
for s in plex.library.sections()
|
||
|
|
]
|
||
|
|
|
||
|
|
def get_library(self, title_or_key: str):
|
||
|
|
plex = self._ensure_connected()
|
||
|
|
sec = plex.library.section(title_or_key)
|
||
|
|
if sec is None:
|
||
|
|
raise ValueError(f"Library '{title_or_key}' not found")
|
||
|
|
return sec
|
||
|
|
|
||
|
|
def list_content_ratings(self, library: str) -> list[dict[str, Any]]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
ratings: dict[str, int] = {}
|
||
|
|
for item in sec.all():
|
||
|
|
r = item.contentRating or "UNRATED"
|
||
|
|
ratings[r] = ratings.get(r, 0) + 1
|
||
|
|
return sorted(
|
||
|
|
[{"rating": k, "count": v} for k, v in ratings.items()],
|
||
|
|
key=lambda x: x["count"],
|
||
|
|
reverse=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_library_stats(self, library: str) -> dict[str, Any]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
all_items = sec.all()
|
||
|
|
total = len(all_items)
|
||
|
|
rated = sum(1 for i in all_items if i.contentRating)
|
||
|
|
unrated = total - rated
|
||
|
|
ratings: dict[str, int] = {}
|
||
|
|
for i in all_items:
|
||
|
|
r = i.contentRating or "UNRATED"
|
||
|
|
ratings[r] = ratings.get(r, 0) + 1
|
||
|
|
return {
|
||
|
|
"library": library,
|
||
|
|
"total_items": total,
|
||
|
|
"rated": rated,
|
||
|
|
"unrated": unrated,
|
||
|
|
"rating_distribution": ratings,
|
||
|
|
}
|
||
|
|
|
||
|
|
# ── movie queries ──────────────────────────────────────────
|
||
|
|
|
||
|
|
def _movie_to_dict(self, movie) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"title": movie.title,
|
||
|
|
"year": getattr(movie, "year", None),
|
||
|
|
"rating_key": movie.ratingKey,
|
||
|
|
"content_rating": movie.contentRating or "",
|
||
|
|
"studio": getattr(movie, "studio", None),
|
||
|
|
"summary": getattr(movie, "summary", "")[:200],
|
||
|
|
"duration_minutes": (
|
||
|
|
getattr(movie, "duration", 0) // 60_000
|
||
|
|
if getattr(movie, "duration", None)
|
||
|
|
else None
|
||
|
|
),
|
||
|
|
"added_at": _parse_plex_datetime(str(movie.addedAt))
|
||
|
|
if hasattr(movie, "addedAt")
|
||
|
|
else None,
|
||
|
|
"updated_at": _parse_plex_datetime(str(movie.updatedAt))
|
||
|
|
if hasattr(movie, "updatedAt")
|
||
|
|
else None,
|
||
|
|
"genres": [g.tag for g in getattr(movie, "genres", [])],
|
||
|
|
"directors": [d.tag for d in getattr(movie, "directors", [])],
|
||
|
|
"roles": [r.tag for r in getattr(movie, "roles", [])][:5],
|
||
|
|
}
|
||
|
|
|
||
|
|
def get_movie(self, library: str, title: str) -> dict[str, Any] | None:
|
||
|
|
self._ensure_connected()
|
||
|
|
try:
|
||
|
|
movie = self.get_library(library).get(title)
|
||
|
|
return self._movie_to_dict(movie)
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_movie_by_key(self, library: str, rating_key: int) -> dict[str, Any] | None:
|
||
|
|
self._ensure_connected()
|
||
|
|
try:
|
||
|
|
movie = self.get_library(library).fetchItem(rating_key)
|
||
|
|
return self._movie_to_dict(movie)
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def search_movies(
|
||
|
|
self, library: str, query: str, limit: int = 20
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
results = sec.search(query)
|
||
|
|
return [self._movie_to_dict(m) for m in results[:limit]]
|
||
|
|
|
||
|
|
def list_all_movies(
|
||
|
|
self, library: str, limit: int = 200, offset: int = 0
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
all_items = sec.all()
|
||
|
|
return [self._movie_to_dict(m) for m in all_items[offset : offset + limit]]
|
||
|
|
|
||
|
|
def get_recently_added(self, library: str, count: int = 10) -> list[dict[str, Any]]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
items = sec.recentlyAdded() if hasattr(sec, "recentlyAdded") else sec.all()
|
||
|
|
return [self._movie_to_dict(m) for m in items[:count]]
|
||
|
|
|
||
|
|
# ── mutation operations ────────────────────────────────────
|
||
|
|
|
||
|
|
def set_content_rating(
|
||
|
|
self, library: str, title: str, rating: str
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
self._ensure_connected()
|
||
|
|
movie = self.get_library(library).get(title)
|
||
|
|
previous = movie.contentRating or ""
|
||
|
|
movie.editContentRating(rating)
|
||
|
|
logger.info("Set %s contentRating from '%s' -> '%s'", title, previous, rating)
|
||
|
|
return {"title": title, "previous_rating": previous, "new_rating": rating}
|
||
|
|
|
||
|
|
def batch_set_content_rating(
|
||
|
|
self, library: str, titles: list[str], rating: str
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
results = []
|
||
|
|
for title in titles:
|
||
|
|
try:
|
||
|
|
result = self.set_content_rating(library, title, rating)
|
||
|
|
results.append({"title": title, "status": "ok", **result})
|
||
|
|
except Exception as e:
|
||
|
|
results.append({"title": title, "status": "error", "error": str(e)})
|
||
|
|
return results
|
||
|
|
|
||
|
|
def update_all_content_ratings(
|
||
|
|
self, library: str, target_rating: str
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
updated = 0
|
||
|
|
skipped = 0
|
||
|
|
errors = []
|
||
|
|
for video in sec.all():
|
||
|
|
try:
|
||
|
|
if video.contentRating != target_rating:
|
||
|
|
video.editContentRating(target_rating)
|
||
|
|
updated += 1
|
||
|
|
else:
|
||
|
|
skipped += 1
|
||
|
|
except Exception as e:
|
||
|
|
errors.append({"title": video.title, "error": str(e)})
|
||
|
|
logger.info(
|
||
|
|
"update_all: library=%s target=%s updated=%d skipped=%d errors=%d",
|
||
|
|
library,
|
||
|
|
target_rating,
|
||
|
|
updated,
|
||
|
|
skipped,
|
||
|
|
len(errors),
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
"library": library,
|
||
|
|
"target_rating": target_rating,
|
||
|
|
"updated": updated,
|
||
|
|
"skipped": skipped,
|
||
|
|
"errors": errors[:50],
|
||
|
|
"total_errors": len(errors),
|
||
|
|
}
|
||
|
|
|
||
|
|
# ── collections ────────────────────────────────────────────
|
||
|
|
|
||
|
|
def list_collections(self, library: str) -> list[dict[str, Any]]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
return [
|
||
|
|
{
|
||
|
|
"title": c.title,
|
||
|
|
"rating_key": c.ratingKey,
|
||
|
|
"child_count": len(c.children),
|
||
|
|
}
|
||
|
|
for c in sec.collections()
|
||
|
|
]
|
||
|
|
|
||
|
|
def get_movie_collections(self, library: str, title: str) -> list[str]:
|
||
|
|
self._ensure_connected()
|
||
|
|
movie = self.get_library(library).get(title)
|
||
|
|
return [c.tag for c in getattr(movie, "collections", [])]
|
||
|
|
|
||
|
|
# ── library maintenance ────────────────────────────────────
|
||
|
|
|
||
|
|
def refresh_library(self, library: str) -> dict[str, Any]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
sec.refresh()
|
||
|
|
logger.info("Refreshed library '%s'", library)
|
||
|
|
return {"library": library, "status": "refresh triggered"}
|
||
|
|
|
||
|
|
def analyze_library(self, library: str) -> dict[str, Any]:
|
||
|
|
self._ensure_connected()
|
||
|
|
sec = self.get_library(library)
|
||
|
|
sec.analyze()
|
||
|
|
logger.info("Analyzed library '%s'", library)
|
||
|
|
return {"library": library, "status": "analysis triggered"}
|
||
|
|
|
||
|
|
# ── run scheduled task ─────────────────────────────────────
|
||
|
|
|
||
|
|
def run_scheduled_task(
|
||
|
|
self, library: str | None = None, target_rating: str | None = None
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
cfg_library = library or "BorneVideo"
|
||
|
|
cfg_target = target_rating or "2"
|
||
|
|
logger.info(
|
||
|
|
"Scheduled task running on library='%s' target='%s'",
|
||
|
|
cfg_library,
|
||
|
|
cfg_target,
|
||
|
|
)
|
||
|
|
stats_before = self.get_library_stats(cfg_library)
|
||
|
|
result = self.update_all_content_ratings(cfg_library, cfg_target)
|
||
|
|
result["stats_before"] = stats_before
|
||
|
|
return result
|
||
|
|
|
||
|
|
# ── MCP tool registry ──────────────────────────────────────
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def tool_registry() -> list[ToolSpec]:
|
||
|
|
return [
|
||
|
|
ToolSpec(
|
||
|
|
"list_libraries",
|
||
|
|
"List all Plex library sections",
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"get_movie",
|
||
|
|
"Get details for a single movie by title",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {
|
||
|
|
"type": "string",
|
||
|
|
"description": "Library section name",
|
||
|
|
},
|
||
|
|
"title": {"type": "string", "description": "Movie title"},
|
||
|
|
},
|
||
|
|
"required": ["library", "title"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"get_movie_by_key",
|
||
|
|
"Get movie details by its Plex rating key",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"rating_key": {"type": "integer"},
|
||
|
|
},
|
||
|
|
"required": ["library", "rating_key"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"search_movies",
|
||
|
|
"Search for movies by title query",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"query": {"type": "string"},
|
||
|
|
"limit": {"type": "integer", "default": 20},
|
||
|
|
},
|
||
|
|
"required": ["library", "query"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"list_all_movies",
|
||
|
|
"List all movies in a library with pagination",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"limit": {"type": "integer", "default": 200},
|
||
|
|
"offset": {"type": "integer", "default": 0},
|
||
|
|
},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"get_recently_added",
|
||
|
|
"Get recently added movies",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"count": {"type": "integer", "default": 10},
|
||
|
|
},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"list_content_ratings",
|
||
|
|
"List all content ratings present in a library with counts",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"library": {"type": "string"}},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"get_library_stats",
|
||
|
|
"Get statistics about a library (total, rated, unrated, distribution)",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"library": {"type": "string"}},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"set_content_rating",
|
||
|
|
"Set the content rating on a single movie",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"title": {"type": "string"},
|
||
|
|
"rating": {
|
||
|
|
"type": "string",
|
||
|
|
"description": "e.g. '2', '7', '11', 'PG-13'",
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"required": ["library", "title", "rating"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"batch_set_content_rating",
|
||
|
|
"Set content rating on multiple movies by title",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"titles": {"type": "array", "items": {"type": "string"}},
|
||
|
|
"rating": {"type": "string"},
|
||
|
|
},
|
||
|
|
"required": ["library", "titles", "rating"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"update_all_content_ratings",
|
||
|
|
"Update all items in a library to a target content rating",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"target_rating": {"type": "string", "default": "2"},
|
||
|
|
},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"list_collections",
|
||
|
|
"List all collections in a library",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"library": {"type": "string"}},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"get_movie_collections",
|
||
|
|
"Get collections a movie belongs to",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string"},
|
||
|
|
"title": {"type": "string"},
|
||
|
|
},
|
||
|
|
"required": ["library", "title"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"refresh_library",
|
||
|
|
"Trigger a library refresh (rescan)",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"library": {"type": "string"}},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"analyze_library",
|
||
|
|
"Trigger a library analysis",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"library": {"type": "string"}},
|
||
|
|
"required": ["library"],
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ToolSpec(
|
||
|
|
"run_scheduled_task",
|
||
|
|
"Run the standard scheduled task: update all items to target rating",
|
||
|
|
inputSchema={
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"library": {"type": "string", "default": "BorneVideo"},
|
||
|
|
"target_rating": {"type": "string", "default": "2"},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
),
|
||
|
|
]
|
||
|
|
|
||
|
|
def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
|
||
|
|
method_map = {
|
||
|
|
"list_libraries": lambda args: self.list_libraries(),
|
||
|
|
"get_movie": self.get_movie,
|
||
|
|
"get_movie_by_key": self.get_movie_by_key,
|
||
|
|
"search_movies": self.search_movies,
|
||
|
|
"list_all_movies": self.list_all_movies,
|
||
|
|
"get_recently_added": self.get_recently_added,
|
||
|
|
"list_content_ratings": self.list_content_ratings,
|
||
|
|
"get_library_stats": self.get_library_stats,
|
||
|
|
"set_content_rating": self.set_content_rating,
|
||
|
|
"batch_set_content_rating": self.batch_set_content_rating,
|
||
|
|
"update_all_content_ratings": self.update_all_content_ratings,
|
||
|
|
"list_collections": self.list_collections,
|
||
|
|
"get_movie_collections": self.get_movie_collections,
|
||
|
|
"refresh_library": self.refresh_library,
|
||
|
|
"analyze_library": self.analyze_library,
|
||
|
|
"run_scheduled_task": self.run_scheduled_task,
|
||
|
|
}
|
||
|
|
fn = method_map.get(name)
|
||
|
|
if fn is None:
|
||
|
|
raise ValueError(f"Unknown tool: {name}")
|
||
|
|
return fn(**arguments)
|