/pixeltable-mcp/audio-index/tools.py

import os
import logging
from typing import Tuple, Dict, Any, Optional

import pixeltable as pxt
from mcp.server.fastmcp import FastMCP
from pixeltable.functions import whisper
from pixeltable.functions.huggingface import sentence_transformer
from pixeltable.iterators.string import StringSplitter
from pixeltable.iterators import AudioSplitter

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('audio_index')

# Initialize MCP server
mcp = FastMCP("Pixeltable Audio Index")

# Constants
DIRECTORY = 'audio_index'
DEFAULT_CHUNK_DURATION = 30.0
DEFAULT_OVERLAP_DURATION = 2.0
DEFAULT_MIN_CHUNK_DURATION = 5.0
DEFAULT_EMBEDDING_MODEL = 'intfloat/e5-large-v2'
DEFAULT_WHISPER_MODEL = 'base.en'

# Registry to hold all audio indexes
# Format: {full_table_name: (audio_index, chunks_view, sentences_view)}
audio_indexes: Dict[str, Tuple[Any, Any, Any]] = {}


def _get_table_names(table_name: str) -> Tuple[str, str, str]:
    """Generate the full table and view names for a given table name.
    
    Args:
        table_name: Base name for the audio index
        
    Returns:
        Tuple of (full_table_name, chunks_view_name, sentences_view_name)
    """
    full_table_name = f'{DIRECTORY}.{table_name}'
    chunks_view_name = f'{DIRECTORY}.{table_name}_chunks'
    sentences_view_name = f'{DIRECTORY}.{table_name}_sentence_chunks'
    return full_table_name, chunks_view_name, sentences_view_name


def _load_existing_index(full_table_name: str, chunks_view_name: str, 
                         sentences_view_name: str) -> bool:
    """Load an existing audio index into the registry.
    
    Args:
        full_table_name: Full name of the audio index table
        chunks_view_name: Name of the chunks view
        sentences_view_name: Name of the sentences view
        
    Returns:
        True if the index was loaded successfully, False otherwise
    """
    try:
        audio_index = pxt.get_table(full_table_name)
        chunks_view = pxt.get_view(chunks_view_name)
        sentences_view = pxt.get_view(sentences_view_name)
        audio_indexes[full_table_name] = (audio_index, chunks_view, sentences_view)
        logger.info(f"Loaded existing audio index '{full_table_name}'")
        return True
    except Exception as e:
        logger.error(f"Failed to load existing audio index '{full_table_name}': {str(e)}")
        return False


def _get_openai_api_key() -> str:
    """Get OpenAI API key from environment variables.
    
    Returns:
        The OpenAI API key
        
    Raises:
        ValueError: If the API key is not found
    """
    api_key = os.getenv('OPENAI_API_KEY')
    if not api_key:
        raise ValueError("OPENAI_API_KEY not found in environment variables")
    return api_key


@mcp.tool()
def setup_audio_index(table_name: str) -> str:
    """Set up an audio index with the provided name and OpenAI API key.

    Args:
        table_name: The name of the audio index (e.g., 'podcasts', 'interviews').

    Returns:
        A message indicating whether the index was created, already exists, or failed.
    """
    global audio_indexes
    
    # Generate table names
    full_table_name, chunks_view_name, sentences_view_name = _get_table_names(table_name)
    
    try:
        # Set the API key
        openai_api_key = _get_openai_api_key()
        os.environ['OPENAI_API_KEY'] = openai_api_key
        logger.info(f"Setting up audio index '{full_table_name}'")

        # Check if the table already exists
        existing_tables = pxt.list_tables()
        if full_table_name in existing_tables:
            if _load_existing_index(full_table_name, chunks_view_name, sentences_view_name):
                return f"Audio index '{full_table_name}' already exists and is ready for use."
            else:
                return f"Audio index '{full_table_name}' exists but could not be loaded."

        # Create directory and table
        pxt.create_dir(DIRECTORY, if_exists='ignore')
        audio_index = pxt.create_table(full_table_name, {'audio_file': pxt.Audio}, if_exists='ignore')
        logger.info(f"Created audio index table '{full_table_name}'")

        # Create view for audio chunks
        chunks_view = pxt.create_view(
            chunks_view_name,
            audio_index,
            iterator=AudioSplitter.create(
                audio=audio_index.audio_file,
                chunk_duration_sec=DEFAULT_CHUNK_DURATION,
                overlap_sec=DEFAULT_OVERLAP_DURATION,
                min_chunk_duration_sec=DEFAULT_MIN_CHUNK_DURATION
            ),
            if_exists='ignore'
        )
        logger.info(f"Created audio chunks view '{chunks_view_name}'")

        # Add transcription to chunks
        chunks_view.add_computed_column(
            transcription=whisper.transcribe(audio=chunks_view.audio_chunk, model=DEFAULT_WHISPER_MODEL)
        )
        logger.info("Added transcription column to chunks view")

        # Create view that chunks transcriptions into sentences
        sentences_view = pxt.create_view(
            sentences_view_name,
            chunks_view,
            iterator=StringSplitter.create(text=chunks_view.transcription.text, separators='sentence'),
            if_exists='ignore'
        )
        logger.info(f"Created sentence chunks view '{sentences_view_name}'")

        # Define the embedding model and create embedding index
        embed_model = sentence_transformer.using(model_id=DEFAULT_EMBEDDING_MODEL)
        sentences_view.add_embedding_index(column='text', string_embed=embed_model)
        logger.info("Added embedding index to sentence chunks view")

        # Store in the registry
        audio_indexes[full_table_name] = (audio_index, chunks_view, sentences_view)
        return f"Audio index '{full_table_name}' created successfully."
    except Exception as e:
        logger.error(f"Error setting up audio index '{full_table_name}': {str(e)}")
        return f"Error setting up audio index '{full_table_name}': {str(e)}"


@mcp.tool()
def insert_audio(table_name: str, audio_location: str) -> str:
    """Insert an audio file into the specified audio index.

    Args:
        table_name: The name of the audio index (e.g., 'podcasts', 'interviews').
        audio_location: The URL or path to the audio file to insert (e.g., local path or S3 URL).

    Returns:
        A confirmation message indicating success or failure.
    """
    full_table_name, _, _ = _get_table_names(table_name)
    
    try:
        if full_table_name not in audio_indexes:
            logger.warning(f"Audio index '{full_table_name}' not set up")
            return f"Error: Audio index '{full_table_name}' not set up. Please call setup_audio_index first."
            
        audio_index, _, _ = audio_indexes[full_table_name]
        audio_index.insert([{'audio_file': audio_location}])
        logger.info(f"Inserted audio file '{audio_location}' into index '{full_table_name}'")
        return f"Audio file '{audio_location}' inserted successfully into index '{full_table_name}'."
    except Exception as e:
        logger.error(f"Error inserting audio file into '{full_table_name}': {str(e)}")
        return f"Error inserting audio file into '{full_table_name}': {str(e)}"


@mcp.tool()
def query_audio(table_name: str, query_text: str, top_n: int = 5) -> str:
    """Query the specified audio index with a text question.

    Args:
        table_name: The name of the audio index (e.g., 'podcasts', 'interviews').
        query_text: The question or text to search for in the audio content.
        top_n: Number of top results to return (default is 5).

    Returns:
        A string containing the top matching sentences and their similarity scores.
    """
    full_table_name, _, _ = _get_table_names(table_name)
    
    try:
        if full_table_name not in audio_indexes:
            logger.warning(f"Audio index '{full_table_name}' not set up")
            return f"Error: Audio index '{full_table_name}' not set up. Please call setup_audio_index first."
            
        _, _, sentences_view = audio_indexes[full_table_name]
        
        # Calculate similarity scores between query and sentences
        logger.info(f"Querying '{full_table_name}' with: '{query_text}'")
        sim = sentences_view.text.similarity(query_text)

        # Get top results
        results = (sentences_view.order_by(sim, asc=False)
                  .select(sentences_view.text, sim=sim, audio_file=sentences_view.audio_file)
                  .limit(top_n)
                  .collect())

        # Format the results
        result_str = f"Query Results for '{query_text}' in '{full_table_name}':\n\n"
        for i, row in enumerate(results.to_pandas().itertuples(), 1):
            result_str += f"{i}. Score: {row.sim:.4f}\n"
            result_str += f"   Text: {row.text}\n"
            result_str += f"   From audio: {row.audio_file}\n\n"
        
        return result_str if len(results) > 0 else "No results found."
    except Exception as e:
        logger.error(f"Error querying audio index '{full_table_name}': {str(e)}")
        return f"Error querying audio index '{full_table_name}': {str(e)}"


@mcp.tool()
def list_audio_tables(random_string: str = "") -> str:
    """List all audio indexes currently available.

    Returns:
        A string listing the current audio indexes.
    """
    try:
        tables = pxt.list_tables()
        audio_tables = [t for t in tables if t.startswith(f'{DIRECTORY}.') and not (
            t.endswith('_chunks') or t.endswith('_sentence_chunks')
        )]
        
        if not audio_tables:
            return "No audio indexes exist."
            
        # Load any tables that exist but aren't in our registry
        for table in audio_tables:
            if table not in audio_indexes:
                table_name = table.split('.')[-1]
                _, chunks_view_name, sentences_view_name = _get_table_names(table_name)
                _load_existing_index(table, chunks_view_name, sentences_view_name)
                
        return f"Current audio indexes: {', '.join(audio_tables)}"
    except Exception as e:
        logger.error(f"Error listing audio indexes: {str(e)}")
        return f"Error listing audio indexes: {str(e)}"