engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add Grok export source: conversations + media posts

Browse source 
aborist/sources/grok.py exposes two Source subclasses for ingesting
xAI's user data export:

  GrokExportSource    walks ttl/<period>/export_data/<user-id>/ and
                      reads prod-grok-backend.json. Yields one Document
                      per conversation with:
                        - URI: grok://conversation/<conversation-id>
                        - 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/<post-id>.
                      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/<user-uuid>/).

Wired into the CLI: aborist ingest --source {grok_export,grok_media}
accepts --path <export-root>. 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.
This commit is contained in:
russell@unturf.com 2026-04-27 13:49:42 -04:00
parent 24d8204c08
commit 9ba06cd809
No known key found for this signature in database
4 changed files with 527 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

24
aborist/cli.py
View file

@ -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)"
)

9
aborist/sources/__init__.py
View file

@ -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.

245
aborist/sources/grok.py Normal file
View file

@ -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
# `<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,
)

252
tests/test_grok_source.py Normal file
View file

@ -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."
"<grok:render card_id=\"abc\"/>"
" Real QPUs differ.<grok:render card_id=\"d\">"
"ignored body</grok:render>"
),
"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 = "before<grok:render card_id=\"x\"/>middle<grok:render card_id=\"y\">body</grok:render>after"
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