From e6fb46560c6bb4f38afbdf23127fe897766d81f0 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 20 Aug 2026 13:33:06 +0100 Subject: [PATCH] fix: minor graph rag improvements --- .../agents/workflows/workflow_engine.py | 43 +++++++- application/core/settings.py | 6 ++ application/llm/openai.py | 70 ++++++++++-- application/parser/file/docling_parser.py | 47 ++++++-- application/retriever/graph_rag.py | 11 ++ tests/agents/test_workflow_input_documents.py | 81 ++++++++++++++ .../agents/test_workflow_node_attachments.py | 6 +- tests/llm/test_openai.py | 78 ++++++++++++++ tests/parser/file/test_docling_parser.py | 101 ++++++++++++++++-- tests/retriever/test_graph_rag.py | 58 +++++++++- 10 files changed, 461 insertions(+), 40 deletions(-) diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index cb362be5..91d58443 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -1,6 +1,7 @@ import json import logging import re +import time import uuid from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING @@ -750,6 +751,12 @@ class WorkflowEngine: supports_images = any(t.startswith("image/") for t in supported) max_files = int(getattr(settings, "WORKFLOW_NODE_NATIVE_MAX_FILES", 5)) extract_max = int(getattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 5)) + # One wall clock for every blocking parse this node issues. The cap + # above bounds how MANY parses run; this bounds how LONG they take in + # total, so N documents cannot serialize N size-scaled windows. + parse_deadline = time.monotonic() + float( + getattr(settings, "WORKFLOW_NODE_EXTRACT_BUDGET_SECONDS", 900) + ) max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024)) # One read-only connection for the whole batch; the resolved-version @@ -835,7 +842,8 @@ class WorkflowEngine: continue extract_count += 1 content = self._extract_attachment_text( - artifact_id, storage_path, mime_type, filename, max_bytes, size=size + artifact_id, storage_path, mime_type, filename, max_bytes, + size=size, deadline=parse_deadline, ) if content is None: logger.warning( @@ -893,8 +901,14 @@ class WorkflowEngine: filename: str, max_bytes: int, size: Optional[int] = None, + deadline: Optional[float] = None, ) -> Optional[str]: - """Get an attachment's text: reuse upload-time extraction, else inline text mimes, else parse.""" + """Get an attachment's text: reuse upload-time extraction, else inline text mimes, else parse. + + Args: + deadline: ``time.monotonic()`` value past which no further blocking + parse may run, shared across the node's documents. + """ from application.parser.document_reader import truncate_text_head_tail from application.storage.storage_creator import StorageCreator @@ -933,10 +947,22 @@ class WorkflowEngine: return truncate_text_head_tail(text) # Non-text mimes parse via the dedicated parsing queue (works on any backend, # no sandbox): the worker re-resolves the artifact run-scoped and reads its bytes. - return self._parse_document_text(artifact_id, size=size) + return self._parse_document_text(artifact_id, size=size, deadline=deadline) - def _parse_document_text(self, artifact_id: str, size: Optional[int] = None) -> Optional[str]: - """Enqueue ``parse_document`` for this run and await the size-scaled markdown; None on failure.""" + def _parse_document_text( + self, + artifact_id: str, + size: Optional[int] = None, + deadline: Optional[float] = None, + ) -> Optional[str]: + """Enqueue ``parse_document`` for this run and await the size-scaled markdown; None on failure. + + Args: + artifact_id: Run-scoped artifact to parse. + size: Stored byte size, used to scale the per-document window. + deadline: ``time.monotonic()`` value the await must not outlive, + shared with the node's other documents. + """ from celery.exceptions import TimeoutError as CeleryTimeoutError from application.api.user.tasks import ( @@ -955,6 +981,13 @@ class WorkflowEngine: # (floored at DOCUMENT_PARSE_TIMEOUT); the task's per-call time limits are # raised to match, else the worker would self-terminate mid-parse. timeout = parse_timeout_for_size(size) + if deadline is not None: + timeout = min(timeout, deadline - time.monotonic()) + if timeout <= 0: + logger.warning( + "Workflow node: parse budget exhausted; skipping %s", artifact_id + ) + return None try: async_result = parse_document.apply_async( args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options], diff --git a/application/core/settings.py b/application/core/settings.py index d89e06d3..c5dc0d2a 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -462,6 +462,12 @@ class Settings(BaseSettings): # avoid serializing dozens of parses; documents past the cap are skipped with # a truncation note instead of extracted. WORKFLOW_NODE_EXTRACT_MAX_FILES: int = 5 + # Total wall clock one node may spend on blocking document parses, shared + # across all of them. The per-document window scales with size (up to + # DOCUMENT_PARSE_TIMEOUT_MAX), so without a shared budget a node could + # serialize WORKFLOW_NODE_EXTRACT_MAX_FILES full windows and hold a web + # threadpool slot for that whole time. + WORKFLOW_NODE_EXTRACT_BUDGET_SECONDS: int = 900 # A workflow run row is pre-created as ``running`` and finalized when its # generator completes; a client disconnect or worker crash can strand it in # ``running`` forever. The beat reaper fails runs still ``running`` past this diff --git a/application/llm/openai.py b/application/llm/openai.py index bad917c6..439c6766 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -18,6 +18,10 @@ NO_API_KEY = "sk-no-key" # when an endpoint turns out not to serve tool calling. _TOOL_REQUEST_KEYS = ("tools", "tool_choice", "parallel_tool_calls") +# Knobs that shape *how* tools are called. A 400 naming one of these is a bad +# argument, so they are dropped on their own before tools are given up. +_TOOL_CHOICE_KEYS = ("tool_choice", "parallel_tool_calls") + # Substrings an OpenAI-compatible server puts in a 400 when it understands the # request but cannot serve tool calling: vLLM without --enable-auto-tool-choice # / --tool-call-parser, llama.cpp, TGI, Ollama and friends. Matched case @@ -26,8 +30,11 @@ _TOOL_REQUEST_KEYS = ("tools", "tool_choice", "parallel_tool_calls") _TOOLS_UNSUPPORTED_MARKERS = ( "enable-auto-tool-choice", "tool-call-parser", - "tool_choice", - "tool choice", + # Deliberately no bare "tool_choice"/"tool choice": a 400 about a malformed + # tool_choice *argument* names the parameter without meaning the endpoint + # lacks tool support. Only unambiguous phrasings belong here. + "unsupported parameter: 'tool_choice'", + "unknown field: tool_choice", "tools are not supported", "tools is not supported", "does not support tools", @@ -42,6 +49,29 @@ _TOOLS_UNSUPPORTED_MARKERS = ( ) +def _provider_message(error: Exception) -> str: + """The server's own message text, without the SDK's whole-body repr. + + ``openai._base_client`` interpolates the entire decoded JSON body into the + exception message, so matching ``str(error)`` would let a *parameter name* + ("param": "tool_choice") read as "this endpoint cannot serve tools". + + Args: + error: The exception raised by the provider SDK. + + Returns: + The server's message when the body carries one, else ``str(error)``. + """ + body = getattr(error, "body", None) + if isinstance(body, dict): + inner = body.get("error") + if isinstance(inner, dict) and inner.get("message"): + return str(inner["message"]) + if body.get("message"): + return str(body["message"]) + return str(error) + + def _is_tools_unsupported_error(error: Exception) -> bool: """Whether a 400 says the endpoint cannot serve tool calling. @@ -51,11 +81,7 @@ def _is_tools_unsupported_error(error: Exception) -> bool: Returns: True when the server's message matches a tool-calling marker. """ - parts = [str(error)] - body = getattr(error, "body", None) - if isinstance(body, dict): - parts.append(str(body.get("message") or "")) - haystack = " ".join(parts).lower() + haystack = _provider_message(error).lower() return any(marker in haystack for marker in _TOOLS_UNSUPPORTED_MARKERS) @@ -648,16 +674,38 @@ class OpenAILLM(BaseLLM): tool-less retry when the endpoint rejected tool calling. Raises: - openai.BadRequestError: Any 400 unrelated to tool calling, and - any 400 on a request that carried no tools. + openai.BadRequestError: Any 400 unrelated to tool calling, any 400 + on a request that carried no tools, and any 400 raised by the + retry itself. """ try: return create(**params) except BadRequestError as error: - if not params.get("tools") or not _is_tools_unsupported_error(error): + if not params.get("tools"): raise + if getattr(error, "param", None) in _TOOL_CHOICE_KEYS: + # A rejected tool_choice is a bad argument, not a missing + # capability -- the v1 translator forwards a client-supplied + # value verbatim. Drop just that knob and keep the tools. + logging.warning( + "Endpoint rejected %r (model=%s): %s. Retrying with tools intact.", + getattr(error, "param", None), + model, + error, + ) + retry = { + key: value for key, value in params.items() + if key not in _TOOL_CHOICE_KEYS + } + return create(**retry) + if not _is_tools_unsupported_error(error): + raise + response = create(**self._params_without_tools(params)) + # Latch only once the tool-less retry has actually worked: a retry + # that also 400s proves nothing about tool support, and must not + # disable tools for the rest of this answer. self._note_tools_rejected(model, error) - return create(**self._params_without_tools(params)) + return response def _raw_gen( self, diff --git a/application/parser/file/docling_parser.py b/application/parser/file/docling_parser.py index 3f1658bb..9d6ee0d9 100644 --- a/application/parser/file/docling_parser.py +++ b/application/parser/file/docling_parser.py @@ -539,11 +539,13 @@ class DoclingParser(BaseParser): pages: Page count reported for the first pass. Returns: - The retry's content, when it clears the chars-per-page floor. + The retry's content: either because it cleared the chars-per-page + floor, or because the document is genuinely text-sparse. Raises: - DocumentParseError: If the retry is near-empty too, so the caller - fails loudly instead of indexing an empty document. + DocumentParseError: Only with positive evidence of a dropout — a + PDF text layer docling should have read without OCR, or a + multi-page document that OCR'd to literally zero characters. """ name = Path(file).name logger.warning( @@ -577,6 +579,7 @@ class DoclingParser(BaseParser): return content detail = "" + layer_pages = layer_chars = 0 if Path(file).suffix.lower() == ".pdf": layer_pages, layer_chars = _pdf_text_layer_probe(file) if layer_chars > 0: @@ -590,11 +593,30 @@ class DoclingParser(BaseParser): f" The PDF has no text layer over {layer_pages} page(s), " "so OCR was the only possible source." ) - raise DocumentParseError( - f"OCR produced {retry_chars} chars over {retry_pages} pages for " - f"{name}; likely OCR pipeline dropout — document not indexed." - + detail + + # Sparse is not the same as dropped. Fail only on positive evidence + # that there was text to find: a text layer docling should have read + # without OCR, or a multi-page document that OCR'd to literally + # nothing (the observed incident). A photo, logo, chart or picture-led + # catalogue is genuinely text-poor — failing it would reject the + # upload permanently, since DocumentParseError skips autoretry. + if layer_chars > 0 or (retry_chars == 0 and retry_pages > 1): + raise DocumentParseError( + f"OCR produced {retry_chars} chars over {retry_pages} pages for " + f"{name}; likely OCR pipeline dropout — document not indexed." + + detail + ) + + logger.warning( + "%s yielded only %d chars over %d page(s) after a full-page-OCR " + "retry; indexing it as genuinely text-sparse rather than treating " + "it as an OCR dropout.%s", + name, + retry_chars, + retry_pages, + detail, ) + return content def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: """Parse file using docling with hybrid OCR. @@ -602,9 +624,11 @@ class DoclingParser(BaseParser): Uses smart OCR approach where the layout model detects text vs bitmap regions. Text is extracted directly, bitmaps are OCR'd only when needed. - When OCR is enabled for a PDF or image, a near-empty result is treated - as a pipeline dropout rather than as the document's content: it is - retried once on a fresh converter, and failing that it raises. + When OCR is enabled for a PDF or image, a near-empty result is retried + once on a fresh converter. It is only rejected as a pipeline dropout if + the retry is near-empty *and* there is evidence text was there to find; + a genuinely text-sparse document (a photo, a picture-led catalogue) is + indexed with a warning. Args: file: Path to the file to parse @@ -615,7 +639,8 @@ class DoclingParser(BaseParser): Raises: DocumentParseError: If docling fails, or if an OCR parse stays - near-empty across both passes. + near-empty across both passes on a document that demonstrably + carried text. """ logger.info(f"parse_file called for: {file}") diff --git a/application/retriever/graph_rag.py b/application/retriever/graph_rag.py index 69eeff09..dafdca46 100644 --- a/application/retriever/graph_rag.py +++ b/application/retriever/graph_rag.py @@ -306,6 +306,17 @@ class GraphRAGRetriever(BaseRetriever): ) failed.append(source_id) + # Every remaining segment is a ClassicRAG fan-out, and each of its legs + # checks out of the *same* per-DSN pool this store is holding. Hand the + # graph connection back first, or concurrent GraphRAG retrievals occupy + # every slot and then block on their own fallbacks until PoolTimeout. + # ``close()`` nulls the connection, so ``_get_data``'s finally stays correct. + if graphless or failed: + try: + store.close() + except Exception as e: + logging.debug("Error releasing GraphRAG store before fallback: %s", e) + if graphless: segments[classic_slot] = self._classic_for_sources(graphless) if failed: diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index 6d221dcd..74182d04 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -493,6 +493,87 @@ def test_extract_parse_opts_out_of_sync_subtask_guard(monkeypatch): assert captured["args"][1] == {"workflow_run_id": "run-extract"} +def _engine_for_parse(monkeypatch, get_impl): + """Engine wired to a fake ``parse_document`` whose ``get`` is ``get_impl``.""" + import application.api.user.tasks as tasks + + agent = _agent(str(uuid.uuid4()), []) + engine = WorkflowEngine.__new__(WorkflowEngine) + engine.agent = agent + engine.workflow_run_id = "run-budget" + + captured: dict = {} + + class _FakeAsyncResult: + def get(self, timeout=None, disable_sync_subtasks=True): + captured.setdefault("timeouts", []).append(timeout) + return get_impl(timeout) + + def _apply_async(args=None, queue=None, **kw): + captured.setdefault("limits", []).append(kw) + return _FakeAsyncResult() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + return engine, captured + + +def test_parse_await_is_clamped_to_the_node_budget(monkeypatch): + """A node's parses share one wall clock; they do not each get a full window. + + WORKFLOW_NODE_EXTRACT_MAX_FILES documents serialized at the full per-document + window would hold a web threadpool slot for max_files x the maximum timeout. + """ + import time + + engine, captured = _engine_for_parse( + monkeypatch, lambda timeout: {"status": "ok", "content": "md"} + ) + + # 100s of a 900s budget already spent by earlier documents in this node. + deadline = time.monotonic() + 100.0 + out = engine._parse_document_text( + "artifact-1", size=25 * 1024 * 1024, deadline=deadline + ) + + assert out == "md" + granted = captured["timeouts"][0] + assert 0 < granted <= 100.0, "the await must not outlive the node's budget" + + +def test_parse_is_skipped_once_the_node_budget_is_spent(monkeypatch, caplog): + """An exhausted budget skips the document instead of starting a new window.""" + import logging + import time + + def _never_called(timeout): + raise AssertionError("no parse may be enqueued past the budget") + + engine, captured = _engine_for_parse(monkeypatch, _never_called) + + with caplog.at_level(logging.WARNING): + out = engine._parse_document_text( + "artifact-2", size=1024, deadline=time.monotonic() - 1.0 + ) + + assert out is None + assert not captured.get("timeouts") + assert "budget" in caplog.text + + +def test_parse_without_a_deadline_keeps_the_size_scaled_window(monkeypatch): + """The per-document scaling is the point of the feature; keep it intact.""" + from application.api.user.tasks import parse_timeout_for_size + + engine, captured = _engine_for_parse( + monkeypatch, lambda timeout: {"status": "ok", "content": "md"} + ) + + size = 13 * 1024 * 1024 + engine._parse_document_text("artifact-3", size=size) + + assert captured["timeouts"][0] == parse_timeout_for_size(size) + + def test_node_extract_path_capped_with_truncation_note(pg_engine, tmp_path, monkeypatch): """A node referencing more docs than the extract cap parses only up to the cap and notes the rest.""" storage = _wire(pg_engine, tmp_path, monkeypatch) diff --git a/tests/agents/test_workflow_node_attachments.py b/tests/agents/test_workflow_node_attachments.py index eebfebe0..089e1329 100644 --- a/tests/agents/test_workflow_node_attachments.py +++ b/tests/agents/test_workflow_node_attachments.py @@ -232,7 +232,7 @@ def test_extract_non_text_uses_parse_worker(monkeypatch): _patch_repo(monkeypatch, {aid: rec}) eng = _engine(monkeypatch) monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED MD" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "EXTRACTED MD" ) cfg = _node_config(input_documents=[aid], file_passing="extract") @@ -296,7 +296,7 @@ def test_native_file_cap_bounds_native_then_extracts(monkeypatch): ) # extract of a non-text image routes through the parsing worker; stub it so it returns text. monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "fallback text" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "fallback text" ) cfg = _node_config(input_documents=ids, file_passing="auto") @@ -572,7 +572,7 @@ def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch): _patch_repo(monkeypatch, {aid: rec}) _patch_capabilities(monkeypatch) monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "EXTRACTED" ) eng = _engine(monkeypatch) diff --git a/tests/llm/test_openai.py b/tests/llm/test_openai.py index d0d0653a..c2f2d589 100644 --- a/tests/llm/test_openai.py +++ b/tests/llm/test_openai.py @@ -2157,6 +2157,84 @@ class TestIsToolsUnsupportedError: assert _is_tools_unsupported_error(error) is True +def _bad_request_param(message, param): + """A 400 that names a request *parameter*, the way OpenAI/Azure do. + + Mirrors ``openai._base_client._make_status_error_from_response``, which + interpolates the whole decoded body into the exception's own message -- + the reason a body-repr match sees tool words that the server never said. + """ + request = httpx.Request("POST", "http://localhost:8000/v1/chat/completions") + response = httpx.Response(400, request=request) + # The SDK unwraps the "error" envelope into ``body`` (so ``error.param`` + # resolves) but keeps the *whole* payload in the message string. + payload = {"error": {"message": message, "param": param}} + return BadRequestError( + f"Error code: 400 - {payload}", response=response, body=payload["error"] + ) + + +@pytest.mark.unit +class TestToolChoiceIsNotACapabilitySignal: + """A rejected ``tool_choice`` is a bad argument, not a missing capability. + + The SDK interpolates the whole JSON body into ``str(error)``, so matching + on that made ``"param": "tool_choice"`` look like "this endpoint cannot do + tools" -- and the v1 translator forwards a client-supplied ``tool_choice`` + verbatim, so any API caller could trigger it. + """ + + TOOLS = [{"type": "function", "function": {"name": "t"}}] + + def test_a_rejected_tool_choice_is_not_a_tools_error(self): + error = _bad_request_param( + "Invalid value: 'require'. Supported values are: 'none', 'auto'.", + "tool_choice", + ) + + assert _is_tools_unsupported_error(error) is False + + def test_the_body_repr_alone_does_not_trip_the_match(self): + # str(error) contains the entire body, including the param name. + error = _bad_request_param("Unsupported value.", "tool_choice") + + assert "tool_choice" in str(error) + assert _is_tools_unsupported_error(error) is False + + def test_retries_without_tool_choice_but_keeps_the_tools(self, llm): + create = _RecordingCreate( + _bad_request_param("Invalid value: 'require'.", "tool_choice"), + _Response(choices=[_Choice(content="answered")]), + ) + llm.client.chat.completions.create = create + + llm._raw_gen( + llm, model="local-model", messages=[{"role": "user", "content": "hi"}], + stream=False, tools=self.TOOLS, tool_choice="require", + ) + + assert len(create.calls) == 2 + assert "tool_choice" not in create.calls[1] + assert create.calls[1]["tools"] == self.TOOLS + assert llm._supports_tools() is True + + def test_a_failed_tool_less_retry_does_not_latch(self, llm): + def _always_400(**kwargs): + raise _bad_request(VLLM_TOOL_ERROR) + + llm.client.chat.completions.create = _always_400 + + with pytest.raises(BadRequestError): + llm._raw_gen( + llm, model="local-model", + messages=[{"role": "user", "content": "hi"}], + stream=False, tools=self.TOOLS, + ) + + # Nothing proved the endpoint is tool-less, so later turns must retry. + assert llm._supports_tools() is True + + @pytest.mark.unit class TestToolCallingFallback: """An OpenAI-compatible endpoint without tool support (vLLM started diff --git a/tests/parser/file/test_docling_parser.py b/tests/parser/file/test_docling_parser.py index 28e612e1..7fa3eba0 100644 --- a/tests/parser/file/test_docling_parser.py +++ b/tests/parser/file/test_docling_parser.py @@ -4,6 +4,7 @@ Covers: DoclingParser (init, _init_parser, _get_ocr_options, _export_content, parse_file), subclass initialization, error handling. """ +import logging import os from pathlib import Path from unittest.mock import MagicMock, patch @@ -1189,9 +1190,8 @@ class TestOCRDropoutGuard: assert create.call_count == 0 assert parser._converter is converter - def test_threshold_is_per_page_not_per_document(self, monkeypatch): - """200 chars is healthy for one page and a dropout for a hundred.""" - from application.parser.file.base_parser import DocumentParseError + def test_threshold_is_per_page_not_per_document(self): + """200 chars is healthy for one page and suspicious for a hundred.""" from application.parser.file.docling_parser import DoclingPDFParser content = "x" * 200 @@ -1203,14 +1203,20 @@ class TestOCRDropoutGuard: assert parser.parse_file(Path("one-page.pdf")) == content assert create.call_count == 0 + # Over a hundred pages the same text is below the floor, so the guard + # spends a full-page-OCR retry on it. It still indexes afterwards -- + # sparse is not the same as dropped (see the dropout tests below). long_doc = DoclingPDFParser(ocr_enabled=True) long_doc._converter = MagicMock() long_doc._converter.convert.return_value = _mock_conversion(content, pages=100) with patch.object( long_doc, "_create_converter", return_value=long_doc._converter + ) as create, patch( + "application.parser.file.docling_parser._pdf_text_layer_probe", + return_value=(100, 0), ): - with pytest.raises(DocumentParseError): - long_doc.parse_file(Path("hundred-page.pdf")) + assert long_doc.parse_file(Path("hundred-page.pdf")) == content + assert create.call_count == 1 def test_guard_does_not_apply_to_non_ocr_formats(self): """DOCX parsers inherit ocr_enabled=True but never OCR anything.""" @@ -1228,7 +1234,28 @@ class TestOCRDropoutGuard: assert create.call_count == 0 def test_guard_applies_to_images(self): - from application.parser.file.base_parser import DocumentParseError + """The retry covers images -- a degraded converter may still recover one.""" + from application.parser.file.docling_parser import DoclingImageParser + + parser = DoclingImageParser(ocr_enabled=True) + degraded = MagicMock() + degraded.convert.return_value = _mock_conversion("") + parser._converter = degraded + + fresh = MagicMock() + fresh.convert.return_value = _mock_conversion("Recovered caption text.") + + with patch.object(parser, "_create_converter", return_value=fresh) as create: + assert parser.parse_file(Path("scan.png")) == "Recovered caption text." + + assert create.call_count == 1 + + def test_a_text_free_image_indexes_rather_than_failing_the_upload(self, caplog): + """A chart, logo or photo has no text to find; that is not a dropout. + + DocumentParseError is in ``dont_autoretry_for``, so raising here fails a + single-file upload permanently. + """ from application.parser.file.docling_parser import DoclingImageParser parser = DoclingImageParser(ocr_enabled=True) @@ -1237,8 +1264,66 @@ class TestOCRDropoutGuard: parser._converter = converter with patch.object(parser, "_create_converter", return_value=converter): - with pytest.raises(DocumentParseError, match="scan.png"): - parser.parse_file(Path("scan.png")) + with caplog.at_level(logging.WARNING): + assert parser.parse_file(Path("chart.png")) == "" + + assert "text-sparse" in caplog.text + + def test_a_text_sparse_scan_indexes_rather_than_failing_the_upload(self): + """20 pages of pictures with a few captions is the document, not a fault.""" + from application.parser.file.docling_parser import DoclingPDFParser + + content = "Fig 1. Fig 2. Fig 3." + parser = DoclingPDFParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion(content, pages=20) + parser._converter = converter + + with patch.object( + parser, "_create_converter", return_value=converter + ), patch( + "application.parser.file.docling_parser._pdf_text_layer_probe", + return_value=(20, 0), + ): + assert parser.parse_file(Path("catalog.pdf")) == content + + def test_a_multi_page_zero_char_parse_is_still_a_dropout(self): + """The incident this guard exists for: every page OCR'd to nothing.""" + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion("", pages=40) + parser._converter = converter + + with patch.object( + parser, "_create_converter", return_value=converter + ), patch( + "application.parser.file.docling_parser._pdf_text_layer_probe", + return_value=(40, 0), + ): + with pytest.raises(DocumentParseError, match="scan.pdf"): + parser.parse_file(Path("scan.pdf")) + + def test_a_pdf_whose_text_layer_was_missed_is_still_a_dropout(self): + """Text docling should have read without OCR at all -- positive evidence.""" + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion("stub", pages=20) + parser._converter = converter + + with patch.object( + parser, "_create_converter", return_value=converter + ), patch( + "application.parser.file.docling_parser._pdf_text_layer_probe", + return_value=(20, 5000), + ): + with pytest.raises(DocumentParseError, match="text layer"): + parser.parse_file(Path("report.pdf")) def test_retry_conversion_failure_is_still_a_parse_error(self): from application.parser.file.base_parser import DocumentParseError diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py index 9e585c79..13a158ed 100644 --- a/tests/retriever/test_graph_rag.py +++ b/tests/retriever/test_graph_rag.py @@ -68,7 +68,9 @@ class TestGraphRAGFallback: assert docs == classic_docs store.search_nodes_by_embedding.assert_not_called() store.get_subgraph.assert_not_called() - store.close.assert_called_once() + # Released early (before the classic fallback checks out of the same + # pool) and again in _get_data's finally; close() is idempotent. + assert store.close.called @patch("application.retriever.graph_rag.GraphStore") @patch("application.retriever.graph_rag.graphrag_available", return_value=False) @@ -112,6 +114,58 @@ def _store_with_graph( return store +@pytest.mark.unit +class TestGraphRAGPoolDiscipline: + """The graph store must not hold a pooled connection across its fallback. + + ``_classic_for_sources`` runs its own per-source fan-out, each leg of which + checks out of the *same* per-DSN pool. Holding the graph store's connection + while recursing into it lets concurrent GraphRAG retrievals occupy every + slot and then block on their own inner fan-outs until PoolTimeout. + """ + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_connection_is_released_before_the_classic_fallback( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = MagicMock() + store.count_nodes_many.return_value = {"src1": 0} + mock_store_cls.return_value = store + + order = [] + store.close.side_effect = lambda: order.append("close") + + rag = _make_retriever() + with patch.object( + rag, "_classic_for_sources", + side_effect=lambda ids: order.append("classic") or [], + ): + rag._get_data() + + assert order[:2] == ["close", "classic"] + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_a_graph_only_retrieval_keeps_its_connection( + self, _avail, mock_store_cls, _patch_llm_creator, _patch_embed + ): + # Nothing falls back, so there is no nested checkout to guard against; + # the store keeps its connection until _get_data's finally. + store = MagicMock() + store.count_nodes_many.return_value = {"src1": 3} + mock_store_cls.return_value = store + + rag = _make_retriever() + with patch.object(rag, "_graph_docs_for_source", return_value=[]): + with patch.object(rag, "_classic_for_sources") as classic: + rag._get_data() + + classic.assert_not_called() + # Exactly one close: _get_data's finally, not an early release. + assert store.close.call_count == 1 + + @pytest.mark.unit class TestGraphRAGHappyPath: @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) @@ -733,7 +787,7 @@ class TestGraphRAGBatching: with pytest.raises(RuntimeError): rag._get_data() - store.close.assert_called_once() + assert store.close.called @patch("application.retriever.graph_rag.GraphStore") @patch("application.retriever.graph_rag.graphrag_available", return_value=True)