plexflow 0.0.64__py3-none-any.whl

plexflow/core/context/torrent/tpb/context.py

@@ -0,0 +1,45 @@
+from plexflow.core.storage.object.plexflow_storage import PlexflowObjectStore
+from plexflow.core.context.plexflow_context import PlexflowContext
+from plexflow.core.context.plexflow_property import PlexflowObjectProperty
+from plexflow.core.torrents.providers.tpb.utils import TPBSearchResult
+from typing import Union, List
+
+class ThePirateBayTorrentContext(PlexflowContext):
+    """A class used to represent a Select Context in Plexflow.
+
+    This class extends PlexflowContext and adds a selected item property.
+
+    Attributes:
+        selected_item (PlexflowObjectProperty): The selected item in the context.
+    """
+
+    def __init__(self, store: PlexflowObjectStore, **kwargs):
+        """Initializes the SelectContext with the given object store and keyword arguments.
+
+        Args:
+            store (PlexflowObjectStore): The object store to be used.
+            **kwargs: Arbitrary keyword arguments.
+        """
+
+        super().__init__(store=store, **kwargs)
+        self._torrents = PlexflowObjectProperty(self.object_store, "torrents", local=True)
+
+    @property
+    def items(self) -> Union[List[TPBSearchResult], None]:
+        """Gets the value of the selected item.
+
+        Returns:
+            Any: The value of the selected item.
+        """
+
+        return self._torrents.value
+    
+    @items.setter
+    def items(self, val: List[TPBSearchResult]):
+        """Sets the value of the selected item.
+
+        Args:
+            val (Any): The value to be set.
+        """
+
+        self._torrents.value = val

plexflow/core/context/torrent/yts/context.py

@@ -0,0 +1,45 @@
+from plexflow.core.storage.object.plexflow_storage import PlexflowObjectStore
+from plexflow.core.context.plexflow_context import PlexflowContext
+from plexflow.core.context.plexflow_property import PlexflowObjectProperty
+from plexflow.core.torrents.providers.yts.utils import YTSSearchResult
+from typing import Union, List
+
+class YTSTorrentContext(PlexflowContext):
+    """A class used to represent a Select Context in Plexflow.
+
+    This class extends PlexflowContext and adds a selected item property.
+
+    Attributes:
+        selected_item (PlexflowObjectProperty): The selected item in the context.
+    """
+
+    def __init__(self, store: PlexflowObjectStore, **kwargs):
+        """Initializes the SelectContext with the given object store and keyword arguments.
+
+        Args:
+            store (PlexflowObjectStore): The object store to be used.
+            **kwargs: Arbitrary keyword arguments.
+        """
+
+        super().__init__(store=store, **kwargs)
+        self._torrents = PlexflowObjectProperty(self.object_store, "torrents", local=True)
+
+    @property
+    def items(self) -> Union[List[YTSSearchResult], None]:
+        """Gets the value of the selected item.
+
+        Returns:
+            Any: The value of the selected item.
+        """
+
+        return self._torrents.value
+    
+    @items.setter
+    def items(self, val: List[YTSSearchResult]):
+        """Sets the value of the selected item.
+
+        Args:
+            val (Any): The value to be set.
+        """
+
+        self._torrents.value = val

plexflow/core/context/watchlist/context.py

