diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index d1397bc38..f34cc42c0 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -445,6 +445,7 @@ def _build_page_dataframe( vlm_model=vlm_model, page_assets_by_page=page_assets_by_page, node_summary_max_pages=page_memory_config.node_summary_max_pages, + node_assembly_concurrency=page_memory_config.node_assembly_concurrency, ) logger.info( "[page_memory] C7 assembled {} node rows (verdict={})", @@ -625,14 +626,15 @@ def _run_hierarchy_scope( fat_leaf_pages = compute_fat_leaf_pages(scope_skeletons, min_pages=fine_min) if fat_leaf_pages: title_pages = sorted(fat_leaf_pages) - title_rendered = render_document_pages( - pdf_path=pdf_path, - page_count=page_count, - output_dir=output_dir, - pages=title_pages, - page_features=page_features, - page_texts=page_texts, - ) + with stage_timer("page_memory.title_render", page_count=len(title_pages)): + title_rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + pages=title_pages, + page_features=page_features, + page_texts=page_texts, + ) title_tags = [ PageTagResult( page_index=page, @@ -649,6 +651,7 @@ def _run_hierarchy_scope( fat_leaf_pages=fat_leaf_pages, budget=None, vlm_model=vlm_model, + max_concurrent=page_memory_config.title_detection_concurrency, ) _record_trace_stage( trace_recorder, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index ff4191a25..7f23d4998 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -25,10 +25,11 @@ import shutil from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, cast from loguru import logger +from app.services.document_parser.support.stage_profiler import stage_timer from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.document_parser.support.parser_rows import serialize_entities from app.services.page_memory.page_assets import ( @@ -416,63 +417,114 @@ def build_node_rows( vlm_model: str | None = None, page_assets_by_page: dict[int, list[PageAsset]] | None = None, node_summary_max_pages: int = _NODE_SUMMARY_MAX_PAGES_DEFAULT, + node_assembly_concurrency: int = 3, ) -> list[dict[str, Any]]: """Assemble one row per leaf section node (node-granularity chunks).""" available_pages = set(raw_text_by_page.keys()) leaves = identify_leaf_nodes(skeletons) views, page_owner = assign_pages_to_leaves(leaves, available_pages=available_pages) page_to_leaves = pages_by_leaf_count(views) + resolved_concurrency = max(1, node_assembly_concurrency) # Resolve body text once per owned page (PyMuPDF, OCR fallback for scanned). resolved_text: dict[int, str] = {} - for view in views: - for page in view.owned_pages: - resolved_text[page] = resolve_page_text( + owned_pages = sorted({page for view in views for page in view.owned_pages}) + if owned_pages: + import gevent + from gevent.pool import Pool as GeventPool + + def _resolve_one(page: int) -> tuple[int, str]: + return page, resolve_page_text( page=page, raw_text=raw_text_by_page.get(page, ""), image_path=image_path_by_page.get(page), vlm_model=vlm_model, ) + with stage_timer( + "page_memory.node_ocr", + page_count=len(owned_pages), + concurrency=resolved_concurrency, + ): + pool = GeventPool(size=min(resolved_concurrency, len(owned_pages))) + greenlets = [pool.spawn(_resolve_one, page) for page in owned_pages] + gevent.joinall(greenlets, raise_error=True) + resolved_pairs = [ + cast(tuple[int, str], greenlet.value) + for greenlet in greenlets + ] + resolved_text = {page: text for page, text in resolved_pairs} + + summaries: dict[int, tuple[str, list[str], list[dict[str, str]]]] = {} + if views: + import gevent + from gevent.pool import Pool as GeventPool + + def _summarize_one( + index: int, + ) -> tuple[int, tuple[str, list[str], list[dict[str, str]]]]: + return index, compute_node_summary( + view=views[index], + page_to_leaves=page_to_leaves, + tag_by_page=tag_by_page, + image_path_by_page=image_path_by_page, + vlm_model=vlm_model, + node_summary_max_pages=node_summary_max_pages, + ) + + with stage_timer( + "page_memory.node_summary", + node_count=len(views), + concurrency=resolved_concurrency, + ): + pool = GeventPool(size=min(resolved_concurrency, len(views))) + greenlets = [ + pool.spawn(_summarize_one, index) + for index in range(len(views)) + ] + gevent.joinall(greenlets, raise_error=True) + summary_pairs = [ + cast( + tuple[int, tuple[str, list[str], list[dict[str, str]]]], + greenlet.value, + ) + for greenlet in greenlets + ] + summaries = {index: result for index, result in summary_pairs} + rows: list[dict[str, Any]] = [] rows_by_path: dict[str, dict[str, Any]] = {} - for view in views: - leaf = view.leaf - content = build_node_content( - view, - page_owner=page_owner, - page_text=resolved_text, - ) - summary, keywords, entities = compute_node_summary( - view=view, - page_to_leaves=page_to_leaves, - tag_by_page=tag_by_page, - image_path_by_page=image_path_by_page, - vlm_model=vlm_model, - node_summary_max_pages=node_summary_max_pages, - ) - know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" - row = { - "content": content, - "path": leaf.section_path, - "type": "page", - "length": len(content), - "keywords": ";".join(keywords), - "summary": summary, - "know_id": know_id, - "tokens": "", - "connectto": "", - "addtime": get_str_time(), - "page_nums": ",".join(str(page) for page in view.pages), - "entities": serialize_entities(entities), - "asset_title": "", - "extra_metadata": _build_page_extra_metadata( - pages=view.pages, - image_path_by_page=image_path_by_page, - ), - } - rows.append(row) - rows_by_path[leaf.section_path] = row + with stage_timer("page_memory.node_rows", node_count=len(views)): + for index, view in enumerate(views): + leaf = view.leaf + content = build_node_content( + view, + page_owner=page_owner, + page_text=resolved_text, + ) + summary, keywords, entities = summaries.get(index, ("", [], [])) + know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" + row = { + "content": content, + "path": leaf.section_path, + "type": "page", + "length": len(content), + "keywords": ";".join(keywords), + "summary": summary, + "know_id": know_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": ",".join(str(page) for page in view.pages), + "entities": serialize_entities(entities), + "asset_title": "", + "extra_metadata": _build_page_extra_metadata( + pages=view.pages, + image_path_by_page=image_path_by_page, + ), + } + rows.append(row) + rows_by_path[leaf.section_path] = row asset_rows: list[dict[str, Any]] = [] if page_assets_by_page: diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index fcee9da5f..4fe72ac23 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -26,6 +26,7 @@ from app.services.page_memory.page_renderer import PageRenderResult from shared.services.ai.prompt_service import build_prompt from shared.services.ai.summary.engine import summarize +from shared.core.exceptions.domain_exceptions import UnavailableException @dataclass @@ -90,6 +91,8 @@ def tag_pages( resolved_max_concurrent = max_concurrent or int( getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 4) ) + if not pages: + return [] def _tag_one(page: PageRenderResult) -> PageTagResult: plan = plan_map.get(page.page_index) @@ -112,9 +115,9 @@ def _tag_one(page: PageRenderResult) -> PageTagResult: pool = GeventPool(size=min(resolved_max_concurrent, len(pages))) greenlets = [pool.spawn(_tag_one, page) for page in pages] - gevent.joinall(greenlets) + gevent.joinall(greenlets, raise_error=True) - results = [g.value for g in greenlets if g.value is not None] + results = [cast(PageTagResult, g.value) for g in greenlets] vlm_calls = sum(1 for r in results if r.strategy_used == "vlm_lite") logger.info( "[page_tagger] tagged {} pages ({} VLM calls, {} text_only, {} skipped, {} failed) concurrency={}", @@ -182,6 +185,8 @@ def _tag_text_only( response_format={"type": "json_object"}, usage_task="page_memory.text_tag", ) + except UnavailableException: + raise except Exception as exc: logger.warning( "[page_tagger] text_only LLM failed for page {}: {}", @@ -262,6 +267,7 @@ def tag_page_titles( fat_leaf_pages: set[int], budget: Any | None = None, vlm_model: str | None = None, + max_concurrent: int | None = None, ) -> list[PageTagResult]: """Run independent VLM title detection on fat-leaf pages. @@ -278,6 +284,8 @@ def tag_page_titles( Deprecated, ignored. Kept for call-site compatibility. vlm_model: VLM model name; falls back to ``$IMAGE_MODEL``. + max_concurrent: + Maximum concurrent title-detection calls. Returns ------- @@ -294,29 +302,57 @@ def tag_page_titles( tag_map = {t.page_index: t for t in tag_results} page_map = {p.page_index: p for p in pages} - vlm_calls = 0 - titles_found = 0 + work_items = [ + (page_idx, page, tag_map[page_idx]) + for page_idx in sorted(fat_leaf_pages) + if (page := page_map.get(page_idx)) is not None + and page_idx in tag_map + and page.image_path + and os.path.exists(page.image_path) + ] + if not work_items: + return tag_results - for page_idx in sorted(fat_leaf_pages): - page = page_map.get(page_idx) - tag = tag_map.get(page_idx) - if page is None or tag is None: - continue + from shared.core.config import settings + + resolved_max_concurrent = max_concurrent or int( + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 3) + ) - # Skip pages without images (text_only / skip) - if not page.image_path or not os.path.exists(page.image_path): - continue + def _detect_one( + page_idx: int, + page: PageRenderResult, + ) -> tuple[int, list[dict[str, Any]]]: + return page_idx, _tag_vlm_titles(page, model=model) - observed = _tag_vlm_titles(page, model=model) + import gevent + from gevent.pool import Pool as GeventPool + + pool = GeventPool(size=min(resolved_max_concurrent, len(work_items))) + greenlets = [ + pool.spawn(_detect_one, page_idx, page) + for page_idx, page, _tag in work_items + ] + gevent.joinall(greenlets, raise_error=True) + + title_pairs = [ + cast(tuple[int, list[dict[str, Any]]], greenlet.value) + for greenlet in greenlets + ] + titles_by_page: dict[int, list[dict[str, Any]]] = dict(title_pairs) + for page_idx, _page, tag in work_items: + observed = titles_by_page.get(page_idx, []) tag.observed_titles = observed - vlm_calls += 1 - titles_found += len(observed) + + vlm_calls = len(work_items) + titles_found = sum(len(titles) for titles in titles_by_page.values()) logger.info( - "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found", + "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found concurrency={}", vlm_calls, len(fat_leaf_pages), titles_found, + resolved_max_concurrent, ) return tag_results @@ -413,6 +449,8 @@ def _tag_vlm_titles( page.page_index, ) return [] + except UnavailableException: + raise except Exception as exc: logger.warning( "[page_tagger] title VLM failed for page {}: {}", diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index c66fc217b..21304323e 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -9,18 +9,15 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") -from app.services.page_memory.node_assembler import ( - SAME_AS_PREFIX, - assign_pages_to_leaves, - build_node_content, - build_node_rows, - identify_leaf_nodes, -) +import pytest + +import app.services.page_memory.node_assembler as node_assembler from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS from app.services.page_memory.page_assets import PageAsset from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.core.exceptions.domain_exceptions import UnavailableException import pandas as pd from PIL import Image @@ -54,14 +51,30 @@ def _same_page_sibling_skeletons() -> list[SectionSkeleton]: return [parent, child_a, child_b] +def _ordered_page_skeletons() -> list[SectionSkeleton]: + return [ + SectionSkeleton( + section_path=f"demo.pdf/{page_index}", + level=1, + start_page=page_index, + end_page=page_index, + title=f"Section {page_index}", + parent_path="demo.pdf", + ) + for page_index in [1, 2, 3] + ] + + def test_identify_leaf_nodes_drops_internal_parents() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) assert [leaf.title for leaf in leaves] == ["3.1 职责", "3.2 管理规定"] def test_page_ownership_first_leaf_owns_shared_page() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) - views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = node_assembler.assign_pages_to_leaves( + leaves, available_pages={231, 232} + ) assert page_owner[231].title == "3.1 职责" assert page_owner[232].title == "3.2 管理规定" @@ -73,26 +86,30 @@ def test_page_ownership_first_leaf_owns_shared_page() -> None: def test_build_node_content_uses_same_as_for_shared_page() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) - views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = node_assembler.assign_pages_to_leaves( + leaves, available_pages={231, 232} + ) by_title = {view.leaf.title: view for view in views} page_text = {231: "text-231", 232: "text-232"} - content_a = build_node_content( + content_a = node_assembler.build_node_content( by_title["3.1 职责"], page_owner=page_owner, page_text=page_text ) - content_b = build_node_content( + content_b = node_assembler.build_node_content( by_title["3.2 管理规定"], page_owner=page_owner, page_text=page_text ) assert content_a == "text-231" - assert content_b.startswith(f"[{SAME_AS_PREFIX} demo.pdf/3 基本规定/3.1 职责 p231]") + assert content_b.startswith( + f"[{node_assembler.SAME_AS_PREFIX} demo.pdf/3 基本规定/3.1 职责 p231]" + ) assert "text-232" in content_b assert "text-231" not in content_b def test_build_node_rows_reuses_tags_without_vlm() -> None: - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, @@ -120,17 +137,154 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: leaf_b = by_path["demo.pdf/3 基本规定/3.2 管理规定"] assert leaf_b["page_nums"] == "231,232" - assert SAME_AS_PREFIX in leaf_b["content"] + assert node_assembler.SAME_AS_PREFIX in leaf_b["content"] assert "text-232" in leaf_b["content"] assert leaf_b["extra_metadata"] == {} +def test_build_node_rows_preserves_order_under_ocr_and_summary_concurrency( + monkeypatch, +) -> None: + import gevent + + def _fake_resolve_page_text(**kwargs) -> str: + page = int(kwargs["page"]) + gevent.sleep(0.01 * (4 - page)) + return f"text-{page}" + + def _fake_compute_node_summary(**kwargs): + view = kwargs["view"] + gevent.sleep(0.01 * view.leaf.start_page) + return f"summary-{view.leaf.start_page}", [f"k{view.leaf.start_page}"], [] + + monkeypatch.setattr( + node_assembler, + "resolve_page_text", + _fake_resolve_page_text, + ) + monkeypatch.setattr( + node_assembler, + "compute_node_summary", + _fake_compute_node_summary, + ) + + rows = node_assembler.build_node_rows( + skeletons=_ordered_page_skeletons(), + raw_text_by_page={1: "", 2: "", 3: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=2, + ) + + assert [row["path"] for row in rows] == [ + "demo.pdf/1", + "demo.pdf/2", + "demo.pdf/3", + ] + assert [row["content"] for row in rows] == ["text-1", "text-2", "text-3"] + assert [row["summary"] for row in rows] == [ + "summary-1", + "summary-2", + "summary-3", + ] + + +def test_build_node_rows_failed_ocr_greenlet_fails_stage(monkeypatch) -> None: + def _fake_resolve_page_text(**kwargs) -> str: + if int(kwargs["page"]) == 2: + raise RuntimeError("ocr failed") + return "ok" + + monkeypatch.setattr( + node_assembler, + "resolve_page_text", + _fake_resolve_page_text, + ) + + with pytest.raises(RuntimeError): + node_assembler.build_node_rows( + skeletons=_ordered_page_skeletons()[:2], + raw_text_by_page={1: "", 2: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=2, + ) + + +def test_build_node_rows_unavailable_propagates_from_ocr(monkeypatch) -> None: + def _fake_resolve_page_text(**kwargs) -> str: + raise UnavailableException( + internal_message="ocr capacity busy", + retry_after=5, + ) + + monkeypatch.setattr( + node_assembler, + "resolve_page_text", + _fake_resolve_page_text, + ) + + with pytest.raises(UnavailableException): + node_assembler.build_node_rows( + skeletons=_ordered_page_skeletons()[:1], + raw_text_by_page={1: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=1, + ) + + +def test_build_node_rows_unavailable_propagates_from_node_summary( + monkeypatch, +) -> None: + def _fake_compute_node_summary(**kwargs): + raise UnavailableException( + internal_message="summary capacity busy", + retry_after=5, + ) + + monkeypatch.setattr( + node_assembler, + "compute_node_summary", + _fake_compute_node_summary, + ) + + with pytest.raises(UnavailableException): + node_assembler.build_node_rows( + skeletons=_ordered_page_skeletons()[:1], + raw_text_by_page={1: "text-1"}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=1, + ) + + def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_path) -> None: page_image = tmp_path / "pages" / "page-231.png" page_image.parent.mkdir() Image.new("RGB", (2, 3), color=(255, 255, 255)).save(page_image) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(page_image)}, @@ -179,7 +333,7 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: parent_path="demo.pdf/4 风险辨识与分级管控", ) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=[parent, child], raw_text_by_page={233: "parent body", 234: "child body"}, image_path_by_page={}, @@ -225,8 +379,8 @@ def test_boundary_page_belongs_to_next_sibling_start() -> None: parent_path="demo.pdf/安全类", ) - leaves = identify_leaf_nodes([coarse, next_sibling]) - views, page_owner = assign_pages_to_leaves( + leaves = node_assembler.identify_leaf_nodes([coarse, next_sibling]) + views, page_owner = node_assembler.assign_pages_to_leaves( leaves, available_pages={301, 302, 303}, ) @@ -262,7 +416,7 @@ def chat_completion_with_usage(self, **kwargs): img = tmp_path / "page-231.png" img.write_bytes(b"\x89PNG\r\n\x1a\n fake") - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(img), 232: str(img)}, @@ -303,7 +457,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: extraction_status="table_html_extracted", ) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, diff --git a/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py new file mode 100644 index 000000000..8eeba6360 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import os + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.page_memory.page_renderer import PageRenderResult +from app.services.page_memory.page_tagger import PageTagResult, tag_page_titles +from shared.core.exceptions.domain_exceptions import UnavailableException + + +def _write_page_image(tmp_path, page_index: int) -> str: + image_path = tmp_path / f"page-{page_index}.png" + image_path.write_bytes(b"png") + return str(image_path) + + +def test_title_detection_preserves_page_index_assignment_under_concurrency( + monkeypatch, + tmp_path, +) -> None: + import gevent + + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + gevent.sleep(0.01 * (4 - page.page_index)) + return [{"text": f"title-{page.page_index}", "prominence": 0.8}] + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + pages = [ + PageRenderResult( + page_index=page_index, + image_path=_write_page_image(tmp_path, page_index), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + for page_index in [1, 2, 3] + ] + tag_results = [ + PageTagResult(page_index=3), + PageTagResult(page_index=1), + PageTagResult(page_index=2), + ] + + results = tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages={1, 2, 3}, + vlm_model="fake-vlm", + max_concurrent=2, + ) + + observed_by_page = { + tag.page_index: tag.observed_titles[0]["text"] + for tag in results + } + assert observed_by_page == { + 1: "title-1", + 2: "title-2", + 3: "title-3", + } + + +def test_title_detection_failed_greenlet_fails_stage(monkeypatch, tmp_path) -> None: + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + if page.page_index == 2: + raise RuntimeError("title detection failed") + return [{"text": f"title-{page.page_index}"}] + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + pages = [ + PageRenderResult( + page_index=page_index, + image_path=_write_page_image(tmp_path, page_index), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + for page_index in [1, 2] + ] + tag_results = [PageTagResult(page_index=1), PageTagResult(page_index=2)] + + with pytest.raises(RuntimeError): + tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages={1, 2}, + vlm_model="fake-vlm", + max_concurrent=2, + ) + + +def test_title_detection_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + page = PageRenderResult( + page_index=1, + image_path=_write_page_image(tmp_path, 1), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + + with pytest.raises(UnavailableException): + tag_page_titles( + pages=[page], + tag_results=[PageTagResult(page_index=1)], + fat_leaf_pages={1}, + vlm_model="fake-vlm", + max_concurrent=1, + ) diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 173992f5f..33d711fd2 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -96,6 +96,34 @@ class AIConfig(BaseModel): default=4, description="Max concurrent local DOCX image summary VLM calls per parse job.", ) + PAGE_MEMORY_VLM_MAX_INFLIGHT: int = Field( + default=16, + description="Global Redis-backed in-flight limit for page-memory VLM calls.", + ) + PAGE_MEMORY_VLM_LEASE_TTL_SECONDS: int = Field( + default=600, + description="Safety TTL for page-memory VLM in-flight leases.", + ) + PAGE_MEMORY_VLM_WAIT_TIMEOUT_SECONDS: int = Field( + default=120, + description="Max wait before page-memory VLM capacity pressure becomes retryable.", + ) + PAGE_MEMORY_SCOPE_CONCURRENCY: int = Field( + default=5, + description="Local per-job hierarchy scope concurrency for page-memory.", + ) + PAGE_MEMORY_TAG_CONCURRENCY: int = Field( + default=4, + description="Local per-job page tagging concurrency for page-memory.", + ) + PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY: int = Field( + default=3, + description="Local per-job title detection concurrency for page-memory.", + ) + PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY: int = Field( + default=3, + description="Local per-job node OCR and summary concurrency for page-memory.", + ) TOKEN_PRICING_TABLE_JSON: str = Field( default="", description=( diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index fd71f49b0..bbac71ed3 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -13,6 +13,8 @@ class PageMemoryConfig: max_pages: int = 1500 scope_concurrency: int = 5 tag_concurrency: int = 4 + title_detection_concurrency: int = 3 + node_assembly_concurrency: int = 3 tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 hierarchy_model: str | None = None @@ -35,7 +37,26 @@ class PageMemoryConfig: @classmethod def default(cls) -> Self: - return cls() + from shared.core.config import settings + + return cls( + scope_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_SCOPE_CONCURRENCY", 5), + 5, + ), + tag_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_TAG_CONCURRENCY", 4), + 4, + ), + title_detection_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 3), + 3, + ), + node_assembly_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY", 3), + 3, + ), + ) @classmethod def from_mapping(cls, value: object) -> Self: @@ -57,6 +78,14 @@ def from_mapping(cls, value: object) -> Self: value.get("tag_concurrency"), default.tag_concurrency, ), + title_detection_concurrency=_as_int( + value.get("title_detection_concurrency"), + default.title_detection_concurrency, + ), + node_assembly_concurrency=_as_int( + value.get("node_assembly_concurrency"), + default.node_assembly_concurrency, + ), tag_mode=resolved_tag_mode, fine_min_pages=_as_int( value.get("fine_min_pages"), diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index f3bf53ebb..6a1cf590a 100644 --- a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -16,11 +16,15 @@ from openai.types.chat import ChatCompletionMessageParam from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.core.exceptions.domain_exceptions import LLMServiceException, UnavailableException from shared.services.http.client_pool import get_sync_client from shared.services.ai.llm_mock import build_mock_chat_completion_response from shared.utils.security_utils import mask_api_key from shared.services.ai.token_tracking import record_tokens +from shared.services.ai.page_memory_vlm_limiter import ( + PageMemoryVlmLease, + get_page_memory_vlm_limiter, +) LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" LLMUsage = dict[str, int] @@ -39,6 +43,10 @@ def _should_mock_llm_calls() -> bool: return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) +def _is_page_memory_usage_task(usage_task: str | None) -> bool: + return bool(usage_task and usage_task.startswith("page_memory.")) + + def _empty_usage() -> LLMUsage: return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} @@ -171,6 +179,19 @@ def _resolve_direct_api_key( # Ali token-pool helpers # ------------------------------------------------------------------ + def _acquire_page_memory_vlm_lease( + self, *, usage_task: str | None + ) -> PageMemoryVlmLease | None: + if not _is_page_memory_usage_task(usage_task): + return None + return get_page_memory_vlm_limiter().acquire(usage_task=usage_task or "") + + def _release_page_memory_vlm_lease( + self, lease: PageMemoryVlmLease | None + ) -> None: + if lease is not None: + get_page_memory_vlm_limiter().release(lease) + def _should_use_ali_pool(self) -> bool: """Whether to route through the AliQuotaManager instead of a fixed key.""" if self._explicit_api_key: @@ -327,23 +348,24 @@ def chat_completion_raw_with_usage( ) return {"mock_content": content}, _empty_usage() - if self._should_use_ali_pool(): - return self._make_ali_pool_raw_call( - model=effective_model, - all_messages=all_messages, - temperature=temperature, - max_tokens=max_tokens, - api_kwargs=api_kwargs, - usage_task=usage_task, - ) - - client = self._client - if client is None: - raise LLMServiceException( - internal_message="OpenAI client is not initialized for direct provider requests", - provider=self.default_model, - ) + lease = self._acquire_page_memory_vlm_lease(usage_task=usage_task) try: + if self._should_use_ali_pool(): + return self._make_ali_pool_raw_call( + model=effective_model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + usage_task=usage_task, + ) + + client = self._client + if client is None: + raise LLMServiceException( + internal_message="OpenAI client is not initialized for direct provider requests", + provider=self.default_model, + ) response = client.chat.completions.create( model=effective_model, messages=all_messages, @@ -361,11 +383,14 @@ def chat_completion_raw_with_usage( return response, usage except LLMServiceException: raise + except UnavailableException: + raise except Exception as exc: + base_url = getattr(self._client, "base_url", None) logger.error( "LLM raw request failed: model={model}, base_url={base_url}, error_chain={error_chain}", model=effective_model, - base_url=client.base_url, + base_url=base_url, error_chain=_summarize_exception_chain(exc), ) raise LLMServiceException( @@ -373,6 +398,8 @@ def chat_completion_raw_with_usage( provider=self.default_model, original_exception=exc, ) from exc + finally: + self._release_page_memory_vlm_lease(lease) def chat_completion_with_usage( self, @@ -428,9 +455,10 @@ def chat_completion_with_usage( model_name=effective_model, ), _empty_usage() - # Route through Ali token pool when applicable - if self._should_use_ali_pool(): - try: + lease = self._acquire_page_memory_vlm_lease(usage_task=usage_task) + try: + # Route through Ali token pool when applicable + if self._should_use_ali_pool(): return self._make_ali_pool_call( model=effective_model, all_messages=all_messages, @@ -439,25 +467,14 @@ def chat_completion_with_usage( api_kwargs=api_kwargs, usage_task=usage_task, ) - except LLMServiceException: - raise - except Exception as exc: - logger.error(f"LLM request failed (Ali pool): model={effective_model}, error={exc}") + + # Non-Ali path: use the single pre-configured client + client = self._client + if client is None: raise LLMServiceException( - internal_message=f"API request failed: {str(exc)}", + internal_message="OpenAI client is not initialized for direct provider requests", provider=self.default_model, - original_exception=exc, - ) from exc - - # Non-Ali path: use the single pre-configured client - client = self._client - if client is None: - raise LLMServiceException( - internal_message="OpenAI client is not initialized for direct provider requests", - provider=self.default_model, - ) - - try: + ) response = client.chat.completions.create( model=effective_model, messages=all_messages, @@ -479,11 +496,14 @@ def chat_completion_with_usage( return content, usage except LLMServiceException: raise + except UnavailableException: + raise except Exception as exc: + base_url = getattr(self._client, "base_url", None) logger.error( "LLM request failed: model={model}, base_url={base_url}, error_chain={error_chain}", model=effective_model, - base_url=client.base_url, + base_url=base_url, error_chain=_summarize_exception_chain(exc), ) raise LLMServiceException( @@ -491,6 +511,8 @@ def chat_completion_with_usage( provider=self.default_model, original_exception=exc, ) from exc + finally: + self._release_page_memory_vlm_lease(lease) def chat_completion( self, diff --git a/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py b/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py new file mode 100644 index 000000000..0c04f6524 --- /dev/null +++ b/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py @@ -0,0 +1,203 @@ +"""Redis-backed global in-flight gate for page-memory VLM calls.""" + +from __future__ import annotations + +import random +import time +from dataclasses import dataclass +from typing import Optional + +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import UnavailableException +from shared.services.redis.redis_sync_service import ( + SyncRedisService, + SyncRedisServiceFactory, +) + + +@dataclass(frozen=True) +class PageMemoryVlmLease: + """A reserved page-memory VLM capacity slot.""" + + usage_task: str + acquired_at: float + current_inflight: int + + +class PageMemoryVlmLimiter: + """Global Redis counter for page-memory VLM in-flight capacity.""" + + INFLIGHT_KEY = "page_memory:vlm:inflight" + user_message = "Document parsing VLM capacity is busy. Please retry shortly." + + _ACQUIRE_SCRIPT: str = """ +local key = KEYS[1] +local max_inflight = tonumber(ARGV[1]) +local ttl_seconds = tonumber(ARGV[2]) + +local current = tonumber(redis.call('GET', key) or '0') +if current >= max_inflight then + redis.call('EXPIRE', key, ttl_seconds) + return {0, current} +end + +local new_count = redis.call('INCR', key) +redis.call('EXPIRE', key, ttl_seconds) +return {1, new_count} +""" + + _RELEASE_SCRIPT: str = """ +local key = KEYS[1] +local current = tonumber(redis.call('GET', key) or '0') +if current <= 1 then + redis.call('DEL', key) + return 0 +end +return redis.call('DECR', key) +""" + + def __init__( + self, + redis_service: SyncRedisService, + *, + max_inflight: int, + lease_ttl_seconds: int, + wait_timeout_seconds: int, + ) -> None: + self.redis = redis_service + self.max_inflight = max(1, max_inflight) + self.lease_ttl_seconds = max(1, lease_ttl_seconds) + self.wait_timeout_seconds = max(1, wait_timeout_seconds) + + @classmethod + def from_settings( + cls, + redis_service: Optional[SyncRedisService] = None, + ) -> "PageMemoryVlmLimiter": + """Create a limiter from internal page-memory settings.""" + return cls( + redis_service or SyncRedisServiceFactory.get_service(), + max_inflight=int(getattr(settings, "PAGE_MEMORY_VLM_MAX_INFLIGHT", 16)), + lease_ttl_seconds=int( + getattr(settings, "PAGE_MEMORY_VLM_LEASE_TTL_SECONDS", 600) + ), + wait_timeout_seconds=int( + getattr(settings, "PAGE_MEMORY_VLM_WAIT_TIMEOUT_SECONDS", 120) + ), + ) + + def acquire(self, *, usage_task: str) -> PageMemoryVlmLease: + """Wait for and reserve one global VLM capacity slot.""" + started_at = time.monotonic() + deadline = started_at + self.wait_timeout_seconds + last_inflight = 0 + + while time.monotonic() < deadline: + acquired, current_inflight = self._try_acquire(usage_task=usage_task) + last_inflight = current_inflight + if acquired: + wait_ms = int((time.monotonic() - started_at) * 1000) + logger.bind( + event="page_memory.vlm_gate.acquired", + usage_task=usage_task, + current_inflight=current_inflight, + max_inflight=self.max_inflight, + wait_ms=wait_ms, + ).info("page_memory.vlm_gate.acquired") + return PageMemoryVlmLease( + usage_task=usage_task, + acquired_at=time.monotonic(), + current_inflight=current_inflight, + ) + + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._sleep(min(self._jittered_backoff(last_inflight), remaining)) + + raise UnavailableException( + internal_message=( + "Page-memory VLM capacity gate timed out " + f"after {self.wait_timeout_seconds}s " + f"({last_inflight}/{self.max_inflight} in-flight)" + ), + retry_after=max(1, min(self.wait_timeout_seconds, 60)), + limit=self.max_inflight, + period="minute", + user_message=self.user_message, + ) + + def release(self, lease: PageMemoryVlmLease) -> None: + """Release a previously acquired page-memory VLM capacity slot.""" + try: + self.redis.eval( + self._RELEASE_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[], + ) + except Exception: + logger.opt(exception=True).warning( + "Failed to release page-memory VLM in-flight slot for {}", + lease.usage_task, + ) + + def get_inflight_count(self) -> int: + """Return the current Redis in-flight count.""" + try: + raw = self.redis.get(self.INFLIGHT_KEY, 0) + return max(0, int(raw or 0)) + except Exception: + return 0 + + def _try_acquire(self, *, usage_task: str) -> tuple[bool, int]: + try: + result = self.redis.eval( + self._ACQUIRE_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[self.max_inflight, self.lease_ttl_seconds], + ) + except Exception as exc: + logger.opt(exception=True).warning( + "Failed to check page-memory VLM in-flight counter for {}", + usage_task, + ) + raise UnavailableException( + internal_message="Failed to check page-memory VLM capacity gate", + retry_after=5, + user_message=self.user_message, + original_exception=exc, + ) from exc + + if not isinstance(result, list) or len(result) < 2: + raise UnavailableException( + internal_message=f"Unexpected page-memory VLM gate result: {result!r}", + retry_after=5, + user_message=self.user_message, + ) + return int(result[0]) == 1, max(0, int(result[1])) + + def _jittered_backoff(self, current_inflight: int) -> float: + pressure = max(current_inflight - self.max_inflight + 1, 1) + base = min(0.25 * pressure, 2.0) + return base + random.uniform(0.0, 0.2) + + def _sleep(self, seconds: float) -> None: + try: + import gevent + + gevent.sleep(seconds) + except ImportError: + time.sleep(seconds) + + +_page_memory_vlm_limiter: Optional[PageMemoryVlmLimiter] = None + + +def get_page_memory_vlm_limiter() -> PageMemoryVlmLimiter: + """Return a singleton page-memory VLM limiter for worker processes.""" + global _page_memory_vlm_limiter + if _page_memory_vlm_limiter is None: + _page_memory_vlm_limiter = PageMemoryVlmLimiter.from_settings() + return _page_memory_vlm_limiter diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index 7a71d1f1b..eaedd05c2 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -29,6 +29,7 @@ from loguru import logger +from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.ai import openai_compatible_client_sync as _client_mod from shared.services.ai.prompt_service import _detect_text_language, build_prompt from shared.services.ai.response_process_service import eval_response @@ -177,6 +178,12 @@ def _call_llm( ) continue return None + except UnavailableException: + if budget is not None: + budget.refund(budget_pool, est=est, stage=budget_stage) + if usage_task.startswith("page_memory."): + raise + return None except Exception as exc: logger.warning("[summary] LLM call failed for {}: {}", usage_task, exc) if budget is not None: diff --git a/packages/shared-python/shared/tests/test_page_memory_config.py b/packages/shared-python/shared/tests/test_page_memory_config.py new file mode 100644 index 000000000..ea97ad773 --- /dev/null +++ b/packages/shared-python/shared/tests/test_page_memory_config.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.models.schemas.page_memory_config import PageMemoryConfig + + +def test_page_memory_config_defaults_resolve_concurrency_settings() -> None: + config = PageMemoryConfig.default() + + assert config.scope_concurrency == 5 + assert config.tag_concurrency == 4 + assert config.title_detection_concurrency == 3 + assert config.node_assembly_concurrency == 3 diff --git a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py new file mode 100644 index 000000000..c69854395 --- /dev/null +++ b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +import os +from typing import Any + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.core.exceptions.domain_exceptions import ( # noqa: E402 + LLMServiceException, + UnavailableException, +) +import shared.services.ai.openai_compatible_client_sync as client_mod # noqa: E402 +from shared.services.ai.page_memory_vlm_limiter import ( # noqa: E402 + PageMemoryVlmLease, + PageMemoryVlmLimiter, +) +import shared.services.ai.summary.engine as summary_engine # noqa: E402 + + +class _FakeRedis: + def __init__(self, *, initial_count: int = 0, fail_eval: bool = False) -> None: + self.count = initial_count + self.fail_eval = fail_eval + + def eval(self, script: str, keys: list[str], args: list[Any]) -> list[int] | int: + if self.fail_eval: + raise RuntimeError("redis unavailable") + if "INCR" in script: + max_inflight = int(args[0]) + if self.count >= max_inflight: + return [0, self.count] + self.count += 1 + return [1, self.count] + self.count = max(self.count - 1, 0) + return self.count + + def get(self, key: str, default: Any = None) -> int: + return self.count + + +class _DelayedCapacityRedis(_FakeRedis): + def __init__(self) -> None: + super().__init__(initial_count=1) + self.acquire_attempts = 0 + + def eval(self, script: str, keys: list[str], args: list[Any]) -> list[int] | int: + if "INCR" not in script: + return super().eval(script, keys, args) + self.acquire_attempts += 1 + if self.acquire_attempts == 1: + return [0, 1] + self.count = 1 + return [1, 1] + + +def _build_limiter(redis: _FakeRedis) -> PageMemoryVlmLimiter: + limiter = PageMemoryVlmLimiter( + redis, + max_inflight=1, + lease_ttl_seconds=60, + wait_timeout_seconds=1, + ) + limiter._jittered_backoff = lambda current_inflight: 0.001 # type: ignore[method-assign] + return limiter + + +def test_page_memory_vlm_limiter_acquire_release_updates_inflight_count() -> None: + redis = _FakeRedis() + limiter = _build_limiter(redis) + + lease = limiter.acquire(usage_task="page_memory.tag") + assert lease.current_inflight == 1 + assert limiter.get_inflight_count() == 1 + + limiter.release(lease) + assert limiter.get_inflight_count() == 0 + + +def test_page_memory_vlm_limiter_waits_then_succeeds() -> None: + redis = _DelayedCapacityRedis() + limiter = _build_limiter(redis) + + lease = limiter.acquire(usage_task="page_memory.title_detection") + + assert lease.current_inflight == 1 + assert redis.acquire_attempts == 2 + + +def test_page_memory_vlm_limiter_timeout_raises_unavailable() -> None: + redis = _FakeRedis(initial_count=1) + limiter = PageMemoryVlmLimiter( + redis, + max_inflight=1, + lease_ttl_seconds=60, + wait_timeout_seconds=1, + ) + limiter._jittered_backoff = lambda current_inflight: 1.1 # type: ignore[method-assign] + + with pytest.raises(UnavailableException): + limiter.acquire(usage_task="page_memory.node_ocr") + + +def test_page_memory_vlm_limiter_redis_failure_raises_unavailable() -> None: + limiter = _build_limiter(_FakeRedis(fail_eval=True)) + + with pytest.raises(UnavailableException): + limiter.acquire(usage_task="page_memory.node_summary") + + +def test_page_memory_provider_exception_still_releases_lease(monkeypatch) -> None: + class _FakeLimiter: + def __init__(self) -> None: + self.release_count = 0 + + def acquire(self, *, usage_task: str) -> PageMemoryVlmLease: + return PageMemoryVlmLease( + usage_task=usage_task, + acquired_at=0.0, + current_inflight=1, + ) + + def release(self, lease: PageMemoryVlmLease) -> None: + self.release_count += 1 + + class _FakeCompletions: + def create(self, **kwargs: Any) -> Any: + raise RuntimeError("provider failed") + + class _FakeChat: + completions = _FakeCompletions() + + class _FakeSdkClient: + chat = _FakeChat() + base_url = "http://provider.example" + + fake_limiter = _FakeLimiter() + monkeypatch.setattr( + client_mod, + "get_page_memory_vlm_limiter", + lambda: fake_limiter, + ) + client = client_mod.OpenAICompatibleClientSync( + api_key="test", + api_url="http://provider.example/v1", + default_model="deepseek-test", + ) + client._client = _FakeSdkClient() # type: ignore[assignment] + + with pytest.raises(LLMServiceException): + client.chat_completion_with_usage( + messages="hello", + usage_task="page_memory.tag", + ) + + assert fake_limiter.release_count == 1 + + +def test_page_memory_summary_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> Any: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + image_path = tmp_path / "page.png" + image_path.write_bytes(b"png") + + monkeypatch.setattr( + summary_engine._client_mod, + "get_openai_client", + lambda model=None: _FakeClient(), + ) + + with pytest.raises(UnavailableException): + summary_engine.summarize( + mode="page", + image_paths=[str(image_path)], + model="fake-vlm", + usage_task="page_memory.node_summary", + ) + + +def test_page_memory_transcription_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> Any: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + image_path = tmp_path / "page.png" + image_path.write_bytes(b"png") + + monkeypatch.setattr( + summary_engine._client_mod, + "get_openai_client", + lambda model=None: _FakeClient(), + ) + + with pytest.raises(UnavailableException): + summary_engine.transcribe( + image_paths=[str(image_path)], + model="fake-vlm", + usage_task="page_memory.node_ocr", + )