feat: parser improvements and fixes

application/agents/base.py

@@ -1139,6 +1139,53 @@ class BaseAgent(ABC):
 
     # ---- LLM generation ----
 
+    def _llm_supports_tools(self) -> bool:
+        """Whether the LLM accepts a ``tools`` payload.
+
+        ``_supports_tools`` is a method on every real provider (the old
+        truthiness test on the *bound method* never gated anything), but
+        test doubles sometimes set it as a plain bool, so both shapes are
+        honored. A provider that never implemented the check — ``BaseLLM``
+        raises — keeps the historical permissive behavior; the provider
+        drops the tools itself.
+
+        Returns:
+            True when the ``tools`` kwarg should be sent.
+        """
+        supports = getattr(self.llm, "_supports_tools", None)
+        if supports is None:
+            return False
+        if not callable(supports):
+            return bool(supports)
+        try:
+            return bool(supports())
+        except Exception:
+            logger.debug(
+                "Tool-support check failed for %s; sending tools anyway",
+                type(self.llm).__name__,
+                exc_info=True,
+            )
+            return True
+
+    def _structured_output_kwarg(self) -> Optional[str]:
+        """Name of the gen kwarg this LLM takes structured output on.
+
+        Read off the LLM *class*, not ``llm_name``: every OpenAI-wire
+        provider (``openai_compatible``, ``groq``, ``novita``,
+        ``openrouter``, ``docsgpt``) dispatches through an ``OpenAILLM``
+        subclass, so a string comparison against "openai" silently
+        dropped guided decoding for all of them. ``BaseLLM`` owns the
+        declaration so the cross-provider fallback adapter can read it too.
+
+        Returns:
+            ``"response_format"``, ``"response_schema"``, or None when the
+            provider has no structured-output kwarg.
+        """
+        # ``type(self.llm)`` — an instance attribute on a test double would
+        # otherwise leak a truthy Mock into the gen kwargs.
+        kwarg = getattr(type(self.llm), "structured_output_kwarg", None)
+        return kwarg if kwarg in ("response_format", "response_schema") else None
+
     def _llm_gen(
         self,
         messages: List[Dict],
@@ -1162,11 +1209,7 @@ class BaseAgent(ABC):
         if self.attachments:
             gen_kwargs["_usage_attachments"] = self.attachments
 
-        if (
-            hasattr(self.llm, "_supports_tools")
-            and self.llm._supports_tools
-            and self.tools
-        ):
+        if self.tools and self._llm_supports_tools():
             gen_kwargs["tools"] = self.tools
         if (
             self.json_schema
@@ -1176,14 +1219,12 @@ class BaseAgent(ABC):
             structured_format = self.llm.prepare_structured_output_format(
                 self.json_schema, strict=getattr(self, "json_schema_strict", True)
             )
-            if structured_format:
-                if self.llm_name == "openai":
-                    gen_kwargs["response_format"] = structured_format
-                elif self.llm_name == "google":
-                    gen_kwargs["response_schema"] = structured_format
+            structured_kwarg = self._structured_output_kwarg()
+            if structured_format and structured_kwarg:
+                gen_kwargs[structured_kwarg] = structured_format
         elif (
             getattr(self, "json_object", False)
-            and self.llm_name == "openai"
+            and self._structured_output_kwarg() == "response_format"
             and hasattr(self.llm, "_supports_structured_output")
             and self.llm._supports_structured_output()
         ):

application/agents/tools/read_document.py

@@ -2,7 +2,8 @@
 
 The ``read_document`` action resolves a parent-scoped input artifact, enqueues a
 ``parse_document`` task on the dedicated ``parsing`` Celery queue, and awaits the
-result with a timeout. The run-scoped authz gate is enforced TWICE — here before
+result with a timeout (inside a worker it parses inline under the same bound — see
+``_dispatch``). The run-scoped authz gate is enforced TWICE — here before
 enqueue (reject cross-tenant) and again in the worker (re-resolve, never trusting a
 raw path). When a ``json_schema`` is supplied the structured payload is validated
 through the existing jsonschema path; the full result may also be persisted as a
@@ -12,7 +13,11 @@ through the existing jsonschema path; the full result may also be persisted as a
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Optional
+import signal
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from typing import Any, Callable, Dict, List, Optional
 
 from celery import current_task
 
@@ -39,6 +44,15 @@ except Exception:  # pragma: no cover - jsonschema is a declared dependency
     jsonschema = None  # type: ignore[assignment]
 
 
+class _InlineParseTimeout(BaseException):
+    """Raised when an inline (in-worker) parse outlives its window.
+
+    Derives from ``BaseException`` on purpose: the SIGALRM interrupt lands deep inside the
+    parser, whose blanket ``except Exception`` would otherwise swallow it and report a
+    generic parse failure instead of the shared timed-out result.
+    """
+
+
 class ReadDocumentTool(Tool):
     """Read Document
     Parse a document (PDF, Word, PowerPoint, ...) to text, markdown, or structured data.
@@ -58,6 +72,8 @@ class ReadDocumentTool(Tool):
         self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
         self.message_id: Optional[str] = self.config.get("message_id")
         self._last_artifact_id: Optional[str] = None
+        # Byte size of the resolved input, used to scale the awaited parse window.
+        self._input_size: Optional[int] = None
 
     # ------------------------------------------------------------------
     # Tool ABC
@@ -157,6 +173,7 @@ class ReadDocumentTool(Tool):
         """Resolve the input run-scoped (reject cross-tenant before enqueue), enqueue+await, validate."""
         input_id = kwargs.get("input")
         json_schema = kwargs.get("json_schema")
+        self._input_size = None
         if not isinstance(input_id, str) or not input_id.strip():
             return {"status": "error", "error": "input artifact id is required."}
         if json_schema is not None:
@@ -199,8 +216,17 @@ class ReadDocumentTool(Tool):
         and blocking on ``get()`` would self-deadlock: concurrent agent tasks each hold a pool
         slot blocked in ``get()`` so ``parse_document`` never gets a free slot. So parse INLINE
         in-process inside a worker; only dispatch+await (degrading on timeout/failure) from web.
+
+        Both branches are bounded by the SAME size-scaled window: the inline parse carries no
+        Celery time limit of its own, so an unbounded one would pin the agent's worker slot
+        until the OUTER task's limit (webhook runs have none) kills the whole agent run.
         """
         parent = self._parent()
+        from application.api.user.tasks import parse_timeout_for_size
+
+        # OCR cost scales with pages, so the parse window grows with the document's size
+        # (floored at DOCUMENT_PARSE_TIMEOUT).
+        timeout = parse_timeout_for_size(self._input_size)
 
         # ``current_task`` is a Celery proxy: truthy only while this runs inside a worker task,
         # falsy in the web process (the bare proxy is NOT identity-None, so test truthiness).
@@ -208,7 +234,12 @@ class ReadDocumentTool(Tool):
             from application.worker import run_parse_document
 
             try:
-                result = run_parse_document(artifact_id, parent, self.user_id, options)
+                result = self._run_inline_bounded(
+                    lambda: run_parse_document(artifact_id, parent, self.user_id, options), timeout
+                )
+            except _InlineParseTimeout:
+                logger.warning("read_document: inline parse timed out after %ss; abandoning it", int(timeout))
+                return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."}
             except Exception as exc:
                 logger.exception("read_document: inline parse failed")
                 return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
@@ -218,12 +249,18 @@ class ReadDocumentTool(Tool):
 
         from celery.exceptions import TimeoutError as CeleryTimeoutError
 
-        from application.api.user.tasks import parse_document
+        from application.api.user.tasks import parse_document, parse_task_time_limits
 
-        timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
+        # The task's per-call time limits are raised to match the awaited window: bound to
+        # the base timeout at import, the worker would otherwise self-terminate a large
+        # parse long before this await gives up.
         queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
         try:
-            async_result = parse_document.apply_async(args=[artifact_id, parent, self.user_id, options], queue=queue)
+            async_result = parse_document.apply_async(
+                args=[artifact_id, parent, self.user_id, options],
+                queue=queue,
+                **parse_task_time_limits(timeout),
+            )
             # The web process (not a worker) awaits here; ``disable_sync_subtasks=False`` keeps
             # the call correct if invoked from a non-prefork (eventlet/gevent) worker where the
             # inline branch above still ran but the blanket guard would otherwise raise.
@@ -237,6 +274,87 @@ class ReadDocumentTool(Tool):
             return {"status": "error", "error": "document parsing produced an unexpected result."}
         return result
 
+    # ------------------------------------------------------------------
+    # Inline parse bound
+    # ------------------------------------------------------------------
+    @classmethod
+    def _run_inline_bounded(cls, fn: Callable[[], Any], timeout: Optional[float]) -> Any:
+        """Run ``fn`` under a wall-clock bound, raising ``_InlineParseTimeout`` on expiry.
+
+        Prefers a SIGALRM interval timer (POSIX, main thread — prefork/solo Celery pools run
+        the task on the child's main thread), which actually interrupts the parse between
+        bytecodes and frees the slot. Otherwise falls back to a helper thread, which can only
+        stop WAITING: the orphaned parse keeps running until the parser finishes.
+
+        Args:
+            fn: Zero-argument callable performing the parse.
+            timeout: Window in seconds; ``None`` or non-positive runs unbounded.
+
+        Returns:
+            Whatever ``fn`` returns.
+        """
+        if not timeout or timeout <= 0:
+            return fn()
+        if cls._sigalrm_usable():
+            return cls._run_with_sigalrm(fn, timeout)
+        return cls._run_in_thread(fn, timeout)
+
+    @staticmethod
+    def _sigalrm_usable() -> bool:
+        """True when this thread may arm SIGALRM without clobbering another user of the timer."""
+        if not hasattr(signal, "SIGALRM") or not hasattr(signal, "setitimer"):
+            return False
+        if threading.current_thread() is not threading.main_thread():
+            return False
+        try:
+            previous = signal.getsignal(signal.SIGALRM)
+            armed = signal.getitimer(signal.ITIMER_REAL)
+        except (ValueError, OSError):  # pragma: no cover - non-main thread / no timer
+            return False
+        if previous not in (None, signal.SIG_DFL):
+            return False
+        return not (armed[0] or armed[1])
+
+    @staticmethod
+    def _run_with_sigalrm(fn: Callable[[], Any], timeout: float) -> Any:
+        """Run ``fn`` with a SIGALRM deadline, always restoring the handler and cancelling the timer."""
+
+        def _on_alarm(signum: int, frame: Any) -> None:
+            raise _InlineParseTimeout()
+
+        previous = signal.signal(signal.SIGALRM, _on_alarm)
+        signal.setitimer(signal.ITIMER_REAL, timeout)
+        try:
+            return fn()
+        finally:
+            # Nested so the handler is restored even if the alarm lands in this very
+            # window (fn returned just as the timer expired).
+            try:
+                signal.setitimer(signal.ITIMER_REAL, 0)
+            finally:
+                signal.signal(signal.SIGALRM, previous)
+
+    @staticmethod
+    def _run_in_thread(fn: Callable[[], Any], timeout: float) -> Any:
+        """Run ``fn`` in a helper thread and stop waiting at ``timeout`` (the parse is NOT killed)."""
+        executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="read-document-inline")
+        try:
+            future = executor.submit(fn)
+            try:
+                return future.result(timeout=timeout)
+            except FutureTimeoutError:
+                if future.done():  # ``fn`` itself raised a TimeoutError; surface it as a failure
+                    raise
+                logger.warning(
+                    "read_document: inline parse exceeded %.0fs off the main thread; the parse thread "
+                    "keeps running until the parser finishes (it cannot be interrupted)",
+                    timeout,
+                )
+                raise _InlineParseTimeout() from None
+        finally:
+            # Never block the agent run on the orphaned parse.
+            executor.shutdown(wait=False)
+
     # ------------------------------------------------------------------
     # Input resolution (run-scoped gate, before enqueue)
     # ------------------------------------------------------------------
@@ -260,6 +378,8 @@ class ReadDocumentTool(Tool):
                     if artifact_id is not None
                     else None
                 )
+                if artifact is not None:
+                    self._input_size = self._version_size(repo, artifact_id, artifact)
         except Exception:
             logger.exception("read_document: failed to resolve input artifact")
             return {"status": "error", "error": f"failed to resolve input artifact {raw_id}."}
@@ -274,6 +394,17 @@ class ReadDocumentTool(Tool):
             return {"status": "error", "error": f"input artifact {raw_id} not found in this conversation/run."}
         return str(artifact_id)
 
+    @staticmethod
+    def _version_size(repo: Any, artifact_id: Any, artifact: Dict[str, Any]) -> Optional[int]:
+        """Best-effort byte size of the artifact's current version (scales the parse window)."""
+        try:
+            version = repo.get_version(artifact_id, artifact.get("current_version"))
+        except Exception:
+            logger.debug("read_document: could not read input size; using the base parse window")
+            return None
+        size = (version or {}).get("size")
+        return int(size) if isinstance(size, (int, float)) else None
+
     def _bridge_chat_attachment(self, raw_id: str) -> Any:
         """Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure."""
         if not self.conversation_id or not self.user_id:
@@ -281,6 +412,10 @@ class ReadDocumentTool(Tool):
         attachment = match_attachment(self.config.get("attachments"), raw_id, self.user_id)
         if attachment is None:
             return None
+        # The attachment row carries the authoritative byte size; the bridged artifact
+        # is a copy of the same bytes, so it sizes the parse window just as well.
+        size = attachment.get("size")
+        self._input_size = int(size) if isinstance(size, (int, float)) else None
         try:
             return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id)
         except AttachmentBridgeError as exc:

application/agents/workflow_agent.py

@@ -387,8 +387,43 @@ class WorkflowAgent(BaseAgent):
                     "mime_type": ref["mime_type"],
                 }
             )
+            self._record_preextracted_text(ref["artifact_id"], attachment)
         return refs, dropped
 
+    def _record_preextracted_text(self, artifact_id: str, attachment: Dict[str, Any]) -> None:
+        """Stage an attachment's upload-time text on the engine so nodes skip re-parsing it."""
+        if self._engine is None:
+            return
+        text = self._usable_attachment_text(attachment)
+        if text:
+            self._engine.preextracted_text[str(artifact_id)] = text
+
+    @staticmethod
+    def _usable_attachment_text(attachment: Dict[str, Any]) -> Optional[str]:
+        """Return an attachment's stored text when it is safe to reuse verbatim, else None.
+
+        The attachment was already parsed (and possibly OCR'd) by ``store_attachment``;
+        reusing that text is what keeps a scanned document from being OCR'd once per
+        node. Reuse requires the stored text to be the WHOLE document: a truncated row
+        was cut head-first at ``ATTACHMENT_MAX_TOKENS`` and would silently lose the tail
+        that the workflow's own head+tail bounding keeps, so it falls back to a parse.
+        A row with no ``extraction`` metadata at all (legacy, pre-dating that field) is
+        trusted when its ``content`` is non-empty: nothing records otherwise, and the
+        non-workflow chat path already feeds that same text to the model.
+        """
+        content = attachment.get("content")
+        if not isinstance(content, str) or not content.strip():
+            return None
+        metadata = attachment.get("metadata")
+        extraction = metadata.get("extraction") if isinstance(metadata, dict) else None
+        if extraction is None:
+            return content
+        if not isinstance(extraction, dict):
+            return None
+        if extraction.get("status") != "ok" or extraction.get("truncated"):
+            return None
+        return content
+
     @staticmethod
     def _read_attachment_bytes(storage: Any, upload_path: str, max_bytes: int) -> Optional[bytes]:
         """Read an attachment with a bounded read and a guaranteed handle close; None on failure."""

application/agents/workflows/workflow_engine.py

@@ -76,6 +76,13 @@ class WorkflowEngine:
         # embedded draft), so run-scoped artifacts must NOT be persisted as orphans.
         # Defaults True; WorkflowAgent flips it off on the draft path.
         self.run_persisted: bool = True
+        # Transient, in-memory ``artifact_id -> text`` for input documents whose text
+        # was already extracted when the attachment was stored, so a node reuses it
+        # instead of re-parsing (re-OCRing) the same bytes. Deliberately NOT part of
+        # ``state``: state is snapshotted into the run row, and the text is large.
+        # Populated once per run by ``WorkflowAgent._bridge_attachments``; lifetime is
+        # this engine instance, i.e. exactly the run.
+        self.preextracted_text: Dict[str, str] = {}
         self.state: WorkflowState = {}
         self.execution_log: List[Dict[str, Any]] = []
         self._condition_result: Optional[str] = None
@@ -805,16 +812,22 @@ class WorkflowEngine:
                 attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
                 native_count += 1
             else:
-                # Inline-text mimes are read directly (cheap); other mimes route
-                # through the parsing worker -- a blocking, ~120s per-document call.
-                # Cap how many documents a single node sends down that path so a
-                # node referencing many non-native documents (e.g. the ``*`` token)
-                # can't serialize dozens of parses. Inline text is not capped.
-                needs_parse = not self._is_inline_text_mime(mime_type)
+                # Inline-text mimes are read directly (cheap), and a document whose
+                # text was already extracted when its attachment was stored is reused
+                # as-is (no parse at all); every other mime routes through the parsing
+                # worker -- a blocking, size-scaled per-document call. Cap how many
+                # documents a single node sends down THAT path so a node referencing
+                # many non-native documents (e.g. the ``*`` token) can't serialize
+                # dozens of parses. Inline text and reused text are not capped: they
+                # cost no blocking parse, so charging them would starve the budget.
+                needs_parse = (
+                    artifact_id not in self.preextracted_text
+                    and not self._is_inline_text_mime(mime_type)
+                )
                 if needs_parse:
                     # Count (and gate on) the parse ATTEMPT, not the success: a
