mycelium-ai 0.5.0__py3-none-any.whl

mycelium/infrastructure/chroma_adapter.py

@@ -0,0 +1,232 @@
+"""ChromaDB integration for storing and searching embeddings."""
+
+import logging
+import re
+from pathlib import Path
+from typing import List, Optional
+
+import chromadb
+from tqdm import tqdm
+
+from ..domain.models import Track, TrackEmbedding, SearchResult, MediaServerType
+from ..domain.repositories import EmbeddingRepository
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaEmbeddingRepository(EmbeddingRepository):
+    """Implementation of EmbeddingRepository using ChromaDB with model-specific collections."""
+
+    def __init__(
+            self,
+            db_path: str,
+            media_server_type: MediaServerType,
+            collection_name: str = "my_music_library",
+            model_id: str = "laion/larger_clap_music_and_speech",
+            batch_size: int = 1000,
+    ):
+        self.db_path = db_path
+        self.base_collection_name = collection_name
+        self.model_id = model_id
+        self.batch_size = batch_size
+        self.media_server_type = media_server_type
+
+        # Initialize ChromaDB client
+        try:
+            Path(db_path).mkdir(parents=True, exist_ok=True)
+        except Exception:
+            logger.error(f"Failed to create database directory at {db_path}. Please check permissions.")
+        self.client = chromadb.PersistentClient(path=db_path)
+
+        # Create model-specific collection name
+        self.collection_name = self._get_collection_name_for_model(model_id)
+
+        # Specify 'cosine' distance metric for normalized embeddings
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name,
+            metadata={"hnsw:space": "cosine", "model_id": model_id}
+        )
+
+        logger.info(
+            f"Collection '{self.collection_name}' ready for model '{model_id}'. Current elements: {self.collection.count()}")
+
+    def _get_collection_name_for_model(self, model_id: str) -> str:
+        """Generate a safe collection name for the given model ID."""
+        # Make model ID safe for collection name (alphanumeric and underscores only)
+        safe_model_id = re.sub(r'\W', '_', model_id.replace('/', '_'))
+        return f"{self.base_collection_name}_{safe_model_id}"
+
+    def save_embeddings(self, embeddings: List[TrackEmbedding]) -> None:
+        """Save track embeddings to ChromaDB."""
+        if not embeddings:
+            return
+
+        # Prepare data for batch insertion
+        ids = []
+        embedding_vectors = []
+        metadatas = []
+
+        for track_embedding in embeddings:
+            track = track_embedding.track
+            ids.append(track.unique_id)
+            embedding_vectors.append(track_embedding.embedding)
+            metadatas.append({
+                "filepath": str(track.filepath),
+                "artist": track.artist,
+                "album": track.album,
+                "title": track.title,
+                "media_server_type": track.media_server_type.value,
+                "media_server_rating_key": track.media_server_rating_key,
+                "model_id": self.model_id
+            })
+
+        # Insert in batches for maximum efficiency
+        for i in tqdm(range(0, len(ids), self.batch_size), desc="Indexing in ChromaDB"):
+            end_idx = min(i + self.batch_size, len(ids))
+            id_batch = ids[i:end_idx]
+            embedding_batch = embedding_vectors[i:end_idx]
+            metadata_batch = metadatas[i:end_idx]
+
+            self.collection.add(
+                ids=id_batch,
+                embeddings=embedding_batch,
+                metadatas=metadata_batch
+            )
+
+        logger.info("Indexing completed!")
+        logger.info(f"Total elements in collection '{self.collection_name}': {self.collection.count()}")
+
+    def search_by_embedding(self, embedding: List[float], n_results: int = 10) -> List[SearchResult]:
+        """Search for similar tracks by embedding."""
+        results = self.collection.query(
+            query_embeddings=[embedding],
+            n_results=n_results
+        )
+
+        return self._parse_search_results(results.copy())
+
+    def get_embedding_count(self) -> int:
+        """Get the total number of embeddings stored."""
+        return self.collection.count()
+
+    @staticmethod
+    def _parse_search_results(results: dict) -> List[SearchResult]:
+        """Parse ChromaDB results into SearchResult objects."""
+        search_results = []
+
+        if not results['ids'] or not results['ids'][0]:
+            return search_results
+
+        for i in range(len(results['ids'][0])):
+            metadata = results['metadatas'][0][i]
+            distance = results['distances'][0][i]
+            unique_id = results['ids'][0][i]
+
+            # Parse unique_id to get media server info
+            media_server_type_str, media_server_rating_key = unique_id.split(':', 1)
+
+            from ..domain.models import MediaServerType
+            try:
+                media_server_type = MediaServerType(media_server_type_str)
+            except ValueError:
+                media_server_type = MediaServerType.PLEX  # Default fallback
+
+            track = Track(
+                artist=metadata['artist'],
+                album=metadata['album'],
+                title=metadata['title'],
+                filepath=Path(metadata['filepath']),
+                media_server_rating_key=media_server_rating_key,
+                media_server_type=media_server_type
+            )
+
+            # Convert distance to similarity score (1 - distance for cosine)
+            similarity_score = 1.0 - distance
+
+            search_results.append(SearchResult(
+                track=track,
+                similarity_score=similarity_score,
+                distance=distance
+            ))
+
+        return search_results
+
+    def has_embedding(self, track_id: str) -> bool:
+        """Check if an embedding exists for a track."""
+        track_id = Track(media_server_type=self.media_server_type,
+              media_server_rating_key=track_id).unique_id
+        logger.debug(
+            f"Checking embedding for track {track_id}: collection_name={self.collection_name}, model_id={self.model_id}"
+        )
+        try:
+            result = self.collection.get(ids=[track_id])
+            exists = len(result['ids']) > 0
+            logger.debug(f"Checking embedding for track {track_id}: exists={exists}")
+            return exists
+        except Exception as e:
+            logger.error(f"Error checking embedding for track {track_id}: {e}")
+            return False
+
+    def save_embedding(self, track_embedding: TrackEmbedding) -> None:
+        """Save a single track embedding to ChromaDB."""
+        track = track_embedding.track
+        track_id = track.unique_id
+
+        logger.info(f"Saving embedding to ChromaDB for track {track_id}: {track.artist} - {track.title}")
+        logger.info(f"Collection count before save: {self.collection.count()}")
+
+        # Check if embedding already exists, if so, update it
+        existing = self.collection.get(ids=[track_id])
+        if existing['ids']:
+            logger.info(f"Updating existing embedding for track {track_id}")
+            self.collection.update(
+                ids=[track_id],
+                embeddings=[track_embedding.embedding],
+                metadatas=[{
+                    "filepath": str(track.filepath),
+                    "artist": track.artist,
+                    "album": track.album,
+                    "title": track.title,
+                    "media_server_type": track.media_server_type.value,
+                    "media_server_rating_key": track.media_server_rating_key,
+                    "model_id": self.model_id
+                }]
+            )
+        else:
+            logger.info(f"Adding new embedding for track {track_id}")
+            self.collection.add(
+                ids=[track_id],
+                embeddings=[track_embedding.embedding],
+                metadatas=[{
+                    "filepath": str(track.filepath),
+                    "artist": track.artist,
+                    "album": track.album,
+                    "title": track.title,
+                    "media_server_type": track.media_server_type.value,
+                    "media_server_rating_key": track.media_server_rating_key,
+                    "model_id": self.model_id
+                }]
+            )
+
+        logger.info(f"Collection count after save: {self.collection.count()}")
+        logger.info(f"Successfully saved embedding to ChromaDB for track {track_id}")
+
+    def get_embedding_by_track_id(self, track_id: str) -> Optional[List[float]]:
+        """Get embedding for a specific track."""
+        track_id = Track(media_server_type=self.media_server_type, media_server_rating_key=track_id).unique_id
+        try:
+            result = self.collection.get(
+                ids=[track_id],
+                include=['embeddings']
+            )
+            if result['embeddings'] is not None and len(result['embeddings']) > 0:
+                embedding = result['embeddings'][0]
+                logger.debug(
+                    f"Retrieved embedding for track {track_id}, size: {len(embedding) if embedding is not None else 0}")
+                return embedding
+            else:
+                logger.debug(f"No embedding found in ChromaDB for track {track_id}")
+                return None
+        except Exception as e:
+            logger.error(f"Error retrieving embedding for track {track_id}: {e}")
+            return None

