New commit

Author: dchiueh

.DS_Store

@@ -0,0 +1,23 @@
+name: Deploy Lambda to AWS
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  prod-project-north-fork:
+    uses: MinneapolisStarTribune/strib-workflows/.github/workflows/
+    with:
+      terraform-generic-repo-name: terraform-generic-lambda
+      terraform-generic-repo-ref: v0.4.0
+      terraform-tfvars-file-path: "./config/terraform/prod-semantic-image-search.tfvars"
+      terraform-backend-bucket: "terragrunt-strib-terraform-state-chad-wagner-us-west-2"
+      aws-account-id: "903601045739"
+      aws-environment: prod
+      aws-region: us-west-2
+    secrets: inherit

.github/workflows/lambda-cicd.yml

@@ -0,0 +1,55 @@
+# Semantic Image Search
+
+This repository contains a prototype for captioning images, embedding text, and serving semantic image search results via an AWS Lambda function. It includes helper scripts for data ingestion, model training, and a minimal Chrome extension.
+
+## Repository Layout
+
+- **`backend/`** – Data preparation and training scripts.
+  - `alttext-embed.py` – Generates captions for images with GPT-4-Vision and writes image embeddings.
+  - `ingest.py` – Fetches ArcXP story revisions to build story–image pairs.
+  - `copy_to_rds.py` – Bulk loads embeddings into a PostgreSQL database using pgvector.
+  - `train.py` – Trains a LightGBM ranking model from the collected data.
+- **`lambda/`** – AWS Lambda code used by the search API.
+  - `handler.py` – Entry point that scores candidate images against a query.
+  - `ranker.py` – Downloads the trained model and manages database connections.
+- **`extension/`** – Simple Chrome extension that queries `/suggest` for lead-art recommendations.
+- **`tests/`** – Unit tests for the handler and ranking helper.
+- **`data/`** – Placeholder directory for generated embeddings and the trained model.
+
+## Getting Started
+
+1. **Install dependencies**
+   ```bash
+   pip install -r backend/requirements.txt
+   ```
+
+2. **Ingest story–image pairs**
+   ```bash
+   scripts/run_ingest.sh --since 2024-01-01T00:00:00Z
+   ```
+   Set `ARC_API_URL` and `ARC_TOKEN` to connect to the Arc API.
+
+3. **Caption and embed images**
+   ```bash
+   scripts/run_alttext.sh <image-id> [...]
+   ```
+   Requires `PHOTO_URL` and `PHOTO_TOKEN` for fetching image bytes. Results are written to `artifacts/images.csv`.
+
+4. **Load embeddings into PostgreSQL**
+   ```bash
+   PG_DSN=postgresql://user:pass@host/db python backend/copy_to_rds.py
+   ```
+
+5. **Train the ranking model**
+   ```bash
+   scripts/run_train.sh
+   ```
+   Produces `artifacts/leadart_ranker.lgb`.
+
+6. **Run tests**
+   ```bash
+   pytest
+   ```
+
+The Lambda function (`lambda/handler.py`) uses the trained model to rank candidates stored in the database. The Chrome extension demonstrates how search suggestions can be integrated into Arc Composer.
+

