lambda_function.py

#!/usr/bin/env python
"""
Serverless-compatible Lambda for embedding videos → OpenSearch
• Idempotent via logical _id
• No delete_by_query
• Bulk with refresh=False, then one final refresh
"""

import os, json, base64, logging, pathlib, boto3, numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List
from dotenv import load_dotenv
from opensearchpy import OpenSearch, helpers, RequestsHttpConnection
from requests_aws4auth import AWS4Auth
from supabase import create_client

load_dotenv()

# ─── 1. CONFIG ─────────────────────────────────────────────────────────
S3_BUCKET_FRAMES  = os.getenv("S3_BUCKET_FRAMES", "oriane-contents")
MODEL_ID          = os.getenv("MODEL_ID", "amazon.titan-embed-image-v1")
AWS_REGION        = os.getenv("AWS_REGION", "us-east-1")
EMB_DIM           = int(os.getenv("EMB_DIM", "1024"))
MAX_BATCH         = int(os.getenv("MAX_FRAMES_PER_BATCH", "20"))
CONCURRENCY_LIMIT = int(os.getenv("CONCURRENCY_LIMIT", "4"))
SUPABASE_URL      = os.getenv("SUPABASE_URL")
SUPABASE_KEY      = os.getenv("SUPABASE_KEY")
OS_ENDPOINT       = os.getenv("OS_ENDPOINT")
LOG_LEVEL         = os.getenv("LOG_LEVEL", "INFO").upper()

for var in ("SUPABASE_URL","SUPABASE_KEY","OS_ENDPOINT"):
    if not globals()[var]:
        raise RuntimeError(f"{var} is required")

# ─── 2. LOGGING & CLIENTS ───────────────────────────────────────────────
logging.basicConfig(level=LOG_LEVEL, format="%(asctime)s %(levelname)s %(message)s")
def get_log(code: str):
    lg = logging.getLogger(f"embed[{code}]")
    if not lg.handlers:
        h = logging.StreamHandler()
        h.setFormatter(logging.Formatter("%(asctime)s %(levelname)s [%(name)s] %(message)s"))
        lg.addHandler(h)
    lg.setLevel(LOG_LEVEL)
    return lg

session = boto3.Session(region_name=AWS_REGION)
s3       = session.client("s3")
bedrock  = session.client("bedrock-runtime")
supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

creds  = session.get_credentials()
auth   = AWS4Auth(creds.access_key, creds.secret_key, AWS_REGION, "aoss", session_token=creds.token)
host   = OS_ENDPOINT.split("://",1)[-1].rstrip("/")
os_client = OpenSearch(
    hosts=[{"host": host, "port": 443}],
    http_auth=auth,
    use_ssl=True,
    verify_certs=True,
    connection_class=RequestsHttpConnection,
)

# ─── 3. HELPERS ─────────────────────────────────────────────────────────
def list_frame_keys(platform: str, code: str) -> List[str]:
    pfx = f"{platform}/{code}/frames/"
    pages = s3.get_paginator("list_objects_v2").paginate(Bucket=S3_BUCKET_FRAMES, Prefix=pfx)
    return [o["Key"] for pg in pages for o in pg.get("Contents",[]) if o["Key"].endswith(".jpg")]

def titan_embed(b64s: List[str]) -> List[List[float]]:
    out = []
    for b64 in b64s:
        body = json.dumps({"inputImage": b64, "embeddingConfig": {"outputEmbeddingLength": EMB_DIM}})
        rsp  = bedrock.invoke_model(modelId=MODEL_ID, body=body,
                                    accept="application/json", contentType="application/json")
        emb  = json.loads(rsp["body"].read())["embedding"]
        out.append(emb)
    return out

def bulk_gen(idx: str, docs: List[dict]):
    for d in docs:
        yield {"index": {"_index": idx, "_id": d["_id"]}}
        yield d

def mark_video(code: str, **fields):
    supabase.table("insta_content").update(fields).eq("code",code).execute()

# ─── 4. PER-VIDEO EMBED ─────────────────────────────────────────────────
def embed_video(platform: str, code: str):
    log = get_log(code)
    log.info("Starting embedding")
    keys = list_frame_keys(platform, code)
    if not keys:
        raise RuntimeError(f"No frames for {code}")

    video_id = f"{platform}.{code}"
    frame_docs, vecs = [], []
    for i in range(0, len(keys), MAX_BATCH):
        chunk = keys[i:i+MAX_BATCH]
        try:
            b64s   = [base64.b64encode(s3.get_object(Bucket=S3_BUCKET_FRAMES,Key=k)["Body"].read()).decode() for k in chunk]
            embeds = titan_embed(b64s)
        except Exception as e:
            for k in chunk:
                idx = int(pathlib.Path(k).stem)
                supabase.table("embedding_errors").insert({"code":code,"frame":idx,"error":str(e)}).execute()
            raise
        for k, v in zip(chunk, embeds):
            fno = int(pathlib.Path(k).stem)
            frame_docs.append({
                "_id":      f"{video_id}#{fno}",
                "video_id": video_id,
                "vector":   v,
                "platform": platform,
                "code":     code,
                "frame":    fno,
            })
            vecs.append(v)

    # upsert frames batch
    helpers.bulk(os_client, bulk_gen("video_frames", frame_docs), refresh=False)

    # upsert summary
    summary = {
        "_id":      video_id,
        "video_id": video_id,
        "vector":   np.mean(vecs, axis=0).tolist(),
        "platform": platform,
        "code":     code,
        "frames":   len(vecs),
    }
    helpers.bulk(os_client, bulk_gen("videos", [summary]), refresh=False)

    mark_video(code, is_embedded=True)
    log.info("Indexed %d frames", len(vecs))

# ─── 5. RECORD HANDLER ───────────────────────────────────────────────────
def already_embedded(code: str) -> bool:
    res = supabase.table("insta_content").select("is_embedded").eq("code",code).maybe_single().execute()
    return bool(res.data and res.data.get("is_embedded"))

def process_record(rec: dict):
    body = json.loads(rec.get("body","{}"))
    platform, code = body.get("platform"), body.get("code")
    if not platform or not code:
        return {"code": None, "status":"invalid"}
    if already_embedded(code):
        return {"code":code, "status":"skipped"}
    embed_video(platform, code)
    return {"code":code, "status":"done"}

# ─── 6. LAMBDA HANDLER ──────────────────────────────────────────────────
def lambda_handler(event, _):
    results = []
    with ThreadPoolExecutor(max_workers=CONCURRENCY_LIMIT) as pool:
        for fut in as_completed([pool.submit(process_record,r) for r in event.get("Records",[])]):
            results.append(fut.result())
    # one final refresh
    os_client.indices.refresh("video_frames")
    os_client.indices.refresh("videos")
    return {"status":"completed","results":results}