mycelium/infrastructure/clap_adapter.py

@@ -0,0 +1,280 @@
+"""CLAP model integration for generating embeddings."""
+
+import logging
+import random
+from pathlib import Path
+from typing import List, Optional
+
+import librosa
+import torch
+from transformers import ClapModel, ClapProcessor
+
+from ..domain.repositories import EmbeddingGenerator
+
+
+class CLAPEmbeddingGenerator(EmbeddingGenerator):
+    """ Implementation of EmbeddingGenerator using LAION's CLAP model. """
+
+    def __init__(
+            self,
+            model_id: str = "laion/larger_clap_music_and_speech",
+            target_sr: int = 48000,
+            chunk_duration_s: int = 10,
+            num_chunks: int = 3,
+            max_load_duration_s: Optional[int] = 120
+    ):
+        self.model_id = model_id
+        self.target_sr = target_sr
+        self.chunk_duration_s = chunk_duration_s
+        self.num_chunks = num_chunks
+        self.max_load_duration_s = max_load_duration_s
+        self.logger = logging.getLogger(__name__)
+
+        self.device = self.get_best_device()
+        self.logger.info(f"Selected device: {self.device}")
+
+        ## Lazy loading. Model is not loaded on instantiation.
+        self.model: Optional[ClapModel] = None
+        self.processor: Optional[ClapProcessor] = None
+
+        self.use_half = self.can_use_half_precision()
+        if self.use_half:
+            self.logger.info("Half precision (FP16) is supported and will be used.")
+        else:
+            self.logger.info("Half precision not supported, using full precision (FP32).")
+
+    def _load_model_if_needed(self):
+        """Loads the model and processor on the first call that needs them."""
+        if self.model is None or self.processor is None:
+            self.logger.info(f"Loading model '{self.model_id}' to device '{self.device}'...")
+
+            self.model = ClapModel.from_pretrained(self.model_id).to(self.device)
+            self.processor = ClapProcessor.from_pretrained(self.model_id)
+
+            if self.use_half and self.device == "cuda":
+                self.logger.info("Applying half precision (FP16) to model for CUDA device.")
+                self.model.half()
+            elif self.use_half and self.device == "mps":
+                self.logger.warning(
+                    "Half precision is supported but disabled on MPS device to prevent potential crashes. Using FP32.")
+
+            self.model.eval()
+            self.logger.info("Model loaded successfully.")
+            try:
+                self.logger.info(f"Model dtype after load: {next(self.model.parameters()).dtype}")
+            except StopIteration:
+                self.logger.debug("Could not determine model dtype (no parameters found).")
+
+    @staticmethod
+    def get_best_device() -> str:
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    def can_use_half_precision(self) -> bool:
+        """Checks once if the device supports half precision."""
+        if self.device == "cuda":
+            # Most modern CUDA devices support FP16.
+            return True
+        if self.device == "mps":
+            # Check for potential runtime errors on some MPS devices.
+            try:
+                torch.tensor([1.0], dtype=torch.half).to(self.device)
+                return True
+            except RuntimeError:
+                self.logger.warning("MPS device does not support half precision, falling back to FP32.")
+                return False
+        return False
+
+    def _get_processor(self) -> ClapProcessor:
+        """Return a ready-to-use processor with a non-optional type."""
+        self._load_model_if_needed()
+        assert self.processor is not None
+        return self.processor
+
+    def _get_model(self) -> ClapModel:
+        """Return a ready-to-use model with a non-optional type."""
+        self._load_model_if_needed()
+        assert self.model is not None
+        return self.model
+
+    def _prepare_inputs(self, inputs: dict) -> dict:
+        """Move inputs to the correct device and cast floating tensors to the model's dtype.
+
+        This prevents dtype mismatches when the model runs in half precision on CUDA.
+        """
+        model = self._get_model()
+        # Determine model parameter dtype (e.g., torch.float32 or torch.float16)
+        model_dtype = next(model.parameters()).dtype
+        prepared = {}
+        for k, v in inputs.items():
+            if isinstance(v, torch.Tensor):
+                if v.is_floating_point():
+                    prepared[k] = v.to(device=self.device, dtype=model_dtype)
+                else:
+                    prepared[k] = v.to(device=self.device)
+            else:
+                prepared[k] = v
+        return prepared
+
+    def generate_embedding(self, filepath: Path) -> Optional[List[float]]:
+        """Generate embedding for a single audio file by delegating to batch method."""
+        results = self.generate_embedding_batch([filepath])
+        return results[0] if results else None
+
+    def generate_embedding_batch(self, filepaths: List[Path]) -> List[Optional[List[float]]]:
+        """Generate embeddings for multiple audio files in a single GPU batch"""
+        if not filepaths:
+            return []
+
+        try:
+            processor = self._get_processor()
+            model = self._get_model()
+
+            all_chunks = []
+            file_chunk_counts = []
+
+            chunk_size_samples = self.chunk_duration_s * self.target_sr
+
+            # Load and prepare all audio files
+            for filepath in filepaths:
+                try:
+                    waveform, _ = librosa.load(
+                        str(filepath),
+                        sr=self.target_sr,
+                        mono=True,
+                        duration=self.max_load_duration_s
+                    )
+
+                    total_samples = len(waveform)
+                    chunks = []
+
+                    # Calculate how many full, non-overlapping chunks can fit.
+                    num_possible_bins = total_samples // chunk_size_samples
+
+                    if num_possible_bins == 0:
+                        self.logger.warning(
+                            f"File {filepath} is too short ({total_samples / self.target_sr:.1f}s) "
+                            f"for even one chunk of {self.chunk_duration_s:.1f}s.")
+                        file_chunk_counts.append(0)
+                        continue
+
+                    # Determine which bin indices to sample from.
+                    if num_possible_bins < self.num_chunks:
+                        self.logger.warning(
+                            f"File {filepath} only has space for {num_possible_bins} non-overlapping chunks, "
+                            f"less than the requested {self.num_chunks}. Using all available chunks."
+                        )
+                        chosen_bin_indices = range(num_possible_bins)
+                    else:
+                        possible_bin_indices = range(num_possible_bins)
+                        chosen_bin_indices = random.sample(possible_bin_indices, k=self.num_chunks)
+
+                    # Create chunks based on the chosen indices.
+                    for bin_index in chosen_bin_indices:
+                        start_idx = bin_index * chunk_size_samples
+                        end_idx = start_idx + chunk_size_samples
+                        chunk = waveform[start_idx:end_idx]
+                        chunks.append(chunk)
+
+                    if not chunks:
+                        self.logger.warning(f"No valid chunks generated for {filepath}.")
+                        file_chunk_counts.append(0)
+                        continue
+
+                    all_chunks.extend(chunks)
+                    file_chunk_counts.append(len(chunks))
+
+                except Exception as e:
+                    self.logger.error(f"Error loading audio file {filepath}: {e}")
+                    file_chunk_counts.append(0)
+
+            if not all_chunks:
+                return [None] * len(filepaths)
+
+            # Process all chunks in a single batch
+            inputs = processor(
+                audios=all_chunks,
+                sampling_rate=self.target_sr,
+                return_tensors="pt",
+                padding=True
+            )
+
+            inputs = self._prepare_inputs(inputs)
+
+            with torch.no_grad():
+                audio_features = model.get_audio_features(**inputs)
+
+            # Split results back to individual files and compute mean embeddings
+            results = []
+            chunk_idx = 0
+
+            for chunk_count in file_chunk_counts:
+                if chunk_count == 0:
+                    results.append(None)
+                else:
+                    file_features = audio_features[chunk_idx:chunk_idx + chunk_count]
+                    mean_embedding = torch.mean(file_features, dim=0)
+                    normalized_embedding = torch.nn.functional.normalize(mean_embedding, p=2, dim=0)
+                    results.append(normalized_embedding.cpu().numpy().tolist())
+                    chunk_idx += chunk_count
+
+            self.logger.info(
+                f"Successfully processed batch of {len(filepaths)} audio files ({len(all_chunks)} total chunks)")
+            return results
+
+        except Exception as e:
+            self.logger.error(f"Error in batch audio embedding generation: {e}", exc_info=True)
+            return [None] * len(filepaths)
+
+    def generate_text_embedding(self, text: str) -> Optional[List[float]]:
+        """Generate embedding for a single text query by delegating to batch method."""
+        results = self.generate_text_embedding_batch([text])
+        return results[0] if results else None
+
+    def generate_text_embedding_batch(self, texts: List[str]) -> List[Optional[List[float]]]:
+        """Generate embeddings for multiple text queries in a single GPU batch for better utilization."""
+        if not texts:
+            return []
+
+        try:
+            processor = self._get_processor()
+            model = self._get_model()
+
+            inputs = processor(
+                text=texts,
+                return_tensors="pt",
+                padding=True
+            )
+
+            inputs = self._prepare_inputs(inputs)
+
+            with torch.no_grad():
+                text_features = model.get_text_features(**inputs)
+                text_embeddings = torch.nn.functional.normalize(text_features, p=2, dim=-1)
+
+            # Convert to list of lists
+            results = text_embeddings.cpu().numpy().tolist()
+            self.logger.info(f"Successfully processed batch of {len(texts)} text queries")
+            return results
+
+        except Exception as e:
+            self.logger.error(f"Error in batch text embedding generation: {e}", exc_info=True)
+            return [None] * len(texts)
+
+    def unload_model(self) -> None:
+        """Unload model to free GPU memory."""
+        if self.model is not None:
+            del self.model
+            del self.processor
+            self.model = None
+            self.processor = None
+
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+            elif self.device == "mps":
+                torch.mps.empty_cache()
+
+            self.logger.info("Model unloaded")