README.md

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import argparse
+import base64
+import csv
+import os
+import json
+from pathlib import Path
+from typing import Iterable, List, Optional, Tuple
+import logging
+
+from common.utils import embed_text
+import requests
+import openai
+
+PHOTO_URL = os.environ.get("PHOTO_URL", "https://photo.example.com")
+PHOTO_TOKEN = os.environ.get("PHOTO_TOKEN")
+
+
+
+def fetch_image(image_id: str) -> Tuple[bytes, str, str]:
+    """Retrieve image bytes and metadata from Merlin/Photo Center."""
+    if requests is None:
+        return os.urandom(10), "", ""
+    url = f"{PHOTO_URL}/{image_id}"
+    headers = {"Authorization": f"Bearer {PHOTO_TOKEN}"} if PHOTO_TOKEN else None
+    try:
+        resp = requests.get(url, headers=headers, timeout=30)
+        resp.raise_for_status()
+        taken_at = resp.headers.get("X-Taken-At", "")
+        uploaded_at = resp.headers.get("X-Uploaded-At", "")
+        return resp.content, taken_at, uploaded_at
+    except Exception as exc:
+        logging.error("Failed to fetch image %s: %s", image_id, exc)
+        return b"", ""
+
+
+def caption_image(image_id: str, image_bytes: Optional[bytes] = None) -> str:
+    """Call a vision-capable model to produce a caption."""
+    try:
+        image_bytes = fetch_image_bytes(image_id)
+        b64 = base64.b64encode(image_bytes).decode()
+        resp = openai.ChatCompletion.create(
+            model="gpt-4-vision-preview",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Describe this image in one sentence."},
+                        {"type": "image_url", "image_url": f"data:image/jpeg;base64,{b64}"},
+                    ],
+                }
+            ],
+        )
+        return resp["choices"][0]["message"]["content"].strip()
+    except Exception as exc:
+        logging.error("Failed to caption image %s: %s", image_id, exc)
+        return f"Image {image_id} caption"
+
+def process(images: Iterable[str]) -> None:
+    Path("artifacts").mkdir(exist_ok=True)
+    out_path = Path("artifacts/images.csv")
+    with out_path.open("w", newline="") as f:
+        writer = csv.writer(f)
+        header = ["image_id", "caption", "embedding", "taken_at", "uploaded_at"]
+        writer.writerow(header)
+        for image_id in images:
+            image_bytes, taken_at, uploaded_at = fetch_image(image_id)
+            caption = caption_image(image_id, image_bytes)
+            vector = embed_text(caption)
+            writer.writerow([
+                image_id,
+                caption,
+                json.dumps(vector),
+                taken_at,
+                uploaded_at,
+            ])
+    print(f"Wrote {out_path}")
+
+
+def main(args: Iterable[str] | None = None) -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("image_ids", nargs="+", help="IDs to caption")
+    ns = parser.parse_args(args)
+    process(ns.image_ids)
+
+
+if __name__ == "__main__":
+    main()

backend/alttext-embed.py

@@ -0,0 +1,27 @@
+"""Bulk load image embeddings into a PostgreSQL database."""
+
+import csv
+import os
+
+import psycopg2
+
+DSN = os.environ.get("PG_DSN")
+TABLE = os.environ.get("PG_TABLE", "images")
+
+
+def copy_csv(path: str) -> None:
+    if not DSN:
+        raise RuntimeError("PG_DSN not configured")
+    conn = psycopg2.connect(DSN)
+    cur = conn.cursor()
+    cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+    with open(path, "r") as f:
+        cur.copy_expert(f"COPY {TABLE} FROM STDIN WITH CSV HEADER", f)
+    conn.commit()
+    cur.close()
+    conn.close()
+    print(f"Copied {path} to {TABLE}")
+
+
+if __name__ == "__main__":
+    copy_csv("artifacts/images.csv")

backend/copy_to_rds.py

@@ -0,0 +1,143 @@
+import argparse
+import csv
+import logging
+import os
+from pathlib import Path
+from typing import Iterable, Set, Dict, Any, Optional
+import requests
+from pymongo import MongoClient
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+API_URL = os.environ.get("ARC_API_URL")
+API_TOKEN = os.environ.get("ARC_TOKEN")
+
+MONGO_URI = os.environ.get("MONGO_URI")
+MONGO_DB = os.environ.get("MONGO_DB")
+MONGO_COLLECTION = os.environ.get("MONGO_COLLECTION")
+
+
+def load_photo_editor_ids(path: str) -> Set[str]:
+    with open(path, newline="") as f:
+        return {row[0] for row in csv.reader(f) if row}
+
+
+def fetch_revisions(story_id: str) -> Iterable[Dict[str, Any]]:
+    if not API_TOKEN:
+        raise RuntimeError("ARC_TOKEN not configured")
+    url = f"{API_URL}/stories/{story_id}/revisions"
+    headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    try:
+        response = requests.get(url, headers=headers, timeout=30)
+        response.raise_for_status()
+        return response.json().get("revisions", [])
+    except Exception as exc:
+        logging.error("Failed to fetch revisions for %s: %s", story_id, exc)
+        return []
+
+
+def story_has_photo_editor(revisions: Iterable[Dict[str, Any]], editor_ids: Set[str]) -> bool:
+    return any(rev.get("user_id") in editor_ids for rev in revisions)
+
+def first_lead_art_revision(
+    revisions: Iterable[Dict[str, Any]],
+    editors: Optional[Set[str]] = None,
+) -> Optional[Dict[str, Any]]:
+    """Return the first revision containing lead art.
+
+    If ``editors`` is provided, only revisions by those users are considered.
+    """
+    for rev in sorted(revisions, key=lambda r: r.get("revision_date", "")):
+        if editors is not None and rev.get("user_id") not in editors:
+            continue
+        promo = rev.get("story", {}).get("promo_items", {})
+        if promo.get("lead_art"):
+            return rev
+    return None
+
+
+def ingest_story(story_id: str) -> Optional[Dict[str, Any]]:
+    """Return metadata for a story including a label based on editor."""
+    revisions = list(fetch_revisions(story_id))
+    rev = first_lead_art_revision(revisions, PHOTO_EDITOR_IDS)
+    label = 2
+    if not rev:
+        rev = first_lead_art_revision(revisions)
+        label = 1
+    if not rev:
+        return None
+    story = rev.get("story", {})
+    image_id = story.get("promo_items", {}).get("lead_art", {}).get("id")
+    text = story.get("text", "")
+    timestamp = rev.get("revision_date")
+    if isinstance(timestamp, str):
+        timestamp = timestamp[:19]
+    return {
+        "story_id": story_id,
+        "text": text,
+        "image_id": image_id,
+        "timestamp": timestamp,
+        "label": label,
+    }
+
+
+def fetch_image_id_from_mongo(story_id: str) -> str:
+    client = MongoClient(MONGO_URI)
+    db = client[MONGO_DB]
+    collection = db[MONGO_COLLECTION]
+    doc = collection.find_one({"_id": story_id}, {"promo_items.lead_art.id": 1})
+    client.close()
+    if doc:
+        return doc.get("promo_items", {}).get("lead_art", {}).get("id", "")
+    return ""
+
+
+def fetch_story_ids_since(since_date: str) -> Iterable[str]:
+    client = MongoClient(MONGO_URI)
+    db = client[MONGO_DB]
+    collection = db[MONGO_COLLECTION]
+    query = {"last_updated_date": {"$gte": since_date}}
+    projection = {"_id": 1}
+    story_ids = [doc["_id"] for doc in collection.find(query, projection)]
+    client.close()
+    return story_ids
+
+
+def process_story(story_id: str, editor_ids: Set[str]) -> Dict[str, str] | None:
+    revisions = list(fetch_revisions(story_id))
+    if story_has_photo_editor(revisions, editor_ids):
+        image_id = fetch_image_id_from_mongo(story_id)
+        return {"story_id": story_id, "image_id": image_id}
+    return None
+
+
+def main(args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--since", required=True, help="ISO date to start from (e.g., 2021-01-01)")
+    parser.add_argument("--editor-csv", default="editor_ids.csv", help="CSV file of photo editor user IDs")
+    parser.add_argument("--workers", type=int, default=8, help="Number of parallel threads")
+    ns = parser.parse_args(args)
+
+    photo_editor_ids = load_photo_editor_ids(ns.editor_csv)
+    story_ids = fetch_story_ids_since(ns.since)
+
+    Path("artifacts").mkdir(exist_ok=True)
+    matched_path = Path("artifacts/story_image_pairs.csv")
+
+    matched_rows = []
+    with ThreadPoolExecutor(max_workers=ns.workers) as executor:
+        future_to_story = {executor.submit(process_story, story_id, photo_editor_ids): story_id for story_id in story_ids}
+        for future in as_completed(future_to_story):
+            result = future.result()
+            if result:
+                matched_rows.append(result)
+
+    with matched_path.open("w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=["story_id", "image_id"])
+        writer.writeheader()
+        writer.writerows(matched_rows)
+
+    print(f"Wrote matched story IDs and image IDs to {matched_path}")
+
+
+if __name__ == "__main__":
+    main()

backend/ingest.py

@@ -0,0 +1,7 @@
+lightgbm
+openai
+psycopg2-binary
+pgvector
+pandas
+tqdm
+requests

backend/requirements.txt

@@ -0,0 +1,10 @@
+CREATE TABLE IF NOT EXISTS images (
+    image_id TEXT PRIMARY KEY,
+    caption  TEXT,
+    embedding VECTOR(768),
+    taken_at TIMESTAMP,
+    uploaded_at TIMESTAMP
+);
+
+-- Create ivfflat index for ANN search
+CREATE INDEX IF NOT EXISTS images_embedding_idx ON images USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);