queue_consumer.py

#!/usr/bin/env python3
"""RabbitMQ consumer for DevKG pipeline jobs.

Listens on the `devkg_jobs` queue, processes Claude Code session transcripts
into RDF Turtle, and uploads them to Fuseki.
"""

import base64
import hashlib
import json
import os
import sys
import tempfile
import time
import traceback
from pathlib import Path

import pika


def _init_vertex_credentials():
    """Decode GOOGLE_APPLICATION_CREDENTIALS_BASE64 to a temp file if present."""
    b64 = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_BASE64")
    if not b64 or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
        return
    decoded = base64.b64decode(b64).decode("utf-8")
    fd, path = tempfile.mkstemp(suffix=".json", prefix="gcp-creds-")
    os.write(fd, decoded.encode())
    os.close(fd)
    os.chmod(path, 0o600)
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = path


_init_vertex_credentials()


RABBITMQ_URL = os.environ.get("RABBITMQ_URL", "amqp://devkg:devkg@localhost:5672/")
FUSEKI_URL = os.environ.get("FUSEKI_URL", "http://localhost:3030")
FUSEKI_DATASET = os.environ.get("FUSEKI_DATASET", "devkg")
FUSEKI_USER = os.environ.get("FUSEKI_USER", "admin")
FUSEKI_PASS = os.environ.get("FUSEKI_PASS", "admin")
FUSEKI_AUTH = (FUSEKI_USER, FUSEKI_PASS)
OUTPUT_DIR = Path(os.environ.get("OUTPUT_DIR", "/app/output/claude"))

QUEUE = "devkg_jobs"
DLX = "devkg_jobs_dlx"
DLQ = "devkg_jobs_failed"
WATERMARK_FILE = OUTPUT_DIR / "watermarks.json"


def resolve_output_dir(container_path: str) -> Path:
    """Pick platform-specific output directory (watermarks live beside TTLs)."""
    base = Path(os.environ.get("OUTPUT_BASE", "/app/output"))
    if "/cursor-sessions/" in container_path or ".cursor/projects" in container_path:
        return base / "cursor"
    if "/pi-sessions/" in container_path or ".pi/agent/sessions" in container_path:
        return base / "pi"
    if "/codex-sessions/" in container_path or ".codex/sessions" in container_path:
        return base / "codex"
    # Honor OUTPUT_DIR for Claude / unknown unless it already points at a platform subdir
    return OUTPUT_DIR


def log(level: str, msg: str):
    print(f"[{level}] {msg}", file=sys.stderr, flush=True)


def write_turtle_atomic(graph, output_file: Path) -> None:
    """Serialize Turtle to a temp file and replace atomically on success."""
    output_file.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_path = tempfile.mkstemp(prefix=f".{output_file.stem}.", suffix=".ttl.tmp", dir=output_file.parent)
    os.close(fd)
    try:
        graph.serialize(destination=tmp_path, format="turtle")
        os.replace(tmp_path, output_file)
    except Exception:
        try:
            os.unlink(tmp_path)
        except FileNotFoundError:
            pass
        raise


