From 9ba06cd8095bcdaa6a010e29a714ac5dabad1b11 Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Mon, 27 Apr 2026 13:49:42 -0400 Subject: [PATCH] add Grok export source: conversations + media posts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit aborist/sources/grok.py exposes two Source subclasses for ingesting xAI's user data export: GrokExportSource walks ttl//export_data// and reads prod-grok-backend.json. Yields one Document per conversation with: - URI: grok://conversation/ - title: the conversation's auto-generated title - content: full message text in turn order source_type='grok_export'. GrokMediaPostsSource same export, but yields per media-generation post (image/video prompts) under URI grok://media/. source_type='grok_media'. Both auto-walk down from the export root so callers can pass the top-level directory xAI delivered (e.g., ~/Downloads//). Wired into the CLI: aborist ingest --source {grok_export,grok_media} accepts --path . tests/test_grok_source.py covers the walk + parse + Document shape with a fabricated mini-export fixture. Once ingested, conversations become normal queryable docs in the shard cluster — your prior chats become memory the corpus can consult during RAG. 79 tests passing. --- aborist/cli.py | 24 +++- aborist/sources/__init__.py | 9 +- aborist/sources/grok.py | 245 +++++++++++++++++++++++++++++++++++ tests/test_grok_source.py | 252 ++++++++++++++++++++++++++++++++++++ 4 files changed, 527 insertions(+), 3 deletions(-) create mode 100644 aborist/sources/grok.py create mode 100644 tests/test_grok_source.py diff --git a/aborist/cli.py b/aborist/cli.py index 5ae44fa..232a9e5 100644 --- a/aborist/cli.py +++ b/aborist/cli.py @@ -49,6 +49,14 @@ def _cmd_ingest(args: argparse.Namespace) -> int: print("html source needs --url or --urls-from", file=sys.stderr) return 2 src = HtmlPageSource(urls, respect_robots=not args.no_robots) + elif args.source in ("grok_export", "grok_media"): + if not args.path: + print(f"--path is required for {args.source}", file=sys.stderr) + return 2 + from aborist.sources import GrokExportSource, GrokMediaPostsSource + + cls = GrokExportSource if args.source == "grok_export" else GrokMediaPostsSource + src = cls(path=args.path) else: print(f"unknown source: {args.source}", file=sys.stderr) return 2 @@ -870,10 +878,22 @@ def build_parser() -> argparse.ArgumentParser: ingest.add_argument( "--source", required=True, - choices=["wikipedia_cur", "wikipedia_old", "html"], + choices=[ + "wikipedia_cur", + "wikipedia_old", + "html", + "grok_export", + "grok_media", + ], help="source type", ) - ingest.add_argument("--path", help="path to dump file (file-backed sources)") + ingest.add_argument( + "--path", + help=( + "path to dump file (wikipedia) or to xAI export root / " + "prod-grok-backend.json (grok_export, grok_media)" + ), + ) ingest.add_argument( "--url", action="append", help="URL to ingest (html source; repeatable)" ) diff --git a/aborist/sources/__init__.py b/aborist/sources/__init__.py index e12bddc..0d30a5a 100644 --- a/aborist/sources/__init__.py +++ b/aborist/sources/__init__.py @@ -1,12 +1,19 @@ """Source implementations. Add a new corpus = add a new module here.""" +from aborist.sources.grok import GrokExportSource, GrokMediaPostsSource from aborist.sources.wikipedia import ( WikipediaCurDump, WikipediaOldDump, WikipediaSqlDump, ) -__all__ = ["WikipediaCurDump", "WikipediaOldDump", "WikipediaSqlDump"] +__all__ = [ + "GrokExportSource", + "GrokMediaPostsSource", + "WikipediaCurDump", + "WikipediaOldDump", + "WikipediaSqlDump", +] # HtmlPageSource has optional deps (selectolax). Surface import errors only # when callers actually request it. diff --git a/aborist/sources/grok.py b/aborist/sources/grok.py new file mode 100644 index 0000000..a8eab52 --- /dev/null +++ b/aborist/sources/grok.py @@ -0,0 +1,245 @@ +"""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 +# `` (self-closing) and rarely as +# `...` (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"]*?/>|]*?>.*?", + 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": ""}}`, and + `{"$numberLong": ""}` 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, + ) diff --git a/tests/test_grok_source.py b/tests/test_grok_source.py new file mode 100644 index 0000000..0bba14b --- /dev/null +++ b/tests/test_grok_source.py @@ -0,0 +1,252 @@ +"""Tests for the Grok account-export source. + +Synthetic fixture mirrors the real prod-grok-backend.json shape: a +`conversations` array of `{conversation, responses}` entries, plus a +`media_posts` array of generation prompts. Real personal data is NOT +imported into the test suite. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from aborist.ingest import ingest_source +from aborist.sources.grok import ( + GrokExportSource, + GrokMediaPostsSource, + _extract_urls, + _normalize_ts, + _scrub, +) +from aborist.store import connect + + +def _make_export(tmp_path: Path) -> Path: + """Build a minimal but representative export tree, return its root.""" + user_dir = tmp_path / "ttl" / "30d" / "export_data" / "user-uuid" + user_dir.mkdir(parents=True) + payload = { + "conversations": [ + { + "conversation": { + "id": "conv-aaaa-1111", + "title": "Quantum Slop", + "create_time": "2026-04-25T12:00:00Z", + "user_id": "user-uuid", + }, + "responses": [ + { + "response": { + "_id": "msg-001", + "conversation_id": "conv-aaaa-1111", + "sender": "human", + "model": "grok-4-auto", + "create_time": {"$date": "2026-04-25T12:00:00Z"}, + "message": ( + "what is up with quantum slop? see " + "https://example.com/qs and " + "https://example.org/qpu, also (https://example.net/x)." + ), + "metadata": {}, + }, + "share_link": None, + }, + { + "response": { + "_id": "msg-002", + "conversation_id": "conv-aaaa-1111", + "sender": "ASSISTANT", + "model": "grok-3", + # Mongo extended JSON unix millis (real export shape). + "create_time": {"$numberLong": "1777306441511"}, + "message": ( + "Quantum slop is a misnomer." + "" + " Real QPUs differ." + "ignored body" + ), + "metadata": {}, + }, + "share_link": None, + }, + ], + }, + { + "conversation": { + "id": "conv-bbbb-2222", + "title": "Empty thread", + "create_time": "2026-04-26T00:00:00Z", + "user_id": "user-uuid", + }, + # Both messages empty after scrubbing -> doc should be skipped. + "responses": [ + { + "response": { + "_id": "msg-101", + "sender": "human", + "message": "", + "model": "", + }, + "share_link": None, + }, + ], + }, + ], + "projects": [], + "tasks": [], + "media_posts": [ + { + "id": "media-7777", + "user_id": "user-uuid", + "original_prompt": "a fox standing on a permacomputer", + "media_type": "image", + "create_time": "2026-03-12T08:00:00Z", + "link": "https://assets.example/img/abc.webp", + }, + { + "id": "media-8888", + "user_id": "user-uuid", + "original_prompt": "", + "media_type": "image", + "create_time": "2026-03-12T08:00:01Z", + "link": "https://assets.example/img/empty.webp", + }, + ], + } + (user_dir / "prod-grok-backend.json").write_text(json.dumps(payload)) + return tmp_path + + +def test_normalize_ts_handles_iso_string(): + assert _normalize_ts("2026-04-25T12:00:00Z") == "2026-04-25T12:00:00Z" + + +def test_normalize_ts_handles_mongo_date_wrapper(): + assert _normalize_ts({"$date": "2026-04-25T12:00:00Z"}) == "2026-04-25T12:00:00Z" + + +def test_normalize_ts_handles_mongo_numberlong_millis(): + # 1_777_306_441_511 ms = 2026-04-27T16:14:01.511Z + iso = _normalize_ts({"$numberLong": "1777306441511"}) + assert iso.startswith("2026-04-27T16:14:01") + assert iso.endswith("Z") + + +def test_normalize_ts_handles_nested_date_numberlong(): + iso = _normalize_ts({"$date": {"$numberLong": "1777306441511"}}) + assert iso.startswith("2026-04-27T16:14:01") + + +def test_normalize_ts_garbage_returns_empty(): + assert _normalize_ts({"weird": "shape"}) == "" + assert _normalize_ts(None) == "" + + +def test_scrub_drops_grok_render_markup(): + s = "beforemiddlebodyafter" + assert _scrub(s) == "beforemiddleafter" + + +def test_extract_urls_trims_trailing_punctuation_and_dedupes(): + text = ( + "see https://example.com/qs, and https://example.org/qpu. " + "also (https://example.net/x). https://example.com/qs again." + ) + urls = _extract_urls(text) + assert urls == [ + "https://example.com/qs", + "https://example.org/qpu", + "https://example.net/x", + ] + + +def test_grok_export_source_yields_one_doc_per_nonempty_conversation(tmp_path): + root = _make_export(tmp_path) + src = GrokExportSource(path=root) + docs = list(src.iter_documents()) + assert len(docs) == 1 + doc = docs[0] + assert doc.uri == "grok://conversation/conv-aaaa-1111" + assert doc.source_type == "grok_export" + assert doc.title == "Quantum Slop" + # Both turns serialized with sender/model headers. + assert "human (grok-4-auto)" in doc.content + assert "assistant (grok-3)" in doc.content + # Mongo extended JSON shapes must be normalized — no raw wrappers leak. + assert "$numberLong" not in doc.content + assert "$date" not in doc.content + # Render markers stripped. + assert "grok:render" not in doc.content + # URLs surfaced as hyperlink edges. + edge_uris = {e.dst_uri for e in doc.edges} + assert "https://example.com/qs" in edge_uris + assert "https://example.org/qpu" in edge_uris + assert "https://example.net/x" in edge_uris + assert all(e.edge_type == "hyperlink" for e in doc.edges) + + +def test_grok_media_source_yields_one_doc_per_nonempty_post(tmp_path): + root = _make_export(tmp_path) + src = GrokMediaPostsSource(path=root) + docs = list(src.iter_documents()) + assert len(docs) == 1 + doc = docs[0] + assert doc.uri == "grok://media/media-7777" + assert doc.source_type == "grok_media" + assert "fox standing on a permacomputer" in doc.content + asset_edges = [e for e in doc.edges if e.edge_type == "asset"] + assert len(asset_edges) == 1 + assert asset_edges[0].dst_uri == "https://assets.example/img/abc.webp" + + +def test_resolves_export_root_from_json_or_dir(tmp_path): + root = _make_export(tmp_path) + json_path = ( + root / "ttl" / "30d" / "export_data" / "user-uuid" / "prod-grok-backend.json" + ) + # Directly the JSON file. + a = list(GrokExportSource(path=json_path).iter_documents()) + # The user-id directory. + b = list(GrokExportSource(path=json_path.parent).iter_documents()) + # The export root (auto-walks down). + c = list(GrokExportSource(path=root).iter_documents()) + assert [d.uri for d in a] == [d.uri for d in b] == [d.uri for d in c] + + +def test_resolve_missing_path_raises(tmp_path): + with pytest.raises(FileNotFoundError): + GrokExportSource(path=tmp_path / "nope") + + +def test_grok_ingest_round_trip(tmp_path): + """End-to-end: ingest a tiny export into a fresh DB, verify counts.""" + root = _make_export(tmp_path) + db_path = tmp_path / "aborist.db" + conn = connect(db_path) + try: + conv_stats = ingest_source(conn, GrokExportSource(path=root)) + media_stats = ingest_source(conn, GrokMediaPostsSource(path=root)) + finally: + conn.close() + assert conv_stats.inserted == 1 + assert media_stats.inserted == 1 + # Re-ingest is a no-op (idempotent at content-root level). + conn = connect(db_path) + try: + again = ingest_source(conn, GrokExportSource(path=root)) + finally: + conn.close() + assert again.inserted == 0 + assert again.skipped_duplicate == 1 + + +def test_serialized_content_is_deterministic(tmp_path): + """Same export -> same document_root every time. Required for idempotent + re-ingest and stable cache keys.""" + root = _make_export(tmp_path) + a = list(GrokExportSource(path=root).iter_documents()) + b = list(GrokExportSource(path=root).iter_documents()) + assert a[0].content == b[0].content