-                    # timed-out/failed parse is the ~120s worst case we must bound,
-                    # so it has to consume cap budget too. Otherwise a degraded
+                    # timed-out/failed parse is the full-window worst case we must
+                    # bound, so it has to consume cap budget too. Otherwise a degraded
                     # parsing backend (every parse fails) never advances the count
                     # and the node keeps issuing blocking parses without limit.
                     if extract_count >= extract_max:
@@ -822,7 +835,7 @@ class WorkflowEngine:
                         continue
                     extract_count += 1
                 content = self._extract_attachment_text(
-                    artifact_id, storage_path, mime_type, filename, max_bytes
+                    artifact_id, storage_path, mime_type, filename, max_bytes, size=size
                 )
                 if content is None:
                     logger.warning(
@@ -873,12 +886,30 @@ class WorkflowEngine:
         return mime_type == "application/pdf" and supports_images
 
     def _extract_attachment_text(
-        self, artifact_id: str, storage_path: str, mime_type: str, filename: str, max_bytes: int
+        self,
+        artifact_id: str,
+        storage_path: str,
+        mime_type: str,
+        filename: str,
+        max_bytes: int,
+        size: Optional[int] = None,
     ) -> Optional[str]:
-        """Get an attachment's text: inline already-text formats, else parse via the parsing worker; None on failure."""
+        """Get an attachment's text: reuse upload-time extraction, else inline text mimes, else parse."""
         from application.parser.document_reader import truncate_text_head_tail
         from application.storage.storage_creator import StorageCreator
 
+        # An uploaded chat attachment was already parsed (and possibly OCR'd) when it
+        # was stored, so re-parsing it here would repeat the dominant cost of the run
+        # for every node that references it. Bounded with the same head+tail window the
+        # inline-text path uses.
+        preextracted = self.preextracted_text.get(artifact_id)
+        if preextracted:
+            logger.info(
+                "Workflow node: reusing upload-time extraction for %s (%s); skipping re-parse",
+                filename,
+                artifact_id,
+            )
+            return truncate_text_head_tail(preextracted)
         if self._is_inline_text_mime(mime_type):
             try:
                 data = StorageCreator.get_storage().get_file(storage_path).read()
@@ -902,13 +933,17 @@ 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)
+        return self._parse_document_text(artifact_id, size=size)
 
-    def _parse_document_text(self, artifact_id: str) -> Optional[str]:
-        """Enqueue ``parse_document`` for this run and await the bounded markdown; None on failure."""
+    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."""
         from celery.exceptions import TimeoutError as CeleryTimeoutError
 
-        from application.api.user.tasks import parse_document
+        from application.api.user.tasks import (
+            parse_document,
+            parse_task_time_limits,
+            parse_timeout_for_size,
+        )
         from application.core.settings import settings
 
         user_id = self._resolve_user_id()
@@ -916,11 +951,15 @@ class WorkflowEngine:
             return None
         options = {"output": "markdown", "include_tables": False, "persist": False}
         queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
-        timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
+        # OCR cost scales with pages, so the window grows with the document's size
+        # (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)
         try:
             async_result = parse_document.apply_async(
                 args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options],
                 queue=queue,
+                **parse_task_time_limits(timeout),
             )
             # A workflow can run inside a Celery worker (scheduled runs / webhooks).
             # In a prefork worker ``task_join_will_block()`` is process-wide, so the

application/api/user/tasks.py

@@ -1,5 +1,6 @@
 import logging
 from datetime import timedelta
+from typing import Dict, Optional
 
 from sqlalchemy.exc import DataError
 
@@ -270,6 +271,47 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None):
     return resp
 
 
+# Seconds the hard (SIGKILL) limit trails the soft limit, giving the soft handler
+# room to unwind and return a terminal result before the worker is killed.
+_PARSE_HARD_LIMIT_GRACE = 30
+
+
+def parse_timeout_for_size(size_bytes: Optional[int]) -> float:
+    """Seconds to allow one ``parse_document``: floored at the base timeout, scaled by size, capped.
+
+    Args:
+        size_bytes: Byte size of the document, or None when unknown (base timeout).
+
+    Returns:
+        The parse window in seconds.
+    """
+    from application.core.settings import settings
+
+    base = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120) or 120)
+    per_mib = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 0) or 0)
+    ceiling = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", base) or base)
+    size = float(size_bytes) if isinstance(size_bytes, (int, float)) else 0.0
+    scaled = base + per_mib * max(size, 0.0) / (1024 * 1024)
+    return min(ceiling, max(base, scaled))
+
+
+def parse_task_time_limits(timeout: float) -> Dict[str, int]:
+    """Per-call Celery limits for a ``parse_document`` the caller awaits for ``timeout`` seconds.
+
+    The task's import-time ``soft_time_limit`` is bound to the BASE timeout, so a caller
+    awaiting a longer, size-scaled window must raise the per-call limits to match or the
+    worker self-terminates the parse first.
+
+    Args:
+        timeout: The awaited parse window in seconds.
+
+    Returns:
+        ``apply_async`` kwargs carrying the soft and hard time limits.
+    """
+    soft = max(1, int(timeout))
+    return {"soft_time_limit": soft, "time_limit": soft + _PARSE_HARD_LIMIT_GRACE}
+
+
 # Not DURABLE: the read_document tool awaits this synchronously with a timeout, so a
 # blind autoretry would double-parse and the caller would already have degraded. The
 # task is routed to the dedicated ``parsing`` queue (celeryconfig task_routes) so a
@@ -286,18 +328,23 @@ def parse_document(self, artifact_id, parent, user_id, options=None):
         # A pathological/malicious document must not pin a parsing-worker slot past the
         # window the caller already abandoned. Return the worker's clean error shape so
         # the slot frees and the Redis result backend still gets a terminal result.
-        limit = getattr(self, "soft_time_limit", None)
+        # ``request.timelimit`` is (hard, soft) and carries the caller's PER-CALL limit,
+        # so prefer it over the import-time default when reporting the window.
+        limit = (getattr(self.request, "timelimit", None) or (None, None))[1]
+        if limit is None:
+            limit = getattr(self, "soft_time_limit", None)
         suffix = f" after {int(limit)}s" if limit else ""
         return {"status": "error", "error": f"document parsing timed out{suffix}."}
 
 
-# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the same window read_document awaits) so
+# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the floor of the window callers await) so
 # the prefork worker self-terminates a runaway parse instead of pinning the slot; the hard
-# limit is the SIGKILL backstop if the soft handler can't unwind in time.
+# limit is the SIGKILL backstop if the soft handler can't unwind in time. Callers awaiting a
+# size-scaled window override both per call via ``parse_task_time_limits``.
 try:
     from application.core.settings import settings as _parse_settings
     parse_document.soft_time_limit = int(_parse_settings.DOCUMENT_PARSE_TIMEOUT)
-    parse_document.time_limit = parse_document.soft_time_limit + 30
+    parse_document.time_limit = parse_document.soft_time_limit + _PARSE_HARD_LIMIT_GRACE
 except Exception:
     pass
 

application/app.py

@@ -31,7 +31,10 @@ from application.api.v1 import v1_bp  # noqa: E402
 from application.celery_init import celery  # noqa: E402
 from application.core.secret_key import resolve_jwt_secret_key  # noqa: E402
 from application.core.settings import settings  # noqa: E402
-from application.storage.db.bootstrap import ensure_database_ready  # noqa: E402
+from application.storage.db.bootstrap import (  # noqa: E402
+    ensure_database_ready,
+    ensure_vector_schema,
+)
 from application.stt.upload_limits import (  # noqa: E402
     build_stt_file_size_limit_message,
     should_reject_stt_request,
@@ -59,6 +62,11 @@ ensure_database_ready(
     logger=logging.getLogger("application.app"),
 )
 
+# Own the vector DB's schema here too, so the retrieval hot path is pure reads
+# instead of re-running DDL for every source of every request.
+if settings.AUTO_VECTOR_SCHEMA:
+    ensure_vector_schema(logger=logging.getLogger("application.app"))
+
 from application.agents.default_tools import (  # noqa: E402
     validate_default_chat_tools,
 )

application/cache.py

@@ -26,6 +26,73 @@ def _cache_default(value):
         return f"<bytes:sha256:{hashlib.sha256(bytes(value)).hexdigest()}>"
     return repr(value)
 
+
+# Generation kwargs that never reach the provider: usage-accounting side
+# channels only. Everything else the caller passes (``response_format``,
+# ``response_schema``, ``tool_choice``, ``reasoning_effort``, sampling
+# params, ...) is part of the request and therefore part of the key —
+# otherwise a workflow node that changed its JSON schema replays the old
+# schema's cached answer for the whole TTL.
+_CACHE_KEY_IGNORED_KWARGS = frozenset({"_usage_attachments", "attachments"})
+
+# Kwargs that make the answer depend on provider-held state no key can
+# capture. ``previous_response_id`` chains a Responses API turn server
+# side, and a cache hit would also skip the ``_last_response_id``
+# bookkeeping the next turn needs — so skip the cache entirely.
+_CACHE_BYPASS_KWARGS = ("previous_response_id",)
+
+
+def _gen_kwargs_fingerprint(extra: dict | None) -> str:
+    """Stable fingerprint of the generation-affecting kwargs.
+
+    Args:
+        extra: Keyword arguments forwarded to the generation call.
+
+    Returns:
+        A sorted JSON dump of the semantic kwargs, or "" when there are none.
+
+    Raises:
+        ValueError: If the kwargs cannot be serialized (callers treat this
+            as "do not cache").
+    """
+    if not extra:
+        return ""
+    filtered = {
+        key: value
+        for key, value in extra.items()
+        if key not in _CACHE_KEY_IGNORED_KWARGS and value is not None
+    }
+    if not filtered:
+        return ""
+    try:
+        return json.dumps(filtered, sort_keys=True, default=_cache_default)
+    except (TypeError, ValueError) as e:
+        raise ValueError(f"Unserializable generation kwargs: {e}") from e
+
+
+def _bypasses_cache(extra: dict | None) -> bool:
+    """Whether a kwarg ties the call to provider-side conversation state."""
+    if not extra:
+        return False
+    return any(extra.get(key) for key in _CACHE_BYPASS_KWARGS)
+
+
+def _is_stream_payload(raw: str) -> bool:
+    """Whether a cached value is a ``stream_cache`` chunk envelope.
+
+    ``gen_cache`` and ``stream_cache`` share one key space, so a
+    non-streaming read can land on a stream entry; without this guard the
+    serialized chunk list would be handed back as the model's answer.
+    """
+    if not raw.startswith("{"):
+        return False
+    try:
+        decoded = json.loads(raw)
+    except ValueError:
+        return False
+    return isinstance(decoded, dict) and isinstance(decoded.get("chunks"), list)
+
+
 _redis_instance = None
 _redis_creation_failed = False
 _instance_lock = Lock()
@@ -116,23 +183,44 @@ def get_pubsub_redis_instance():
     return _pubsub_redis_instance
 
 
-def gen_cache_key(messages, model="docgpt", tools=None):
+def gen_cache_key(messages, model="docgpt", tools=None, extra=None):
+    """Build the Redis key for one generation call.
+
+    Args:
+        messages: Chat messages for the call.
+        model: Model identifier.
+        tools: Tool schemas, when the call carries any.
+        extra: Remaining generation kwargs (``response_format``,
+            ``response_schema``, sampling params, ...). Non-semantic keys
+            are dropped before hashing; the suffix is omitted entirely when
+            nothing semantic remains, so keys for plain calls are unchanged.
+
+    Returns:
+        Hex digest used as the cache key.
+
+    Raises:
+        ValueError: If ``messages`` holds a non-dict entry, or ``extra``
+            cannot be serialized.
+    """
     if not all(isinstance(msg, dict) for msg in messages):
         raise ValueError("All messages must be dictionaries.")
     messages_str = json.dumps(messages, default=_cache_default)
     tools_str = json.dumps(str(tools)) if tools else ""
     combined = f"{model}_{messages_str}_{tools_str}"
+    extra_str = _gen_kwargs_fingerprint(extra)
+    if extra_str:
+        combined = f"{combined}_{extra_str}"
     cache_key = get_hash(combined)
     return cache_key
 
 
 def gen_cache(func):
     def wrapper(self, model, messages, stream, tools=None, *args, **kwargs):
-        if tools is not None:
+        if tools is not None or _bypasses_cache(kwargs):
             return func(self, model, messages, stream, tools, *args, **kwargs)
-        
+
         try:
-            cache_key = gen_cache_key(messages, model, tools)
+            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
         except ValueError as e:
             logger.error(f"Cache key generation failed: {e}")
             return func(self, model, messages, stream, tools, *args, **kwargs)
@@ -142,7 +230,9 @@ def gen_cache(func):
             try:
                 cached_response = redis_client.get(cache_key)
                 if cached_response:
-                    return cached_response.decode("utf-8")
+                    decoded = cached_response.decode("utf-8")
+                    if not _is_stream_payload(decoded):
+                        return decoded
             except Exception as e:
                 logger.error(f"Error getting cached response: {e}", exc_info=True)
 
@@ -160,12 +250,12 @@ def gen_cache(func):
 
 def stream_cache(func):
     def wrapper(self, model, messages, stream, tools=None, *args, **kwargs):
-        if tools is not None:
+        if tools is not None or _bypasses_cache(kwargs):
             yield from func(self, model, messages, stream, tools, *args, **kwargs)
             return
-        
+
         try:
-            cache_key = gen_cache_key(messages, model, tools)
+            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
         except ValueError as e:
             logger.error(f"Cache key generation failed: {e}")
             yield from func(self, model, messages, stream, tools, *args, **kwargs)

application/core/settings.py

@@ -75,6 +75,9 @@ class Settings(BaseSettings):
     AUTO_MIGRATE: bool = True
     # On app startup, create the target Postgres database if it's missing (requires CREATEDB privilege). Dev-friendly default.
     AUTO_CREATE_DB: bool = True
+    # On app startup, create the pgvector/graph tables and verify the embedding dimension. Set False to manage them
+    # out-of-band — there is no Alembic migration for the vector DB because it may be a separate cluster.
+    AUTO_VECTOR_SCHEMA: bool = True
     LLM_PATH: str = os.path.join(current_dir, "models/docsgpt-7b-f16.gguf")
     DEFAULT_MAX_HISTORY: int = 150
     DEFAULT_LLM_TOKEN_LIMIT: int = 128000  # Fallback when model not found in registry
@@ -108,6 +111,12 @@ class Settings(BaseSettings):
     DOCLING_COMPILE_TORCH_MODELS: bool = False
     DOCLING_TABULAR_MAX_BYTES: int = 2_000_000
     DOCLING_MARKUP_MAX_BYTES: int = 8_000_000
+    # Chars-per-page floor below which an OCR'd PDF/image parse is treated as a docling
+    # pipeline dropout (long-running workers were observed returning zero characters for
+    # every scanned page after a long scanned PDF, with no error) rather than as content.
+    # Such a parse is retried once on a fresh full-page-OCR converter and then fails
+    # loudly instead of indexing an empty document. 0 disables the guard.
+    DOCLING_OCR_MIN_CHARS_PER_PAGE: int = 20
     # Read PDF *attachments* via their embedded text layer (pypdfium2) instead
     # of docling, falling back to docling when there is no text layer to read.
     # Attachments go into a prompt, so docling's structural markdown earns far
@@ -127,6 +136,9 @@ class Settings(BaseSettings):
     # ``RetrieverCreator.retrievers`` registry keys (``classic`` / ``default``),
     # NOT the legacy ``classic_rag`` label which never matched the registry.
     RETRIEVERS_ENABLED: list = ["classic", "default"]
+    # Concurrent per-source vector searches within one retrieval (multi-source chats);
+    # the query is embedded once and shared across sources.
+    RETRIEVAL_MAX_PARALLEL_SOURCES: int = 4
     # Kill-switch for per-source retrieval dispatch. When False the retrieval
     # path collapses to today's single-retriever behavior (consumed by the
     # Dispatcher in a later change; defined here so the flag exists up front).
@@ -225,6 +237,9 @@ class Settings(BaseSettings):
     # dialect form (``postgresql+psycopg://``) are all accepted and
     # normalized internally for ``psycopg.connect()``.
     PGVECTOR_CONNECTION_STRING: Optional[str] = None
+    # Per-process psycopg connection pool for the vector store; 0 = one direct
+    # connection per store instance, the legacy behaviour.
+    PGVECTOR_POOL_MAX_SIZE: int = 8
     # IVFFlat probes for vector search. ``None`` derives sqrt(lists) from the
     # index itself; set an integer to pin it. Higher = better recall, more scan.
     PGVECTOR_IVFFLAT_PROBES: Optional[int] = None
@@ -428,6 +443,12 @@ class Settings(BaseSettings):
     # ``read_document`` parsing on a dedicated Celery ``parsing`` queue (backend parser).
     DOCUMENT_PARSE_QUEUE: str = "parsing"  # queue the parse_document task is routed to
     DOCUMENT_PARSE_TIMEOUT: int = 120  # seconds the tool awaits the enqueued parse before degrading
+    # The base timeout is a FLOOR: the awaited window (and the task's per-call Celery
+    # time limits) grow with the document's size, because OCR cost scales with pages
+    # -- a 30-page scan needs ~60s, a 100-page scan several minutes. Without this a
+    # large scan is silently dropped from the node/tool at the base window.
+    DOCUMENT_PARSE_TIMEOUT_PER_MB: int = 60  # extra seconds of parse window per MiB of input
+    DOCUMENT_PARSE_TIMEOUT_MAX: int = 900  # absolute ceiling on the size-scaled parse window
     DOCUMENT_PARSE_MAX_BYTES: int = 0  # cap on a parsed document's bytes (0 = reuse SANDBOX_MAX_INPUT_BYTES)
     DOCUMENT_MAX_DECOMPRESSED_BYTES: int = 300 * 1024 * 1024
     DOCUMENT_MAX_ARCHIVE_ENTRIES: int = 10000

application/graphrag/extraction.py

@@ -26,7 +26,9 @@ from typing import Any, Callable, Dict, List, Optional
 from application.core.settings import settings
 from application.llm.llm_creator import LLMCreator
 from application.storage.db.source_config import SourceConfig
-from application.vectorstore.base import EmbeddingsSingleton
+# ``EmbeddingsSingleton`` is re-exported here so callers and tests can reach the
+# shared instance cache from this module.
+from application.vectorstore.base import EmbeddingsSingleton, get_embeddings  # noqa: F401
 
 logger = logging.getLogger(__name__)
 
@@ -193,9 +195,7 @@ def extract_graph_for_source(
     skipped_over_cap = max(0, len(pending) - cap)
     to_process = pending[:cap]
 
-    embedding = EmbeddingsSingleton.get_instance(
-        settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-    )
+    embedding = get_embeddings()
 
     llm = _build_extraction_llm(
         _resolve_extraction_model(config), user, request_id

application/graphrag/store.py

@@ -1,12 +1,14 @@
 """Per-source knowledge-graph store co-located with the pgvector ``documents`` table.
 
