cameron.stream
[READ-ONLY] Mirror of https://github.com/just-cameron/central. Autonomous AI building collective intelligence on ATProtocol. The central node of the comind network.
central.comind.network/
1"""Backfill historical records from ATProto into the indexer.
2
3Fetches all records for each agent's relevant collections from their PDS,
4embeds them, and stores in the database. Skips records already indexed.
5
6Usage:
7 uv run python -m indexer.backfill # All agents
8 uv run python -m indexer.backfill --did did:plc:... # Specific agent
9 uv run python -m indexer.backfill --dry-run # Count only
10"""
11
12import argparse
13import logging
14import time
15from datetime import datetime
16from typing import Optional
17
18import httpx
19
20from indexer import db, embeddings
21from indexer.worker import BASE_COLLECTIONS, SEED_DIDS
22
23logging.basicConfig(
24 level=logging.INFO,
25 format="%(asctime)s [%(levelname)s] %(message)s",
26)
27logger = logging.getLogger(__name__)
28
29
30def resolve_pds(did: str) -> Optional[str]:
31 """Resolve DID to PDS endpoint."""
32 try:
33 resp = httpx.get(f"https://plc.directory/{did}", timeout=10)
34 resp.raise_for_status()
35 doc = resp.json()
36 return doc["service"][0]["serviceEndpoint"]
37 except Exception as e:
38 logger.error(f"Failed to resolve PDS for {did}: {e}")
39 return None
40
41
42def resolve_handle(did: str) -> Optional[str]:
43 """Resolve DID to handle."""
44 try:
45 resp = httpx.get(f"https://plc.directory/{did}", timeout=10)
46 resp.raise_for_status()
47 doc = resp.json()
48 for alias in doc.get("alsoKnownAs", []):
49 if alias.startswith("at://"):
50 return alias[5:]
51 except Exception:
52 pass
53 return None
54
55
56def list_records(pds: str, did: str, collection: str):
57 """List all records in a collection with pagination."""
58 cursor = None
59 while True:
60 params = {
61 "repo": did,
62 "collection": collection,
63 "limit": 100,
64 }
65 if cursor:
66 params["cursor"] = cursor
67
68 resp = httpx.get(
69 f"{pds}/xrpc/com.atproto.repo.listRecords",
70 params=params,
71 timeout=30,
72 )
73
74 if resp.status_code == 400:
75 # Collection doesn't exist for this repo
76 return
77 resp.raise_for_status()
78 data = resp.json()
79
80 records = data.get("records", [])
81 if not records:
82 return
83
84 for record in records:
85 yield record
86
87 cursor = data.get("cursor")
88 if not cursor:
89 break
90
91
92def get_existing_uris(engine) -> set:
93 """Get set of all URIs already in the database."""
94 session = db.get_session(engine)
95 try:
96 from sqlalchemy import text
97 result = session.execute(text("SELECT uri FROM cognition_records"))
98 return {row[0] for row in result}
99 finally:
100 session.close()
101
102
103def backfill_agent(engine, did: str, handle: str, pds: str, collections: list[str], dry_run: bool = False):
104 """Backfill all collections for a single agent."""
105 logger.info(f"Backfilling @{handle} ({did}) from {pds}")
106
107 existing = get_existing_uris(engine)
108 total_new = 0
109 total_skipped = 0
110
111 for collection in collections:
112 new = 0
113 skipped = 0
114 batch_texts = []
115 batch_records = []
116
117 for record in list_records(pds, did, collection):
118 uri = record["uri"]
119 if uri in existing:
120 skipped += 1
121 continue
122
123 value = record.get("value", {})
124 content = embeddings.extract_content(value)
125 if not content:
126 skipped += 1
127 continue
128
129 rkey = uri.split("/")[-1]
130 created_at = None
131 if created_str := value.get("createdAt"):
132 try:
133 created_at = datetime.fromisoformat(
134 created_str.replace("Z", "+00:00")
135 )
136 except ValueError:
137 pass
138
139 batch_texts.append(content)
140 batch_records.append({
141 "uri": uri,
142 "did": did,
143 "collection": collection,
144 "rkey": rkey,
145 "content": content,
146 "created_at": created_at,
147 "handle": handle,
148 })
149
150 # Process in batches of 100
151 if len(batch_texts) >= 100:
152 if not dry_run:
153 _store_batch(engine, batch_texts, batch_records)
154 new += len(batch_texts)
155 batch_texts = []
156 batch_records = []
157 # Rate limit
158 time.sleep(0.5)
159
160 # Process remaining
161 if batch_texts:
162 if not dry_run:
163 _store_batch(engine, batch_texts, batch_records)
164 new += len(batch_texts)
165
166 if new > 0 or skipped > 0:
167 logger.info(f" {collection}: {new} new, {skipped} skipped")
168 total_new += new
169 total_skipped += skipped
170
171 logger.info(f" Total: {total_new} new, {total_skipped} skipped")
172 return total_new
173
174
175def _store_batch(engine, texts: list[str], records: list[dict]):
176 """Embed and store a batch of records."""
177 try:
178 embs = embeddings.embed_batch(texts)
179 except Exception as e:
180 logger.error(f"Embedding failed: {e}")
181 return
182
183 session = db.get_session(engine)
184 try:
185 for rec, emb in zip(records, embs):
186 db.upsert_record(
187 session,
188 uri=rec["uri"],
189 did=rec["did"],
190 collection=rec["collection"],
191 rkey=rec["rkey"],
192 content=rec["content"],
193 embedding=emb,
194 created_at=rec["created_at"],
195 handle=rec["handle"],
196 )
197 except Exception as e:
198 logger.error(f"Store failed: {e}")
199 session.rollback()
200 finally:
201 session.close()
202
203
204def main():
205 parser = argparse.ArgumentParser(description="Backfill ATProto records into indexer")
206 parser.add_argument("--did", help="Backfill specific DID only")
207 parser.add_argument("--dry-run", action="store_true", help="Count records without indexing")
208 args = parser.parse_args()
209
210 engine = db.get_engine()
211 db.init_db(engine)
212
213 dids = [args.did] if args.did else sorted(SEED_DIDS)
214 collections = sorted(BASE_COLLECTIONS)
215
216 total = 0
217 for did in dids:
218 pds = resolve_pds(did)
219 if not pds:
220 continue
221 handle = resolve_handle(did) or did[:30]
222 count = backfill_agent(engine, did, handle, pds, collections, dry_run=args.dry_run)
223 total += count
224
225 logger.info(f"Backfill complete. {total} new records indexed.")
226
227
228if __name__ == "__main__":
229 main()