@@ -0,0 +1,46 @@
+from plexflow.core.storage.object.plexflow_storage import PlexflowObjectStore
+from plexflow.core.context.plexflow_context import PlexflowContext
+from plexflow.core.context.plexflow_property import PlexflowObjectProperty
+from plexflow.core.plex.watchlist.datatypes import MediaContainer
+from typing import Union
+
+class WatchlistContext(PlexflowContext):
+    """
+    A class used to represent a context for managing a watchlist in Plexflow.
+
+    This class extends PlexflowContext and adds a property for the watchlist.
+
+    Attributes:
+        watchlist (PlexflowObjectProperty): The watchlist in the context.
+    """
+
+    def __init__(self, store: PlexflowObjectStore, **kwargs):
+        """
+        Initializes the WatchlistContext with the given object store and optional keyword arguments.
+
+        Args:
+            store (PlexflowObjectStore): The object store to be used.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(store=store, **kwargs)
+        self.watchlist = PlexflowObjectProperty(self.object_store, "watchlist", local=True)
+
+    @property
+    def list(self) -> Union[MediaContainer, None]:
+        """
+        Gets the value of the watchlist.
+
+        Returns:
+            Any: The value of the selected watchlist item.
+        """
+        return self.watchlist.value
+    
+    @list.setter
+    def list(self, val: MediaContainer):
+        """
+        Sets the value of the watchlist.
+
+        Args:
+            val (Any): The value to be set.
+        """
+        self.watchlist.value = val

plexflow/core/downloads/candidates/download_candidate.py

@@ -0,0 +1,210 @@
+import math
+from typing import List, Optional, Set
+from plexflow.core.torrents.results.universal import UniversalTorrent
+from plexflow.core.subtitles.results.subtitle import Subtitle
+from plexflow.utils.imdb.imdb_codes import IMDbCode
+
+class DownloadCandidate:
+    """
+    Represents a download candidate for a torrent with associated subtitles.
+    """
+
+    def __init__(self, torrent: UniversalTorrent, subtitles: List[Subtitle]):
+        """
+        Initializes a new instance of the DownloadCandidate class.
+
+        Args:
+            torrent (UniversalTorrent): The torrent associated with the download candidate.
+            subtitles (List[Subtitle]): The list of subtitles associated with the download candidate.
+        """
+        self.torrent = torrent
+        self.subtitles = subtitles
+
+    @property
+    def imdb_code(self) -> IMDbCode:
+        """
+        Get the IMDb code of the torrent.
+
+        Returns:
+            IMDbCode: The IMDb code of the torrent.
+        """
+        return self.torrent.imdb_code
+
+    @property
+    def is_season_pack(self) -> bool:
+        """
+        Check if the torrent is a season pack.
+
+        Returns:
+            bool: True if the torrent is a season pack, False otherwise.
+        """
+        return self.torrent.is_season_pack
+
+    @property
+    def season(self) -> Optional[int]:
+        """
+        Get the season number of the torrent.
+
+        Returns:
+            Optional[int]: The season number of the torrent, or None if not applicable.
+        """
+        return self.torrent.season
+
+    @property
+    def episode(self) -> Optional[int]:
+        """
+        Get the episode number of the torrent.
+
+        Returns:
+            Optional[int]: The episode number of the torrent, or None if not applicable.
+        """
+        return self.torrent.episode
+
+    @property
+    def max_peers(self) -> int:
+        """
+        Get the maximum number of peers for the torrent.
+
+        Returns:
+            int: The maximum number of peers for the torrent.
+        """
+        return self.torrent.max_peers
+
+    @property
+    def max_seeds(self) -> int:
+        """
+        Get the maximum number of seeds for the torrent.
+
+        Returns:
+            int: The maximum number of seeds for the torrent.
+        """
+        return self.torrent.max_seeds
+
+    @property
+    def min_seeds(self) -> int:
+        """
+        Get the minimum number of seeds for the torrent.
+
+        Returns:
+            int: The minimum number of seeds for the torrent.
+        """
+        return self.torrent.min_seeds
+
+    @property
+    def min_peers(self) -> int:
+        """
+        Get the minimum number of peers for the torrent.
+
+        Returns:
+            int: The minimum number of peers for the torrent.
+        """
+        return self.torrent.min_peers
+
+    @property
+    def sources(self) -> Set:
+        """
+        Get the set of sources for the torrent.
+
+        Returns:
+            Set: The set of sources for the torrent.
+        """
+        return self.torrent.sources
+
+    @property
+    def max_size_bytes(self) -> int:
+        """
+        Get the maximum size of the torrent in bytes.
+
+        Returns:
+            int: The maximum size of the torrent in bytes.
+        """
+        return self.torrent.max_size_bytes
+
+    @property
+    def min_size_bytes(self) -> int:
+        """
+        Get the minimum size of the torrent in bytes.
+
+        Returns:
+            int: The minimum size of the torrent in bytes.
+        """
+        return self.torrent.min_size_bytes
+
+    @property
+    def has_native_subtitles(self) -> bool:
+        """
+        Check if the torrent has native subtitles.
+
+        Returns:
+            bool: True if the torrent has native subtitles, False otherwise.
+        """
+        return self.torrent.has_native_subtitles
+
+    @property
+    def has_native_dutch_subtitles(self) -> bool:
+        """
+        Check if the torrent has native Dutch subtitles.
+
+        Returns:
+            bool: True if the torrent has native Dutch subtitles, False otherwise.
+        """
+        return self.torrent.has_native_dutch_subtitles
+
+    @property
+    def has_native_english_subtitles(self) -> bool:
+        """
+        Check if the torrent has native English subtitles.
+
+        Returns:
+            bool: True if the torrent has native English subtitles, False otherwise.
+        """
+        return self.torrent.has_native_english_subtitles
+
+    @property
+    def fitness(self) -> float:
+        """
+        Calculate the fitness score of the download candidate.
+
+        Returns:
+            float: The fitness score of the download candidate.
+        """
+        return self.max_seeds / max(math.sqrt(self.min_size_bytes), 1)
+
+    @property
+    def has_subtitles(self) -> bool:
+        """
+        Check if the download candidate has subtitles.
+
+        Returns:
+            bool: True if the download candidate has subtitles, False otherwise.
+        """
+        return any([
+            self.has_native_subtitles,
+            len(self.subtitles) > 0
+        ])
+
+    @property
+    def has_dutch_subtitles(self) -> bool:
+        """
+        Check if the download candidate has Dutch subtitles.
+
+        Returns:
+            bool: True if the download candidate has Dutch subtitles, False otherwise.
+        """
+        return any([
+            self.has_native_dutch_subtitles,
+            any(s.language == "nl" for s in self.subtitles),
+        ])
+
+    @property
+    def has_english_subtitles(self) -> bool:
+        """
+        Check if the download candidate has English subtitles.
+
+        Returns:
+            bool: True if the download candidate has English subtitles, False otherwise.
+        """
+        return any([
+            self.has_native_english_subtitles,
+            any(s.language == "en" for s in self.subtitles),
+        ])

plexflow/core/downloads/candidates/filtered.py

@@ -0,0 +1,51 @@
+from plexflow.core.downloads.candidates.download_candidate import DownloadCandidate
+from typing import List
+
+class FilteredCandidates:
+    """
+    Represents a collection of filtered download candidates based on specific criteria.
+
+    Args:
+        imdb_id (str): The IMDb ID of the movie.
+        year (int): The release year of the movie.
+        titles (List[str]): A list of movie titles.
+        candidates (List[DownloadCandidate]): A list of download candidates.
+
+    Attributes:
+        imdb_id (str): The IMDb ID of the movie.
+        year (int): The release year of the movie.
+        titles (List[str]): A list of movie titles.
+        candidates (List[DownloadCandidate]): A list of download candidates.
+    """
+
+    def __init__(self, imdb_id: str, year: int, titles: List[str], candidates: List[DownloadCandidate]):
+        self.imdb_id = imdb_id
+        self.year = year
+        self.titles = titles
+        self.candidates = candidates
+
+    def __iter__(self):
+        """
+        Returns an iterator that yields candidates from the `candidates` list
+        if they match the criteria specified by the `_is_match` method.
+
+        Yields:
+            candidate: A candidate that matches the criteria.
+        """
+        return (candidate for candidate in self.candidates if self._is_match(candidate))
+
+    def _is_match(self, candidate: DownloadCandidate) -> bool:
+        """
+        Checks if the given candidate is a match for the current filter.
+
+        Args:
+            candidate (DownloadCandidate): The candidate to check.
+
+        Returns:
+            bool: True if the candidate is a match, False otherwise.
+        """
+        if candidate.imdb_code == self.imdb_id:
+            return True
+        if candidate.year == self.year and candidate.title in self.titles:
+            return True
+        return False

plexflow/core/downloads/candidates/utils.py

@@ -0,0 +1,39 @@
+from typing import Iterable, List
+from plexflow.core.torrents.results.universal import UniversalTorrent
+from plexflow.core.subtitles.results.subtitle import Subtitle
+from plexflow.core.downloads.candidates.download_candidate import DownloadCandidate
+
+def create_download_candidates(torrents: Iterable[UniversalTorrent], subtitles: Iterable[Subtitle]) -> List[DownloadCandidate]:
+    """
+    This function creates a list of DownloadCandidate objects from an iterable of UniversalTorrent and an iterable of Subtitle.
+    All subtitles that are compatible with a torrent will become together a DownloadCandidate.
+
+    Parameters:
+    torrents (Iterable[UniversalTorrent]): An iterable of UniversalTorrent objects.
+    subtitles (Iterable[Subtitle]): An iterable of Subtitle objects.
+
+    Returns:
+    List[DownloadCandidate]: A list of DownloadCandidate objects.
+
+    Raises:
+    ValueError: If either torrents or subtitles is not an iterable.
+    """
+
+    # Check if inputs are iterables
+    if not isinstance(torrents, Iterable):
+        raise ValueError("torrents should be an iterable of UniversalTorrent objects")
+    if not isinstance(subtitles, Iterable):
+        raise ValueError("subtitles should be an iterable of Subtitle objects")
+
+    download_candidates = []
+
+    for torrent in torrents:
+        # Check if torrent is an instance of UniversalTorrent
+        if not isinstance(torrent, UniversalTorrent):
+            raise ValueError("Each torrent should be an instance of UniversalTorrent")
+
+        compatible_subtitles = [subtitle for subtitle in subtitles if torrent.is_compatible_with(subtitle)]
+        download_candidate = DownloadCandidate(torrent, compatible_subtitles)
+        download_candidates.append(download_candidate)
+
+    return download_candidates

plexflow/core/env/env.py

@@ -0,0 +1,31 @@
+import os
+from airflow.models import Variable
+from airflow.exceptions import AirflowException
+
+def get_var(var_name: str) -> str:
+    """
+    Function to retrieve a variable from Airflow or OS environment.
+
+    This function first attempts to retrieve the variable from Airflow's variables.
+    If the variable is not found in Airflow or an error occurs, it then attempts to retrieve it from the OS environment.
+
+    Args:
+        var_name (str): The name of the variable to retrieve.
+
+    Returns:
+        str: The value of the variable. If the variable is not found, returns None.
+
+    Examples:
+        >>> get_var('HOME')
+        '/home/user'
+        >>> get_var('UNKNOWN_VAR')
+        None
+    """
+    try:
+        # Try to get the variable from Airflow
+        var_value = Variable.get(var_name)
+    except AirflowException:
+        # If it fails, get the variable from os environment
+        var_value = os.getenv(var_name)
+
+    return var_value

plexflow/core/genai/bot.py

@@ -0,0 +1,9 @@
+import os
+import cohere
+from plexflow.core.genai.utils.loader import load_preamble
+
+class CohereBot:
+    def __init__(self, preamble_id: str) -> None:
+        self.co = cohere.Client(api_key=os.getenv("COHERE_API_KEY"))
+        self.chat_history = []
+        self.preamble = load_preamble(preamble_id)

plexflow/core/genai/plexa.py

@@ -0,0 +1,54 @@
+
+import json
+import os
+import cohere
+from plexflow.core.genai.utils.loader import load_preamble
+
+class Plexa:
+    def __init__(self) -> None:
+        self.bot_name = "Plexa"
+        self.co = cohere.Client(api_key=os.getenv("COHERE_API_KEY"))
+        self.chat_history = []
+        self.preamble = load_preamble("plexa_v2")
+        self.hint = "BE BRIEF IN YOUR RESPONSE, SPLIT TEXT IN MESSAGES AND FORMAT AS VALID JSON"
+    
+    def add_user_message(self, message: str, context: dict):
+        self.chat_history.append({
+            "role": "USER", "text":
+                json.dumps({
+                    "movie": context.get("title"),
+                    "release_date": context.get("release_date"),
+                    "user_question": message,
+                    "hint": self.hint,
+                })
+            })
+    
+    def add_bot_message(self, message: str, context: dict):
+        self.chat_history.append({
+            "role": "CHATBOT", "text":
+                json.dumps({
+                    "response": message,
+                    "parts": [
+                        message
+                    ]
+                })
+            })
+
+    def react(self, question: str, context: dict):
+        message = json.dumps({
+            "movie": context.get("title"),
+            "release_date": context.get("release_date"),
+            "user_question": question,
+            "hint": self.hint,
+        })
+
+        return self.co.chat(
+                    message=message,
+                    temperature=1,
+                    connectors=[{"id": "web-search"}],
+                    model="command-r-plus",
+                    chat_history=self.chat_history[-10:],
+                    preamble=self.preamble,
+                )
+    
+    

plexflow/core/genai/torrent/imdb_verify.py

@@ -0,0 +1,65 @@
+from plexflow.core.genai.bot import CohereBot
+import json
+import requests
+from bs4 import BeautifulSoup
+
+class TorrentImdbMatcher(CohereBot):
+    def __init__(self) -> None:
+        super().__init__(preamble_id="torrent_imdb_matcher")
+    
+    def parse(self, content: str):
+        response = self.co.chat(
+            message=content,
+            temperature=1,
+            model="command-r-plus",
+            preamble=self.preamble,
+        )
+        
+#        print(response.json())
+        content = response.text
+        
+        if content.startswith('```json'):
+            content = content.lstrip("`json")
+        
+        if content.endswith('`'):
+            content = content.rstrip("`")
+
+        return json.loads(content)
+
+def get_page_text(url):
+    r = requests.get(
+        url,
+        headers={
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36",
+            "Cookie": r"dom3ic8zudi28v8lr6fgphwffqoz0j6c=b7c768c9-98eb-49e8-8ec5-2c9e5c5828c9%3A2%3A1; sb_main_755b5f8e271690d6cb76076f459e9c82=1; PHPSESSID=g2g6grs8ejkoqisi81h1tomtk6; sb_count_755b5f8e271690d6cb76076f459e9c82=4; hu8935j4i9fq3hpuj9q39=true; fencekey=8ad2dbbf2b3ba78729048b3d823574ee; AdskeeperStorage=%7B%220%22%3A%7B%22svspr%22%3A%22%22%2C%22svsds%22%3A9%7D%2C%22C1543068%22%3A%7B%22page%22%3A5%2C%22time%22%3A%221718527871490%22%7D%2C%22C1343686%22%3A%7B%22page%22%3A5%2C%22time%22%3A%221718527871482%22%7D%2C%22C385455%22%3A%7B%22page%22%3A5%2C%22time%22%3A%221718527871553%22%7D%7D"
+        }
+    )
+    
+    r.raise_for_status()
+    
+    soup = BeautifulSoup(r.text, 'html.parser')
+    
+    return soup.get_text()
+
+def matches_with_imdb(imdb_id: str, torrent_url: str):
+    parser = TorrentImdbMatcher()
+
+    imdb_info = get_page_text(f"https://www.imdb.com/title/{imdb_id}")
+    torrent_info = get_page_text(torrent_url)
+    
+    prompt = f"""
+    IMDB INFO:
+    {imdb_info}
+    
+    TORRENT INFO:
+    {torrent_info}
+    """
+
+    print("prompt has", len(prompt), "chars")
+
+    data = parser.parse(prompt)
+    
+    matched = data.get("match")
+    reasons = data.get("reasons", [])
+    
+    return matched, reasons

plexflow/core/genai/torrent/movie.py

@@ -0,0 +1,25 @@
+from plexflow.core.genai.bot import CohereBot
+import json
+
+class MovieTorrentParser(CohereBot):
+    def __init__(self) -> None:
+        super().__init__(preamble_id="movie_torrent_parser")
+    
+    def parse(self, content: str):
+        response = self.co.chat(
+            message=content,
+            temperature=1,
+            model="command-r-plus",
+            preamble=self.preamble,
+        )
+        
+#        print(response.json())
+        content = response.text
+        
+        if content.startswith('```json'):
+            content = content.lstrip("`json")
+        
+        if content.endswith('`'):
+            content = content.rstrip("`")
+
+        return json.loads(content)

plexflow/core/genai/utils/loader.py

@@ -0,0 +1,5 @@
+
+
+def load_preamble(name: str):
+    with open(f"config/genai/{name}.txt", "r") as file:
+        return file.read()