埋め込みからインサイトへ：TwelveLabs MarengoとS3ベクトルを活用した実践クロスモーダル検索

import boto3
import json
import numpy as np
import uuid
import time
import os
from typing import List, Dict
from botocore.exceptions import ClientError

# AWS Configuration
AWS_REGION = "us-east-1"  # Update to your region
AWS_PROFILE = "default"   # Update to your profile

# S3 Vectors Configuration
VECTOR_BUCKET_NAME = "marengo-vectors-" + str(uuid.uuid4())[:8]
VECTOR_INDEX_NAME = "embeddings-index"
VECTOR_DIMENSION = 1024  # Marengo embedding dimension

# Marengo Model
MODEL_ID = 'twelvelabs.marengo-embed-2-7-v1:0'

# Temporary S3 bucket for Bedrock output (required by Bedrock API)
TEMP_S3_BUCKET = "<YOUR_S3_BUCKET>"  # TODO: Replace with your S3 bucket name

print(f"Vector Bucket: {VECTOR_BUCKET_NAME}")
print(f"Vector Index: {VECTOR_INDEX_NAME}")
print(f"Model: {MODEL_ID}")

# Initialize AWS clients
session = boto3.Session(profile_name=AWS_PROFILE)
bedrock_client = session.client('bedrock-runtime', region_name=AWS_REGION)
s3_client = session.client('s3')
s3vectors_client = session.client('s3vectors', region_name=AWS_REGION)
print("AWS clients initialized")

# Create Vector Bucket
try:
    s3vectors_client.create_vector_bucket(
        vectorBucketName=VECTOR_BUCKET_NAME,
        encryptionConfiguration={'sseType': 'AES256'}
    )
    print(f"Vector bucket '{VECTOR_BUCKET_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Vector bucket already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

def upload_file_to_s3(local_path: str, bucket: str, key: str) -> str:
    """
    Upload a local file to S3
    
    Args:
        local_path: Path to local file
        bucket: S3 bucket name
        key: S3 object key (path in bucket)
    
    Returns:
        S3 URI of uploaded file
    """
    try:
        s3_client.upload_file(local_path, bucket, key)
        s3_uri = f"s3://{bucket}/{key}"
        print(f"✅ Uploaded {os.path.basename(local_path)} to {s3_uri}")
        return s3_uri
    except ClientError as e:
        print(f"❌ Error uploading file: {e}")
        raise

def generate_text_embedding(text: str) -> List[float]:
    """
    Generate embedding for text using Marengo on Bedrock
    """
    # Create unique output path
    output_prefix = f'embeddings/{uuid.uuid4()}'
    
    # Start async embedding generation
    response = bedrock_client.start_async_invoke(
        modelId=MODEL_ID,
        modelInput={
            "inputType": "text",
            "inputText": text
        },
        outputDataConfig={
            "s3OutputDataConfig": {
                "s3Uri": f's3://{TEMP_S3_BUCKET}/{output_prefix}'
            }
        }
    )
    
    invocation_arn = response["invocationArn"]
    print(f"Generating text embedding for: '{text[:50]}...'")
    
    # Wait for completion
    status = None
    while status not in ["Completed", "Failed"]:
        response = bedrock_client.get_async_invoke(invocationArn=invocation_arn)
        status = response['status']
        time.sleep(2)
    
    if status != "Completed":
        raise Exception(f"Embedding generation failed")
    
    # Retrieve embedding from S3
    response = s3_client.list_objects_v2(Bucket=TEMP_S3_BUCKET, Prefix=output_prefix)
    for obj in response.get('Contents', []):
        if obj['Key'].endswith('output.json'):
            result = s3_client.get_object(Bucket=TEMP_S3_BUCKET, Key=obj['Key'])
            data = json.loads(result['Body'].read())
            return data['data'][0]['embedding']
    
    raise Exception("No embedding output found")

def generate_video_embedding(video_path: str, start_sec: float = 0, length_sec: float = None) -> List[float]:
    """
    Generate embedding for video using Marengo on Bedrock
    """
    # Read video file and encode to base64
    with open(video_path, 'rb') as video_file:
        video_base64 = base64.b64encode(video_file.read()).decode('utf-8')
    
    # Build model input with optional time segments
    model_input = {
        "inputType": "video",
        "mediaSource": {"base64String": video_base64},
        "embeddingOption": ["visual-text", "audio"]  # Capture both visual and audio
    }
    
    if start_sec is not None:
        model_input["startSec"] = start_sec
    if length_sec is not None:
        model_input["lengthSec"] = length_sec
    
    # ... rest follows the same async pattern

def generate_video_embedding_from_s3(s3_uri: str, start_sec: float = 0, 
                                    length_sec: float = None) -> List[float]:
    """Generate embedding for video from S3 using Marengo on Bedrock"""
    output_prefix = f'{OUTPUT_PREFIX}video/{uuid.uuid4()}'
    
    model_input = {
        "inputType": "video",
        "mediaSource": {
            "s3Location": {
                "uri": s3_uri,
                "bucketOwner": ACCOUNT_ID
            }
        },
        "embeddingOption": ["visual-text", "audio"]
    }
    
    response = bedrock_client.start_async_invoke(
        modelId=MODEL_ID,
        modelInput=model_input,
        outputDataConfig={
            "s3OutputDataConfig": {
                "s3Uri": f's3://{TEMP_S3_BUCKET}/{output_prefix}',
                "bucketOwner": ACCOUNT_ID
            }
        }
    )
    
    print(f"🎬 Generating video embedding from S3: {s3_uri}")
    # Enhanced status monitoring with retry logic...

# Upload video to S3
video_s3_key = f"{MEDIA_PREFIX}videos/sample_video.mp4"
video_s3_uri = upload_file_to_s3('video.mp4', TEMP_S3_BUCKET, video_s3_key)

# Generate embedding from S3
video_embedding_s3 = generate_video_embedding_from_s3(video_s3_uri, start_sec=0, length_sec=10)

# Compare VIDEO embeddings
print("🎬 VIDEO EMBEDDING VERIFICATION")
print("-" * 50)

# Generate using Base64
print("Generating via Base64...")
video_base64_emb = generate_video_embedding('video.mp4', start_sec=0, length_sec=10)

# Upload and generate using S3 URI
print("Generating via S3 URI...")
video_s3_key = f"{MEDIA_PREFIX}verification/video.mp4"
video_s3_uri = upload_file_to_s3('video.mp4', TEMP_S3_BUCKET, video_s3_key)
video_s3_emb = generate_video_embedding_from_s3(video_s3_uri, start_sec=0, length_sec=10)

# Calculate similarity
similarity = cosine_similarity(video_base64_emb, video_s3_emb)
print(f"📊 Results:")
print(f"  Cosine similarity: {similarity:.6f}")
print(f"  Are they identical? {'✅ YES' if similarity > 0.9999 else '❌ NO'}")

# Visualize text embeddings
if text_embeddings:
    plt.figure(figsize=(12, 3))
    for i, (text, emb, _) in enumerate(text_embeddings[:3]):
        plt.subplot(1, 3, i+1)
        plt.imshow([emb[:100]], aspect='auto', cmap='coolwarm')
        plt.title(f'Text {i+1}', fontsize=10)
        plt.xlabel('Dimensions (first 100)')
        plt.colorbar(orientation='horizontal', pad=0.1)
    plt.suptitle('Text Embeddings Visualization', fontsize=12)
    plt.tight_layout()
    plt.show()

def store_embeddings(embeddings_data: List[tuple]) -> bool:
    """
    Store embeddings in S3 Vectors with metadata
    """
    vectors_to_insert = []
    
    for i, item in enumerate(embeddings_data):
        # Support (text, embedding) and (text, embedding, media_type)
        if len(item) == 3:
            text, embedding, media_type = item
        else:
            text, embedding = item
            media_type = "text"
        
        vector_entry = {
            'key': f'vector_{i:04d}',           # deterministic, idempotent upserts
            'data': {
                'float32': [float(v) for v in embedding]  # convert to float32-compatible list
            },
            'metadata': {
                'text': text,                   # human-readable description
                'media_type': media_type,       # useful for UI badges and filtering
                'id': i                         # numeric id for quick reference
            }
        }
        vectors_to_insert.append(vector_entry)

    try:
        s3vectors_client.put_vectors(
            vectorBucketName=VECTOR_BUCKET_NAME,
            indexName=VECTOR_INDEX_NAME,
            vectors=vectors_to_insert
        )
        print(f"Stored {len(vectors_to_insert)} vectors in S3 Vectors")
        
        # Optional: quick summary
        media_counts = {}
        for item in embeddings_data:
            mtype = item[2] if len(item) == 3 else "text"
            media_counts[mtype] = media_counts.get(mtype, 0) + 1
        print("Media types stored:")
        for m, c in media_counts.items():
            print(f"  - {m}: {c}")
        return True
    except ClientError as e:
        print(f"Error storing vectors: {e}")
        return False

def search_similar(query_text: str = None, query_embedding: List[float] = None, top_k: int = 3) -> List[Dict]:
    """
    Search for similar vectors using either a text query or a pre-computed embedding
    """
    # Generate embedding for query if text is provided
    if query_text and not query_embedding:
        print(f"\nSearching for: '{query_text}'")
        query_embedding = generate_text_embedding(query_text)
    elif query_embedding:
        print(f"\nSearching with provided embedding...")
    else:
        raise ValueError("Either query_text or query_embedding must be provided")
    
    # Query S3 Vectors
    try:
        response = s3vectors_client.query_vectors(
            vectorBucketName=VECTOR_BUCKET_NAME,
            indexName=VECTOR_INDEX_NAME,
            topK=top_k,
            queryVector={'float32': [float(v) for v in query_embedding]},
            returnMetadata=True,
            returnDistance=True
        )
        
        results = []
        for vector in response.get('vectors', []):
            # Handle both old vectors (without media_type) and new vectors (with media_type)
            results.append({
                'text': vector['metadata'].get('text', 'No description'),
                'media_type': vector['metadata'].get('media_type', 'text'),  # Default to 'text' for old vectors
                'similarity': 1.0 - vector.get('distance', 0),  # Convert distance to similarity
                'key': vector.get('key')
            })
        
        return results
        
    except ClientError as e:
        print(f"Query failed: {e}")
        return []