-GraphRAG is pgvector-only: the graph tables live in the same DB as the
-pgvector store and are created on-demand (``CREATE TABLE IF NOT EXISTS`` +
-``CREATE EXTENSION IF NOT EXISTS vector``), mirroring
-``PGVectorStore._ensure_table_exists`` rather than going through app-DB Alembic.
-That DB may be a separate cluster (e.g. Neon) from the app DB where ``sources``
-lives, so ``source_id`` is a plain indexed UUID column with no cross-DB FK and
-all ids are generated in Python.
+GraphRAG is pgvector-only: the graph tables live in the same DB as the pgvector
+store and are created at boot by ``ensure_vector_schema`` (``CREATE TABLE IF NOT
+EXISTS`` + ``CREATE EXTENSION IF NOT EXISTS vector``), mirroring
+``PGVectorStore.create_schema`` rather than going through app-DB Alembic.
+Constructing a store runs no DDL and opens no connection; the write methods
+re-check the schema once per instance as a safety net for a process that never
+ran the boot hook. That DB may be a separate cluster (e.g. Neon) from the app DB
+where ``sources`` lives, so ``source_id`` is a plain indexed UUID column with no
+cross-DB FK and all ids are generated in Python.
 """
 
 from __future__ import annotations
@@ -18,9 +20,15 @@ from typing import Any, Dict, List, Optional
 from psycopg.types.json import Jsonb
 
 from application.core.settings import settings
+from application.vectorstore import pgconn
 
 DEFAULT_NAME_EMBEDDING_DIM = 768
 
+# Mirrors ``PGVectorStore.DEFAULT_POOL_MAX_SIZE``; duplicated rather than
+# imported so the graph store never pulls the vector store (and its embeddings
+# stack) in at import time. 0 disables pooling.
+DEFAULT_POOL_MAX_SIZE = 8
+
 MAX_SUBGRAPH_NODES = 500
 MAX_SUBGRAPH_EDGES = 2000
 
@@ -30,6 +38,14 @@ GRAPH_OVERVIEW_MAX_LIMIT = 250
 PGVECTOR_SOURCE_COLUMN = "source_id"
 
 
+def _resolve_pool_max_size() -> int:
+    """Pool size from settings, defensively — 0 means one direct connection."""
+    value = getattr(settings, "PGVECTOR_POOL_MAX_SIZE", DEFAULT_POOL_MAX_SIZE)
+    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+        return value
+    return DEFAULT_POOL_MAX_SIZE
+
+
 def _safe_identifier(name: str) -> str:
     """Return ``name`` if it is a bare SQL identifier, else raise.
 
@@ -97,35 +113,74 @@ class GraphStore:
         self._psycopg = psycopg
         self._register_vector = register_vector
         self._connection = None
-        self._ensure_tables()
+        self._pooled = False
+        self._pool_max_size = _resolve_pool_max_size()
+        # No DDL here: the graph tables are created at boot alongside the
+        # pgvector table, and every write re-checks once (see _ensure_tables_once).
+        self._tables_ensured = False
 
     def _get_connection(self):
+        """Get or create this store's connection, pooled unless pooling is off.
+
+        Shares :mod:`application.vectorstore.pgconn`'s per-DSN pool with
+        ``PGVectorStore``, so a retrieval that touches both pays one checkout
+        each instead of two fresh connect handshakes.
+        """
         if self._connection is None or self._connection.closed:
-            self._connection = self._psycopg.connect(self._connection_string)
-            self._register_vector(self._connection)
+            if self._pool_max_size > 0:
+                self._connection = pgconn.pool_for(
+                    self._connection_string, self._pool_max_size
+                ).getconn()
+                self._pooled = True
+                # No _register_pgvector_types here: the pool's ``configure``
+                # hook already ran it on this physical connection, and repeating
+                # it would cost a catalog lookup on every checkout. The DDL path
+                # (_ensure_tables) still re-registers after CREATE EXTENSION.
+            else:
+                self._connection = self._psycopg.connect(self._connection_string)
+                self._register_pgvector_types(self._connection)
+                self._pooled = False
         return self._connection
 
+    def _register_pgvector_types(self, conn) -> None:
+        """Register pgvector's adapters, tolerating a not-yet-created extension.
+
+        ``register_vector`` looks the ``vector`` type up in the catalog and
+        raises when it is absent — which is exactly the state of a brand-new
+        database, before ``CREATE EXTENSION`` has run. Swallow that so the
+        schema bootstrap can proceed; it re-registers once the type exists.
+        """
+        try:
+            self._register_vector(conn)
+        except Exception as e:
+            logging.debug("pgvector types not registered yet: %s", e)
+
     def _embedding_dim(self) -> int:
         """Dimension of the configured embeddings model, matching ``PGVectorStore``.
 
         Falls back to ``DEFAULT_NAME_EMBEDDING_DIM`` so the graph table and the
         pgvector ``documents`` table always agree on the configured model.
         """
-        from application.vectorstore.base import EmbeddingsSingleton
+        from application.vectorstore.base import get_embeddings
 
-        embedding = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-        )
+        embedding = get_embeddings()
         return getattr(embedding, "dimension", DEFAULT_NAME_EMBEDDING_DIM)
 
-    def _ensure_tables(self):
-        conn = self._get_connection()
+    @staticmethod
+    def create_schema(conn, *, dimension: int = DEFAULT_NAME_EMBEDDING_DIM) -> None:
+        """Create the graph tables and indexes on ``conn`` without committing.
+
+        Shared by the boot hook (``ensure_vector_schema``) and the store's own
+        write-path safety net; the caller owns the transaction.
+
+        Args:
+            conn: Open psycopg connection to the pgvector database.
+            dimension: Width of the node name-embedding vectors.
+        """
         cursor = conn.cursor()
         try:
             cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
 
-            embedding_dim = self._embedding_dim()
-
             cursor.execute(
                 f"""
                 CREATE TABLE IF NOT EXISTS graph_nodes (
@@ -137,7 +192,7 @@ class GraphStore:
                     description TEXT,
                     degree INT DEFAULT 0,
                     doc_freq INT DEFAULT 0,
-                    name_embedding vector({embedding_dim}),
+                    name_embedding vector({dimension}),
                     UNIQUE (source_id, normalized_name)
                 );
                 """
@@ -206,14 +261,44 @@ class GraphStore:
             # post-filter in search_nodes_by_embedding that silently returns
             # zero nodes, and graph_rag does not fall back when the source has
             # nodes. Add an index deliberately once a graph is large enough.
+        finally:
+            cursor.close()
 
+    def _ensure_tables(self):
+        """Create the graph schema under an advisory lock, then commit."""
+        # Same key as the pgvector store and the boot hook: one lock guards all
+        # DDL in this database, so concurrent workers never race each other.
+        from application.vectorstore.pgvector import SCHEMA_LOCK_KEY
+
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+                )
+            finally:
+                cursor.close()
+            self.create_schema(conn, dimension=self._embedding_dim())
             conn.commit()
+            # The extension may have just been created; pick up its adapters.
+            self._register_pgvector_types(conn)
         except Exception as e:
             conn.rollback()
             logging.error(f"Error creating graph tables: {e}")
             raise
-        finally:
-            cursor.close()
+
+    def _ensure_tables_once(self) -> None:
+        """Create the graph schema on this instance's first write, at most once.
+
+        Readers never create tables — boot owns the schema. This is the safety
+        net for a process that never ran the boot hook (scripts, tests, the
+        first extraction on a brand-new deployment).
+        """
+        if getattr(self, "_tables_ensured", False):
+            return
+        self._ensure_tables()
+        self._tables_ensured = True
 
     def _upsert_node(
         self,
@@ -287,6 +372,7 @@ class GraphStore:
         incremented, the type is refreshed if previously empty, and the
         embedding is refreshed when provided. Returns the node id either way.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -351,6 +437,7 @@ class GraphStore:
         source_chunk_ids: Optional[List[str]] = None,
     ) -> str:
         """Insert an edge and bump the degree of both endpoints. Returns its id."""
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -384,6 +471,7 @@ class GraphStore:
         )
 
     def link_node_chunk(self, source_id: str, node_id: str, chunk_id: str):
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -415,6 +503,7 @@ class GraphStore:
         are not bumped here — the caller runs ``set_node_degrees`` once at the
         end. Returns ``(nodes_upserted, edges_added)``.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         node_ids: Dict[str, str] = {}
@@ -540,6 +629,53 @@ class GraphStore:
             cursor.close()
             conn.rollback()
 
+    def count_nodes_many(self, source_ids: List[str]) -> Dict[str, int]:
+        """Node counts for several sources in one round trip.
+
+        Replaces the retriever's per-source ``count_nodes`` fan-out: N sources
+        used to cost N queries (each on its own fresh connection). Ids with no
+        rows are filled in with 0 in Python, so the caller always gets an entry
+        for every id it asked about.
+
+        Args:
+            source_ids: Source ids to count; empty/falsy entries are ignored.
+
+        Returns:
+            dict: ``{source_id: node_count}``, zero-filled. All zeros when the
+            query fails, which drives the ClassicRAG fallback exactly as a
+            failing ``count_nodes`` does.
+        """
+        ids = [str(s) for s in source_ids if s]
+        if not ids:
+            return {}
+        counts: Dict[str, int] = {source_id: 0 for source_id in ids}
+        # ``source_id`` is a UUID column, so Postgres hands back the canonical
+        # lowercase text. Map it to the exact string the caller passed, or a
+        # differently-cased id would land under a second key and read as 0.
+        by_canonical = {source_id.lower(): source_id for source_id in ids}
+        conn = self._get_connection()
+        cursor = conn.cursor()
+        try:
+            cursor.execute(
+                """
+                SELECT source_id, count(*)
+                FROM graph_nodes
+                WHERE source_id = ANY(%s)
+                GROUP BY source_id;
+                """,
+                (ids,),
+            )
+            for row in cursor.fetchall():
+                returned = str(row[0])
+                counts[by_canonical.get(returned.lower(), returned)] = int(row[1])
+            return counts
+        except Exception as e:
+            logging.error(f"Error counting nodes: {e}")
+            return counts
+        finally:
+            cursor.close()
+            conn.rollback()
+
     def search_nodes_by_embedding(
         self, source_id: str, query_embedding: List[float], k: int = 10
     ) -> List[Dict[str, Any]]:
@@ -860,6 +996,7 @@ class GraphStore:
         (``WHERE id IN (src, dst)`` bumps the endpoint a single time when
         ``src == dst``). ``UNION`` deduplicates the two endpoints of each edge.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -891,6 +1028,7 @@ class GraphStore:
             cursor.close()
 
     def mark_chunk(self, source_id: str, chunk_id: str, status: str):
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -952,6 +1090,7 @@ class GraphStore:
 
     def delete_by_source(self, source_id: str):
         """Remove every graph row for a source (no FK cascade across clusters)."""
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -972,10 +1111,26 @@ class GraphStore:
         finally:
             cursor.close()
 
+    def close(self) -> None:
+        """Release this store's connection: back to the pool, or closed outright.
+
+        A pooled connection is rolled back first when it is still in a
+        transaction, so the next borrower gets a clean session.
+        """
+        conn = getattr(self, "_connection", None)
+        if conn is None:
+            return
+        self._connection = None
+        pgconn.release(
+            getattr(self, "_connection_string", ""),
+            conn,
+            getattr(self, "_pooled", False),
+        )
+
     def __del__(self):
-        if (
-            hasattr(self, "_connection")
-            and self._connection
-            and not self._connection.closed
-        ):
-            self._connection.close()
+        """Release the connection when the object is destroyed. Never raises."""
+        try:
+            self.close()
+        except Exception:
+            # Interpreter teardown can null out module globals; never raise here.
+            pass

application/llm/base.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import ClassVar
+from typing import ClassVar, Dict, Optional, Tuple
 
 import httpx
 import openai
@@ -38,6 +38,24 @@ class BaseLLM(ABC):
     # calls by vendor. Subclasses override.
     provider_name: ClassVar[str] = "unknown"
 
