arborist/aborist/sources/grok.py

"""Grok account-export source.

xAI's "Export Account Data" downloads ship as a directory tree containing
`prod-grok-backend.json` (conversations + media posts) plus binary assets.

Grain choice: one Document = one conversation. A conversation has a stable
id, a human-authored title, and a tree of human/assistant messages. The
serialized turn sequence becomes Document.content; the chunker splits long
conversations across multiple 512-token chunks naturally. Per-message
provenance lives inline in each turn header.

A second source (GrokMediaPostsSource) yields Grok's image/video generation
prompts as small documents, each linking to its asset URL via an edge.
"""

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator

from aborist.document import Document, Edge
from aborist.source import Source


# Strip Grok-internal render markers from message text. These appear as
# `<grok:render card_id="..."/>` (self-closing) and rarely as
# `<grok:render ...>...</grok:render>` (block). Drop the markup; it has no
# meaning outside the live Grok UI and would otherwise pollute embeddings
# and full-text search.
_GROK_RENDER_RE = re.compile(
    r"<grok:render\b[^>]*?/>|<grok:render\b[^>]*?>.*?</grok:render>",
    re.DOTALL,
)

# Liberal URL extractor: http(s) + non-whitespace + bracket characters.
# Trailing punctuation (.,;:!?)] gets trimmed.
_URL_RE = re.compile(r"https?://[^\s<>()\[\]'\"`]+", re.IGNORECASE)
_URL_TRAIL = ".,;:!?)]"

# Reasonable upper bound on conversation count to scan in a single export.
_HARD_DOC_CAP = 1_000_000


def _scrub(text: str) -> str:
    """Drop Grok-internal markup but keep the conversational text intact."""
    if not text:
        return ""
    return _GROK_RENDER_RE.sub("", text)


def _extract_urls(text: str) -> list[str]:
    """Return distinct URLs from message body, in first-seen order."""
    out: list[str] = []
    seen: set[str] = set()
    for m in _URL_RE.finditer(text):
        u = m.group(0).rstrip(_URL_TRAIL)
        if u and u not in seen:
            seen.add(u)
            out.append(u)
    return out


def _normalize_sender(sender: str | None) -> str:
    """Collapse Grok's sender labels into {human, assistant}."""
    if not sender:
        return "unknown"
    s = sender.strip().lower()
    if s in ("assistant", "model"):
        return "assistant"
    return s  # 'human' stays as 'human'; anything else preserved.


def _normalize_ts(raw) -> str:
    """Coerce Grok's varied timestamp shapes into an ISO-8601 string.

    Real exports interleave: bare ISO strings, Mongo extended JSON
    `{"$date": "..."}`, `{"$date": {"$numberLong": "<millis>"}}`, and
    `{"$numberLong": "<millis>"}` directly. Anything unparseable returns "".
    """
    if raw is None:
        return ""
    if isinstance(raw, str):
        return raw
    if isinstance(raw, (int, float)):
        # Heuristic: a value > 10^12 is plausibly unix millis, else seconds.
        # Grok exports use millis. Trust that.
        return _millis_to_iso(int(raw))
    if isinstance(raw, dict):
        if "$date" in raw:
            return _normalize_ts(raw["$date"])
        if "$numberLong" in raw:
            try:
                return _millis_to_iso(int(raw["$numberLong"]))
            except (TypeError, ValueError):
                return ""
    return ""


def _millis_to_iso(ms: int) -> str:
    """Unix millis -> ISO-8601 UTC string (best-effort)."""
    try:
        return (
            datetime.fromtimestamp(ms / 1000, tz=timezone.utc)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z")
        )
    except (OverflowError, OSError, ValueError):
        return ""