def file_hash(path: str) -> str:
    """Compute SHA256 hash of a file's contents."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def load_watermarks() -> dict:
    """Load watermark file mapping session_id -> content hash."""
    if not WATERMARK_FILE.exists():
        return {}
    with open(WATERMARK_FILE) as f:
        return json.load(f)


def save_watermarks(watermarks: dict) -> None:
    """Save watermark state."""
    WATERMARK_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(WATERMARK_FILE, "w") as f:
        json.dump(watermarks, f, indent=2)


def connect_with_retry(url: str, max_retries: int = 10) -> pika.BlockingConnection:
    """Connect to RabbitMQ with exponential backoff."""
    delay = 1.0
    for attempt in range(1, max_retries + 1):
        try:
            params = pika.URLParameters(url)
            params.heartbeat = 600
            params.blocked_connection_timeout = 300
            conn = pika.BlockingConnection(params)
            log("INFO", f"Connected to RabbitMQ (attempt {attempt})")
            return conn
        except pika.exceptions.AMQPConnectionError as e:
            if attempt == max_retries:
                raise
            log("WARN", f"Connection failed (attempt {attempt}/{max_retries}): {e}")
            time.sleep(delay)
            delay = min(delay * 2, 30.0)


def setup_queues(channel):
    """Declare the main queue with dead-letter exchange."""
    # Dead-letter exchange + queue
    channel.exchange_declare(exchange=DLX, exchange_type="fanout", durable=True)
    channel.queue_declare(queue=DLQ, durable=True)
    channel.queue_bind(queue=DLQ, exchange=DLX)

    # Main queue with DLX
    channel.queue_declare(
        queue=QUEUE,
        durable=True,
        arguments={"x-dead-letter-exchange": DLX},
    )
    channel.basic_qos(prefetch_count=1)


def translate_path(host_path: str) -> str:
    """Translate host path to container path.

    Host: ~/.cursor/projects/{slug}/agent-transcripts/{uuid}/{uuid}.jsonl
    Container: /cursor-sessions/projects/{slug}/agent-transcripts/{uuid}/{uuid}.jsonl

    Host: ~/.claude/projects/{slug}/{session}.jsonl
    Container: /claude-sessions/{slug}/{session}.jsonl

    Host: ~/.pi/agent/sessions/{slug}/{session}.jsonl
    Container: /pi-sessions/{slug}/{session}.jsonl

    Host: ~/.codex/sessions/{yyyy}/{mm}/{dd}/{session}.jsonl
    Container: /codex-sessions/{yyyy}/{mm}/{dd}/{session}.jsonl
    """
    # Cursor MUST be checked before Claude's "/projects/" match — paths contain
    # ".cursor/projects/" which would otherwise remap into /claude-sessions.
    idx = host_path.find("/.cursor/projects/")
    if idx != -1:
        return "/cursor-sessions" + host_path[idx + len("/.cursor") :]

    marker_claude = "/.claude/projects/"
    idx = host_path.find(marker_claude)
    if idx != -1:
        return "/claude-sessions" + host_path[idx + len("/.claude/projects") :]

    # Legacy fallback for already-translated or unusual Claude paths
    marker_projects = "/projects/"
    idx = host_path.find(marker_projects)
    if idx != -1 and ".cursor" not in host_path:
        return "/claude-sessions" + host_path[idx + len("/projects") :]

    marker_pi = "/sessions/"
    idx = host_path.find(marker_pi)
    if idx != -1 and ".pi" in host_path:
        return "/pi-sessions" + host_path[idx + len("/sessions") :]

    idx = host_path.find("/.codex/sessions/")
    if idx != -1:
        return "/codex-sessions" + host_path[idx + len("/.codex/sessions") :]

    return host_path  # can't translate, return as-is


def _derive_link_metadata(container_path: str, session_id: str) -> dict:
    """Derive provenance metadata for Wikidata linking runs from a transcript path."""
    if "/cursor-sessions/" in container_path or ".cursor/projects" in container_path:
        platform = "cursor"
    elif "/pi-sessions/" in container_path or ".pi/agent/sessions" in container_path:
        platform = "pi"
    elif "/codex-sessions/" in container_path or ".codex/sessions" in container_path:
        platform = "codex"
    else:
        platform = "claude"

    project = ""
    p = Path(container_path)
    if ".cursor" in p.parts:
        idx = p.parts.index("projects")
        if idx + 1 < len(p.parts):
            project = p.parts[idx + 1]
    elif "projects" in p.parts:
        idx = p.parts.index("projects")
        if idx + 1 < len(p.parts):
            project = p.parts[idx + 1]

    return {
        "source_platform": platform,
        "session_id": session_id or p.stem,
        "source_file": str(p.resolve()),
        "project": project,
    }


def process_message(body: bytes) -> None:
    """Process a single pipeline job."""
    msg = json.loads(body)
    transcript_path = msg.get("transcript_path", "")
    session_id = msg.get("session_id", "")

    if not transcript_path:
        log("WARN", "Message missing transcript_path, skipping")
        return

    # Skip subagent sessions
    if "/subagents/" in transcript_path:
        log("INFO", f"Skipping subagent session: {session_id}")
        return

    # Translate host path to container path
    container_path = translate_path(transcript_path)

    if not os.path.exists(container_path):
        raise FileNotFoundError(f"Transcript not found: {container_path} (host: {transcript_path})")

    # Derive output filename (platform-specific dir avoids watermark collisions)
    basename = session_id or Path(container_path).stem
    output_dir = resolve_output_dir(container_path)
    watermark_file = output_dir / "watermarks.json"
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / f"{basename}.ttl"

    # SHA256 watermark: skip if file content hasn't changed since last processing
    watermarks = {}
    if watermark_file.exists():
        with open(watermark_file) as f:
            watermarks = json.load(f)
    current_hash = file_hash(container_path)
    if watermarks.get(basename) == current_hash:
        log("SKIP", f"{basename} (unchanged since last processing)")
        return

    log("INFO", f"Processing: {basename}")

    # Import pipeline modules (deferred to avoid import errors during setup)
    from pipeline.llm_providers import get_extraction_model

    model = get_extraction_model()

    devkg_triple_count = 0

    def _run_devkg():
        nonlocal devkg_triple_count

        if "/cursor-sessions/" in container_path or ".cursor/projects" in container_path:
            from pipeline.cursor_to_rdf import build_graph
        elif "/pi-sessions/" in container_path or ".pi/agent/sessions" in container_path:
            from pipeline.pi_to_rdf import build_graph
        elif "/codex-sessions/" in container_path or ".codex/sessions" in container_path:
            from pipeline.codex_to_rdf import build_graph
        else:
            from pipeline.jsonl_to_rdf import build_graph

        from pipeline.load_fuseki import ensure_dataset, upload_turtle

        graph = build_graph(container_path, skip_extraction=False, model=model)

        if os.environ.get("DEVKG_SKIP_LINKING", "").lower() not in ("1", "true", "yes"):
            try:
                from pipeline.link_entities import link_entities_into_graph
                link_trace_metadata = _derive_link_metadata(container_path, session_id)
                link_stats = link_entities_into_graph(
                    graph, trace_metadata=link_trace_metadata
                )
                log(
                    "INFO",
                    "  wikidata: "
                    f"{link_stats['linked']} linked "
                    f"({link_stats['cache_hits']} cache, "
                    f"{link_stats['agentic_calls']} agentic, "
                    f"{link_stats['negative_cache_hits']} negative, "
                    f"{link_stats['skipped']} skipped)",
                )
            except Exception as e:
                log("WARN", f"  wikidata linking failed (continuing): {e}")

        write_turtle_atomic(graph, output_file)
        devkg_triple_count = len(graph)

        ensure_dataset(FUSEKI_URL, FUSEKI_DATASET, auth=FUSEKI_AUTH)
        upload_turtle(FUSEKI_URL, FUSEKI_DATASET, str(output_file), auth=FUSEKI_AUTH)
        log("DONE", f"  devkg: {devkg_triple_count} triples -> {output_file}")

    _run_devkg()  # raises on failure — devkg failure is fatal for the job

    # Update watermark after successful processing
    watermarks[basename] = current_hash
    watermark_file.parent.mkdir(parents=True, exist_ok=True)
    with open(watermark_file, "w") as f:
        json.dump(watermarks, f, indent=2)

    log("DONE", f"{basename} ({devkg_triple_count} devkg triples)")


def on_message(channel, method, properties, body):
    """RabbitMQ message callback."""
    try:
        process_message(body)
        channel.basic_ack(delivery_tag=method.delivery_tag)
    except Exception as e:
        try:
            msg = json.loads(body) if body else {}
            session_id = msg.get("session_id", "unknown")
        except Exception:
            session_id = body[:80].decode(errors="replace") if body else "unknown"
        log("ERROR", f"{session_id}: {e}")
        traceback.print_exc(file=sys.stderr)
        channel.basic_nack(delivery_tag=method.delivery_tag, requeue=False)


def main():
    connection = connect_with_retry(RABBITMQ_URL)
    channel = connection.channel()
    setup_queues(channel)

    channel.basic_consume(queue=QUEUE, on_message_callback=on_message)

    log("READY", f"Waiting for jobs on {QUEUE}")
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        log("INFO", "Shutting down")
        channel.stop_consuming()
    finally:
        connection.close()


if __name__ == "__main__":
    main()