semantic_search.py

"""ChromaDB semantic search module for music memories app."""

import chromadb
from sentence_transformers import SentenceTransformer

# Initialize the embedding model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# Initialize ChromaDB client with persistent storage
chroma_client = chromadb.PersistentClient(path="./chroma_db")

# Create collections for different semantic search types
song_vibes_collection = chroma_client.get_or_create_collection(name="song_vibes")
memory_vibes_collection = chroma_client.get_or_create_collection(name="memory_vibes")
context_vibes_collection = chroma_client.get_or_create_collection(name="context_vibes")
playlist_journeys_collection = chroma_client.get_or_create_collection(name="playlist_journeys")


# ============== SONG VIBES ==============

def add_song_vibe(song_id: int, title: str, artist: str, lyrics: str = "") -> None:
    """Add song embedding based on lyrics/title/artist."""
    text = f"{title} by {artist} - {lyrics}".strip()
    song_vibes_collection.add(
        ids=[f"song_{song_id}"],
        documents=[text],
        metadatas=[{"id": song_id, "title": title, "artist": artist}],
    )


def search_song_vibes(query: str, n_results: int = 5) -> list[dict]:
    """Search songs by vibe/lyrics."""
    results = song_vibes_collection.query(
        query_texts=[query],
        n_results=n_results,
        include=["documents", "metadatas", "distances"],
    )
    if not results["metadatas"] or not results["metadatas"][0]:
        return []
    return [
        {"id": m["id"], "title": m["title"], "artist": m["artist"],
         "document": results["documents"][0][i], "distance": results["distances"][0][i]}
        for i, m in enumerate(results["metadatas"][0])
    ]


def remove_song_vibe(song_id: int) -> None:
    """Remove a song vibe."""
    song_vibes_collection.delete(ids=[f"song_{song_id}"])


# ============== MEMORY VIBES ==============

def add_memory_vibe(memory_id: int, user_id: int, description: str) -> None:
    """Add memory embedding based on description."""
    memory_vibes_collection.add(
        ids=[f"memory_{memory_id}"],
        documents=[description],
        metadatas=[{"id": memory_id, "user_id": user_id}],
    )


def search_memory_vibes(query: str, n_results: int = 5) -> list[dict]:
    """Search memories by description."""
    results = memory_vibes_collection.query(
        query_texts=[query],
        n_results=n_results,
        include=["documents", "metadatas", "distances"],
    )
    if not results["metadatas"] or not results["metadatas"][0]:
        return []
    return [
        {"id": m["id"], "user_id": m["user_id"],
         "document": results["documents"][0][i], "distance": results["distances"][0][i]}
        for i, m in enumerate(results["metadatas"][0])
    ]


def remove_memory_vibe(memory_id: int) -> None:
    """Remove a memory vibe."""
    memory_vibes_collection.delete(ids=[f"memory_{memory_id}"])


# ============== CONTEXT VIBES ==============

def add_context_vibe(context_id: int, weather: str, time_of_day: str, location_type: str) -> None:
    """Add context embedding."""
    text = f"Weather: {weather}, Time: {time_of_day}, Location: {location_type}".strip()
    context_vibes_collection.add(
        ids=[f"context_{context_id}"],
        documents=[text],
        metadatas=[{"id": context_id, "weather": weather, "time_of_day": time_of_day, "location_type": location_type}],
    )


def search_context_vibes(query: str, n_results: int = 5) -> list[dict]:
    """Search contexts by description."""
    results = context_vibes_collection.query(
        query_texts=[query],
        n_results=n_results,
        include=["documents", "metadatas", "distances"],
    )
    if not results["metadatas"] or not results["metadatas"][0]:
        return []
    return [
        {"id": m["id"], "weather": m["weather"], "time_of_day": m["time_of_day"],
         "location_type": m["location_type"], "document": results["documents"][0][i],
         "distance": results["distances"][0][i]}
        for i, m in enumerate(results["metadatas"][0])
    ]


def remove_context_vibe(context_id: int) -> None:
    """Remove a context vibe."""
    context_vibes_collection.delete(ids=[f"context_{context_id}"])


# ============== PLAYLIST JOURNEYS ==============

def add_playlist_journey(playlist_id: int, name: str, mood_description: str) -> None:
    """Add playlist journey embedding for mood transitions."""
    playlist_journeys_collection.add(
        ids=[f"playlist_{playlist_id}"],
        documents=[f"{name}: {mood_description}"],
        metadatas=[{"id": playlist_id, "name": name}],
    )


def search_playlist_journeys(query: str, n_results: int = 5) -> list[dict]:
    """Search playlists by mood/vibe."""
    results = playlist_journeys_collection.query(
        query_texts=[query],
        n_results=n_results,
        include=["documents", "metadatas", "distances"],
    )
    if not results["metadatas"] or not results["metadatas"][0]:
        return []
    return [
        {"id": m["id"], "name": m["name"],
         "document": results["documents"][0][i], "distance": results["distances"][0][i]}
        for i, m in enumerate(results["metadatas"][0])
    ]


def remove_playlist_journey(playlist_id: int) -> None:
    """Remove a playlist journey."""
    playlist_journeys_collection.delete(ids=[f"playlist_{playlist_id}"])