def _format_conversation(conv: dict, responses: list[dict]) -> str:
    """Serialize a conversation tree into chunkable plain text.

    Returns "" if no message has any body after scrubbing — the title alone
    is not substantive enough to commit. Format is otherwise deterministic so
    the same export re-ingested produces the same document_root.
    """
    turns: list[str] = []
    for wrapper in responses:
        r = wrapper.get("response") or {}
        sender = _normalize_sender(r.get("sender"))
        model = r.get("model") or ""
        ts = _normalize_ts(r.get("create_time"))
        body = _scrub(r.get("message") or "").strip()
        if not body:
            continue
        head = f"[{ts}] {sender}"
        if model:
            head += f" ({model})"
        head += ":"
        turns.append(f"{head}\n{body}")
    if not turns:
        return ""
    title = conv.get("title") or ""
    parts: list[str] = []
    if title:
        parts.append(f"# {title}")
    parts.extend(turns)
    return "\n\n".join(parts)


def _resolve_export_root(path: str | Path) -> Path:
    """Find the directory containing prod-grok-backend.json.

    Accepts either the JSON file itself, the user-id directory containing it,
    or any ancestor down to the export root.
    """
    p = Path(path)
    if p.is_file() and p.name == "prod-grok-backend.json":
        return p.parent
    if p.is_dir():
        if (p / "prod-grok-backend.json").is_file():
            return p
        # Walk down at most a few levels to find it.
        matches = list(p.rglob("prod-grok-backend.json"))
        if matches:
            return matches[0].parent
    raise FileNotFoundError(
        f"prod-grok-backend.json not found at or under {path}"
    )


class GrokExportSource(Source):
    """One Document per Grok conversation."""

    source_type = "grok_export"

    def __init__(self, path: str | Path):
        self.export_root = _resolve_export_root(path)
        self._json_path = self.export_root / "prod-grok-backend.json"

    def iter_documents(self) -> Iterator[Document]:
        with self._json_path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        convs = data.get("conversations") or []
        for entry in convs[:_HARD_DOC_CAP]:
            conv = entry.get("conversation") or {}
            responses = entry.get("responses") or []
            conv_id = conv.get("id")
            if not conv_id:
                continue
            content = _format_conversation(conv, responses)
            if not content:
                continue
            edges: list[Edge] = []
            seen_urls: set[str] = set()
            for wrapper in responses:
                r = wrapper.get("response") or {}
                msg = _scrub(r.get("message") or "")
                for u in _extract_urls(msg):
                    if u in seen_urls:
                        continue
                    seen_urls.add(u)
                    edges.append(Edge(edge_type="hyperlink", dst_uri=u))
            yield Document(
                uri=f"grok://conversation/{conv_id}",
                content=content,
                source_type=self.source_type,
                title=conv.get("title") or None,
                edges=edges,
            )


class GrokMediaPostsSource(Source):
    """One Document per Grok media-generation prompt (images/video)."""

    source_type = "grok_media"

    def __init__(self, path: str | Path):
        self.export_root = _resolve_export_root(path)
        self._json_path = self.export_root / "prod-grok-backend.json"

    def iter_documents(self) -> Iterator[Document]:
        with self._json_path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        media = data.get("media_posts") or []
        for post in media:
            mp_id = post.get("id")
            prompt = (post.get("original_prompt") or "").strip()
            if not mp_id or not prompt:
                continue
            asset_url = post.get("link")
            ts = post.get("create_time") or ""
            media_type = post.get("media_type") or ""
            header_bits: list[str] = []
            if ts:
                header_bits.append(f"[{ts}]")
            if media_type:
                header_bits.append(media_type)
            header = " ".join(header_bits)
            content = (header + "\n\n" + prompt).strip() if header else prompt
            edges: list[Edge] = []
            if asset_url:
                edges.append(Edge(edge_type="asset", dst_uri=asset_url))
            title = prompt.splitlines()[0][:80] if prompt else None
            yield Document(
                uri=f"grok://media/{mp_id}",
                content=content,
                source_type=self.source_type,
                title=title,
                edges=edges,
            )