wehub-resource-sync
184 lines
7.1 KiB
Python
184 lines
7.1 KiB
Python
import pixeltable as pxt
|
|
import os
|
|
from mcp.server.fastmcp import FastMCP
|
|
from pixeltable.functions import openai
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
from pixeltable.functions.video import extract_audio
|
|
from pixeltable.iterators import AudioSplitter
|
|
from pixeltable.iterators.string import StringSplitter
|
|
from datetime import datetime
|
|
|
|
mcp = FastMCP("Pixeltable")
|
|
|
|
# Base directory for all indexes
|
|
DIRECTORY = 'video_index'
|
|
|
|
# Registry to hold all video indexes
|
|
video_indexes = {}
|
|
|
|
|
|
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_video_index(table_name: str) -> str:
|
|
"""Set up a video index with the provided name and OpenAI API key.
|
|
|
|
Args:
|
|
table_name: The name of the video index (e.g., 'lectures', 'interviews').
|
|
|
|
Returns:
|
|
A message indicating whether the index was created, already exists, or failed.
|
|
"""
|
|
global video_indexes
|
|
|
|
# Construct full table and view names
|
|
full_table_name = f'{DIRECTORY}.{table_name}'
|
|
chunks_view_name = f'{DIRECTORY}.{table_name}_chunks'
|
|
sentences_view_name = f'{DIRECTORY}.{table_name}_sentence_chunks'
|
|
|
|
try:
|
|
# Set the API key
|
|
openai_api_key = _get_openai_api_key()
|
|
os.environ['OPENAI_API_KEY'] = openai_api_key
|
|
|
|
# Check if the table already exists
|
|
existing_tables = pxt.list_tables()
|
|
if full_table_name in existing_tables:
|
|
video_index = pxt.get_table(full_table_name)
|
|
chunks_view = pxt.get_table(chunks_view_name)
|
|
sentences_view = pxt.get_table(sentences_view_name)
|
|
video_indexes[full_table_name] = (video_index, chunks_view, sentences_view)
|
|
return f"Video index '{full_table_name}' already exists and is ready for use."
|
|
|
|
# Create directory and table
|
|
pxt.create_dir(DIRECTORY, if_exists='ignore')
|
|
video_index = pxt.create_table(
|
|
full_table_name,
|
|
{'video_file': pxt.Video, 'uploaded_at': pxt.Timestamp},
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Extract audio from video
|
|
video_index.add_computed_column(
|
|
audio_extract=extract_audio(video_index.video_file, format='mp3')
|
|
)
|
|
|
|
# Create view for audio chunks
|
|
chunks_view = pxt.create_view(
|
|
chunks_view_name,
|
|
video_index,
|
|
iterator=AudioSplitter.create(
|
|
audio=video_index.audio_extract,
|
|
chunk_duration_sec=30.0,
|
|
overlap_sec=2.0,
|
|
min_chunk_duration_sec=5.0
|
|
),
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Add transcription to chunks
|
|
chunks_view.add_computed_column(
|
|
transcription=openai.transcriptions(audio=chunks_view.audio_chunk, model='whisper-1')
|
|
)
|
|
|
|
# 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'
|
|
)
|
|
|
|
# Define the embedding model and create embedding index
|
|
embed_model = sentence_transformer.using(model_id='intfloat/e5-large-v2')
|
|
sentences_view.add_embedding_index(column='text', string_embed=embed_model)
|
|
|
|
# Store in the registry
|
|
video_indexes[full_table_name] = (video_index, chunks_view, sentences_view)
|
|
return f"Video index '{full_table_name}' created successfully."
|
|
except Exception as e:
|
|
return f"Error setting up video index '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def insert_video(table_name: str, video_location: str) -> str:
|
|
"""Insert a video file into the specified video index.
|
|
|
|
Args:
|
|
table_name: The name of the video index (e.g., 'lectures', 'interviews').
|
|
video_location: The URL or path to the video file to insert (e.g., local path or S3 URL).
|
|
|
|
Returns:
|
|
A confirmation message indicating success or failure.
|
|
"""
|
|
full_table_name = f'{DIRECTORY}.{table_name}'
|
|
try:
|
|
if full_table_name not in video_indexes:
|
|
return f"Error: Video index '{full_table_name}' not set up. Please call setup_video_index first."
|
|
video_index, _, _ = video_indexes[full_table_name]
|
|
video_index.insert([{'video_file': video_location, 'uploaded_at': datetime.now()}])
|
|
return f"Video file '{video_location}' inserted successfully into index '{full_table_name}'."
|
|
except Exception as e:
|
|
return f"Error inserting video file into '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def query_video(table_name: str, query_text: str, top_n: int = 5) -> str:
|
|
"""Query the specified video index with a text question.
|
|
|
|
Args:
|
|
table_name: The name of the video index (e.g., 'lectures', 'interviews').
|
|
query_text: The question or text to search for in the video 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 = f'{DIRECTORY}.{table_name}'
|
|
try:
|
|
if full_table_name not in video_indexes:
|
|
return f"Error: Video index '{full_table_name}' not set up. Please call setup_video_index first."
|
|
_, _, sentences_view = video_indexes[full_table_name]
|
|
# Calculate similarity scores between query and sentences
|
|
sim = sentences_view.text.similarity(query_text)
|
|
|
|
# Get top results
|
|
results = (sentences_view.order_by(sim, asc=False)
|
|
.select(sentences_view.text, sim=sim, video_file=sentences_view.video_file,
|
|
uploaded_at=sentences_view.uploaded_at)
|
|
.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 video: {row.video_file}\n"
|
|
result_str += f" Uploaded: {row.uploaded_at}\n\n"
|
|
|
|
return result_str if result_str else "No results found."
|
|
except Exception as e:
|
|
return f"Error querying video index '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def list_video_tables() -> str:
|
|
"""List all video indexes currently available.
|
|
|
|
Returns:
|
|
A string listing the current video indexes.
|
|
"""
|
|
tables = pxt.list_tables()
|
|
video_tables = [t for t in tables if t.startswith(f'{DIRECTORY}.')]
|
|
return f"Current video indexes: {', '.join(video_tables)}" if video_tables else "No video indexes exist." |