mycelium/infrastructure/plex_adapter.py

@@ -0,0 +1,145 @@
+"""Plex integration for accessing music library."""
+
+import logging
+from pathlib import Path
+from typing import List, Optional
+from datetime import datetime
+
+from plexapi.audio import Artist
+from plexapi.server import PlexServer
+from tqdm import tqdm
+
+from ..domain.models import Track, Playlist, MediaServerType
+from ..domain.repositories import MediaServerRepository
+
+
+class PlexMusicRepository(MediaServerRepository):
+    """Implementation of MediaServerRepository for accessing Plex music library."""
+
+    def __init__(
+        self,
+        plex_url: str = None,
+        plex_token: str = None,
+        music_library_name: str = "Music"
+    ):
+        self.plex_url = plex_url
+        self.plex_token = plex_token
+        self.music_library_name = music_library_name
+        self.logger = logging.getLogger(__name__)
+
+    def get_all_tracks(self) -> List[Track]:
+        """Get all tracks from the Plex music library."""
+        try:
+            plex = PlexServer(self.plex_url, self.plex_token, timeout=3600)
+            music_lib = plex.library.section(self.music_library_name)
+            self.logger.info(f"Connected to Plex. Scanning library '{self.music_library_name}'...")
+        except Exception as e:
+            raise ConnectionError(f"Error connecting to Plex server: {e}")
+
+        all_tracks = []
+
+        # Hierarchical iteration for better robustness and memory efficiency
+        artists = music_lib.all(libtype='artist')
+        artists: List[Artist]
+        for artist in tqdm(artists, desc="Processing Artists"):
+            try:
+                for album in artist.albums():
+                    for track in album.tracks():
+                        for part in track.iterParts():
+                            filepath = Path(part.file)
+                            if filepath.exists():
+                                track_obj = Track(
+                                    artist=artist.title,
+                                    album=album.title,
+                                    title=track.title,
+                                    filepath=filepath,
+                                    media_server_rating_key=str(track.ratingKey),
+                                    media_server_type=MediaServerType.PLEX
+                                )
+                                all_tracks.append(track_obj)
+                            else:
+                                self.logger.warning(f"File not found, skipping: {filepath}")
+            except Exception as e:
+                self.logger.error(f"Error processing artist {artist.title}: {e}. Continuing...", exc_info=True)
+
+        return all_tracks
+    
+    def get_track_by_id(self, track_id: str) -> Optional[Track]:
+        """Get a specific track by Plex rating key."""
+        try:
+            plex = PlexServer(self.plex_url, self.plex_token)
+            track = plex.fetchItem(int(track_id))
+            
+            # Get the first available part of the track
+            for part in track.iterParts():
+                filepath = Path(part.file)
+                if filepath.exists():
+                    return Track(
+                        artist=track.grandparentTitle or "Unknown Artist",
+                        album=track.parentTitle or "Unknown Album",
+                        title=track.title,
+                        filepath=filepath,
+                        media_server_rating_key=str(track.ratingKey),
+                        media_server_type=MediaServerType.PLEX
+                    )
+            
+            return None
+            
+        except Exception as e:
+            self.logger.error(f"Error getting track {track_id}: {e}", exc_info=True)
+            return None
+
+    def create_playlist(self, playlist: Playlist, batch_size: int = 100) -> Playlist:
+        """Create a playlist on the Plex server using batch processing for large playlists.
+        
+        Args:
+            playlist: The playlist to create
+            batch_size: Number of tracks to add per batch (default: 100)
+        """
+        try:
+            plex = PlexServer(self.plex_url, self.plex_token)
+            
+            # Get Plex track objects for all tracks in the playlist
+            plex_tracks = []
+            for track in playlist.tracks:
+                try:
+                    plex_track = plex.fetchItem(int(track.media_server_rating_key))
+                    plex_tracks.append(plex_track)
+                except Exception as e:
+                    self.logger.warning(f"Could not fetch track {track.media_server_rating_key}: {e}")
+                    continue
+            
+            if not plex_tracks:
+                raise ValueError("No valid tracks found for playlist creation")
+            
+            total_tracks = len(plex_tracks)
+            self.logger.info(f"Creating playlist '{playlist.name}' with {total_tracks} tracks")
+            
+            # Create the playlist with the first batch
+            first_batch = plex_tracks[:batch_size]
+            created_playlist = plex.createPlaylist(title=playlist.name, items=first_batch)
+            self.logger.info(f"Created playlist '{playlist.name}' with initial batch of {len(first_batch)} tracks")
+            
+            # Add remaining tracks in batches
+            remaining_tracks = plex_tracks[batch_size:]
+            if remaining_tracks:
+                self.logger.info(f"Adding {len(remaining_tracks)} remaining tracks in batches of {batch_size}")
+                
+                for i in range(0, len(remaining_tracks), batch_size):
+                    batch = remaining_tracks[i:i + batch_size]
+                    created_playlist.addItems(batch)
+                    self.logger.debug(f"Added batch {i//batch_size + 1}: {len(batch)} tracks")
+                
+                self.logger.info(f"Successfully completed playlist creation with all {total_tracks} tracks")
+            
+            # Return the playlist with server ID and creation time
+            return Playlist(
+                name=playlist.name,
+                tracks=playlist.tracks,
+                created_at=datetime.now(),
+                server_id=str(created_playlist.ratingKey)
+            )
+            
+        except Exception as e:
+            self.logger.error(f"Error creating playlist '{playlist.name}': {e}", exc_info=True)
+            raise