rohitsar567 Claude Opus 4.7 (1M context) commited on
Commit
380491c
·
1 Parent(s): 119579e

fix(rag): KI-117 — boot-time cleanup of dangling profile chunks

Browse files

The HF Space carries its own Chroma DB which still contains the legacy
`profile_anonymous` row from KI-102's earliest deploy — a chunk written
without a `session_id` metadata field. That row poisons every retrieval
whose `where` clause references session_id, because Chroma raises when a
filtered row is missing the filtered key. KI-111's wrap catches the
raise so the bot doesn't crash, but retrieval returns empty on HF until
the bad row is purged.

KI-112 already prevents future bad writes (input guards on
upsert_profile_chunk), and the local DB was cleaned manually. This
patch adds a one-shot startup task in backend/main.py that, on app
boot, scans the Chroma collection for any `doc_type='profile'` chunks
lacking a non-empty `session_id` metadata field and deletes them. Runs
idempotently — if there are no bad rows, it's a no-op. After HF
rebuilds with this code, the boot task self-heals HF's DB on first
request.

- New helper `_startup_purge_dangling_profile_chunks()` registered via a
third `@app.on_event("startup")` hook (sits alongside the existing
admin-overrides loader + llm_health probe loop).
- Uses `coll.get(where={"doc_type":"profile"}, ...)` then filters for
rows whose `session_id` metadata is missing / non-str / empty.
- Logs `KI-117: purged N dangling profile chunks at boot` on success
and `KI-117: total chunks after cleanup: N` for HF Space log
verification.
- Wrapped in try/except at two levels — boot must never crash even if
Chroma is wedged. Inner failures log WARNING + continue.
- Heavy work runs in `asyncio.to_thread` so we don't block the event
loop during boot.

Local verification: function runs, reports `0 dangling profile chunks`
(local was already cleaned by KI-112) + `1832 total chunks after
cleanup`. Full pytest suite (90 tests) still green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Files changed (1) hide show
  1. backend/main.py +85 -0
backend/main.py CHANGED
@@ -211,6 +211,91 @@ async def _startup_llm_health_probe():
211
  asyncio.create_task(llm_health.background_probe_loop())
212
 
213
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
214
  @app.get("/api/health", response_model=HealthResponse)
215
  async def health():
216
  missing = settings.validate()
 
211
  asyncio.create_task(llm_health.background_probe_loop())
212
 
213
 
214
+ async def _startup_purge_dangling_profile_chunks():
215
+ """KI-117 — boot-time self-heal of dangling `doc_type='profile'` chunks.
216
+
217
+ Background: KI-102's earliest deploy wrote a `profile_anonymous` chunk
218
+ WITHOUT a `session_id` metadata field. That legacy row poisoned every
219
+ subsequent retrieval whose `where` clause referenced session_id, because
220
+ Chroma raises when a filtered row is missing the filtered key. KI-112
221
+ added input guards so no new bad rows can be written, and the local DB
222
+ was cleaned manually. But the HF Space carries its OWN copy of the
223
+ Chroma DB and still contains the dangling row.
224
+
225
+ This handler scans the collection for any `doc_type='profile'` chunks
226
+ whose metadata lacks a non-empty `session_id` and deletes them. Runs
227
+ idempotently — if there are no bad rows, it's a no-op. After HF rebuilds
228
+ with this code, the boot task self-heals HF's DB on first request.
229
+
230
+ Wrapped in try/except so a Chroma hiccup never crashes boot.
231
+ """
232
+ def _do_purge() -> None:
233
+ from rag.retrieve import get_collection
234
+
235
+ coll = get_collection()
236
+ try:
237
+ res = coll.get(
238
+ where={"doc_type": "profile"},
239
+ limit=10000,
240
+ include=["metadatas"],
241
+ )
242
+ except Exception as e:
243
+ logging.warning(
244
+ "KI-117: profile-chunk scan failed (%s: %s) — skipping cleanup",
245
+ type(e).__name__, e,
246
+ )
247
+ return
248
+
249
+ ids = res.get("ids") or []
250
+ metas = res.get("metadatas") or []
251
+ bad_ids: list[str] = []
252
+ for cid, meta in zip(ids, metas):
253
+ sid = (meta or {}).get("session_id")
254
+ if not (isinstance(sid, str) and sid.strip()):
255
+ bad_ids.append(cid)
256
+
257
+ if bad_ids:
258
+ try:
259
+ coll.delete(ids=bad_ids)
260
+ logging.info(
261
+ "KI-117: purged %d dangling profile chunks at boot (ids=%s)",
262
+ len(bad_ids),
263
+ bad_ids[:10] + (["..."] if len(bad_ids) > 10 else []),
264
+ )
265
+ except Exception as e:
266
+ logging.warning(
267
+ "KI-117: delete(ids=...) failed (%s: %s) — bad rows remain",
268
+ type(e).__name__, e,
269
+ )
270
+ return
271
+ else:
272
+ logging.info("KI-117: no dangling profile chunks found (DB clean)")
273
+
274
+ try:
275
+ total = coll.count()
276
+ logging.info("KI-117: total chunks after cleanup: %d", total)
277
+ except Exception as e:
278
+ logging.warning(
279
+ "KI-117: post-cleanup count() failed (%s: %s)",
280
+ type(e).__name__, e,
281
+ )
282
+
283
+ try:
284
+ await asyncio.to_thread(_do_purge)
285
+ except Exception as e:
286
+ # Belt + suspenders — boot must never crash.
287
+ logging.warning(
288
+ "KI-117: boot cleanup raised at top level (%s: %s) — continuing boot",
289
+ type(e).__name__, e,
290
+ )
291
+
292
+
293
+ @app.on_event("startup")
294
+ async def _startup_purge_dangling_profile_chunks_handler():
295
+ """KI-117 — register the boot-time cleanup as a FastAPI startup hook."""
296
+ await _startup_purge_dangling_profile_chunks()
297
+
298
+
299
  @app.get("/api/health", response_model=HealthResponse)
300
  async def health():
301
  missing = settings.validate()