+    # Name of the gen kwarg this provider takes structured output on
+    # ("response_format" for OpenAI-wire classes, "response_schema" for
+    # Google); None = the provider has no structured-output kwarg.
+    structured_output_kwarg: ClassVar[Optional[str]] = None
+
+    # (json_schema, strict) last passed to ``prepare_structured_output_format``;
+    # lets a cross-provider fallback re-prepare the schema in the backup's own
+    # format instead of forwarding a kwarg the backup cannot read.
+    _structured_output_source: Optional[Tuple[Dict, bool]] = None
+
+    # Gen kwargs that carry structured output, across providers. Kept here
+    # (not on the provider subclasses) because the adapter has to recognize
+    # every provider's kwarg, not just its own.
+    _STRUCTURED_OUTPUT_KWARGS: ClassVar[Tuple[str, ...]] = (
+        "response_format",
+        "response_schema",
+    )
+
     def __init__(
         self,
         decoded_token=None,
@@ -300,6 +318,132 @@ class BaseLLM(ABC):
                 prepared.append({**message, "content": parts})
         return prepared
 
+    @staticmethod
+    def _fallback_enforces_structured_output(fallback) -> bool:
+        """Whether the fallback's capabilities still allow schema enforcement."""
+        supports = getattr(fallback, "_supports_structured_output", None)
+        if supports is None:
+            return True
+        if not callable(supports):
+            return bool(supports)
+        try:
+            return bool(supports())
+        except Exception:
+            logger.debug(
+                "Structured-output capability check failed for %s; assuming supported",
+                type(fallback).__name__,
+                exc_info=True,
+            )
+            return True
+
+    @staticmethod
+    def _recover_structured_output_source(present: Dict) -> Optional[Tuple[Dict, bool]]:
+        """Recover a raw (schema, strict) pair from already-prepared kwargs.
+
+        Callers that hand-build an OpenAI ``response_format`` never went through
+        ``prepare_structured_output_format``, so nothing was recorded; the raw
+        schema is still readable out of the envelope. A Google
+        ``response_schema`` is a lossy, type-mapped conversion — not reversible —
+        so it yields nothing.
+        """
+        response_format = present.get("response_format")
+        if not isinstance(response_format, dict):
+            return None
+        json_schema = response_format.get("json_schema")
+        if not isinstance(json_schema, dict):
+            return None
+        schema = json_schema.get("schema")
+        if not isinstance(schema, dict) or not schema:
+            return None
+        return schema, bool(json_schema.get("strict", True))
+
+    def _adapt_structured_output_kwargs(self, fallback, kwargs: Dict) -> Dict:
+        """Re-express the primary's structured-output kwargs for ``fallback``.
+
+        Structured output is provider-specific: OpenAI-wire classes take
+        ``response_format``, Google takes ``response_schema``. Forwarding the
+        primary's kwarg verbatim to a different-family backup either loses
+        enforcement silently (Google swallows ``response_format`` in
+        ``**kwargs``) or raises ``TypeError`` inside the OpenAI SDK
+        (``response_schema`` is not a Chat-Completions param), which turns the
+        fallback into no fallback at all.
+
+        Args:
+            fallback: Backup LLM the request is about to be re-sent to.
+            kwargs: The primary's generation kwargs.
+
+        Returns:
+            A copy of ``kwargs`` carrying the backup's own structured-output
+            kwarg, or neither when the backup cannot enforce a schema.
+        """
+        adapted = dict(kwargs)
+        present = {
+            name: adapted.pop(name)
+            for name in self._STRUCTURED_OUTPUT_KWARGS
+            if name in adapted
+        }
+        if not present:
+            return adapted
+
+        target = getattr(fallback, "structured_output_kwarg", None)
+        if not target or not self._fallback_enforces_structured_output(fallback):
+            logger.warning(
+                f"Fallback {fallback.model_id} cannot enforce structured "
+                f"output; continuing unstructured"
+            )
+            return adapted
+
+        response_format = present.get("response_format")
+        if (
+            isinstance(response_format, dict)
+            and response_format.get("type") == "json_object"
+        ):
+            # json_object mode is an OpenAI-wire shape with no wired equivalent
+            # elsewhere; keep it only within the same family.
+            if target == "response_format":
+                adapted[target] = response_format
+            else:
+                logger.warning(
+                    f"Fallback {fallback.model_id} has no json_object mode; "
+                    f"continuing unstructured"
+                )
+            return adapted
+
+        if target in present:
+            # Same wire family (OpenAI -> openai_compatible, Google -> Google):
+            # the prepared value is already in the backup's format.
+            adapted[target] = present[target]
+            return adapted
+
+        source = self._structured_output_source
+        if not source:
+            source = self._recover_structured_output_source(present)
+        if not source:
+            logger.warning(
+                f"Fallback {fallback.model_id} needs {target} but the source "
+                f"schema is unavailable; continuing unstructured"
+            )
+            return adapted
+
+        schema, strict = source
+        try:
+            prepared = fallback.prepare_structured_output_format(schema, strict=strict)
+        except Exception:
+            logger.warning(
+                f"Failed to prepare structured output for fallback "
+                f"{fallback.model_id}; continuing unstructured",
+                exc_info=True,
+            )
+            prepared = None
+        if prepared:
+            adapted[target] = prepared
+        else:
+            logger.warning(
+                f"Fallback {fallback.model_id} cannot enforce structured "
+                f"output; continuing unstructured"
+            )
+        return adapted
+
     def _execute_with_fallback(
         self, method_name: str, decorators: list, *args, **kwargs
     ):
@@ -371,6 +515,9 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            fallback_kwargs = self._adapt_structured_output_kwargs(
+                fallback, fallback_kwargs
+            )
             if fallback_kwargs.get("messages"):
                 fallback_kwargs["messages"] = self._prepare_fallback_messages(
                     fallback,
@@ -502,6 +649,9 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            fallback_kwargs = self._adapt_structured_output_kwargs(
+                fallback, fallback_kwargs
+            )
             if fallback_kwargs.get("messages"):
                 fallback_kwargs["messages"] = self._prepare_fallback_messages(
                     fallback,
@@ -673,9 +823,13 @@ class BaseLLM(ABC):
     def _supports_structured_output(self):
         return False
 
-    def prepare_structured_output_format(self, json_schema):
-        """Prepare structured output format specific to the LLM provider"""
-        _ = json_schema
+    def prepare_structured_output_format(self, json_schema, strict=True):
+        """Prepare structured output format specific to the LLM provider.
+
+        Overrides must record ``self._structured_output_source`` so a
+        cross-provider fallback can re-prepare the schema in its own format.
+        """
+        _ = (json_schema, strict)
         return None
 
     def get_supported_attachment_types(self):

application/llm/google_ai.py

@@ -12,6 +12,7 @@ from application.storage.storage_creator import StorageCreator
 
 class GoogleLLM(BaseLLM):
     provider_name = "google"
+    structured_output_kwarg = "response_schema"
 
     def __init__(
         self, api_key=None, user_api_key=None, decoded_token=None, *args, **kwargs
@@ -674,6 +675,8 @@ class GoogleLLM(BaseLLM):
         ``strict`` is accepted for signature parity with the OpenAI provider;
         Google enforces the schema natively via ``response_schema``.
         """
+        # Recorded so a cross-provider fallback can re-prepare the raw schema.
+        self._structured_output_source = (json_schema, strict) if json_schema else None
         if not json_schema:
             return None
         type_map = {

application/llm/openai.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os.path
 
-from openai import OpenAI
+from openai import BadRequestError, OpenAI
 
 from application.core.settings import settings
 from application.llm.base import BaseLLM
@@ -14,6 +14,50 @@ from application.storage.storage_creator import StorageCreator
 # Placeholder sent to OpenAI-compatible backends that require no credentials.
 NO_API_KEY = "sk-no-key"
 
+# Request keys that only make sense alongside ``tools``; all dropped together
+# when an endpoint turns out not to serve tool calling.
+_TOOL_REQUEST_KEYS = ("tools", "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
+# insensitively, and only for requests that actually carried tools — so a false
+# positive degrades to a tool-less answer rather than a failed chat.
+_TOOLS_UNSUPPORTED_MARKERS = (
+    "enable-auto-tool-choice",
+    "tool-call-parser",
+    "tool_choice",
+    "tool choice",
+    "tools are not supported",
+    "tools is not supported",
+    "does not support tools",
+    "doesn't support tools",
+    "tool calling",
+    "tool call is not supported",
+    "function calling",
+    "functions are not supported",
+    "does not support function",
+    "unsupported parameter: 'tools'",
+    "unknown field: tools",
+)
+
+
+def _is_tools_unsupported_error(error: Exception) -> bool:
+    """Whether a 400 says the endpoint cannot serve tool calling.
+
+    Args:
+        error: The exception raised by the provider SDK.
+
+    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()
+    return any(marker in haystack for marker in _TOOLS_UNSUPPORTED_MARKERS)
+
 
 def _truncate_base64_for_logging(messages):
     """
@@ -115,6 +159,12 @@ class _RespChoice:
 
 class OpenAILLM(BaseLLM):
     provider_name = "openai"
+    structured_output_kwarg = "response_format"
+
+    # Flipped once this endpoint answered 400 on a request carrying tools.
+    # Instance-scoped (the assignment shadows this class default), so an
+    # agent loop pays for the discovery round trip at most once.
+    _tools_rejected = False
 
     def __init__(
         self,
@@ -562,6 +612,53 @@ class OpenAILLM(BaseLLM):
                 return normalized
         return ""
 
+    def _note_tools_rejected(self, model, error) -> None:
+        """Record that this endpoint refuses tool calling, warning once."""
+        if not self._tools_rejected:
+            logging.warning(
+                "Endpoint rejected tool calling (model=%s, base_url=%s): %s. "
+                "Retrying this request without tools and skipping tools for "
+                "the rest of this session — set supports_tools: false in the "
+                "model catalog for %s to avoid the discovery round trip.",
+                model,
+                self._effective_base_url,
+                error,
+                model,
+            )
+        self._tools_rejected = True
+
+    @staticmethod
+    def _params_without_tools(params: dict) -> dict:
+        """Copy of ``params`` with every tool-calling key removed."""
+        return {
+            key: value for key, value in params.items()
+            if key not in _TOOL_REQUEST_KEYS
+        }
+
+    def _create_with_tool_fallback(self, create, params: dict, model):
+        """Send a request, retrying once without tools on a tools-unsupported 400.
+
+        Args:
+            create: The provider ``create`` callable for this endpoint.
+            params: Request kwargs, possibly including ``tools``.
+            model: Model id, used for the warning.
+
+        Returns:
+            The provider response — from the first call, or from the
+            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.
+        """
+        try:
+            return create(**params)
+        except BadRequestError as error:
+            if not params.get("tools") or not _is_tools_unsupported_error(error):
+                raise
+            self._note_tools_rejected(model, error)
+            return create(**self._params_without_tools(params))
+
     def _raw_gen(
         self,
         baseself,
@@ -620,7 +717,9 @@ class OpenAILLM(BaseLLM):
             request_params["response_format"] = response_format
         self._last_usage = None
         self._stream_reached_finish = False
-        response = self.client.chat.completions.create(**request_params)
+        response = self._create_with_tool_fallback(
+            self.client.chat.completions.create, request_params, model
+        )
         logging.debug("OpenAI request completed")
         self._record_chat_usage(getattr(response, "usage", None))
         if tools:
@@ -694,7 +793,12 @@ class OpenAILLM(BaseLLM):
         self._last_usage = None
         self._stream_reached_finish = False
         self._last_finish_reason = None
-        response = self.client.chat.completions.create(**request_params)
+        # The request is issued before the loop, so a tools-unsupported 400
+        # lands here with nothing yielded yet — the retry inside
+        # ``_create_with_tool_fallback`` can never duplicate delivered output.
+        response = self._create_with_tool_fallback(
+            self.client.chat.completions.create, request_params, model
+        )
 
         try:
             for line in response:
@@ -1298,7 +1402,9 @@ class OpenAILLM(BaseLLM):
             stream=False,
             kwargs=kwargs,
         )
-        response = self.client.responses.create(**params)
+        response = self._create_with_tool_fallback(
+            self.client.responses.create, params, model
+        )
         if response is None:
             raise RuntimeError("Responses API returned no response object")
         logging.debug(
@@ -1371,7 +1477,11 @@ class OpenAILLM(BaseLLM):
             stream=True,
             kwargs=kwargs,
         )
-        response = self.client.responses.create(**params)
+        # Issued before the event loop below, so the tool-less retry cannot
+        # duplicate output already yielded.
+        response = self._create_with_tool_fallback(
+            self.client.responses.create, params, model
+        )
 
         func_calls = {}
         reasoning_items = []
@@ -1485,6 +1595,12 @@ class OpenAILLM(BaseLLM):
                 response.close()
 
     def _supports_tools(self):
+        # A 400 already told us this endpoint can't serve tool calling
+        # (vLLM without --enable-auto-tool-choice, and friends). Answer
+        # False from then on so later turns skip tools outright instead of
+        # paying for another rejected round trip.
+        if self._tools_rejected:
+            return False
         # When the LLM was constructed via LLMCreator with a registered
         # AvailableModel, ``self.capabilities`` is the per-model record.
         # BYOM users can disable tool support; respect that. Otherwise
@@ -1515,6 +1631,8 @@ class OpenAILLM(BaseLLM):
             kwargs["reasoning_effort"] = effort
 
     def prepare_structured_output_format(self, json_schema, strict=True):
+        # Recorded so a cross-provider fallback can re-prepare the raw schema.
+        self._structured_output_source = (json_schema, strict) if json_schema else None
         if not json_schema:
             return None
         try:

application/parser/chunking_strategies.py

@@ -288,12 +288,9 @@ class SemanticChunker(_BaseStrategyChunker):
         sentences = self._split_sentences(text)
         if len(sentences) < 2:
             raise ValueError("too few sentences for semantic chunking")
-        from application.vectorstore.base import EmbeddingsSingleton
-        from application.core.settings import settings
+        from application.vectorstore.base import get_embeddings
 
-        embeddings = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME
-        ).embed_documents(sentences)
+        embeddings = get_embeddings().embed_documents(sentences)
         breakpoints = self._breakpoints(embeddings)
         groups = self._group(sentences, breakpoints)
         return [g for g in self._enforce_tokens(groups) if g.strip()]

application/parser/file/docling_parser.py

@@ -9,10 +9,11 @@ images (PNG, JPEG, TIFF, BMP, WEBP), WebVTT, and specialized XML formats.
 import importlib.util
 import logging
 import os
+import re
 import tempfile
 import zipfile
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from application.parser.file.base_parser import BaseParser, DocumentParseError
 from application.utils import truncate_to_line_boundary
@@ -213,6 +214,124 @@ def _parse_markup_bounded(
             pass
 
 
+# Suffixes whose text comes (wholly or partly) from OCR, and are therefore
+# covered by the near-empty-output dropout guard. Everything else docling
+# handles (docx/xlsx/html/vtt/...) has a native text layer and legitimately
+# short files, so the guard would only produce false alarms there.
+_IMAGE_SUFFIXES = frozenset(
+    {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp"}
+)
+_OCR_GUARDED_SUFFIXES = _IMAGE_SUFFIXES | {".pdf"}
+
+# Chars-per-page floor below which an OCR parse is treated as a pipeline
+# dropout rather than as the document's real content. Scanned pages that
+# carry any text at all clear this comfortably; the observed failure mode
+# returns literally zero characters per page.
+_DEFAULT_OCR_MIN_CHARS_PER_PAGE = 20
+
+# Docling emits this placeholder for a picture it did not read; it is markup,
+# not document text, so it must not count towards the chars-per-page ratio.
+_IMAGE_PLACEHOLDER_RE = re.compile(r"<!--\s*image\s*-->")
+
+
+def _ocr_min_chars_per_page() -> int:
+    """Chars-per-page floor for the OCR dropout guard; 0 disables it."""
+    from application.core.settings import settings
+
+    try:
+        return int(
+            getattr(
+                settings,
+                "DOCLING_OCR_MIN_CHARS_PER_PAGE",
+                _DEFAULT_OCR_MIN_CHARS_PER_PAGE,
+            )
+        )
+    except (TypeError, ValueError):
+        return _DEFAULT_OCR_MIN_CHARS_PER_PAGE
+
+
+def _text_char_count(content: Optional[str]) -> int:
+    """Count real text characters in exported content, ignoring placeholders."""
+    if not content:
+        return 0
+    return len(_IMAGE_PLACEHOLDER_RE.sub("", content).strip())
+
+
+def _positive_len(value: object) -> int:
+    """``len(value)`` when it is a positive int, 0 when it is neither."""
+    try:
+        length = len(value)
+    except TypeError:
+        return 0
+    return length if isinstance(length, int) and length > 0 else 0
+
+
+def _result_page_count(result: object, file: Path) -> int:
+    """Page count of a conversion result, never below 1.
+
+    Args:
+        result: Docling ``ConversionResult``.
+        file: Path of the file that was converted.
+
+    Returns:
+        Number of pages; 1 for images and whenever docling reports none.
+    """
+    if Path(file).suffix.lower() in _IMAGE_SUFFIXES:
+        return 1
+    document = getattr(result, "document", None)
+    pages = _positive_len(getattr(document, "pages", None))
+    if pages:
+        return pages
+    num_pages = getattr(document, "num_pages", None)
+    if callable(num_pages):
+        try:
+            reported = num_pages()
+        except Exception:  # pragma: no cover - defensive, docling-version drift
+            reported = None
+        if isinstance(reported, int) and reported > 0:
+            return reported
+    return _positive_len(getattr(result, "pages", None)) or 1
+
+
+def _pdf_text_layer_probe(file: Path) -> Tuple[int, int]:
+    """Measure a PDF's embedded text layer with pypdfium2.
+
+    Used only to explain a suspected OCR dropout: a near-empty parse of a PDF
+    that *does* carry a text layer means docling failed to read text it never
+    needed OCR for, while an absent text layer points at the OCR stage itself.
+
+    Args:
+        file: Path to the PDF.
+
+    Returns:
+        Tuple of (page count, characters in the text layer); (0, 0) when the
+        file cannot be probed.
+    """
+    try:
+        import pypdfium2 as pdfium
+    except ImportError:
+        return 0, 0
+    try:
+        pdf = pdfium.PdfDocument(str(file))
+    except Exception:
+        return 0, 0
+    try:
+        total = 0
+        for index in range(len(pdf)):
+            page = pdf[index]
+            textpage = page.get_textpage()
+            try:
+                total += max(0, textpage.count_chars())
+            finally:
+                textpage.close()
+                page.close()
+        return len(pdf), total
+    except Exception:
+        return 0, 0
+    finally:
+        pdf.close()
+
+
 class DoclingParser(BaseParser):
     """Parser using docling for advanced document processing.
 
@@ -291,6 +410,13 @@ class DoclingParser(BaseParser):
             ocr_options = self._get_ocr_options()
             if ocr_options is not None:
                 pipeline_options.ocr_options = ocr_options
+            # Docling's *default* OCR options carry their own flag, so without
+            # this the setting was silently dropped whenever `_get_ocr_options`
+            # returned None (use_rapidocr=False) — including the dropout retry,
+            # whose whole point is forcing full-page OCR.
+            active_ocr_options = getattr(pipeline_options, "ocr_options", None)
+            if hasattr(active_ocr_options, "force_full_page_ocr"):
+                active_ocr_options.force_full_page_ocr = self.force_full_page_ocr
 
         return DocumentConverter(
             format_options={
@@ -383,18 +509,113 @@ class DoclingParser(BaseParser):
 
         return content
 
+    def _ocr_guard_applies(self, file: Path) -> bool:
+        """Whether the near-empty-output dropout guard covers this parse."""
+        return (
+            bool(self.ocr_enabled)
+            and _ocr_min_chars_per_page() > 0
+            and Path(file).suffix.lower() in _OCR_GUARDED_SUFFIXES
+        )
+
+    def _recover_from_ocr_dropout(
+        self, file: Path, first_pass_chars: int, pages: int
+    ) -> str:
+        """Retry a near-empty OCR parse once on a fresh full-page-OCR converter.
+
+        Docling caches its pipeline (and the threaded page queue behind it) on
+        the ``DocumentConverter`` instance, so a converter that has degraded
+        mid-worker keeps returning empty pages while a fresh one starts clean —
+        hence a brand new converter rather than a second `convert` call.
+
+        ``self._converter`` is dropped either way instead of being replaced by
+        the retry converter: the degraded instance must not survive, but the
+        retry one is configured for full-page OCR, and keeping it would impose
+        that cost (and its worse results on text PDFs) on every later file in
+        the worker. The next parse lazily rebuilds one with normal options.
+
+        Args:
+            file: Path to the file being parsed.
+            first_pass_chars: Text characters the first pass produced.
+            pages: Page count reported for the first pass.
+
+        Returns:
+            The retry's content, when it clears the chars-per-page floor.
+
+        Raises:
+            DocumentParseError: If the retry is near-empty too, so the caller
+                fails loudly instead of indexing an empty document.
+        """
+        name = Path(file).name
+        logger.warning(
+            "OCR output for %s is near-empty (%d chars over %d page(s)); "
+            "retrying once on a fresh converter with full-page OCR",
+            name,
+            first_pass_chars,
+            pages,
+        )
+        original_force_full_page_ocr = self.force_full_page_ocr
+        try:
+            self.force_full_page_ocr = True
+            result = self._create_converter().convert(str(file))
+            content = self._export_content(result.document)
+        finally:
+            self.force_full_page_ocr = original_force_full_page_ocr
+            self._converter = None
+
+        retry_chars = _text_char_count(content)
+        retry_pages = max(pages, _result_page_count(result, file))
+        if retry_chars >= _ocr_min_chars_per_page() * retry_pages:
+            logger.warning(
+                "Recovered %s on retry: %d chars over %d page(s) after a "
+                "near-empty first pass (%d chars); the degraded converter has "
+                "been discarded",
+                name,
+                retry_chars,
+                retry_pages,
+                first_pass_chars,
+            )
+            return content
+
+        detail = ""
+        if Path(file).suffix.lower() == ".pdf":
+            layer_pages, layer_chars = _pdf_text_layer_probe(file)
+            if layer_chars > 0:
+                detail = (
+                    f" The PDF carries a {layer_chars}-char text layer over "
+                    f"{layer_pages} page(s) that docling should have read "
+                    "without OCR at all."
+                )
+            elif layer_pages > 0:
+                detail = (
+                    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
+        )
+
     def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
         """Parse file using docling with hybrid OCR.
 
         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.
+
         Args:
             file: Path to the file to parse
             errors: Error handling mode (ignored, docling handles internally)
 
         Returns:
             Parsed document content as markdown string
+
+        Raises:
+            DocumentParseError: If docling fails, or if an OCR parse stays
+                near-empty across both passes.
         """
         logger.info(f"parse_file called for: {file}")
 
@@ -405,10 +626,29 @@ class DoclingParser(BaseParser):
             logger.info(f"Converting file with hybrid OCR: {file}")
             result = self._converter.convert(str(file))
             content = self._export_content(result.document)
-            logger.info(f"Parse complete, content length: {len(content)} chars")
+            pages = _result_page_count(result, file)
+            chars = _text_char_count(content)
+            logger.info(
+                "Parse complete for %s: %d chars (%d text) over %d page(s), "
+                "%.1f chars/page",
+                Path(file).name,
+                len(content),
+                chars,
+                pages,
+                chars / pages,
+            )
+
+            if self._ocr_guard_applies(file) and chars < (
+                _ocr_min_chars_per_page() * pages
+            ):
+                return self._recover_from_ocr_dropout(file, chars, pages)
 
             return content
 
+        except DocumentParseError:
+            # Already the loud, actionable failure — do not re-wrap it into a
+            # generic "Failed to parse ..." and lose the dropout diagnosis.
+            raise
         except Exception as e:
             logger.error(f"Error parsing file with docling: {e}", exc_info=True)
             # ``errors`` governs *decoding* leniency, not whether a total

application/retriever/classic_rag.py

@@ -1,8 +1,17 @@
 import logging
+from typing import Any, Dict, List, Optional, Tuple
 
 from application.core.settings import settings
 from application.llm.llm_creator import LLMCreator
 from application.retriever.base import BaseRetriever
+from application.retriever.fanout import (  # noqa: F401  (re-exported)
+    DEFAULT_MAX_PARALLEL_SOURCES,
+    EMBEDDING_ATTRS as _EMBEDDING_ATTRS,
+    embed_questions,
+    max_parallel_sources,
+    run_source_jobs,
+    store_embeddings,
+)
 from application.retriever.labels import labels_from_metadata
 from application.utils import num_tokens_from_string
 from application.vectorstore.vector_creator import VectorCreator
@@ -10,6 +19,11 @@ from application.vectorstore.vector_creator import VectorCreator
 logger = logging.getLogger(__name__)
 
 
+def _max_parallel_sources(n_sources: int) -> int:
+    """Worker count for the per-source fan-out, bounded by the source count."""
+    return max_parallel_sources(n_sources, settings)
+
+
 class ClassicRAG(BaseRetriever):
     # The group's real top-k, set by the Dispatcher when it inflates ``chunks``
     # for a prescreen fetch. None → ``chunks`` is already the top-k.
@@ -152,13 +166,25 @@ class ClassicRAG(BaseRetriever):
             logger.error(f"Error rephrasing query: {e}", exc_info=True)
             return self.original_question
 
-    def _fetch_candidates(self, docsearch, question, src_k, score_threshold):
+    def _fetch_candidates(
+        self,
+        docsearch,
+        question: str,
+        src_k: int,
+        score_threshold: Optional[float],
+        query_vector: Optional[List[float]] = None,
+    ):
         """Fetch candidate hits for one vector store (vector search).
 
         Returns plain hits, or ``(hit, score)`` pairs when ``include_scores`` is
         set. Subclasses override this to change candidate sourcing (e.g. RRF
         fusion) while inheriting the surrounding per-source resolution and
         budgeting.
+
+        Args:
+            query_vector: Query embedding computed once for the whole
+                retrieval. Forwarded so the store skips embedding the query
+                again; stores that don't support it ignore the kwarg.
         """
         # ``score_threshold`` is honoured by pgvector/mongodb and safely ignored
         # by stores whose ``search`` swallows kwargs. The candidate count is
@@ -167,6 +193,8 @@ class ClassicRAG(BaseRetriever):
         search_kwargs = {"k": k}
         if score_threshold is not None:
             search_kwargs["score_threshold"] = score_threshold
+        if query_vector is not None:
+            search_kwargs["query_vector"] = query_vector
         if self.include_scores:
             return docsearch.search_with_scores(question, **search_kwargs)
         return docsearch.search(question, **search_kwargs)
@@ -175,6 +203,152 @@ class ClassicRAG(BaseRetriever):
         """Label for the scores ``_fetch_candidates`` attaches (None if unscored)."""
         return getattr(docsearch, "score_kind", None)
 
+    def _resolve_source(
+        self, vectorstore_id: str, chunks_per_source: int
+    ) -> Dict[str, Any]:
+        """Resolve one source's fetch parameters (top-k, threshold, query).
+
+        Per-source overrides come from the Dispatcher; absent, the source gets
+        the global behaviour — byte-identical to the pre-override path.
+        """
+        src_cfg = self.per_source_retrieval.get(vectorstore_id)
+        if src_cfg is None:
+            # No per-source override → the effective rephrase_query defaults to
+            # True, so use the (lazily-cached) rephrased question. In the
+            # non-deferred path the cache is already populated.
+            return {
+                "id": vectorstore_id,
+                "src_k": chunks_per_source,
+                "score_threshold": None,
+                "question": self._get_rephrased_question(),
+            }
+        src_k = max(1, int(src_cfg.chunks))
+        # Prescreen fetches a larger candidate set up front; the Dispatcher's
+        # prescreen stage trims back to max_keep afterwards. Raise the fetch
+        # size to candidate_k here.
+        ps_cfg = (
+            src_cfg.prescreen_config() if hasattr(src_cfg, "prescreen_config") else None
+        )
+        if ps_cfg is not None:
+            src_k = max(src_k, int(ps_cfg.candidate_k))
+        return {
+            "id": vectorstore_id,
+            "src_k": src_k,
+            "score_threshold": src_cfg.score_threshold,
+            "question": (
+                self._get_rephrased_question()
+                if src_cfg.rephrase_query
+                else self.original_question
+            ),
+        }
+
+    def _plan_sources(self, chunks_per_source: int) -> List[Dict[str, Any]]:
+        """Resolve every source's fetch parameters, in source order.
+
+        Runs on the calling thread: the lazy rephrase behind it is an LLM
+        side-call that must happen once, not once per worker.
+        """
+        plans = []
+        for vectorstore_id in self.vectorstores:
+            if not vectorstore_id:
+                continue
+            try:
+                plans.append(self._resolve_source(vectorstore_id, chunks_per_source))
+            except Exception as e:
+                logger.error(
+                    f"Error searching vectorstore {vectorstore_id}: {e}", exc_info=True
+                )
+        return plans
+
+    @staticmethod
+    def _store_embeddings(docsearch):
+        """Return the embeddings object a vector store searches with, if any."""
+        return store_embeddings(docsearch)
+
+    def _embed_questions(
+        self, docsearch, questions: List[str]
+    ) -> Dict[str, List[float]]:
+        """Embed each distinct query once for the whole retrieval."""
+        return embed_questions(docsearch, questions)
+
+    def _search_source(
+        self,
+        plan: Dict[str, Any],
+        docsearch=None,
+        query_vector: Optional[List[float]] = None,
+    ) -> Optional[Tuple[Any, Optional[str]]]:
+        """Search one source, returning ``(candidates, score_kind)``.
+
+        Builds the vector store when not supplied, so each worker thread owns
+        its own store instance (and therefore its own DB connection). Errors are
+        logged and reported as ``None`` so one bad source cannot take the rest
+        of the retrieval down with it.
+        """
+        try:
+            if docsearch is None:
+                docsearch = VectorCreator.create_vectorstore(
+                    settings.VECTOR_STORE, plan["id"], settings.EMBEDDINGS_KEY
+                )
+            docs_temp = self._fetch_candidates(
+                docsearch,
+                plan["question"],
+                plan["src_k"],
+                plan["score_threshold"],
+                query_vector=query_vector,
+            )
+            score_kind = self._score_kind(docsearch) if self.include_scores else None
+            return docs_temp, score_kind
+        except Exception as e:
+            logger.error(
+                f"Error searching vectorstore {plan['id']}: {e}", exc_info=True
+            )
+            return None
+
+    def _fetch_all(
+        self, plans: List[Dict[str, Any]]
+    ) -> List[Optional[Tuple[Any, Optional[str]]]]:
+        """Fetch every source's candidates, one embedding and one fan-out.
+
+        The first store is built on the calling thread because its embeddings
+        object supplies the shared query vector — and priming the embeddings
+        singleton there keeps the workers off a concurrent model load. The
+        searches then run on a bounded pool, and results come back in the
+        original source order so the merge stays deterministic.
+        """
+        first_store = None
+        try:
+            first_store = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, plans[0]["id"], settings.EMBEDDINGS_KEY
+            )
+        except Exception as e:
+            logger.error(
+                f"Error searching vectorstore {plans[0]['id']}: {e}", exc_info=True
+            )
+
+        if first_store is None:
+            # The first source is already a logged failure; the rest still run,
+            # each embedding its own query (there is no store to borrow one from).
+            return [None] + self._run_jobs([(plan, None, None) for plan in plans[1:]])
+
+        questions = list(dict.fromkeys(plan["question"] for plan in plans))
+        vectors = self._embed_questions(first_store, questions)
+        return self._run_jobs(
+            [
+                (plan, first_store if idx == 0 else None, vectors.get(plan["question"]))
+                for idx, plan in enumerate(plans)
+            ]
+        )
+
+    def _run_jobs(self, jobs: List[Tuple]) -> List[Optional[Tuple[Any, Optional[str]]]]:
+        """Run the per-source search jobs, concurrently when there are several."""
+        if not jobs:
+            return []
+        return run_source_jobs(
+            lambda job: self._search_source(*job),
+            jobs,
+            workers=_max_parallel_sources(len(jobs)),
+        )
+
     def _get_data(self):
         if self.chunks == 0 or not self.vectorstores:
             logger.info(
@@ -194,96 +368,64 @@ class ClassicRAG(BaseRetriever):
         token_budget = max(int(self.doc_token_limit * 0.9), 100)
         cumulative_tokens = 0
 
-        for vectorstore_id in self.vectorstores:
-            if vectorstore_id:
-                try:
-                    # Per-source overrides (set by the Dispatcher). Absent →
-                    # global behaviour, byte-identical to before.
-                    src_cfg = self.per_source_retrieval.get(vectorstore_id)
-                    if src_cfg is not None:
-                        src_k = max(1, int(src_cfg.chunks))
-                        # Prescreen fetches a larger candidate set up front; the
-                        # Dispatcher's prescreen stage trims back to max_keep
-                        # afterwards. Raise the fetch size to candidate_k here.
-                        ps_cfg = (
-                            src_cfg.prescreen_config()
-                            if hasattr(src_cfg, "prescreen_config")
-                            else None
-                        )
-                        if ps_cfg is not None:
-                            src_k = max(src_k, int(ps_cfg.candidate_k))
-                        score_threshold = src_cfg.score_threshold
-                        question = (
-                            self._get_rephrased_question()
-                            if src_cfg.rephrase_query
-                            else self.original_question
-                        )
-                    else:
-                        src_k = chunks_per_source
-                        score_threshold = None
-                        # No per-source override → the effective rephrase_query
-                        # defaults to True, so use the (lazily-cached) rephrased
-                        # question. In the non-deferred path the cache is already
-                        # populated, so this matches today's behaviour exactly.
-                        question = self._get_rephrased_question()
+        # Resolve every source, then fetch them all (one query embedding, one
+        # bounded fan-out). The merge below stays serial and in source order, so
+        # dedupe/budget/trim semantics are exactly what they were.
+        plans = self._plan_sources(chunks_per_source)
+        results = self._fetch_all(plans) if plans else []
 
-                    docsearch = VectorCreator.create_vectorstore(
-                        settings.VECTOR_STORE, vectorstore_id, settings.EMBEDDINGS_KEY
-                    )
-                    docs_temp = self._fetch_candidates(
-                        docsearch, question, src_k, score_threshold
-                    )
-                    score_kind = (
-                        self._score_kind(docsearch) if self.include_scores else None
-                    )
+        for plan, result in zip(plans, results):
+            if result is None:
+                continue
+            vectorstore_id = plan["id"]
+            src_k = plan["src_k"]
+            docs_temp, score_kind = result
+            try:
+                # ``_fetch_candidates`` over-fetches (k >= 20) so a prescreen
+                # stage has candidates to filter; trim back to src_k so
+                # ``chunks`` is the final top-k it claims to be. With
+                # prescreen on, src_k is already raised to candidate_k above,
+                # so the stage still sees its full candidate set.
+                kept = 0
 
-                    # ``_fetch_candidates`` over-fetches (k >= 20) so a prescreen
-                    # stage has candidates to filter; trim back to src_k so
-                    # ``chunks`` is the final top-k it claims to be. With
-                    # prescreen on, src_k is already raised to candidate_k above,
-                    # so the stage still sees its full candidate set.
-                    kept = 0
-
-                    for doc in docs_temp:
-                        if kept >= src_k or cumulative_tokens >= token_budget:
-                            break
-
-                        score = None
-                        if isinstance(doc, tuple):
-                            doc, score = doc
-
-                        if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
-                            page_content = doc.page_content
-                            metadata = doc.metadata
-                        else:
-                            page_content = doc.get("text", doc.get("page_content", ""))
-                            metadata = doc.get("metadata", {})
-
-                        labels = labels_from_metadata(
-                            metadata, page_content, vectorstore_id
-                        )
-
-                        doc_text_with_header = f"{labels['filename']}\n{page_content}"
-                        doc_tokens = num_tokens_from_string(doc_text_with_header)
-
-                        if cumulative_tokens + doc_tokens < token_budget:
-                            entry = {"text": page_content, **labels}
-                            if self.include_scores:
-                                entry["score"] = score
-                                entry["score_kind"] = score_kind
-                            all_docs.append(entry)
-                            cumulative_tokens += doc_tokens
-                            kept += 1
-
-                    if cumulative_tokens >= token_budget:
+                for doc in docs_temp:
+                    if kept >= src_k or cumulative_tokens >= token_budget:
                         break
 
-                except Exception as e:
-                    logger.error(
-                        f"Error searching vectorstore {vectorstore_id}: {e}",
-                        exc_info=True,
-                    )
-                    continue
+                    score = None
+                    if isinstance(doc, tuple):
+                        doc, score = doc
+
+                    if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
+                        page_content = doc.page_content
+                        metadata = doc.metadata
+                    else:
+                        page_content = doc.get("text", doc.get("page_content", ""))
+                        metadata = doc.get("metadata", {})
+
+                    labels = labels_from_metadata(metadata, page_content, vectorstore_id)
+
+                    doc_text_with_header = f"{labels['filename']}\n{page_content}"
+                    doc_tokens = num_tokens_from_string(doc_text_with_header)
+
+                    if cumulative_tokens + doc_tokens < token_budget:
+                        entry = {"text": page_content, **labels}
+                        if self.include_scores:
+                            entry["score"] = score
+                            entry["score_kind"] = score_kind
+                        all_docs.append(entry)
+                        cumulative_tokens += doc_tokens
+                        kept += 1
+
+                if cumulative_tokens >= token_budget:
+                    break
+
+            except Exception as e:
+                logger.error(
+                    f"Error searching vectorstore {vectorstore_id}: {e}",
+                    exc_info=True,
+                )
+                continue
 
         # ``chunks_per_source`` has a floor of 1 so no attached source is
         # starved, which means N sources always yield at least N documents —

application/retriever/fanout.py

@@ -0,0 +1,127 @@
+"""Shared helpers for fanning one query out across several vector stores.
+
+Attaching N sources used to cost N query embeddings and N serial round trips.
+The pieces here collapse that to one embedding plus one bounded, order-
+preserving thread pool, and are reused by :mod:`application.retriever.classic_rag`
+and :mod:`application.services.search_service` so both paths behave the same.
+"""
+
+from __future__ import annotations
+
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar
+
+from application.core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+# Fallback fan-out width when ``RETRIEVAL_MAX_PARALLEL_SOURCES`` is unset.
+DEFAULT_MAX_PARALLEL_SOURCES = 4
+
+# Attributes the vector stores keep their embeddings object on, in the order
+# they are probed. Stores that build embeddings inline (elasticsearch, lancedb)
+# expose none of them and are handled by the ``_get_embeddings`` fallback.
+EMBEDDING_ATTRS = ("_embedding", "_embeddings", "embeddings")
+
+Job = TypeVar("Job")
+Result = TypeVar("Result")
+
+
+def max_parallel_sources(n_sources: int, settings_obj: Any = None) -> int:
+    """Worker count for the per-source fan-out, bounded by the source count.
+
+    Args:
+        n_sources: Number of sources about to be searched.
+        settings_obj: Settings to read ``RETRIEVAL_MAX_PARALLEL_SOURCES`` from;
+            defaults to the global settings.
+
+    Returns:
+        int: Worker count, at least 1 and never above ``n_sources``.
+    """
+    config = settings if settings_obj is None else settings_obj
+    configured = getattr(
+        config, "RETRIEVAL_MAX_PARALLEL_SOURCES", DEFAULT_MAX_PARALLEL_SOURCES
+    )
+    try:
+        configured = int(configured)
+    except (TypeError, ValueError):
+        configured = DEFAULT_MAX_PARALLEL_SOURCES
+    return max(1, min(configured, n_sources))
+
+
+def store_embeddings(docsearch) -> Optional[Any]:
+    """Return the embeddings object a vector store searches with, if any."""
+    for attr in EMBEDDING_ATTRS:
+        embedder = getattr(docsearch, attr, None)
+        if embedder is not None and hasattr(embedder, "embed_query"):
+            return embedder
+    getter = getattr(docsearch, "_get_embeddings", None)
+    if callable(getter):
+        try:
+            return getter(settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY)
+        except Exception as e:
+            logger.debug("Could not resolve store embeddings: %s", e)
+    return None
+
+
+def embed_questions(docsearch, questions: Iterable[str]) -> Dict[str, List[float]]:
+    """Embed each distinct query once for the whole retrieval.
+
+    Every source of one retrieval shares a single embeddings config
+    (``settings.EMBEDDINGS_NAME`` is global — nothing per-source overrides it),
+    so a vector computed with one store's embedder is valid for all of them.
+    That turns N sources from N query embeddings into one.
+
+    Args:
+        docsearch: Any store of the retrieval; only its embedder is used.
+        questions: Queries to embed (duplicates are embedded once).
+
+    Returns:
+        dict: ``{question: vector}``, or ``{}`` when no embedder is reachable
+        or embedding fails — in which case every store embeds its own query,
+        exactly as before.
+    """
+    embedder = store_embeddings(docsearch)
+    if embedder is None:
+        return {}
+    try:
+        return {
+            question: embedder.embed_query(question)
+            for question in dict.fromkeys(questions)
+        }
+    except Exception as e:
+        logger.warning(
+            "Query embedding failed (%s); each store will embed its own.", e
+        )
+        return {}
+
+
+def run_source_jobs(
+    fn: Callable[[Job], Result],
+    jobs: Iterable[Job],
+    workers: Optional[int] = None,
+) -> List[Result]:
+    """Run one job per source, concurrently when there are several.
+
+    Args:
+        fn: Per-source callable; it must swallow its own errors, since a raise
+            here aborts the whole fan-out.
+        jobs: Job payloads, in source order.
+        workers: Pool size; defaults to :func:`max_parallel_sources`.
+
+    Returns:
+        list: ``fn`` results in the original job order, so the caller's merge
+        stays deterministic no matter which source finishes first.
+    """
+    jobs = list(jobs)
+    if not jobs:
+        return []
+    if workers is None:
+        workers = max_parallel_sources(len(jobs))
+    if workers == 1:
+        return [fn(job) for job in jobs]
+    with ThreadPoolExecutor(
+        max_workers=workers, thread_name_prefix="rag-source"
+    ) as pool:
+        return list(pool.map(fn, jobs))

application/retriever/graph_rag.py

@@ -7,7 +7,11 @@ reused) rephrase.
 
 Composes :class:`ClassicRAG` rather than subclassing: PPR doesn't fit the
 ``_fetch_candidates`` hook, but the composed instance supplies the rephrase, the
-token-budget loop, and the per-source fallback when a source has no graph.
+token-budget loop, and the fallback for sources that have no graph.
+
+Per request the whole source group costs one node-count query, one query
+embedding, and one ClassicRAG run for all the graphless sources together — all
+over a single pooled connection shared with the vector store.
 """
 
 from __future__ import annotations
@@ -25,7 +29,7 @@ from application.retriever.base import BaseRetriever
 from application.retriever.classic_rag import ClassicRAG
 from application.retriever.labels import labels_from_metadata
 from application.utils import num_tokens_from_string
-from application.vectorstore.base import EmbeddingsSingleton
+from application.vectorstore.base import get_embeddings
 
 SEED_NODES = 10
 SUBGRAPH_HOPS = 1
@@ -88,9 +92,7 @@ class GraphRAGRetriever(BaseRetriever):
         self.per_source_retrieval = {}
 
     def _embed_query(self, question: str) -> List[float]:
-        embedding = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-        )
+        embedding = get_embeddings()
         return embedding.embed_query(question)
 
     def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]:
@@ -157,10 +159,17 @@ class GraphRAGRetriever(BaseRetriever):
         base = self.base_chunks if self.base_chunks is not None else self.chunks
         return max(1, base // max(1, len(self.vectorstores)))
 
-    def _graph_docs_for_source(self, store, source_id) -> List[Dict[str, Any]]:
-        """Local PPR retrieval for one source (caller guarantees it has a graph)."""
-        question = self._classic._get_rephrased_question()
-        query_embedding = self._embed_query(question)
+    def _graph_docs_for_source(
+        self, store, source_id, query_embedding: List[float]
+    ) -> List[Dict[str, Any]]:
+        """Local PPR retrieval for one source (caller guarantees it has a graph).
+
+        Args:
+            store: Open :class:`GraphStore` shared by every source of this run.
+            source_id: Source to retrieve from.
+            query_embedding: Embedding of the rephrased question, computed once
+                by the caller for the whole retrieval.
+        """
         seed_rows = store.search_nodes_by_embedding(
             source_id, query_embedding, k=SEED_NODES
         )
@@ -204,15 +213,26 @@ class GraphRAGRetriever(BaseRetriever):
             cumulative_tokens += doc_tokens
         return docs
 
-    def _classic_for_source(self, source_id) -> List[Dict[str, Any]]:
-        """Reuse the composed ClassicRAG to retrieve one source's chunks."""
+    def _classic_for_sources(self, source_ids) -> List[Dict[str, Any]]:
+        """Reuse the composed ClassicRAG to retrieve a whole batch of sources.
+
+        One inner run for every graphless source instead of one run per source:
+        ClassicRAG then embeds the query once and fans the sources out itself.
+        The trade is that the per-source chunk split and the shared doc-token
+        budget apply across the batch — i.e. exactly plain ClassicRAG semantics
+        over those sources, rather than each source getting its own full budget.
+        """
+        source_ids = [source_id for source_id in source_ids if source_id]
+        if not source_ids:
+            return []
+        wanted = set(source_ids)
         original = self._classic.vectorstores
         original_overrides = self._classic.per_source_retrieval
         original_base = self._classic.base_chunks
         try:
-            self._classic.vectorstores = [source_id]
+            self._classic.vectorstores = list(source_ids)
             self._classic.per_source_retrieval = {
-                k: v for k, v in self.per_source_retrieval.items() if k == source_id
+                k: v for k, v in self.per_source_retrieval.items() if k in wanted
             }
             # The Dispatcher sets these on *this* object; the inner retriever is
             # the one that reads them.
@@ -223,8 +243,79 @@ class GraphRAGRetriever(BaseRetriever):
             self._classic.per_source_retrieval = original_overrides
             self._classic.base_chunks = original_base
 
+    def _classic_for_source(self, source_id) -> List[Dict[str, Any]]:
+        """Retrieve one source through the batched classic path."""
+        return self._classic_for_sources([source_id])
+
+    def _retrieve_with_store(self, store, sources) -> List[Dict[str, Any]]:
+        """Split ``sources`` by graph presence, then batch each half.
+
+        Graph sources keep their own slot in source order; every graphless
+        source collapses into a single ClassicRAG run that occupies the slot of
+        the first graphless source. Sources whose PPR retrieval raises are
+        collected and retried as one more classic batch, appended at the end.
+        """
+        try:
+            counts = store.count_nodes_many(sources)
+        except Exception as e:
+            logging.error(f"GraphRAG count_nodes failed for {sources}: {e}")
+            counts = {}
+
+        segments: List[List[Dict[str, Any]]] = []
+        graph_slots: Dict[str, int] = {}
+        graphed: List[str] = []
+        graphless: List[str] = []
+        classic_slot = None
+        for source_id in sources:
+            if counts.get(source_id, 0) > 0:
+                if source_id in graph_slots:
+                    continue
+                graph_slots[source_id] = len(segments)
+                segments.append([])
+                graphed.append(source_id)
+            else:
+                if classic_slot is None:
+                    classic_slot = len(segments)
+                    segments.append([])
+                graphless.append(source_id)
+
+        failed: List[str] = []
+        query_embedding = None
+        if graphed:
+            # Embedded once for the whole retrieval, not once per graph source.
+            try:
+                query_embedding = self._embed_query(
+                    self._classic._get_rephrased_question()
+                )
+            except Exception as e:
+                logging.error(
+                    f"GraphRAG query embedding failed, falling back: {e}",
+                    exc_info=True,
+                )
+                failed, graphed = list(graphed), []
+
+        for source_id in graphed:
+            try:
+                segments[graph_slots[source_id]] = self._graph_docs_for_source(
+                    store, source_id, query_embedding
+                )
+            except Exception as e:
+                logging.error(
+                    f"GraphRAG retrieval failed for {source_id}, falling back: {e}",
+                    exc_info=True,
+                )
+                failed.append(source_id)
+
+        if graphless:
+            segments[classic_slot] = self._classic_for_sources(graphless)
+        if failed:
+            segments.append(self._classic_for_sources(failed))
+
+        return [doc for segment in segments for doc in segment]
+
     def _get_data(self) -> List[Dict[str, Any]]:
-        if not self.vectorstores:
+        sources = [source_id for source_id in self.vectorstores if source_id]
+        if not sources:
             return []
 
         store = None
@@ -235,31 +326,17 @@ class GraphRAGRetriever(BaseRetriever):
                 logging.error(f"GraphRAG store unavailable, falling back: {e}")
                 store = None
 
-        all_docs: List[Dict[str, Any]] = []
-        for source_id in self.vectorstores:
-            if not source_id:
-                continue
-            has_graph = False
-            if store is not None:
-                try:
-                    has_graph = store.count_nodes(source_id) > 0
-                except Exception as e:
-                    logging.error(f"GraphRAG count_nodes failed for {source_id}: {e}")
-                    has_graph = False
-
-            if not has_graph:
-                all_docs.extend(self._classic_for_source(source_id))
-                continue
+        if store is None:
+            return self._classic_for_sources(sources)
 
+        try:
+            return self._retrieve_with_store(store, sources)
+        finally:
+            # Hand the pooled connection back; the store is per-request.
             try:
-                all_docs.extend(self._graph_docs_for_source(store, source_id))
+                store.close()
             except Exception as e:
-                logging.error(
-                    f"GraphRAG retrieval failed for {source_id}, falling back: {e}",
-                    exc_info=True,
-                )
-                all_docs.extend(self._classic_for_source(source_id))
-        return all_docs
+                logging.debug("Error closing GraphRAG store: %s", e)
 
     def search(self, query: str = "") -> List[Dict[str, Any]]:
         if query:

application/retriever/hybrid_rag.py

@@ -55,16 +55,22 @@ class HybridRetriever(ClassicRAG):
         carry their own label."""
         return "rrf"
 
-    def _fetch_candidates(self, docsearch, question, src_k, score_threshold):
+    def _fetch_candidates(
+        self, docsearch, question, src_k, score_threshold, query_vector=None
+    ):
         """Return RRF-fused vector+keyword hits for one vector store.
 
         Inherits the per-source resolution and budgeting from
         :meth:`ClassicRAG._get_data`; only candidate sourcing differs.
         RRF scores are not cosine similarities, so ``score_threshold`` is
-        intentionally not applied to the fused list.
+        intentionally not applied to the fused list. ``query_vector`` is the
+        retrieval's single query embedding — the keyword half never needs one.
         """
         candidate_k = min(max(src_k * 2, 20), 500)
-        vector_hits = docsearch.search(question, k=candidate_k)
+        vector_kwargs = {"k": candidate_k}
+        if query_vector is not None:
+            vector_kwargs["query_vector"] = query_vector
+        vector_hits = docsearch.search(question, **vector_kwargs)
         keyword_hits = docsearch.keyword_search(question, k=candidate_k)
         fused = fuse_with_scores(vector_hits, keyword_hits)
         if self.include_scores:

application/services/search_service.py

@@ -8,9 +8,10 @@ MCP error responses, etc.).
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from application.core.settings import settings
+from application.retriever.fanout import embed_questions, run_source_jobs
 from application.storage.db.repositories.agents import AgentsRepository
 from application.storage.db.session import db_readonly
 from application.vectorstore.vector_creator import VectorCreator
@@ -89,6 +90,79 @@ def _authorized_source_ids(conn, agent: Dict[str, Any], source_ids: List[str]) -
     return allowed
 
 
+def _search_one(
+    source_id: str,
+    docsearch: Any,
+    query: str,
+    k: int,
+    query_vector: Optional[List[float]],
+) -> Optional[List[Any]]:
+    """Search one source, returning its hits or ``None`` when it fails.
+
+    Builds the vector store when not supplied, so each worker thread owns its
+    own store instance (and therefore its own DB connection). Errors are logged
+    and reported as ``None`` so one broken index cannot take the rest down.
+    """
+    try:
+        if docsearch is None:
+            docsearch = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+            )
+        search_kwargs: Dict[str, Any] = {"k": k}
+        if query_vector is not None:
+            search_kwargs["query_vector"] = query_vector
+        return docsearch.search(query, **search_kwargs)
+    except Exception as e:
+        logger.error(
+            f"Error searching vectorstore {source_id}: {e}",
+            exc_info=True,
+        )
+        return None
+
+
+def _fetch_sources(
+    query: str, source_ids: List[str], k: int
+) -> List[Optional[List[Any]]]:
+    """Fetch every source's hits: one query embedding, one bounded fan-out.
+
+    The first store is built on the calling thread because its embeddings
+    object supplies the shared query vector — and priming the embeddings
+    singleton there keeps the workers off a concurrent model load. Results come
+    back in the original source order so the merge stays deterministic.
+    """
+
+    def _job(job) -> Optional[List[Any]]:
+        source_id, docsearch, query_vector = job
+        return _search_one(source_id, docsearch, query, k, query_vector)
+
+    first_store = None
+    try:
+        first_store = VectorCreator.create_vectorstore(
+            settings.VECTOR_STORE, source_ids[0], settings.EMBEDDINGS_KEY
+        )
+    except Exception as e:
+        logger.error(
+            f"Error searching vectorstore {source_ids[0]}: {e}",
+            exc_info=True,
+        )
+
+    if first_store is None:
+        # The first source is already a logged failure; the rest still run,
+        # each embedding its own query (there is no store to borrow one from).
+        return [None] + run_source_jobs(
+            _job, [(sid, None, None) for sid in source_ids[1:]]
+        )
+
+    vector = embed_questions(first_store, [query]).get(query)
+    return run_source_jobs(
+        _job,
+        [
+            (sid, first_store if idx == 0 else None, vector)
+            for idx, sid in enumerate(source_ids)
+        ],
+    )
+
+
 def _search_sources(
     query: str, source_ids: List[str], chunks: int
 ) -> List[Dict[str, Any]]:
@@ -101,19 +175,24 @@ def _search_sources(
         return []
 
     results: List[Dict[str, Any]] = []
+    # Blank ids build no store but still count here, so the per-source budget
+    # matches what the serial implementation handed each real source.
     chunks_per_source = max(1, chunks // len(source_ids))
     seen_texts: set[int] = set()
 
-    for source_id in source_ids:
-        if not source_id or not source_id.strip():
+    active_ids = [sid for sid in source_ids if sid and sid.strip()]
+    if not active_ids:
+        return []
+
+    # Fetch every source up front, then merge serially in source order so the
+    # dedupe / cap semantics are exactly what they were.
+    fetched = _fetch_sources(query, active_ids, chunks_per_source * 2)
+
+    for source_id, docs in zip(active_ids, fetched):
+        if docs is None:
             continue
 
         try:
-            docsearch = VectorCreator.create_vectorstore(
-                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
-            )
-            docs = docsearch.search(query, k=chunks_per_source * 2)
-
             for doc in docs:
                 if len(results) >= chunks:
                     break

application/storage/db/bootstrap.py

@@ -7,9 +7,15 @@ On app startup the Flask factory (and Celery worker init) can call
    configured role to have ``CREATEDB`` privilege).
 2. Apply every pending Alembic migration up to ``head``.
 
-Both steps are gated by settings that default ON for dev convenience and
-can be turned off in prod (``AUTO_CREATE_DB`` / ``AUTO_MIGRATE``) where
-schema is managed out-of-band by a deploy pipeline.
+:func:`ensure_vector_schema` does the same job for the *vector* database
+(``pgvector``): it owns the ``documents`` table and, when GraphRAG is on,
+the graph tables, so the retrieval hot path never runs DDL. That database
+may be a separate cluster, which is why it has no Alembic migration.
+
+Every step is gated by a setting that defaults ON for dev convenience and
+can be turned off in prod (``AUTO_CREATE_DB`` / ``AUTO_MIGRATE`` /
+``AUTO_VECTOR_SCHEMA``) where schema is managed out-of-band by a deploy
+pipeline.
 
 All heavy imports (alembic, psycopg, sqlalchemy.exc sub-symbols) are
 deferred to inside the function so merely importing this module has no
@@ -19,6 +25,7 @@ side effects and is cheap for test collection.
 from __future__ import annotations
 
 import logging
+import time
 from typing import Optional
 
 
@@ -69,6 +76,134 @@ def ensure_database_ready(
         _run_migrations(log)
 
 
+def ensure_vector_schema(*, logger: Optional[logging.Logger] = None) -> None:
+    """Create the pgvector schema once at boot and verify its dimension.
+
+    Owning the DDL here is what lets ``PGVectorStore`` construction stay free:
+    the retriever builds one store per source per request, and each used to run
+    ``CREATE EXTENSION`` / ``CREATE TABLE`` / two ``CREATE INDEX`` statements
+    before its first SELECT. A dimension mismatch between the table and the
+    configured embeddings model is fatal here rather than silent garbage
+    retrieval at query time.
+
+    Args:
+        logger: Optional logger. Defaults to this module's logger.
+
+    Raises:
+        RuntimeError: The existing ``documents`` table was built for a
+            different embedding width than the configured model produces.
+        Exception: Any connection or DDL failure is re-raised so the app fails
+            fast, matching :func:`ensure_database_ready`.
+    """
+    log = logger or logging.getLogger(__name__)
+
+    from application.core.settings import settings
+
+    store_kind = (settings.VECTOR_STORE or "").lower()
+    if store_kind != "pgvector":
+        log.debug(
+            "ensure_vector_schema: VECTOR_STORE is %r, not pgvector; nothing to do.",
+            settings.VECTOR_STORE,
+        )
+        return
+
+    dsn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
+    if not dsn and getattr(settings, "POSTGRES_URI", None):
+        from application.core.db_uri import normalize_pgvector_connection_string
+
+        dsn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
+    if not dsn:
+        log.info(
+            "ensure_vector_schema: no pgvector connection string configured "
+            "(PGVECTOR_CONNECTION_STRING / POSTGRES_URI); skipping."
+        )
+        return
+
+    import psycopg
+
+    from application.vectorstore.pgvector import (
+        DEFAULT_EMBEDDING_DIM,
+        SCHEMA_LOCK_KEY,
+        PGVectorStore,
+    )
+
+    # Loading the model here is deliberate: the process loads it on first
+    # retrieval anyway, and EmbeddingsSingleton caches it.
+    dim: Optional[int] = None
+    try:
+        from application.vectorstore.base import get_embeddings
+
+        dim = getattr(get_embeddings(), "dimension", None)
+    except Exception as exc:  # noqa: BLE001 — never block boot on the model
+        log.warning(
+            "ensure_vector_schema: could not load the embeddings model (%s); "
+            "creating the table with %d dimensions and skipping the dimension "
+            "check.",
+            exc,
+            DEFAULT_EMBEDDING_DIM,
+        )
+    if dim is None:
+        log.warning(
+            "ensure_vector_schema: the embeddings model exposes no dimension; "
+            "using %d and skipping the dimension check.",
+            DEFAULT_EMBEDDING_DIM,
+        )
+
+    graph_enabled = bool(getattr(settings, "GRAPHRAG_ENABLED", False))
+    started = time.monotonic()
+    # A plain connection, never the store's pool: this can run pre-fork under
+    # ``gunicorn --preload``, and an inherited pooled socket is a broken one.
+    conn = psycopg.connect(dsn)
+    try:
+        cursor = conn.cursor()
+        try:
+            # Serialize concurrent workers; released when this transaction ends.
+            cursor.execute(
+                "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+            )
+        finally:
+            cursor.close()
+
+        PGVectorStore.create_schema(conn, dimension=dim or DEFAULT_EMBEDDING_DIM)
+        if graph_enabled:
+            from application.graphrag.store import (
+                DEFAULT_NAME_EMBEDDING_DIM,
+                GraphStore,
+            )
+
+            GraphStore.create_schema(
+                conn, dimension=dim or DEFAULT_NAME_EMBEDDING_DIM
+            )
+        conn.commit()
+
+        actual = PGVectorStore.table_dimension(conn)
+        if dim and actual and actual != dim:
+            raise RuntimeError(
+                f"pgvector table 'documents' has vector({actual}) but "
+                f"EMBEDDINGS_NAME={settings.EMBEDDINGS_NAME} produces {dim}-dim "
+                "vectors; re-ingest into a fresh table or point "
+                "PGVECTOR_CONNECTION_STRING at the matching database."
+            )
+
+        log.info(
+            "ensure_vector_schema: table 'documents' ready "
+            "(dimension=%s, graph tables=%s) in %d ms.",
+            actual or dim or DEFAULT_EMBEDDING_DIM,
+            "yes" if graph_enabled else "no",
+            int((time.monotonic() - started) * 1000),
+        )
+    except Exception as exc:
+        try:
+            conn.rollback()
+        except Exception:
+            # Connection already gone; nothing left to roll back.
+            pass
+        log.error("ensure_vector_schema: %s", exc)
+        raise
+    finally:
+        conn.close()
+
+
 def _ensure_database_exists(uri: str, log: logging.Logger) -> None:
     """Create the target database if a connection reveals it's missing.
 

application/vectorstore/base.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import requests
 
@@ -196,18 +197,111 @@ class EmbeddingsSingleton:
         }
 
         if embeddings_name in embeddings_factory:
-            return embeddings_factory[embeddings_name](*args, **kwargs)
+            if args or kwargs:
+                logging.debug(
+                    "Dropping %d positional and %d keyword argument(s) for pinned "
+                    "embeddings model %s: its factory takes none.",
+                    len(args),
+                    len(kwargs),
+                    embeddings_name,
+                )
+            return embeddings_factory[embeddings_name]()
         else:
             return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
 
 
+def _azure_configured() -> bool:
+    """True when the Azure OpenAI deployment settings are all present."""
+    return bool(
+        settings.OPENAI_API_BASE
+        and settings.OPENAI_API_VERSION
+        and settings.AZURE_DEPLOYMENT_NAME
+    )
+
+
+def get_embeddings(
+    embeddings_name: Optional[str] = None, embeddings_key: Optional[str] = None
+):
+    """Resolve the configured embeddings instance. The single entry point.
+
+    Callers that reach for :meth:`EmbeddingsSingleton.get_instance` directly
+    reproduce neither the bundled local-model path (the Docker image ships
+    ``/app/models/all-mpnet-base-v2``) nor the OpenAI/Azure key handling: they
+    download a second copy of the model from the hub, and passing the key
+    positionally to the pinned HuggingFace names raises ``TypeError`` because
+    those factories take no arguments. Route every caller through here.
+
+    Args:
+        embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
+        embeddings_key: API key; defaults to ``settings.EMBEDDINGS_KEY``.
+
+    Returns:
+        The shared embeddings instance for the resolved model.
+    """
+    embeddings_name = embeddings_name or settings.EMBEDDINGS_NAME
+    embeddings_key = (
+        embeddings_key if embeddings_key is not None else settings.EMBEDDINGS_KEY
+    )
+
+    # Check for remote embeddings first
+    if settings.EMBEDDINGS_BASE_URL:
+        logging.info(
+            f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}"
+        )
+        return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
+
+    if embeddings_name == "openai_text-embedding-ada-002":
+        if _azure_configured():
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
+            )
+        else:
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name, openai_api_key=embeddings_key
+            )
+    elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
+        possible_paths = [
+            "/app/models/all-mpnet-base-v2",  # Docker absolute path
+            "./models/all-mpnet-base-v2",  # Relative path
+        ]
+        local_model_path = None
+        for path in possible_paths:
+            if os.path.exists(path):
+                local_model_path = path
+                logging.info(f"Found local model at path: {path}")
+                break
+            else:
+                logging.info(f"Path does not exist: {path}")
+        if local_model_path:
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                local_model_path,
+            )
+        else:
+            logging.warning(
+                f"Local model not found in any of the paths: {possible_paths}. Falling back to HuggingFace download."
+            )
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name,
+            )
+    else:
+        embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
+    return embedding_instance
+
+
 class BaseVectorStore(ABC):
     def __init__(self):
         pass
 
     @abstractmethod
     def search(self, *args, **kwargs):
