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