fix: minor graph rag improvements

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],

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

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,

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}")
 

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: