russell@unturf.com
9ba06cd809
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.
252 lines
9 KiB
Python
252 lines
9 KiB
Python
"""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
|