-        """Search for similar documents/chunks in the vectorstore"""
+        """Search for similar documents/chunks in the vectorstore.
+
+        Implementations accept an optional ``query_vector`` kwarg: the query
+        already embedded by the caller, so a multi-source retrieval embeds once
+        instead of once per store. A store that cannot use it must still swallow
+        the kwarg (every signature here ends in ``**kwargs``) and embed the
+        question itself.
+        """
         pass
 
     def keyword_search(self, question, k=10):
@@ -279,53 +373,9 @@ class BaseVectorStore(ABC):
         return deleted
 
     def is_azure_configured(self):
-        return (
-            settings.OPENAI_API_BASE
-            and settings.OPENAI_API_VERSION
-            and settings.AZURE_DEPLOYMENT_NAME
-        )
+        """Kept for compatibility; delegates to the module-level check."""
+        return _azure_configured()
 
     def _get_embeddings(self, embeddings_name, embeddings_key=None):
-        # Check for remote embeddings first
-        if settings.EMBEDDINGS_BASE_URL:
-            logging.info(
-                f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}"
-            )
-            return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
-
-        if embeddings_name == "openai_text-embedding-ada-002":
-            if self.is_azure_configured():
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
-                )
-            else:
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name, openai_api_key=embeddings_key
-                )
-        elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
-            possible_paths = [
-                "/app/models/all-mpnet-base-v2",  # Docker absolute path
-                "./models/all-mpnet-base-v2",  # Relative path
-            ]
-            local_model_path = None
-            for path in possible_paths:
-                if os.path.exists(path):
-                    local_model_path = path
-                    logging.info(f"Found local model at path: {path}")
-                    break
-                else:
-                    logging.info(f"Path does not exist: {path}")
-            if local_model_path:
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    local_model_path,
-                )
-            else:
-                logging.warning(
-                    f"Local model not found in any of the paths: {possible_paths}. Falling back to HuggingFace download."
-                )
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name,
-                )
-        else:
-            embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
-        return embedding_instance
+        """Resolve embeddings for this store; see :func:`get_embeddings`."""
+        return get_embeddings(embeddings_name, embeddings_key)

application/vectorstore/elasticsearch.py

@@ -73,9 +73,27 @@ class ElasticsearchStore(BaseVectorStore):
 
         return es_client
 
-    def search(self, question, k=2, index_name=settings.ELASTIC_INDEX, *args, **kwargs):
-        embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
-        vector = embeddings.embed_query(question)
+    def search(
+        self,
+        question,
+        k=2,
+        index_name=settings.ELASTIC_INDEX,
+        *args,
+        query_vector=None,
+        **kwargs,
+    ):
+        """Search by kNN + full text, fused with RRF.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        vector = query_vector
+        if vector is None:
+            embeddings = self._get_embeddings(
+                settings.EMBEDDINGS_NAME, self.embeddings_key
+            )
+            vector = embeddings.embed_query(question)
         knn = {
             "filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
             "field": "vector",

application/vectorstore/faiss.py

@@ -174,16 +174,28 @@ class FaissStore(BaseVectorStore):
         return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)]
 
     def search_with_scores(
-        self, question: str, k: int = 4, *args, **kwargs
+        self,
+        question: str,
+        k: int = 4,
+        *args,
+        query_vector: Optional[List[float]] = None,
+        **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Same search as :meth:`search`, pairing each hit with its L2 distance."""
+        """Same search as :meth:`search`, pairing each hit with its L2 distance.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
         # FAISS has no relevance-threshold knob; drop it so the per-source
         # score_threshold is safely ignored rather than crashing the forward.
         kwargs.pop("score_threshold", None)
         if self.index is None or self.index.ntotal == 0:
             return []
 
-        vector = np.array([self.embeddings.embed_query(question)], dtype=np.float32)
+        if query_vector is None:
+            query_vector = self.embeddings.embed_query(question)
+        vector = np.array([query_vector], dtype=np.float32)
         distances, rows = self.index.search(vector, min(k, self.index.ntotal))
 
         results = []

application/vectorstore/lancedb.py

@@ -80,10 +80,19 @@ class LanceDBVectorStore(BaseVectorStore):
         self.ensure_table_exists()
         self.docsearch.add(vectors)
 
-    def search(self, query: str, k: int = 2, *args, **kwargs):
-        """Search LanceDB for the top k most similar vectors."""
+    def search(self, query: str, k: int = 2, *args, query_vector=None, **kwargs):
+        """Search LanceDB for the top k most similar vectors.
+
+        Args:
+            query_vector: Precomputed embedding of ``query``; when given the
+                store skips embedding the query itself.
+        """
         self.ensure_table_exists()
-        query_embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key).embed_query(query)
+        query_embedding = query_vector
+        if query_embedding is None:
+            query_embedding = self._get_embeddings(
+                settings.EMBEDDINGS_NAME, self.embeddings_key
+            ).embed_query(query)
         results = self.docsearch.search(query_embedding).limit(k).to_list()
         return [(result["_distance"], result["text"], result["metadata"]) for result in results]
 

application/vectorstore/milvus.py

@@ -110,12 +110,20 @@ class MilvusStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: Optional[float] = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Search, pairing each hit with its cosine similarity."""
+        """Search, pairing each hit with its cosine similarity.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        if query_vector is None:
+            query_vector = self._embeddings.embed_query(question)
         results = self._client.search(
             collection_name=self._collection,
-            data=[self._embeddings.embed_query(question)],
+            data=[query_vector],
             filter=self._filter,
             limit=k,
             output_fields=["text", "metadata"],

application/vectorstore/mongodb.py

@@ -54,7 +54,9 @@ class MongoDBVectorStore(BaseVectorStore):
 
     score_kind = "cosine_similarity"
 
-    def search(self, question, k=2, *args, score_threshold=None, **kwargs):
+    def search(
+        self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
+    ):
         """Search via Atlas ``$vectorSearch``.
 
         Args:
@@ -62,21 +64,32 @@ class MongoDBVectorStore(BaseVectorStore):
             k: Maximum number of results.
             score_threshold: Optional ``vectorSearchScore`` floor in ``[0, 1]``;
                 results scoring below it are dropped.
+            query_vector: Precomputed embedding of ``question``, so a caller
+                searching several sources embeds the query only once.
         """
         return [
             doc
             for doc, _ in self.search_with_scores(
-                question, k, *args, score_threshold=score_threshold, **kwargs
+                question,
+                k,
+                *args,
+                score_threshold=score_threshold,
+                query_vector=query_vector,
+                **kwargs,
             )
         ]
 
-    def search_with_scores(self, question, k=2, *args, score_threshold=None, **kwargs):
+    def search_with_scores(
+        self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
+    ):
         """Same search as :meth:`search`, pairing each hit with its score.
 
         The score is Atlas' ``vectorSearchScore`` — the same quantity
-        ``score_threshold`` is compared against.
+        ``score_threshold`` is compared against. ``query_vector`` skips the
+        per-store query embedding when the caller already has one.
         """
-        query_vector = self._embedding.embed_query(question)
+        if query_vector is None:
+            query_vector = self._embedding.embed_query(question)
 
         pipeline = [
             {

application/vectorstore/pgconn.py

@@ -0,0 +1,112 @@
+"""Process-wide psycopg connection pools keyed by DSN.
+
+Every component that talks to the pgvector database — the vector store and the
+GraphRAG graph store — checks out of the *same* pool for a given DSN, so a
+request that touches N sources costs pooled checkouts instead of N fresh TCP +
+TLS + auth handshakes. Keeping the registry in its own module (rather than in
+``pgvector``) is what lets ``application.graphrag.store`` join it without
+importing the vector store.
+
+The boot-time schema hook deliberately does not use this: pools are built
+lazily on first checkout, i.e. after the fork in a Celery/gunicorn worker, so a
+preloading parent never hands a live socket to a child.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Dict
+
+# Seconds a caller waits for a free pooled connection before giving up.
+POOL_TIMEOUT_SECONDS = 30.0
+
+# dsn -> psycopg_pool.ConnectionPool, one per process.
+_POOLS: Dict[str, Any] = {}
+_POOLS_LOCK = threading.Lock()
+
+
+def configure_pooled_connection(conn) -> None:
+    """Register pgvector's type adapters once per physical pooled connection.
+
+    Tolerates a database where ``CREATE EXTENSION vector`` has not run yet:
+    ``register_vector`` looks the type up in the catalog and raises when it is
+    absent, which would leave the pool unable to open the very connection the
+    schema bootstrap needs. The write path re-registers once the type exists.
+    """
+    from pgvector.psycopg import register_vector
+
+    try:
+        register_vector(conn)
+    except Exception as e:
+        logging.debug("pgvector types not registered yet: %s", e)
+
+
+def pool_for(dsn: str, max_size: int):
+    """Return this process's connection pool for ``dsn``, creating it lazily.
+
+    The pool is keyed by DSN alone and is built on first checkout — i.e. after
+    the fork in a Celery/gunicorn worker, so a preloading parent never hands a
+    live socket to a child. The boot-time schema hook must therefore use a
+    direct connection rather than this pool.
+
+    Args:
+        dsn: libpq connection string; the pool's only identity.
+        max_size: Upper bound on physical connections, used on creation only.
+
+    Returns:
+        psycopg_pool.ConnectionPool: The shared pool for ``dsn``.
+    """
+    pool = _POOLS.get(dsn)
+    if pool is not None:
+        return pool
+    with _POOLS_LOCK:
+        pool = _POOLS.get(dsn)
+        if pool is None:
+            from psycopg_pool import ConnectionPool
+
+            pool = ConnectionPool(
+                conninfo=dsn,
+                min_size=1,
+                max_size=max_size,
+                open=True,
+                timeout=POOL_TIMEOUT_SECONDS,
+                name="docsgpt-pgvector",
+                configure=configure_pooled_connection,
+                check=ConnectionPool.check_connection,
+            )
+            _POOLS[dsn] = pool
+    return pool
+
+
+def release(dsn: str, conn, pooled: bool) -> None:
+    """Hand ``conn`` back: to ``dsn``'s pool when pooled, else close it.
+
+    A pooled connection is rolled back first when it is still in a transaction,
+    so the next borrower gets a clean session. Never raises — releasing is
+    always best-effort cleanup on a path that is often a destructor.
+
+    Args:
+        dsn: Connection string the connection was opened for.
+        conn: The psycopg connection being released; ``None`` is a no-op.
+        pooled: Whether the connection came from :func:`pool_for`.
+    """
+    if conn is None:
+        return
+    try:
+        if pooled:
+            try:
+                if conn.info.transaction_status.name != "IDLE":
+                    conn.rollback()
+            except Exception:
+                # Connection already broken; the pool discards it.
+                pass
+            pool = _POOLS.get(dsn)
+            if pool is not None:
+                pool.putconn(conn)
+            elif not conn.closed:
+                conn.close()
+        elif not conn.closed:
+            conn.close()
+    except Exception as e:
+        logging.debug("Error releasing pooled connection: %s", e)

application/vectorstore/pgvector.py

@@ -6,12 +6,28 @@ from typing import List, Optional, Any, Dict
 from psycopg.types.json import Jsonb
 
 from application.core.settings import settings
+from application.vectorstore import pgconn
 from application.vectorstore.base import BaseVectorStore
 from application.vectorstore.document_class import Document
 
 # table name -> IVFFlat ``lists`` (None when the table has no such index)
 _IVFFLAT_LISTS_CACHE: Dict[str, Optional[int]] = {}
 
+DEFAULT_EMBEDDING_DIM = 768
+DEFAULT_POOL_MAX_SIZE = 8
+# Advisory-lock key shared with the boot hook so concurrent workers serialize DDL.
+SCHEMA_LOCK_KEY = "docsgpt:vectors:ddl"
+
+# The connection pools moved to ``application.vectorstore.pgconn`` so the graph
+# store can share them without importing this module. These names stay bound
+# here — same objects, not copies — because callers and tests reach for
+# ``pgvector._POOLS`` / patch ``pgvector._pool_for``.
+POOL_TIMEOUT_SECONDS = pgconn.POOL_TIMEOUT_SECONDS
+_POOLS = pgconn._POOLS
+_POOLS_LOCK = pgconn._POOLS_LOCK
+_configure_pooled_connection = pgconn.configure_pooled_connection
+_pool_for = pgconn.pool_for
+
 
 class PGVectorStore(BaseVectorStore):
     def __init__(
@@ -63,17 +79,64 @@ class PGVectorStore(BaseVectorStore):
         self._psycopg = psycopg
         self._register_vector = register_vector
         self._connection = None
-        self._ensure_table_exists()
+        self._pooled = False
+        self._pool_max_size = self._resolve_pool_max_size()
+        # No DDL here. The retriever builds one store per source per request, so
+        # construction must stay free: schema is owned at boot by
+        # ``ensure_vector_schema``, and the write path below re-checks it once as
+        # a safety net for processes that never ran the boot hook.
+        self._schema_ensured = False
+
+    @staticmethod
+    def _resolve_pool_max_size() -> int:
+        """Pool size from settings, defensively — 0 means one direct connection."""
+        value = getattr(settings, "PGVECTOR_POOL_MAX_SIZE", DEFAULT_POOL_MAX_SIZE)
+        if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+            return value
+        return DEFAULT_POOL_MAX_SIZE
 
     def _get_connection(self):
-        """Get or create database connection"""
+        """Get or create this store's connection, pooled unless pooling is off."""
         if self._connection is None or self._connection.closed:
-            self._connection = self._psycopg.connect(self._connection_string)
-            # Register pgvector types
-            self._register_vector(self._connection)
-            self._apply_ivfflat_probes(self._connection)
+            if self._pool_max_size > 0:
+                self._connection = _pool_for(
+                    self._connection_string, self._pool_max_size
+                ).getconn()
+                self._pooled = True
+            else:
+                self._connection = self._psycopg.connect(self._connection_string)
+                # Register pgvector types
+                self._register_pgvector_types(self._connection)
+                self._pooled = False
+            self._apply_probes_once(self._connection)
         return self._connection
 
+    def _register_pgvector_types(self, conn) -> None:
+        """Register pgvector's adapters, tolerating a not-yet-created extension."""
+        try:
+            self._register_vector(conn)
+        except Exception as e:
+            logging.debug("pgvector types not registered yet: %s", e)
+
+    def _apply_probes_once(self, conn) -> None:
+        """Set ``ivfflat.probes`` once per physical connection and table.
+
+        A pooled connection outlives the store that checked it out, and the
+        ``SET`` is session-level, so the probe lookup must not repeat on every
+        checkout. The marker lives on the connection object itself because that
+        is what the pool recycles.
+        """
+        if getattr(conn, "_docsgpt_probes_table", None) == self._table_name:
+            return
+        self._apply_ivfflat_probes(conn)
+        try:
+            # ``SET`` (not SET LOCAL) is undone by a rollback, so make it stick
+            # before any query runs on this connection.
+            conn.commit()
+            conn._docsgpt_probes_table = self._table_name
+        except Exception as e:  # never let session tuning break a query
+            logging.debug("Could not persist ivfflat.probes: %s", e)
+
     def _ivfflat_lists(self, conn) -> Optional[int]:
         """Return the ``lists`` value of this table's IVFFlat index, if any.
 
@@ -125,30 +188,48 @@ class PGVectorStore(BaseVectorStore):
         except Exception as e:  # older pgvector / no index — search still works
             logging.debug("Could not set ivfflat.probes: %s", e)
 
-    def _ensure_table_exists(self):
-        """Create table and enable pgvector extension if they don't exist"""
-        conn = self._get_connection()
+    @staticmethod
+    def create_schema(
+        conn,
+        *,
+        table_name: str = "documents",
+        vector_column: str = "embedding",
+        text_column: str = "text",
+        metadata_column: str = "metadata",
+        dimension: int = DEFAULT_EMBEDDING_DIM,
+    ) -> None:
+        """Create the extension, table and indexes on ``conn`` without committing.
+
+        Shared by the boot hook (``ensure_vector_schema``) and the store's own
+        write-path safety net; the caller owns the transaction.
+
+        Args:
+            conn: Open psycopg connection.
+            table_name: Documents table to create.
+            vector_column: Embedding column name.
+            text_column: Chunk-text column name.
+            metadata_column: JSONB metadata column name.
+            dimension: Width of the embedding vectors.
+        """
         cursor = conn.cursor()
-        
+
         try:
             # Enable pgvector extension
             cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
-            
-            embedding_dim = getattr(self._embedding, 'dimension', 768)
-            
+
             # Create table with vector column
             create_table_query = f"""
-            CREATE TABLE IF NOT EXISTS {self._table_name} (
+            CREATE TABLE IF NOT EXISTS {table_name} (
                 id SERIAL PRIMARY KEY,
-                {self._text_column} TEXT NOT NULL,
-                {self._vector_column} vector({embedding_dim}),
-                {self._metadata_column} JSONB,
+                {text_column} TEXT NOT NULL,
+                {vector_column} vector({dimension}),
+                {metadata_column} JSONB,
                 source_id TEXT NOT NULL,
                 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
             );
             """
             cursor.execute(create_table_query)
-            
+
             # NO vector index is created here, deliberately.
             #
             # This runs when the table is first created, i.e. on an EMPTY
@@ -166,25 +247,88 @@ class PGVectorStore(BaseVectorStore):
             # data, once a corpus is large enough to need one.
             # Create index for source_id filtering
             source_index_query = f"""
-            CREATE INDEX IF NOT EXISTS {self._table_name}_source_id_idx
-            ON {self._table_name} (source_id);
+            CREATE INDEX IF NOT EXISTS {table_name}_source_id_idx
+            ON {table_name} (source_id);
             """
             cursor.execute(source_index_query)
 
             # Functional GIN index backing keyword_search full-text queries.
             fts_index_query = f"""
-            CREATE INDEX IF NOT EXISTS {self._table_name}_text_fts_idx
-            ON {self._table_name} USING gin(to_tsvector('english', {self._text_column}));
+            CREATE INDEX IF NOT EXISTS {table_name}_text_fts_idx
+            ON {table_name} USING gin(to_tsvector('english', {text_column}));
             """
             cursor.execute(fts_index_query)
+        finally:
+            cursor.close()
 
+    @staticmethod
+    def table_dimension(
+        conn, table_name: str = "documents", vector_column: str = "embedding"
+    ) -> Optional[int]:
+        """Width declared by the table's vector column, or ``None`` if unknown.
+
+        ``None`` means the table is absent or the column is not a ``vector``.
+        """
+        cursor = conn.cursor()
+        try:
+            cursor.execute("SELECT to_regclass(%s);", (table_name,))
+            row = cursor.fetchone()
+            if not row or row[0] is None:
+                return None
+            cursor.execute(
+                "SELECT format_type(a.atttypid, a.atttypmod) FROM pg_attribute a "
+                "WHERE a.attrelid = %s::regclass AND a.attname = %s "
+                "AND NOT a.attisdropped",
+                (table_name, vector_column),
+            )
+            row = cursor.fetchone()
+            if not row or not row[0]:
+                return None
+            match = re.search(r"vector\((\d+)\)", str(row[0]))
+            return int(match.group(1)) if match else None
+        finally:
+            cursor.close()
+
+    def _ensure_table_exists(self) -> None:
+        """Create this store's schema under an advisory lock, then commit."""
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+                )
+            finally:
+                cursor.close()
+            self.create_schema(
+                conn,
+                table_name=self._table_name,
+                vector_column=self._vector_column,
+                text_column=self._text_column,
+                metadata_column=self._metadata_column,
+                dimension=getattr(self._embedding, "dimension", DEFAULT_EMBEDDING_DIM),
+            )
             conn.commit()
+            # The extension may have just been created; pick up its adapters so
+            # the insert that follows can bind a vector.
+            self._register_pgvector_types(conn)
         except Exception as e:
             conn.rollback()
             logging.error(f"Error creating table: {e}")
             raise
-        finally:
-            cursor.close()
+
+    def _ensure_schema_once(self) -> None:
+        """Create the schema on this instance's first write, at most once.
+
+        Readers never create schema — boot owns it. This is the safety net for a
+        process that never ran the boot hook (scripts, tests, the first ingest on
+        a brand-new deployment); a read against a missing table surfaces
+        psycopg's own error, like every other store.
+        """
+        if self._schema_ensured:
+            return
+        self._ensure_table_exists()
+        self._schema_ensured = True
 
     score_kind = "cosine_similarity"
 
@@ -194,6 +338,7 @@ class PGVectorStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: float = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Document]:
         """Search for similar documents using vector similarity.
@@ -204,11 +349,19 @@ class PGVectorStore(BaseVectorStore):
             score_threshold: Optional cosine-similarity floor in ``[0, 1]``.
                 Cosine distance = ``1 - similarity``; rows with similarity below
                 the threshold (distance above ``1 - threshold``) are dropped.
+            query_vector: Precomputed embedding of ``question``. Supplied by a
+                caller searching several sources with one query, so the query is
+                embedded once instead of once per store.
         """
         return [
             doc
             for doc, _ in self.search_with_scores(
-                question, k, *args, score_threshold=score_threshold, **kwargs
+                question,
+                k,
+                *args,
+                score_threshold=score_threshold,
+                query_vector=query_vector,
+                **kwargs,
             )
         ]
 
@@ -291,6 +444,7 @@ class PGVectorStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: float = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[tuple]:
         """Same search as :meth:`search`, pairing each hit with its similarity.
@@ -298,8 +452,13 @@ class PGVectorStore(BaseVectorStore):
         The score is the cosine similarity (``1 - cosine_distance``) — the exact
         quantity ``score_threshold`` is compared against, so a caller can read a
         result's score and pick a threshold from it directly.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
         """
-        query_vector = self._embedding.embed_query(question)
+        if query_vector is None:
+            query_vector = self._embedding.embed_query(question)
 
         conn = self._get_connection()
         cursor = conn.cursor()
@@ -334,6 +493,11 @@ class PGVectorStore(BaseVectorStore):
                     (Document(page_content=text, metadata=metadata), score)
                 )
 
+            # End the read transaction. On a persistent (pooled) connection an
+            # uncommitted SELECT leaves the backend "idle in transaction",
+            # pinning a snapshot and blocking VACUUM for as long as the store
+            # lives.
+            conn.commit()
             return documents
 
         except Exception as e:
@@ -381,10 +545,16 @@ class PGVectorStore(BaseVectorStore):
                 metadata = metadata or {}
                 documents.append(Document(page_content=text, metadata=metadata))
 
+            conn.commit()
             return documents
 
         except Exception as e:
             logging.error(f"Error in keyword search: {e}", exc_info=True)
+            try:
+                conn.rollback()
+            except Exception:
+                # Connection already gone; nothing left to roll back.
+                pass
             return []
         finally:
             cursor.close()
@@ -402,7 +572,8 @@ class PGVectorStore(BaseVectorStore):
 
         embeddings = self._embedding.embed_documents(texts)
         metadatas = metadatas or [{}] * len(texts)
-        
+
+        self._ensure_schema_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         
@@ -474,11 +645,17 @@ class PGVectorStore(BaseVectorStore):
                     "text": text,
                     "metadata": metadata or {}
                 })
-            
+
+            conn.commit()
             return chunks
-            
+
         except Exception as e:
             logging.error(f"Error getting chunks: {e}")
+            try:
+                conn.rollback()
+            except Exception:
+                # Connection already gone; nothing left to roll back.
+                pass
             return []
         finally:
             cursor.close()
@@ -495,7 +672,8 @@ class PGVectorStore(BaseVectorStore):
 
         if not embeddings:
             raise ValueError("Could not generate embedding for chunk")
-        
+
+        self._ensure_schema_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         
@@ -570,7 +748,22 @@ class PGVectorStore(BaseVectorStore):
         finally:
             cursor.close()
 
+    def close(self) -> None:
+        """Release this store's connection: back to the pool, or closed outright.
+
+        A pooled connection is rolled back first when it is still in a
+        transaction, so the next borrower gets a clean session.
+        """
+        conn = getattr(self, "_connection", None)
+        if conn is None:
+            return
+        self._connection = None
+        pgconn.release(self._connection_string, conn, getattr(self, "_pooled", False))
+
     def __del__(self):
-        """Close database connection when object is destroyed"""
-        if hasattr(self, '_connection') and self._connection and not self._connection.closed:
-            self._connection.close()
+        """Release the connection when the object is destroyed. Never raises."""
+        try:
+            self.close()
+        except Exception:
+            # Interpreter teardown can null out module globals; never raise here.
+            pass

application/vectorstore/qdrant.py

@@ -109,12 +109,20 @@ class QdrantStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: Optional[float] = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Search, pairing each hit with its cosine similarity."""
+        """Search, pairing each hit with its cosine similarity.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        if query_vector is None:
+            query_vector = self._embeddings.embed_query(question)
         hits = self._client.query_points(
             collection_name=self._collection,
-            query=self._embeddings.embed_query(question),
+            query=query_vector,
             query_filter=self._filter,
             limit=k,
             with_payload=True,