feat: Enhance document processing by caching image groups and document summaries

This commit is contained in:
Ross
2025-10-18 13:13:42 +01:00
parent 90ac197cbe
commit fbf4ab926b
+198 -126
View File
@@ -31,6 +31,8 @@ from loguru import logger
import sys
IMAGE_GROUPS = {}
CAPTURE_INPUT_DIR = None
DOCUMENT_SUMMARYS = {}
def text_of(node) -> str:
"""Return the text content of a node, stripping extra whitespace."""
@@ -191,40 +193,40 @@ def slugify(s: str, max_len: int = 80) -> str:
return s
def find_article_name_in_json(obj) -> str | None:
"""Look for an article name in known keys or nested values.
Common keys include: aname, articleName, name, title, documentTitle
"""
if not isinstance(obj, dict):
return None
candidates = [
"aname",
"articleName",
"article_name",
"name",
"title",
"documentTitle",
"document_title",
]
for key in candidates:
if key in obj:
v = obj.get(key)
if isinstance(v, str) and v.strip():
return v.strip()
# try shallow nested search for non-empty strings
for v in obj.values():
if isinstance(v, str) and v.strip():
# skip values that look like html (we prefer explicit name keys)
if "<" in v or ">" in v:
continue
# short string looks promising
if len(v.strip()) <= 200:
return v.strip()
return None
#def find_article_name_in_json(obj) -> str | None:
# """Look for an article name in known keys or nested values.
#
# Common keys include: aname, articleName, name, title, documentTitle
# """
# if not isinstance(obj, dict):
# return None
#
# candidates = [
# "aname",
# "articleName",
# "article_name",
# "name",
# "title",
# "documentTitle",
# "document_title",
# ]
# for key in candidates:
# if key in obj:
# v = obj.get(key)
# if isinstance(v, str) and v.strip():
# return v.strip()
#
# # try shallow nested search for non-empty strings
# for v in obj.values():
# if isinstance(v, str) and v.strip():
# # skip values that look like html (we prefer explicit name keys)
# if "<" in v or ">" in v:
# continue
# # short string looks promising
# if len(v.strip()) <= 200:
# return v.strip()
#
# return None
def extract_title_from_html(html: str) -> str | None:
@@ -297,68 +299,6 @@ def recursive_search_for_value(obj, value: str) -> bool:
return False
def find_title_for_docid(docid: str, search_dir: str) -> str | None:
"""Look for files in search_dir that contain the docid and try to extract a title.
This checks nearby summary/media JSONs which often contain a `title` field.
"""
if not docid:
return None
# Phase 1: scan all JSON files for breadcrumb/list-style captures and prefer
# their enhancedDocumentName/name when present. Breadcrumb captures often
# don't include the docid in their filename, so scan every JSON once.
all_paths = sorted(glob.glob(os.path.join(search_dir, "*.json")))
for path in all_paths:
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception:
continue
if isinstance(data, list):
for item in data:
try:
if isinstance(item, dict) and item.get("documentId") == docid:
ed = item.get("enhancedDocumentName") or item.get("name")
if isinstance(ed, str) and ed.strip():
return html.unescape(ed.strip())
except Exception:
continue
# Phase 2: look only at files whose filename contains the docid for other
# title signals (top-level names, recursive title keys, searchResults).
pattern = os.path.join(search_dir, f"*{docid}*.json")
paths = sorted(glob.glob(pattern))
for path in paths:
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception:
continue
# top-level keys commonly used
if isinstance(data, dict):
for key in ("enhancedDocumentName", "documentName", "name"):
if key in data:
v = data.get(key)
if isinstance(v, str) and v.strip():
return html.unescape(v.strip())
# fallback: direct title anywhere in the JSON
title = recursive_search_for_key(data, "title")
if title and isinstance(title, str) and title.strip():
return html.unescape(title.strip())
# sometimes summary objects contain results with titles
sr = data.get("searchResults") if isinstance(data, dict) else None
if isinstance(sr, dict):
results = sr.get("results")
if isinstance(results, list):
for r in results:
if isinstance(r, dict) and r.get("id") == docid and r.get("title"):
return html.unescape(r.get("title"))
return None
def find_breadcrumbs_for_docid(docid: str, search_dir: str) -> list:
@@ -531,6 +471,8 @@ def process_file(path: str, out_dir: str, overwrite: bool = False) -> tuple[bool
base = os.path.basename(path)
name, _ = os.path.splitext(base)
logger.debug(f"Processing file: {path}")
# attempt to extract article name for nicer filenames
try:
with open(path, "r", encoding="utf-8") as f:
@@ -541,33 +483,22 @@ def process_file(path: str, out_dir: str, overwrite: bool = False) -> tuple[bool
# deterministic doc id (if present in filename)
docid = None
article_name = None
if isinstance(data_peek, dict):
article_name = find_article_name_in_json(data_peek)
docid = base.split("_document_content_")[1].split("_")[0]
# if no explicit article name, try to infer from docid present in filename
if not article_name:
# try to parse a uuid-like docid from filename
m = re.search(r"([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})", base)
if m:
docid = m.group(1)
# Prefer nearby breadcrumb/document captures (which include
# enhancedDocumentName/name) over titles that appear in
# capture_index.jsonl. This makes section-level names
# like "Brain Tumor in Newborn/Infant" the canonical title.
title_from_nearby = find_title_for_docid(docid, os.path.dirname(path) or ".")
if not title_from_nearby:
title_from_nearby = find_title_in_capture_index(docid, os.path.dirname(path) or ".")
if title_from_nearby:
article_name = title_from_nearby
logger.debug(f"Extracted docid: {docid}")
#logger.debug(f"DOCUMENT_SUMMARYS keys: {DOCUMENT_SUMMARYS.get(docid)}")
article_name = DOCUMENT_SUMMARYS.get(docid).get("title")
# if still not found, try to extract from inline HTML
if not article_name and isinstance(data_peek, dict):
html_candidate = find_html_in_json(data_peek)
if html_candidate:
title_from_html = extract_title_from_html(html_candidate)
if title_from_html:
article_name = title_from_html
#if not article_name and isinstance(data_peek, dict):
# html_candidate = find_html_in_json(data_peek)
# if html_candidate:
# title_from_html = extract_title_from_html(html_candidate)
# if title_from_html:
# article_name = title_from_html
logger.debug(f"Determined article name: {article_name}")
if article_name:
slug = slugify(article_name)
@@ -584,6 +515,7 @@ def process_file(path: str, out_dir: str, overwrite: bool = False) -> tuple[bool
target_dir = os.path.join(out_dir, "external")
else:
# per request: only save files that have titles
logger.error(f"No article name found for {path}; skipping.")
return False, "no-title"
out_path = os.path.join(target_dir, out_base + ".md")
@@ -647,9 +579,9 @@ def process_file(path: str, out_dir: str, overwrite: bool = False) -> tuple[bool
image_data_to_add = {}
if "imageGroups" in data_peek:
if isinstance(data_peek, dict) and "imageGroups" in data_peek:
for img_group in data_peek.get("imageGroups", []):
logger.debug(f"Image group: {pformat(img_group)}")
#logger.debug(f"Image group: {pformat(img_group)}")
group_id = img_group.get("imageGroupId")
try:
@@ -659,7 +591,133 @@ def process_file(path: str, out_dir: str, overwrite: bool = False) -> tuple[bool
except KeyError:
logger.warning(f"Image group {group_id} not found in IMAGE_GROUPS.")
logger.debug(f"Extracted image data to add: {pformat(image_data_to_add)}")
#logger.debug(f"Extracted image data to add: {pformat(image_data_to_add)}")
# If we have image groups to add, copy matching files and append a formatted
# Images section to the markdown. We search the capture input directory
# recursively for filenames containing the imageId and copy any matches.
try:
if image_data_to_add:
images_dir = os.path.join(target_dir, "images")
os.makedirs(images_dir, exist_ok=True)
appended_images = []
images_md_lines = ["\n\n## Images\n"]
# map imageId -> (rel_path, caption) for replacement in-md
image_map: dict = {}
search_root = CAPTURE_INPUT_DIR or os.path.dirname(path) or "."
for group_name, images in image_data_to_add.items():
images_md_lines.append(f"\n### {group_name}\n")
# render each image in the group
for img in images:
try:
iid = img.get('imageId') or img.get('id')
caption = img.get('caption') or img.get('imageTitle') or img.get('title') or ""
thumb = img.get('thumbnailUrl') or ""
# find matching files under the capture input dir recursively
matches = []
if iid:
matches = glob.glob(os.path.join(search_root, "**", f"*{iid}*"), recursive=True)
# also consider thumbnail URL id in case different naming
if not matches and isinstance(thumb, str) and thumb:
m = re.search(r"/thumbnail/([0-9a-fA-F-\-]+)", thumb)
if m:
tid = m.group(1)
matches = glob.glob(os.path.join(search_root, "**", f"*{tid}*"), recursive=True)
rel_paths = []
for match in sorted(set(matches)):
# skip JSON metadata files; prefer binary images
b = os.path.basename(match)
if b.endswith('.json') or b.endswith('.meta.json'):
# still copy metadata alongside image when present
# we'll copy meta files later if we also find an image
continue
try:
dst = os.path.join(images_dir, os.path.basename(match))
if os.path.abspath(match) != os.path.abspath(dst):
shutil.copy2(match, dst)
rel = os.path.join('images', os.path.basename(match))
rel_paths.append(rel)
appended_images.append(rel)
# associate imageId -> rel + caption for later replacement
if iid:
# prefer explicit caption if available
cval = caption or img.get('caption') or img.get('imageTitle') or img.get('title') or ''
image_map[iid] = (rel, cval)
except Exception:
continue
# If we found images, render them; otherwise, render the thumbnail URL if present
if rel_paths:
for rp in rel_paths:
cap = caption or os.path.basename(rp)
# render image and then visible caption on the next line
images_md_lines.append(f"![{cap}]({rp})")
# ensure caption is visible (italicized)
images_md_lines.append(f"*{cap}*")
images_md_lines.append("")
else:
# try to render a thumbnail path (as-is) if no local file
if thumb:
# normalize thumb path to a relative filename when possible
thumb_name = os.path.basename(thumb.split('?')[0])
images_md_lines.append(f"![{caption}]({thumb_name})")
images_md_lines.append(f"*{caption}*")
images_md_lines.append("")
else:
# at minimum show caption text
if caption:
images_md_lines.append(f"*{caption}*")
images_md_lines.append("")
except Exception:
continue
# attach Images section to markdown
# Before attaching, replace inline references to these images so
# their alt text shows the caption (not original alt).
try:
for iid, (rel, cap) in image_map.items():
# replace markdown image links that include the imageId or filename
bn = os.path.basename(rel)
# 1) replace markdown image links where the URL contains the imageId
md = re.sub(r'!\[[^\]]*\]\([^)]*' + re.escape(iid) + r'[^)]*\)', f'![{cap}]({rel})\n\n*{cap}*', md)
# 2) replace markdown image links where URL contains the basename
md = re.sub(r'!\[[^\]]*\]\([^)]*' + re.escape(bn) + r'[^)]*\)', f'![{cap}]({rel})\n\n*{cap}*', md)
# 3) replace HTML <img ... src="...basename..."> with markdown image
md = re.sub(r'<img[^>]+src=["\"][^"\']*' + re.escape(bn) + r'[^"\']*["\"][^>]*>', f'![{cap}]({rel})\n\n*{cap}*', md)
# 4) replace full URLs ending with basename
md = re.sub(r'https?://[^)\s]*' + re.escape(bn), rel, md)
# 5) replace URL-encoded variants
try:
q = urllib.parse.quote(bn)
md = md.replace(q, rel)
except Exception:
pass
except Exception:
pass
md = md.rstrip() + "\n" + "\n".join(images_md_lines) + "\n"
# update a top-level images index so we can see associations across runs
try:
images_index_path = os.path.join(out_dir, "_images_index.json")
images_index = {}
if os.path.exists(images_index_path):
with open(images_index_path, 'r', encoding='utf-8') as ixf:
try:
images_index = json.load(ixf)
except Exception:
images_index = {}
images_index[out_path] = appended_images
with open(images_index_path, 'w', encoding='utf-8') as ixf:
json.dump(images_index, ixf, indent=2, ensure_ascii=False)
except Exception:
pass
except Exception:
# non-fatal; continue without images
pass
# Before writing, compute a normalized content hash to detect duplicates
def normalize_for_hash(s: str) -> str:
@@ -743,10 +801,9 @@ def main(argv: Iterable[str] | None = None) -> int:
for path in files:
base = os.path.basename(path)
# only process files that match the desired prefix
if not base.startswith("app.statdx.com_document_") or base.endswith("meta.json") or not "_media_" in base:
if not base.startswith("app.statdx.com_document_") or base.endswith("meta.json") or ("_media_" not in base):
continue
logger.debug(f"Caching image groups from media file: {path}")
with open(path, "r", encoding="utf-8") as f:
image_group_data = json.load(f)
@@ -760,14 +817,29 @@ def main(argv: Iterable[str] | None = None) -> int:
"images": images,
}
#logger.debug(f"Image group {n}: {pformat(group)}")
# Then cache document summary data
for path in files:
base = os.path.basename(path)
# only process files that match the desired prefix
if "document_summary" not in base or "meta" in base:
continue
with open(path, "r", encoding="utf-8") as f:
summary_data = json.load(f)
doc_id = base.split("_document_summary_")[1].split("_")[0]
DOCUMENT_SUMMARYS[doc_id] = summary_data #logger.debug(f"Document summary {n}: {pformat(doc)}")
logger.debug(f"Cached {len(IMAGE_GROUPS)} image groups from media files.")
logger.debug(f"Cached {len(DOCUMENT_SUMMARYS)} document summaries from summary files.")
ok = 0
for path in files:
base = os.path.basename(path)
# only process files that match the desired prefix
if not base.startswith("app.statdx.com_document_") or base.endswith("meta.json"):
if not base.startswith("app.statdx.com_document_") or base.endswith("meta.json") or ("_media_" in base) or ("breadcrumbs" in base) or ("_summary_" in base) or ("_tables_" in base):
#if args.verbose:
# print(f"SKIP (not matching prefix): {path}")
continue