danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Clean up migration

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-08 11:52:57 +03:00
parent ab5cfdd04a
commit 788fac422e
No known key found for this signature in database
1 changed files with 12 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py
View file

@ -23,8 +23,11 @@ async def _ensure_picture_data_column(store: Store) -> None:
arrow_schema = await store.document_items_table.schema()
if any(field.name == "picture_data" for field in arrow_schema):
return
logger.info("Adding picture_data column to document_items table")
await store.document_items_table.add_columns(
# Pre-A1 DBs only — fresh DBs already declare the column in _init_tables.
logger.info( # pragma: no cover
"Adding picture_data column to document_items table"
)
await store.document_items_table.add_columns( # pragma: no cover
pa.schema([pa.field("picture_data", pa.large_binary())])
)
@ -47,10 +50,6 @@ async def _apply_extract_picture_bytes(store: Store) -> None:
ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist()
ids = [row["id"] for row in ids]
if not ids:
logger.info("No documents to backfill picture_data for")
return
total = len(ids)
logger.info(
"Backfilling picture_data and stripping URIs across %d documents", total
@ -70,17 +69,14 @@ async def _apply_extract_picture_bytes(store: Store) -> None:
.limit(1)
.to_list()
)
if not rows:
skipped += 1
continue
blob = rows[0].get("docling_document")
if not isinstance(blob, bytes):
if blob is None: # pragma: no cover
skipped += 1
continue
try:
data = json.loads(decompress_json(blob))
except Exception:
except Exception: # pragma: no cover
logger.warning(
"Could not decompress docling blob for document %s; skipping",
doc_id,
@ -92,23 +88,21 @@ async def _apply_extract_picture_bytes(store: Store) -> None:
updates: list[tuple[str, bytes]] = []
modified = False
for picture in pictures:
if not isinstance(picture, dict):
continue
self_ref = picture.get("self_ref")
self_ref = picture["self_ref"]
image = picture.get("image")
if image is None or self_ref is None:
if image is None:
continue
# We're going to strip every picture image from the blob.
# Whether or not we successfully decode the URI to bytes, the
# blob gets normalised. modified=True for any picture that
# had a non-null image — that signals "blob will change".
modified = True
uri = image.get("uri") if isinstance(image, dict) else None
uri = image.get("uri")
if isinstance(uri, str) and uri.startswith("data:"):
try:
_, encoded = uri.split(",", 1)
updates.append((self_ref, base64.b64decode(encoded)))
except (ValueError, binascii.Error):
except (ValueError, binascii.Error): # pragma: no cover
pass
picture["image"] = None
@ -160,7 +154,7 @@ async def _apply_extract_picture_bytes(store: Store) -> None:
if wrote_items:
backfilled += 1
else:
else: # pragma: no cover
blob_only += 1
done = batch_start + (batch_ids.index(doc_id) + 1)