diff --git a/application/agents/base.py b/application/agents/base.py index 4b1e63bb..003f2c01 100644 --- a/application/agents/base.py +++ b/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() ): diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index 15a8b13e..bff32d45 100644 --- a/application/agents/tools/read_document.py +++ b/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: diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index 20b1c6e5..9633cba6 100644 --- a/application/agents/workflow_agent.py +++ b/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.""" diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index d19b2053..cb362be5 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/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 diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 2229f804..bc6f55fd 100644 --- a/application/api/user/tasks.py +++ b/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 diff --git a/application/app.py b/application/app.py index ecfa03cf..719b67dd 100644 --- a/application/app.py +++ b/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, ) diff --git a/application/cache.py b/application/cache.py index 4daee8d4..73efb53b 100644 --- a/application/cache.py +++ b/application/cache.py @@ -26,6 +26,73 @@ def _cache_default(value): return f"" 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) diff --git a/application/core/settings.py b/application/core/settings.py index b8135fe6..d89e06d3 100644 --- a/application/core/settings.py +++ b/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 diff --git a/application/graphrag/extraction.py b/application/graphrag/extraction.py index 6d76d79e..de5d8b42 100644 --- a/application/graphrag/extraction.py +++ b/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 diff --git a/application/graphrag/store.py b/application/graphrag/store.py index be0168fa..120035e8 100644 --- a/application/graphrag/store.py +++ b/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 diff --git a/application/llm/base.py b/application/llm/base.py index 395a2d71..8e1aabc6 100644 --- a/application/llm/base.py +++ b/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): diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 21bd1ed3..2ee72d8b 100644 --- a/application/llm/google_ai.py +++ b/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 = { diff --git a/application/llm/openai.py b/application/llm/openai.py index 8c2043e3..bad917c6 100644 --- a/application/llm/openai.py +++ b/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: diff --git a/application/parser/chunking_strategies.py b/application/parser/chunking_strategies.py index 933ec7f0..69efb8de 100644 --- a/application/parser/chunking_strategies.py +++ b/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()] diff --git a/application/parser/file/docling_parser.py b/application/parser/file/docling_parser.py index 9272ff8c..3f1658bb 100644 --- a/application/parser/file/docling_parser.py +++ b/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"") + + +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 diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py index f02ec558..21c855a3 100644 --- a/application/retriever/classic_rag.py +++ b/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 — diff --git a/application/retriever/fanout.py b/application/retriever/fanout.py new file mode 100644 index 00000000..81b0e1e5 --- /dev/null +++ b/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)) diff --git a/application/retriever/graph_rag.py b/application/retriever/graph_rag.py index 240d0f19..69eeff09 100644 --- a/application/retriever/graph_rag.py +++ b/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: diff --git a/application/retriever/hybrid_rag.py b/application/retriever/hybrid_rag.py index cf739253..f37fca1a 100644 --- a/application/retriever/hybrid_rag.py +++ b/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: diff --git a/application/services/search_service.py b/application/services/search_service.py index 6f2e5b4c..c739bf1c 100644 --- a/application/services/search_service.py +++ b/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 diff --git a/application/storage/db/bootstrap.py b/application/storage/db/bootstrap.py index 87470f3a..34042ef9 100644 --- a/application/storage/db/bootstrap.py +++ b/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. diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index 1a999f26..cd4e04c4 100644 --- a/application/vectorstore/base.py +++ b/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) diff --git a/application/vectorstore/elasticsearch.py b/application/vectorstore/elasticsearch.py index dfa0314f..fd9666fe 100644 --- a/application/vectorstore/elasticsearch.py +++ b/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", diff --git a/application/vectorstore/faiss.py b/application/vectorstore/faiss.py index 18eac04d..48ba6af9 100644 --- a/application/vectorstore/faiss.py +++ b/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 = [] diff --git a/application/vectorstore/lancedb.py b/application/vectorstore/lancedb.py index 25d62318..33ebcf3f 100644 --- a/application/vectorstore/lancedb.py +++ b/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] diff --git a/application/vectorstore/milvus.py b/application/vectorstore/milvus.py index 3050a877..14a63b5e 100644 --- a/application/vectorstore/milvus.py +++ b/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"], diff --git a/application/vectorstore/mongodb.py b/application/vectorstore/mongodb.py index 64e8591a..d0aeb51a 100644 --- a/application/vectorstore/mongodb.py +++ b/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 = [ { diff --git a/application/vectorstore/pgconn.py b/application/vectorstore/pgconn.py new file mode 100644 index 00000000..4538972a --- /dev/null +++ b/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) diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index 2b3640e9..c49e53fd 100644 --- a/application/vectorstore/pgvector.py +++ b/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() \ No newline at end of file + """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 \ No newline at end of file diff --git a/application/vectorstore/qdrant.py b/application/vectorstore/qdrant.py index fed977dc..a55ec25c 100644 --- a/application/vectorstore/qdrant.py +++ b/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, diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index c9a07052..ae9f3040 100644 --- a/tests/agents/test_base_agent.py +++ b/tests/agents/test_base_agent.py @@ -3,6 +3,13 @@ from unittest.mock import Mock, patch import pytest from application.agents.classic_agent import ClassicAgent +from application.llm.anthropic import AnthropicLLM +from application.llm.docsgpt_provider import DocsGPTAPILLM +from application.llm.google_ai import GoogleLLM +from application.llm.groq import GroqLLM +from application.llm.novita import NovitaLLM +from application.llm.open_router import OpenRouterLLM +from application.llm.openai import OpenAILLM @pytest.mark.unit @@ -620,24 +627,23 @@ class TestBaseAgentLLMGeneration: def test_llm_gen_with_json_schema( self, agent_base_params, - mock_llm, - mock_llm_creator, mock_llm_handler_creator, log_context, ): - mock_llm._supports_structured_output = Mock(return_value=True) - mock_llm.prepare_structured_output_format = Mock( - return_value={"schema": "test"} - ) + # A real OpenAILLM: the format kwarg is chosen from the LLM class, + # so a bare Mock deliberately gets no structured output. + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + llm.gen_stream = Mock() agent_base_params["json_schema"] = {"type": "object"} agent_base_params["llm_name"] = "openai" + agent_base_params["llm"] = llm agent = ClassicAgent(**agent_base_params) messages = [{"role": "user", "content": "test"}] agent._llm_gen(messages, log_context) - call_args = mock_llm.gen_stream.call_args[1] + call_args = llm.gen_stream.call_args[1] assert "response_format" in call_args @@ -1163,24 +1169,21 @@ class TestLLMGenAdvanced: def test_llm_gen_google_structured_output( self, agent_base_params, - mock_llm, - mock_llm_creator, mock_llm_handler_creator, log_context, ): - mock_llm._supports_structured_output = Mock(return_value=True) - mock_llm.prepare_structured_output_format = Mock( - return_value={"schema": "test"} - ) + llm = GoogleLLM(api_key="g-test", user_api_key=None) + llm.gen_stream = Mock() agent_base_params["json_schema"] = {"type": "object"} agent_base_params["llm_name"] = "google" + agent_base_params["llm"] = llm agent = ClassicAgent(**agent_base_params) messages = [{"role": "user", "content": "test"}] agent._llm_gen(messages, log_context) - call_kwargs = mock_llm.gen_stream.call_args[1] + call_kwargs = llm.gen_stream.call_args[1] assert "response_schema" in call_kwargs def test_llm_gen_no_tools_when_unsupported( @@ -1529,3 +1532,295 @@ class TestBaseAgentDocumentBudgetOrdering: assert "" in user, "documents must survive a long question" assert agent.retrieved_docs, "documents must not all be shed" assert len(user) < len(huge_question), "question must be truncated" + + +# --------------------------------------------------------------------------- +# _llm_gen — structured-output provider gating +# --------------------------------------------------------------------------- + + +SCHEMA = {"type": "object", "properties": {"answer": {"type": "string"}}} + + +def _build_agent_with_llm(agent_base_params, llm, llm_name, **overrides): + """Agent wired to a real LLM instance whose ``gen_stream`` is captured.""" + llm.gen_stream = Mock() + agent_base_params["llm"] = llm + agent_base_params["llm_name"] = llm_name + agent_base_params.update(overrides) + return ClassicAgent(**agent_base_params) + + +@pytest.mark.unit +class TestLLMGenStructuredOutputProviderGating: + """Structured output must key off the LLM class, not ``llm_name``. + + Catalog models loaded from MODELS_CONFIG_DIR dispatch under + ``openai_compatible`` (and ``groq`` / ``novita`` / ``openrouter`` / + ``docsgpt``), all of which instantiate an ``OpenAILLM`` subclass. A + string comparison against "openai" dropped guided decoding for every + one of them. + """ + + @pytest.mark.parametrize( + "llm_name,llm_factory", + [ + ("openai", lambda: OpenAILLM(api_key="sk-test", user_api_key=None)), + ( + "openai_compatible", + lambda: OpenAILLM( + api_key="sk-test", + user_api_key=None, + base_url="https://vllm.example.invalid/v1", + ), + ), + ("groq", lambda: GroqLLM(api_key="gsk-test", user_api_key=None)), + ("novita", lambda: NovitaLLM(api_key="nv-test", user_api_key=None)), + ( + "openrouter", + lambda: OpenRouterLLM(api_key="or-test", user_api_key=None), + ), + ("docsgpt", lambda: DocsGPTAPILLM(user_api_key=None)), + ], + ) + def test_openai_wire_providers_get_response_format( + self, + agent_base_params, + mock_llm_handler_creator, + log_context, + llm_name, + llm_factory, + ): + llm = llm_factory() + agent = _build_agent_with_llm( + agent_base_params, llm, llm_name, json_schema=SCHEMA + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert call_kwargs["response_format"]["type"] == "json_schema" + assert "response_schema" not in call_kwargs + + @pytest.mark.parametrize( + "llm_name", + ["openai", "openai_compatible", "groq", "novita", "openrouter", "docsgpt"], + ) + def test_openai_wire_providers_get_json_object_mode( + self, agent_base_params, mock_llm_handler_creator, log_context, llm_name + ): + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + agent = _build_agent_with_llm( + agent_base_params, llm, llm_name, json_object=True + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert call_kwargs["response_format"] == {"type": "json_object"} + + def test_google_llm_gets_response_schema( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + llm = GoogleLLM(api_key="g-test", user_api_key=None) + agent = _build_agent_with_llm( + agent_base_params, llm, "google", json_schema=SCHEMA + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert "response_schema" in call_kwargs + assert "response_format" not in call_kwargs + + def test_google_llm_gets_no_json_object_mode( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + """``{"type": "json_object"}`` is an OpenAI-wire shape only.""" + llm = GoogleLLM(api_key="g-test", user_api_key=None) + agent = _build_agent_with_llm( + agent_base_params, llm, "google", json_object=True + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert "response_format" not in call_kwargs + assert "response_schema" not in call_kwargs + + def test_non_openai_wire_llm_gets_neither( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + """Even a capability-claiming Anthropic LLM takes no format kwarg.""" + llm = AnthropicLLM(api_key="ant-test", user_api_key=None) + llm._supports_structured_output = Mock(return_value=True) + llm.prepare_structured_output_format = Mock(return_value={"schema": "x"}) + agent = _build_agent_with_llm( + agent_base_params, llm, "anthropic", json_schema=SCHEMA + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert "response_format" not in call_kwargs + assert "response_schema" not in call_kwargs + + def test_non_openai_wire_llm_gets_no_json_object_mode( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + llm = AnthropicLLM(api_key="ant-test", user_api_key=None) + llm._supports_structured_output = Mock(return_value=True) + agent = _build_agent_with_llm( + agent_base_params, llm, "anthropic", json_object=True + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + assert "response_format" not in llm.gen_stream.call_args[1] + + def test_llm_double_without_class_declaration_gets_neither( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + """The kwarg name is read off the LLM *class*. + + A ``Mock`` answers every attribute lookup with a fresh truthy Mock, so + reading it off the instance would smuggle a Mock-named kwarg into the + gen call. Its class declares nothing, so it stays unstructured. + """ + llm = Mock() + llm._supports_structured_output = Mock(return_value=True) + llm.prepare_structured_output_format = Mock(return_value={"schema": "x"}) + agent = _build_agent_with_llm( + agent_base_params, llm, "custom", json_schema=SCHEMA + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert "response_format" not in call_kwargs + assert "response_schema" not in call_kwargs + assert all(isinstance(key, str) for key in call_kwargs) + + def test_declaration_comes_from_the_llm_class_attribute( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + """``BaseLLM.structured_output_kwarg`` is the single source of truth — + the agent no longer isinstance-checks provider classes.""" + llm = AnthropicLLM(api_key="ant-test", user_api_key=None) + llm._supports_structured_output = Mock(return_value=True) + llm.prepare_structured_output_format = Mock(return_value={"schema": "x"}) + agent = _build_agent_with_llm( + agent_base_params, llm, "anthropic", json_schema=SCHEMA + ) + assert agent._structured_output_kwarg() is None + + with patch.object( + type(llm), "structured_output_kwarg", "response_format", create=True + ): + assert agent._structured_output_kwarg() == "response_format" + + def test_capability_flag_still_vetoes_openai_compatible( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + """A model whose registry caps deny structured output stays clean.""" + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + llm._supports_structured_output = Mock(return_value=False) + agent = _build_agent_with_llm( + agent_base_params, llm, "openai_compatible", json_schema=SCHEMA + ) + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + call_kwargs = llm.gen_stream.call_args[1] + assert "response_format" not in call_kwargs + assert "response_schema" not in call_kwargs + + +# --------------------------------------------------------------------------- +# _llm_gen — tools gate +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestLLMGenToolsGate: + """``_supports_tools`` must be *called*, not merely looked up. + + Testing the bound method for truthiness always passed, so the + agent-level gate was dead and only the provider-level drop kept an + unsupported ``tools`` payload off the wire. + """ + + def test_tools_dropped_when_llm_reports_unsupported( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + llm.gen_stream = Mock() + llm._supports_tools = Mock(return_value=False) + agent_base_params["llm"] = llm + agent = ClassicAgent(**agent_base_params) + agent.tools = [{"type": "function", "function": {"name": "test"}}] + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + assert "tools" not in llm.gen_stream.call_args[1] + + def test_tools_attached_when_llm_reports_supported( + self, agent_base_params, mock_llm_handler_creator, log_context + ): + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + llm.gen_stream = Mock() + llm._supports_tools = Mock(return_value=True) + agent_base_params["llm"] = llm + agent = ClassicAgent(**agent_base_params) + agent.tools = [{"type": "function", "function": {"name": "test"}}] + + agent._llm_gen([{"role": "user", "content": "test"}], log_context) + + assert llm.gen_stream.call_args[1]["tools"] == agent.tools + + def test_bool_attribute_double_is_honored( + self, + agent_base_params, + mock_llm, + mock_llm_creator, + mock_llm_handler_creator, + ): + """Test doubles set ``_supports_tools`` as a plain bool.""" + mock_llm._supports_tools = True + agent = ClassicAgent(**agent_base_params) + agent.tools = [{"type": "function", "function": {"name": "test"}}] + + agent._llm_gen([{"role": "user", "content": "test"}]) + + assert mock_llm.gen_stream.call_args[1]["tools"] == agent.tools + + def test_missing_capability_check_drops_tools( + self, agent_base_params, mock_llm_handler_creator + ): + llm = Mock(spec=["gen_stream", "model_id"]) + llm.gen_stream = Mock() + llm.model_id = "gpt-4" + agent_base_params["llm"] = llm + agent = ClassicAgent(**agent_base_params) + agent.tools = [{"type": "function", "function": {"name": "test"}}] + + agent._llm_gen([{"role": "user", "content": "test"}]) + + assert "tools" not in llm.gen_stream.call_args[1] + + def test_unimplemented_capability_check_still_sends_tools( + self, agent_base_params, mock_llm_handler_creator + ): + """``BaseLLM._supports_tools`` raises; the provider drops them.""" + llm = OpenAILLM(api_key="sk-test", user_api_key=None) + llm.gen_stream = Mock() + llm._supports_tools = Mock( + side_effect=NotImplementedError("Subclass must implement") + ) + agent_base_params["llm"] = llm + agent = ClassicAgent(**agent_base_params) + agent.tools = [{"type": "function", "function": {"name": "test"}}] + + agent._llm_gen([{"role": "user", "content": "test"}]) + + assert llm.gen_stream.call_args[1]["tools"] == agent.tools diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index 0c1f5918..6d221dcd 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -592,3 +592,218 @@ def test_node_extract_cap_bounds_parse_attempts_even_when_every_parse_times_out( assert parse_calls["n"] == 2 # not 4 -- failed parses still consume cap budget notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID] assert len(notes) == 1 + + +# --------------------------------------------------------------------------- +# Reuse of the text ``store_attachment`` already extracted (no double OCR) +# --------------------------------------------------------------------------- + +DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" +_OK_EXTRACTION = { + "status": "ok", + "parser": "DoclingParser", + "truncated": False, + "original_tokens": 12, + "stored_tokens": 12, +} + + +def _with_extraction(attachment: dict, content, extraction=_OK_EXTRACTION) -> dict: + """Add the ``content`` + ``metadata.extraction`` shape ``store_attachment`` persists.""" + attachment["content"] = content + attachment["metadata"] = {} if extraction is None else {"extraction": dict(extraction)} + return attachment + + +def _forbid_parse(monkeypatch) -> None: + """Make any enqueue of the parsing worker a hard failure.""" + import application.api.user.tasks as tasks + + def _apply_async(*a, **k): + raise AssertionError("the document must not be re-parsed") + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + +def _count_parses(monkeypatch, calls: dict) -> None: + """Stub the parsing worker so each blocking parse succeeds and is counted.""" + import application.api.user.tasks as tasks + + class _R: + def get(self, timeout=None, disable_sync_subtasks=True): + return {"status": "ok", "content": "PARSED"} + + def _apply_async(args=None, queue=None, **kw): + calls["n"] = calls.get("n", 0) + 1 + return _R() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + +def _bridged_engine(pg_engine, tmp_path, monkeypatch, attachments): + """Run the bridge for ``attachments`` and return (engine, refs).""" + wf_id = _make_workflow(pg_engine) + agent = _agent(wf_id, attachments) + _patch_engine(monkeypatch) + list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + return engine, engine.captured_inputs["input_documents"] + + +def test_preextracted_attachment_text_is_reused_without_reparsing( + pg_engine, tmp_path, monkeypatch +): + """An attachment already parsed at upload time is inlined verbatim; the worker is never called.""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + att = _with_extraction( + _stage_attachment(storage, b"PK-docx-bytes", "scan.docx", DOCX_MIME), + "PRE-EXTRACTED OCR TEXT", + ) + engine, refs = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att]) + + # The text rides on the engine only -- never in the persisted run state/refs. + aid = refs[0]["artifact_id"] + assert engine.preextracted_text == {aid: "PRE-EXTRACTED OCR TEXT"} + assert set(refs[0]) == {"artifact_id", "ref", "filename", "mime_type"} + assert "content" not in str(engine.state["input_documents"]) + + _forbid_parse(monkeypatch) + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + + assert out == [{"id": aid, "mime_type": "text/plain", "content": "PRE-EXTRACTED OCR TEXT"}] + + +def test_preextracted_text_is_bounded_head_and_tail(pg_engine, tmp_path, monkeypatch): + """Reused text goes through the same head+tail window as the inline-text path.""" + from application.parser.document_reader import _TEXT_MAX_BYTES + + storage = _wire(pg_engine, tmp_path, monkeypatch) + big = "A" * (_TEXT_MAX_BYTES * 3) + att = _with_extraction( + _stage_attachment(storage, b"PK-docx-bytes", "huge.docx", DOCX_MIME), big + ) + engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att]) + + _forbid_parse(monkeypatch) + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + + content = out[0]["content"] + assert "...[truncated" in content + assert len(content.encode("utf-8")) < _TEXT_MAX_BYTES + 200 + + +def test_truncated_attachment_text_falls_through_to_parse(pg_engine, tmp_path, monkeypatch): + """A row cut at ATTACHMENT_MAX_TOKENS lost its tail, so the node re-parses instead.""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + att = _with_extraction( + _stage_attachment(storage, b"PK-docx-bytes", "long.docx", DOCX_MIME), + "HEAD ONLY", + extraction={**_OK_EXTRACTION, "truncated": True}, + ) + engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att]) + + assert engine.preextracted_text == {} + calls: dict = {} + _count_parses(monkeypatch, calls) + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + + assert calls["n"] == 1 + assert out[0]["content"] == "PARSED" + + +def test_failed_or_empty_attachment_text_falls_through_to_parse( + pg_engine, tmp_path, monkeypatch +): + """An empty/whitespace ``content`` (or a failed extraction) is never reused.""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + empty = _with_extraction( + _stage_attachment(storage, b"PK-a", "empty.docx", DOCX_MIME), " " + ) + failed = _with_extraction( + _stage_attachment(storage, b"PK-b", "failed.docx", DOCX_MIME), + None, + extraction={"status": "failed", "parser": None, "truncated": False, "error": "boom"}, + ) + engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [empty, failed]) + + assert engine.preextracted_text == {} + calls: dict = {} + _count_parses(monkeypatch, calls) + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + + assert calls["n"] == 2 + assert [a["content"] for a in out] == ["PARSED", "PARSED"] + + +def test_legacy_row_without_extraction_metadata_is_trusted(pg_engine, tmp_path, monkeypatch): + """A pre-``extraction`` row with non-empty content is reused (nothing records otherwise).""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + att = _stage_attachment(storage, b"PK-docx-bytes", "legacy.docx", DOCX_MIME) + att["content"] = "LEGACY TEXT" # no ``metadata`` key at all + engine, refs = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att]) + + assert engine.preextracted_text == {refs[0]["artifact_id"]: "LEGACY TEXT"} + _forbid_parse(monkeypatch) + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + assert out[0]["content"] == "LEGACY TEXT" + + +def test_preextracted_docs_do_not_consume_the_extract_cap(pg_engine, tmp_path, monkeypatch): + """Reused text costs no blocking parse, so it must not spend the per-node parse budget.""" + from application.core.settings import settings + + storage = _wire(pg_engine, tmp_path, monkeypatch) + attachments = [ + _with_extraction( + _stage_attachment(storage, b"PK-0", "pre0.docx", DOCX_MIME), "REUSED 0" + ), + _with_extraction( + _stage_attachment(storage, b"PK-1", "pre1.docx", DOCX_MIME), "REUSED 1" + ), + _stage_attachment(storage, b"PK-2", "raw2.docx", DOCX_MIME), + _stage_attachment(storage, b"PK-3", "raw3.docx", DOCX_MIME), + ] + engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, attachments) + + monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False) + calls: dict = {} + _count_parses(monkeypatch, calls) + + out = engine._materialize_node_attachments( + AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[] + ) + + # All four documents made it in: the two reused ones did not spend cap budget, + # so both remaining documents still fit the cap of 2 blocking parses. + assert [a["content"] for a in out] == ["REUSED 0", "REUSED 1", "PARSED", "PARSED"] + assert calls["n"] == 2 + assert not [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID] + + +@pytest.mark.parametrize( + "attachment,expected", + [ + ({"content": "text", "metadata": {"extraction": {"status": "ok", "truncated": False}}}, "text"), + ({"content": "text", "metadata": {"extraction": {"status": "ok", "truncated": True}}}, None), + ({"content": "text", "metadata": {"extraction": {"status": "failed"}}}, None), + ({"content": "text", "metadata": {}}, "text"), + ({"content": "text"}, "text"), + ({"content": "text", "metadata": {"extraction": "nope"}}, None), + ({"content": "", "metadata": {"extraction": {"status": "ok", "truncated": False}}}, None), + ({"content": None}, None), + ({}, None), + ], +) +def test_usable_attachment_text_matrix(attachment, expected): + """Only whole, successfully-extracted, non-empty text is eligible for reuse.""" + assert WorkflowAgent._usable_attachment_text(attachment) == expected diff --git a/tests/agents/test_workflow_node_attachments.py b/tests/agents/test_workflow_node_attachments.py index e9412f88..eebfebe0 100644 --- a/tests/agents/test_workflow_node_attachments.py +++ b/tests/agents/test_workflow_node_attachments.py @@ -23,6 +23,7 @@ from application.agents.workflows.schemas import ( WorkflowNode, ) from application.agents.workflows.workflow_engine import WorkflowEngine +from application.core.settings import settings from application.storage.db.repositories.artifacts import ArtifactsRepository RUN_ID = "11111111-1111-1111-1111-111111111111" @@ -231,7 +232,7 @@ def test_extract_non_text_uses_parse_worker(monkeypatch): _patch_repo(monkeypatch, {aid: rec}) eng = _engine(monkeypatch) monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED MD" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED MD" ) cfg = _node_config(input_documents=[aid], file_passing="extract") @@ -295,7 +296,7 @@ def test_native_file_cap_bounds_native_then_extracts(monkeypatch): ) # extract of a non-text image routes through the parsing worker; stub it so it returns text. monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "fallback text" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "fallback text" ) cfg = _node_config(input_documents=ids, file_passing="auto") @@ -383,6 +384,93 @@ def test_duplicate_refs_attach_once(monkeypatch): assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}] +# --------------------------------------------------------------------------- +# Size-scaled parse window +# --------------------------------------------------------------------------- + + +def _capture_parse(monkeypatch) -> dict: + """Patch parse_document.apply_async so nothing touches a broker; capture the call.""" + import application.api.user.tasks as tasks + + captured: Dict[str, Any] = {} + + class _R: + def get(self, timeout=None, disable_sync_subtasks=True): + captured["timeout"] = timeout + return {"status": "ok", "content": "PARSED"} + + def _apply_async(args=None, queue=None, **kw): + captured["kwargs"] = kw + return _R() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + return captured + + +@pytest.mark.unit +def test_parse_window_scales_with_the_document_size(monkeypatch): + """The version's ``size`` reaches the parse: a longer await AND matching Celery limits.""" + from application.api.user.tasks import parse_task_time_limits, parse_timeout_for_size + + size = 8 * 1024 * 1024 + aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=size) + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + captured = _capture_parse(monkeypatch) + + out = eng._materialize_node_attachments( + _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES + ) + + expected = parse_timeout_for_size(size) + assert out == [{"id": aid, "mime_type": "text/plain", "content": "PARSED"}] + assert captured["timeout"] == expected + # Without the per-call limits the worker self-terminates at the base timeout, + # so the longer await would buy nothing. + assert captured["kwargs"] == parse_task_time_limits(expected) + assert captured["kwargs"]["soft_time_limit"] > settings.DOCUMENT_PARSE_TIMEOUT + + +@pytest.mark.unit +def test_parse_window_falls_back_to_the_base_timeout_without_a_size(monkeypatch): + """A NULL/missing version size keeps the configured base window as the floor.""" + aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=None) + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + captured = _capture_parse(monkeypatch) + + eng._materialize_node_attachments( + _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES + ) + + assert captured["timeout"] == float(settings.DOCUMENT_PARSE_TIMEOUT) + assert captured["kwargs"]["soft_time_limit"] == settings.DOCUMENT_PARSE_TIMEOUT + + +@pytest.mark.unit +def test_preextracted_text_skips_the_parse_worker(monkeypatch): + """Text already extracted at upload time is inlined; the parsing worker is never called.""" + aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + eng.preextracted_text[aid] = "ALREADY OCRED" + + import application.api.user.tasks as tasks + + monkeypatch.setattr( + tasks.parse_document, + "apply_async", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not re-parse")), + ) + + out = eng._materialize_node_attachments( + _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES + ) + + assert out == [{"id": aid, "mime_type": "text/plain", "content": "ALREADY OCRED"}] + + # --------------------------------------------------------------------------- # _execute_agent_node wiring # --------------------------------------------------------------------------- @@ -484,7 +572,7 @@ def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch): _patch_repo(monkeypatch, {aid: rec}) _patch_capabilities(monkeypatch) monkeypatch.setattr( - WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED" + WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED" ) eng = _engine(monkeypatch) diff --git a/tests/agents/tools/test_read_document_tool.py b/tests/agents/tools/test_read_document_tool.py index 31192687..ed9808dc 100644 --- a/tests/agents/tools/test_read_document_tool.py +++ b/tests/agents/tools/test_read_document_tool.py @@ -8,6 +8,10 @@ validation, and the surfaced action params. from __future__ import annotations +import logging +import signal +import threading +import time import uuid from typing import Any, Dict, Optional @@ -22,7 +26,10 @@ _ART_ID = str(uuid.uuid4()) # --------------------------------------------------------------------------- # Run-scoped input resolution (mocks the repo gate) # --------------------------------------------------------------------------- -def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str]): +def _stub_repo( + monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str], + size: Optional[int] = None, +): class _Repo: def __init__(self, conn): pass @@ -45,6 +52,9 @@ def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[s return None return {"id": artifact_id, "current_version": 1, "title": "statement.pdf"} + def get_version(self, artifact_id, version): + return {"storage_path": "s/1", "mime_type": "application/pdf", "size": size} + class _Conn: def __enter__(self): return object() @@ -83,6 +93,7 @@ def _patch_task(monkeypatch, *, payload=None, exc=None): def _apply_async(args=None, queue=None, **kw): captured["args"] = args captured["queue"] = queue + captured["limits"] = kw result = _FakeAsyncResult(payload=payload, exc=exc) captured["result"] = result return result @@ -208,6 +219,42 @@ def test_artifact_ref_sets_last_artifact_id(monkeypatch): assert tool.get_artifact_id("read_document") == "new-art" +# --------------------------------------------------------------------------- +# Size-scaled parse window +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_parse_window_scales_with_the_input_size(monkeypatch): + """A large input widens the await AND the task's per-call Celery time limits.""" + from application.api.user.tasks import parse_task_time_limits, parse_timeout_for_size + + size = 8 * 1024 * 1024 + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None, size=size) + captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "x"}) + + out = _tool().execute_action("read_document", input=_ART_ID, persist=False) + + expected = parse_timeout_for_size(size) + assert out["status"] == "ok" + assert captured["result"].get_kwargs["timeout"] == expected + # The task's import-time soft limit is the BASE timeout, so the per-call limits + # must be raised too or the worker kills the parse the caller is still awaiting. + assert captured["limits"] == parse_task_time_limits(expected) + assert captured["limits"]["soft_time_limit"] > rd.settings.DOCUMENT_PARSE_TIMEOUT + + +@pytest.mark.unit +def test_parse_window_falls_back_to_the_base_timeout_without_a_size(monkeypatch): + """An unknown input size keeps the configured base timeout as the floor.""" + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None, size=None) + captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "x"}) + + _tool().execute_action("read_document", input=_ART_ID, persist=False) + + base = float(rd.settings.DOCUMENT_PARSE_TIMEOUT) + assert captured["result"].get_kwargs["timeout"] == base + assert captured["limits"]["soft_time_limit"] == int(base) + + # --------------------------------------------------------------------------- # Cross-tenant: rejected BEFORE enqueue # --------------------------------------------------------------------------- @@ -354,3 +401,216 @@ def test_dispatch_enqueues_when_not_in_worker(monkeypatch): assert out["status"] == "ok" and out["content"] == "queued" assert captured["args"][0] == _ART_ID assert captured["args"][1] == {"conversation_id": "conv-1"} + + +# --------------------------------------------------------------------------- +# The inline (in-worker) parse is time-bounded too +# +# The inline branch has no Celery time limit of its own, so an unbounded parse +# would pin the agent's worker slot until the OUTER task's limit (if any) kills +# the whole run. It must degrade on the SAME window the dispatch branch awaits. +# --------------------------------------------------------------------------- +_TIMED_OUT = "document parsing timed out after" + + +def _inline(monkeypatch, run_parse, *, timeout=0.2) -> ReadDocumentTool: + """Drive the inline branch (current_task truthy) with a patched parse window.""" + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) + monkeypatch.setattr(rd, "current_task", object()) + + import application.api.user.tasks as tasks + monkeypatch.setattr( + tasks.parse_document, "apply_async", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue inside a worker")), + ) + monkeypatch.setattr(tasks, "parse_timeout_for_size", lambda size: timeout) + + import application.worker as worker + monkeypatch.setattr(worker, "run_parse_document", run_parse) + return _tool() + + +@pytest.mark.unit +def test_inline_parse_within_the_window_returns_its_result(monkeypatch): + seen: Dict[str, Any] = {} + + def _run(artifact_id, parent, user_id, options): + seen["thread"] = threading.current_thread() + seen["armed"] = signal.getitimer(signal.ITIMER_REAL) + time.sleep(0.05) + return {"status": "ok", "content": "inline", "truncated": False} + + before = signal.getsignal(signal.SIGALRM) + tool = _inline(monkeypatch, _run, timeout=0.5) + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + + assert out["status"] == "ok" and out["content"] == "inline" + # Signal strategy: the parse runs on the calling (main) thread under an armed timer, + assert seen["thread"] is threading.main_thread() + assert seen["armed"][0] > 0 + # and the handler + timer are restored even though the parse returned early. + assert signal.getsignal(signal.SIGALRM) is before + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + + +@pytest.mark.unit +def test_inline_parse_times_out_on_the_signal_path(monkeypatch, caplog): + seen: Dict[str, Any] = {} + + def _run(*args, **kwargs): + seen["thread"] = threading.current_thread() + time.sleep(0.5) + seen["finished"] = True + return {"status": "ok", "content": "too late"} + + before = signal.getsignal(signal.SIGALRM) + tool = _inline(monkeypatch, _run, timeout=0.2) + started = time.monotonic() + with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"): + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + elapsed = time.monotonic() - started + + # Same error shape the worker returns on its soft limit. + assert out == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."} + # SIGALRM actually INTERRUPTS the parse, freeing the slot (it does not merely abandon it). + assert elapsed < 0.45 + assert seen["thread"] is threading.main_thread() + assert "finished" not in seen + assert signal.getsignal(signal.SIGALRM) is before + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + assert any("timed out" in rec.getMessage() for rec in caplog.records if rec.levelname == "WARNING") + + +@pytest.mark.unit +def test_inline_parse_times_out_on_the_thread_path(monkeypatch, caplog): + """Off the main thread (gevent/threads pools, Windows) the timer is unusable.""" + seen: Dict[str, Any] = {} + + def _run(*args, **kwargs): + seen["thread"] = threading.current_thread() + time.sleep(0.5) + seen["finished"] = True + return {"status": "ok", "content": "too late"} + + tool = _inline(monkeypatch, _run, timeout=0.2) + main_handler = signal.getsignal(signal.SIGALRM) + box: Dict[str, Any] = {} + + def _call(): + box["out"] = tool.execute_action("read_document", input=_ART_ID, persist=False) + + with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"): + caller = threading.Thread(target=_call, name="fake-worker-pool-thread") + caller.start() + caller.join(2.0) + + assert not caller.is_alive() + assert box["out"] == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."} + # The parse ran in the helper executor, not on the caller's thread... + assert seen["thread"].name.startswith("read-document-inline") + assert "finished" not in seen + # ...and the orphan cannot be interrupted, which the warning must say. + assert any("keeps running" in rec.getMessage() for rec in caplog.records if rec.levelname == "WARNING") + # The main thread's signal state was never touched from the worker thread. + assert signal.getsignal(signal.SIGALRM) is main_handler + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + + +@pytest.mark.unit +def test_inline_parse_falls_back_to_a_thread_when_sigalrm_is_taken(monkeypatch): + """A foreign SIGALRM handler must not be clobbered; use the thread strategy instead.""" + seen: Dict[str, Any] = {} + + def _run(*args, **kwargs): + seen["thread"] = threading.current_thread() + return {"status": "ok", "content": "threaded", "truncated": False} + + def _foreign(signum, frame): # pragma: no cover - must never be invoked + raise AssertionError("foreign SIGALRM handler fired") + + tool = _inline(monkeypatch, _run, timeout=0.5) + previous = signal.signal(signal.SIGALRM, _foreign) + try: + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + still_installed = signal.getsignal(signal.SIGALRM) + finally: + signal.signal(signal.SIGALRM, previous) + + assert out["status"] == "ok" and out["content"] == "threaded" + assert seen["thread"].name.startswith("read-document-inline") + assert still_installed is _foreign + + +@pytest.mark.unit +def test_inline_signal_timer_does_not_fire_after_a_fast_parse(monkeypatch): + before = signal.getsignal(signal.SIGALRM) + tool = _inline(monkeypatch, lambda *a, **k: {"status": "ok", "content": "fast"}, timeout=0.2) + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + + assert out["status"] == "ok" and out["content"] == "fast" + # A leaked timer would raise inside this sleep, well past the window. + time.sleep(0.3) + assert signal.getsignal(signal.SIGALRM) is before + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + + +@pytest.mark.unit +def test_inline_parse_failure_still_degrades_to_error(monkeypatch): + def _boom(*args, **kwargs): + raise RuntimeError("docling blew up") + + tool = _inline(monkeypatch, _boom, timeout=0.5) + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + + assert out["status"] == "error" + assert "document parsing failed" in out["error"] and "docling blew up" in out["error"] + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + + +@pytest.mark.unit +@pytest.mark.parametrize("window", [0, None]) +def test_inline_parse_runs_unbounded_without_a_positive_window(monkeypatch, window): + seen: Dict[str, Any] = {} + + def _run(*args, **kwargs): + seen["thread"] = threading.current_thread() + seen["armed"] = signal.getitimer(signal.ITIMER_REAL) + seen["handler"] = signal.getsignal(signal.SIGALRM) + return {"status": "ok", "content": "unbounded", "truncated": False} + + before = signal.getsignal(signal.SIGALRM) + tool = _inline(monkeypatch, _run, timeout=window) + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + + assert out["status"] == "ok" and out["content"] == "unbounded" + # Called directly: no timer armed, no handler swapped, no helper thread. + assert seen["armed"] == (0.0, 0.0) + assert seen["handler"] is before + assert seen["thread"] is threading.main_thread() + + +@pytest.mark.unit +def test_inline_timeout_is_not_swallowed_by_the_parser_catch_all(monkeypatch, caplog): + """``parse_document_bytes`` wraps the parse in ``except Exception`` and returns an error dict. + + The interrupt must fly straight through that catch-all (and still run its ``finally`` + cleanup) so the caller reports the shared timed-out shape, not "parsing failed: ...". + """ + seen: Dict[str, Any] = {} + + def _run(*args, **kwargs): + try: + time.sleep(0.5) + return {"status": "ok", "content": "too late"} + except Exception as exc: # mirrors parser/document_reader.py's catch-all + return {"status": "error", "error": f"parsing failed: {type(exc).__name__}: {exc}"} + finally: + seen["cleaned_up"] = True + + tool = _inline(monkeypatch, _run, timeout=0.2) + with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"): + out = tool.execute_action("read_document", input=_ART_ID, persist=False) + + assert out == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."} + assert seen["cleaned_up"] is True + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index 0785a552..832d8183 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -830,3 +830,67 @@ def test_bare_worker_consumes_app_and_parsing_queues(): names = {queue.name for queue in celeryconfig.task_queues} assert "docsgpt" in names assert settings.DOCUMENT_PARSE_QUEUE in names + + +class TestParseTimeoutForSize: + """The awaited parse window is floored at the base timeout, scales with size, and is capped.""" + + @pytest.mark.unit + def test_unknown_size_uses_the_base_timeout(self): + from application.api.user.tasks import parse_timeout_for_size + from application.core.settings import settings + + base = float(settings.DOCUMENT_PARSE_TIMEOUT) + assert parse_timeout_for_size(None) == base + assert parse_timeout_for_size(0) == base + # A negative/garbage size can never shrink the window below the floor. + assert parse_timeout_for_size(-1) == base + assert parse_timeout_for_size("nope") == base + + @pytest.mark.unit + def test_window_grows_with_the_document_size(self, monkeypatch): + from application.api.user import tasks + from application.core.settings import settings + + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False) + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 60, raising=False) + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", 900, raising=False) + + assert tasks.parse_timeout_for_size(1024 * 1024) == 180.0 + assert tasks.parse_timeout_for_size(5 * 1024 * 1024) == 420.0 + # Fractional sizes scale proportionally, never rounded down to the floor. + assert tasks.parse_timeout_for_size(512 * 1024) == 150.0 + + @pytest.mark.unit + def test_window_is_capped(self, monkeypatch): + from application.api.user import tasks + from application.core.settings import settings + + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False) + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 60, raising=False) + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", 900, raising=False) + + # A huge document cannot pin a parsing-worker slot indefinitely. + assert tasks.parse_timeout_for_size(500 * 1024 * 1024) == 900.0 + + @pytest.mark.unit + def test_scaling_disabled_by_zero_per_mb(self, monkeypatch): + from application.api.user import tasks + from application.core.settings import settings + + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False) + monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 0, raising=False) + + assert tasks.parse_timeout_for_size(50 * 1024 * 1024) == 120.0 + + @pytest.mark.unit + def test_task_time_limits_track_the_awaited_window(self): + from application.api.user.tasks import parse_document, parse_task_time_limits + + limits = parse_task_time_limits(420.0) + assert limits == {"soft_time_limit": 420, "time_limit": 450} + # Same grace as the import-time binding, so per-call limits stay comparable. + grace = parse_document.time_limit - parse_document.soft_time_limit + assert limits["time_limit"] - limits["soft_time_limit"] == grace + # Never zero/negative, whatever the caller computed. + assert parse_task_time_limits(0)["soft_time_limit"] == 1 diff --git a/tests/conftest.py b/tests/conftest.py index 74bd3ff8..5682e183 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -43,6 +43,7 @@ import os # opt back in by setting the env var explicitly. os.environ.setdefault("AUTO_MIGRATE", "false") os.environ.setdefault("AUTO_CREATE_DB", "false") +os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false") import subprocess import sys diff --git a/tests/graphrag/test_extraction.py b/tests/graphrag/test_extraction.py index 89c7362c..9abd87c1 100644 --- a/tests/graphrag/test_extraction.py +++ b/tests/graphrag/test_extraction.py @@ -2,8 +2,8 @@ The LLM and the embeddings model are mocked in every test so the suite makes no real model or network calls. A live ``GraphStore`` is exercised against the -pgvector DB with a unique temp ``source_id`` and torn down via -``delete_by_source``; if no DB is reachable the live tests skip. +ephemeral pytest-postgresql cluster (never the operator's dev DB) with a unique +temp ``source_id``; if pgvector is unavailable there the live tests skip. """ from __future__ import annotations @@ -16,54 +16,55 @@ import pytest import application.graphrag.extraction as extraction_module from application.graphrag.store import GraphStore from application.storage.db.source_config import SourceConfig +from application.vectorstore import pgconn extract_graph_for_source = extraction_module.extract_graph_for_source TEST_EMBEDDING_DIM = 8 -def _resolve_connection_string(): - from application.core.settings import settings - - conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None) - if not conn and getattr(settings, "POSTGRES_URI", None): - from application.core.db_uri import normalize_pgvector_connection_string - - conn = normalize_pgvector_connection_string(settings.POSTGRES_URI) - return conn +@pytest.fixture(autouse=True) +def _close_pools(): + """Never leak a pool into another test; an ephemeral DSN dies with its DB.""" + yield + for dsn, pool in list(pgconn._POOLS.items()): + try: + pool.close() + except Exception: + pass + pgconn._POOLS.pop(dsn, None) -_GRAPH_TABLES = ( - "graph_node_chunks", - "graph_edges", - "graph_nodes", - "graph_ingest_progress", -) +def _ephemeral_dsn(info) -> str: + """libpq DSN for the ephemeral pytest-postgresql database.""" + password = f":{info.password}" if info.password else "" + return ( + f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}" + ) -def _drop_graph_tables(conn_string): - import psycopg +def _live_store(monkeypatch, info): + """Graph store on a fresh ephemeral database, schema created up front. - with psycopg.connect(conn_string) as conn: - with conn.cursor() as cursor: - cursor.execute( - f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;" - ) - conn.commit() - - -def _live_store(monkeypatch): + Construction runs no DDL any more (boot owns the schema), so the tables are + created explicitly here — what ``ensure_vector_schema`` does in production. + """ monkeypatch.setattr( GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM ) - conn = _resolve_connection_string() - if not conn: - pytest.skip("No pgvector connection string configured") + dsn = _ephemeral_dsn(info) + # The pipeline builds its own GraphStore() from settings, so point those at + # the ephemeral cluster too — never at the operator's configured DB. + from application.core import settings as settings_module + + monkeypatch.setattr( + settings_module.settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False + ) + store = GraphStore(connection_string=dsn) try: - _drop_graph_tables(conn) - store = GraphStore(connection_string=conn) + store._ensure_tables() except Exception as exc: - pytest.skip(f"pgvector DB not reachable: {exc}") + pytest.skip(f"pgvector extension unavailable: {exc}") return store @@ -102,6 +103,11 @@ class _StubEmbedding: @pytest.fixture def stub_embedding(monkeypatch): + from application.core.settings import settings + + # The resolver short-circuits to the remote API when this is configured, + # which would bypass the stub on a dev machine that sets it. + monkeypatch.setattr(settings, "EMBEDDINGS_BASE_URL", None) embedding = _StubEmbedding() monkeypatch.setattr( extraction_module.EmbeddingsSingleton, @@ -135,10 +141,10 @@ def _extraction_json(entities, relationships): @pytest.mark.integration class TestExtractionLive: @pytest.fixture - def store(self, monkeypatch): - store = _live_store(monkeypatch) + def store(self, monkeypatch, postgresql): + store = _live_store(monkeypatch, postgresql.info) yield store - _drop_graph_tables(store._connection_string) + store.close() @pytest.fixture def source_id(self): @@ -419,3 +425,37 @@ class TestParsing: assert extraction_module._chunk_id({"chunk_id": "abc"}) == "abc" assert extraction_module._chunk_id({"id": 9}) == "9" assert extraction_module._chunk_id({"text": "no id"}) is None + + +@pytest.mark.unit +class TestEmbeddingsResolution: + def test_extraction_uses_shared_resolver(self, monkeypatch): + """Extraction must resolve embeddings through ``get_embeddings``.""" + from unittest.mock import MagicMock + + fake_store = MagicMock() + fake_store.pending_chunks.return_value = [] + monkeypatch.setattr( + "application.graphrag.store.GraphStore", lambda *a, **k: fake_store + ) + _install_stub_llm(monkeypatch, _StubLLM([])) + + calls = [] + fake_embedding = MagicMock() + + def _resolver(*args, **kwargs): + calls.append((args, kwargs)) + return fake_embedding + + monkeypatch.setattr(extraction_module, "get_embeddings", _resolver) + + summary = extract_graph_for_source( + str(uuid.uuid4()), + user="owner-1", + chunks=[], + config=SourceConfig(), + request_id="req-1", + ) + + assert calls == [((), {})] + assert summary["chunks_processed"] == 0 diff --git a/tests/graphrag/test_retriever_live.py b/tests/graphrag/test_retriever_live.py new file mode 100644 index 00000000..58f99c30 --- /dev/null +++ b/tests/graphrag/test_retriever_live.py @@ -0,0 +1,213 @@ +"""Live GraphRAG retrieval: real graph store, real pool, batched classic fallback. + +Runs against the ephemeral pytest-postgresql cluster (never the operator's dev +database) with an 8-dimension stub embeddings model, so the whole +``GraphRAGRetriever._get_data`` path — count batching, PPR over a real graph, +one pooled connection, one ClassicRAG call for the graphless sources — is +exercised end to end in a few seconds. Skips when the cluster's Postgres has no +pgvector build installed. +""" + +from __future__ import annotations + +import uuid +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from application.retriever.graph_rag import GraphRAGRetriever +from application.vectorstore import pgconn + +pytestmark = pytest.mark.integration + +STUB_DIM = 8 + + +class _StubEmbeddings: + """Deterministic tiny embeddings; no model download, no network.""" + + dimension = STUB_DIM + + def embed_documents(self, texts): + return [self._vector(1.0) for _ in texts] + + def embed_query(self, query): + return self._vector(1.0) + + @staticmethod + def _vector(seed: float): + vector = [0.0] * STUB_DIM + vector[0] = seed + return vector + + +def _dsn(info) -> str: + password = f":{info.password}" if info.password else "" + return ( + f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}" + ) + + +@pytest.fixture(autouse=True) +def _close_pools(): + """Never leak a pool into another test; the ephemeral DSN dies with its DB.""" + yield + for dsn, pool in list(pgconn._POOLS.items()): + try: + pool.close() + except Exception: + pass + pgconn._POOLS.pop(dsn, None) + + +@pytest.fixture +def stub_embeddings(): + stub = _StubEmbeddings() + with patch( + "application.vectorstore.base.get_embeddings", return_value=stub + ), patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=stub, + ), patch( + "application.retriever.graph_rag.get_embeddings", return_value=stub + ): + yield stub + + +@pytest.fixture +def live_dsn(postgresql, monkeypatch): + """DSN of an empty ephemeral DB with GraphRAG enabled and pooling on.""" + try: + with postgresql.cursor() as cursor: + cursor.execute("CREATE EXTENSION vector;") + postgresql.commit() + except Exception as exc: + postgresql.rollback() + pytest.skip(f"pgvector extension unavailable: {exc}") + + dsn = _dsn(postgresql.info) + from application.core import settings as settings_module + + settings = settings_module.settings + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False) + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True, raising=False) + monkeypatch.setattr(settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False) + monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False) + monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False) + return dsn + + +def _seed(dsn, graph_source_id): + """Ingest one chunk and one graph node for ``graph_source_id``.""" + from application.graphrag.store import GraphStore + from application.vectorstore.pgvector import PGVectorStore + + vector_store = PGVectorStore( + source_id=graph_source_id, connection_string=dsn + ) + try: + chunk_ids = vector_store.add_texts( + ["Ada Lovelace wrote the first algorithm."], + [{"title": "Ada", "source": "/docs/ada.txt"}], + ) + finally: + vector_store.close() + + store = GraphStore(connection_string=dsn) + try: + store._ensure_tables() + node_id = store.upsert_node( + graph_source_id, + "Ada Lovelace", + "ada lovelace", + "person", + "A mathematician.", + _StubEmbeddings._vector(1.0), + ) + store.link_node_chunk(graph_source_id, node_id, chunk_ids[0]) + finally: + store.close() + return chunk_ids[0] + + +def _retriever(sources): + with patch( + "application.retriever.classic_rag.LLMCreator.create_llm", + Mock(return_value=MagicMock()), + ): + return GraphRAGRetriever( + source={"question": "who wrote the first algorithm?", + "active_docs": list(sources)}, + chat_history=None, + prompt="", + chunks=2, + doc_token_limit=50000, + model_id="test-model", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "user1"}, + ) + + +class TestGraphRAGRetrieverLive: + def test_one_pool_one_classic_batch_and_the_connection_goes_back( + self, live_dsn, stub_embeddings + ): + graph_source = str(uuid.uuid4()) + graphless_source = str(uuid.uuid4()) + _seed(live_dsn, graph_source) + + rag = _retriever([graphless_source, graph_source]) + seen = [] + + def _classic(): + seen.append(list(rag._classic.vectorstores)) + return [ + { + "title": "classic", + "text": "classic chunk", + "source": graphless_source, + "filename": "classic", + } + ] + + rag._classic._get_data = Mock(side_effect=_classic) + + docs = rag._get_data() + + # The graphless source is retrieved by exactly one ClassicRAG run. + assert rag._classic._get_data.call_count == 1 + assert seen == [[graphless_source]] + + # The graph source came back through PPR, not the fallback. + texts = [doc["text"] for doc in docs] + assert texts == ["classic chunk", "Ada Lovelace wrote the first algorithm."] + assert docs[1]["source"] == "/docs/ada.txt" + + # One pool for the DSN, shared with the vector store that seeded it, + # and the retriever handed its connection back when it finished. + assert list(pgconn._POOLS) == [live_dsn] + stats = pgconn._POOLS[live_dsn].get_stats() + assert stats["pool_size"] <= 4 + assert stats["pool_available"] == stats["pool_size"] + + def test_all_graphless_sources_take_one_classic_call( + self, live_dsn, stub_embeddings + ): + from application.graphrag.store import GraphStore + + store = GraphStore(connection_string=live_dsn) + store._ensure_tables() # tables exist, but no source has a graph + store.close() + + sources = [str(uuid.uuid4()) for _ in range(3)] + rag = _retriever(sources) + seen = [] + rag._classic._get_data = Mock( + side_effect=lambda: seen.append(list(rag._classic.vectorstores)) or [] + ) + + assert rag._get_data() == [] + assert seen == [sources] + stats = pgconn._POOLS[live_dsn].get_stats() + assert stats["pool_available"] == stats["pool_size"] diff --git a/tests/graphrag/test_store.py b/tests/graphrag/test_store.py index abc0cc84..f0467f11 100644 --- a/tests/graphrag/test_store.py +++ b/tests/graphrag/test_store.py @@ -16,16 +16,20 @@ and the helpers build matching ones. from __future__ import annotations import uuid -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import application.graphrag.store as store_module +from application.vectorstore import pgconn +from application.vectorstore import pgvector as pgvector_module GraphStore = store_module.GraphStore TEST_EMBEDDING_DIM = 8 +POOL_DSN = "postgresql://u:p@localhost/graphpool" + _REAL_EMBEDDING_DIM = GraphStore._embedding_dim @@ -36,46 +40,28 @@ def _mock_embedding_dim(monkeypatch): ) -def _resolve_connection_string(): - from application.core.settings import settings - - conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None) - if not conn and getattr(settings, "POSTGRES_URI", None): - from application.core.db_uri import normalize_pgvector_connection_string - - conn = normalize_pgvector_connection_string(settings.POSTGRES_URI) - return conn +@pytest.fixture(autouse=True) +def _close_pools(): + """Never leak a pool into another test; an ephemeral DSN dies with its DB.""" + yield + for dsn, pool in list(pgconn._POOLS.items()): + try: + pool.close() + except Exception: + pass + pgconn._POOLS.pop(dsn, None) -_GRAPH_TABLES = ( - "graph_node_chunks", - "graph_edges", - "graph_nodes", - "graph_ingest_progress", -) +def _ephemeral_dsn(info) -> str: + """libpq DSN for the ephemeral pytest-postgresql database. - -def _drop_graph_tables(conn_string): - import psycopg - - with psycopg.connect(conn_string) as conn: - with conn.cursor() as cursor: - cursor.execute( - f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;" - ) - conn.commit() - - -def _live_store(): - conn = _resolve_connection_string() - if not conn: - pytest.skip("No pgvector connection string configured") - try: - _drop_graph_tables(conn) - store = GraphStore(connection_string=conn) - except Exception as exc: - pytest.skip(f"pgvector DB not reachable: {exc}") - return store + Deliberately not the operator's ``POSTGRES_URI``: these tests create and + drop graph tables, and the dev database is not theirs to rewrite. + """ + password = f":{info.password}" if info.password else "" + return ( + f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}" + ) def _embedding(seed: float) -> list: @@ -87,10 +73,20 @@ def _embedding(seed: float) -> list: @pytest.mark.integration class TestGraphStoreLive: @pytest.fixture - def store(self): - store = _live_store() + def store(self, postgresql): + """Graph store on a fresh ephemeral database. + + Construction no longer creates tables, so the fixture calls + ``_ensure_tables`` explicitly — exactly what ``ensure_vector_schema`` + does at boot in production — and read-before-write tests still pass. + """ + store = GraphStore(connection_string=_ephemeral_dsn(postgresql.info)) + try: + store._ensure_tables() + except Exception as exc: + pytest.skip(f"pgvector extension unavailable: {exc}") yield store - _drop_graph_tables(store._connection_string) + store.close() @pytest.fixture def source_id(self): @@ -344,6 +340,61 @@ class TestGraphStoreLive: finally: store.delete_by_source(source_id) + def test_count_nodes_many_batches_and_zero_fills(self, store): + """One query for N sources; a source with no graph still gets an entry.""" + a, b, c = (str(uuid.uuid4()) for _ in range(3)) + try: + store.upsert_node(a, "A1", "a1") + store.upsert_node(a, "A2", "a2") + store.upsert_node(b, "B1", "b1") + + counts = store.count_nodes_many([a, b, c]) + + assert counts == {a: 2, b: 1, c: 0} + # Agrees with the per-source query it replaces. + assert [store.count_nodes(s) for s in (a, b, c)] == [2, 1, 0] + assert store.count_nodes_many([]) == {} + finally: + store.delete_by_source(a) + store.delete_by_source(b) + + def test_pooled_connection_is_returned_to_the_shared_pool(self, store): + """The live store borrows from the shared pool and gives the socket back.""" + source_id = str(uuid.uuid4()) + assert store.count_nodes_many([source_id]) == {source_id: 0} + assert store._pooled is True + assert list(pgconn._POOLS) == [store._connection_string] + + pool = pgconn._POOLS[store._connection_string] + store.close() + + assert store._connection is None + stats = pool.get_stats() + assert stats["pool_available"] == stats["pool_size"] + + def test_the_vector_store_reuses_the_graph_store_pool(self, store, postgresql): + """Same DSN, one pool: the graph store does not double the connections.""" + from application.vectorstore.pgvector import PGVectorStore + + stub = MagicMock() + stub.dimension = TEST_EMBEDDING_DIM + stub.embed_query.return_value = [0.0] * TEST_EMBEDDING_DIM + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=stub, + ): + vector_store = PGVectorStore( + source_id="live-source", connection_string=store._connection_string + ) + try: + store.count_nodes_many([str(uuid.uuid4())]) + vector_store._get_connection() + + assert list(pgconn._POOLS) == [store._connection_string] + assert vector_store._pooled is True + finally: + vector_store.close() + def test_delete_by_source_isolation(self, store): keep = str(uuid.uuid4()) drop = str(uuid.uuid4()) @@ -381,6 +432,8 @@ class TestGraphStoreParameterization: conn.cursor.return_value = cursor store._connection = conn store._get_connection = lambda: conn + # Boot owns the schema; the write-path safety net has its own tests. + store._tables_ensured = True return store, cursor def test_delete_by_source_binds_source_id(self): @@ -448,6 +501,7 @@ class TestEmbeddingDim: def test_uses_configured_model_dimension(self, monkeypatch): from application.vectorstore import base as base_module + monkeypatch.setattr(base_module.settings, "EMBEDDINGS_BASE_URL", None) fake_embedding = MagicMock() fake_embedding.dimension = 1536 monkeypatch.setattr( @@ -463,6 +517,7 @@ class TestEmbeddingDim: def test_falls_back_to_default_dimension(self, monkeypatch): from application.vectorstore import base as base_module + monkeypatch.setattr(base_module.settings, "EMBEDDINGS_BASE_URL", None) fake_embedding = object() monkeypatch.setattr( base_module.EmbeddingsSingleton, @@ -473,3 +528,321 @@ class TestEmbeddingDim: store = GraphStore.__new__(GraphStore) assert store._embedding_dim() == store_module.DEFAULT_NAME_EMBEDDING_DIM + + +@pytest.mark.unit +class TestEmbeddingDimResolution: + def test_uses_shared_resolver(self, monkeypatch): + """The dimension probe must not build its own embeddings instance.""" + from unittest.mock import patch + + fake_embedding = MagicMock() + fake_embedding.dimension = 1536 + monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM) + + with patch( + "application.vectorstore.base.get_embeddings", + return_value=fake_embedding, + ) as mock_resolver: + store = GraphStore.__new__(GraphStore) + assert store._embedding_dim() == 1536 + + mock_resolver.assert_called_once_with() + + +@pytest.mark.unit +class TestGraphSchemaIsBootOwned: + """Construction must not run DDL: reads happen once per query, per source.""" + + def _mock_store(self): + store = GraphStore.__new__(GraphStore) + cursor = MagicMock() + cursor.fetchone.return_value = [str(uuid.uuid4())] + cursor.fetchall.return_value = [] + conn = MagicMock() + conn.cursor.return_value = cursor + store._connection = conn + store._get_connection = lambda: conn + store._tables_ensured = False + store._ensure_tables = MagicMock() + return store, cursor + + def test_init_opens_no_connection_and_creates_no_tables(self): + from unittest.mock import patch + + with patch.dict( + "sys.modules", + { + "psycopg": MagicMock(), + "pgvector": MagicMock(), + "pgvector.psycopg": MagicMock(), + }, + ), patch.object(GraphStore, "_ensure_tables") as ensure, patch.object( + GraphStore, "_get_connection" + ) as get_conn: + store = GraphStore(connection_string="postgresql://u:p@localhost/db") + + ensure.assert_not_called() + get_conn.assert_not_called() + assert store._tables_ensured is False + + @pytest.mark.parametrize( + "call", + [ + lambda s: s.upsert_node("sid", "N", "n"), + lambda s: s.add_edge("sid", "a", "b"), + lambda s: s.link_node_chunk("sid", "n", "c1"), + lambda s: s.apply_chunk("sid", "c1", [], [], {}), + lambda s: s.set_node_degrees("sid"), + lambda s: s.mark_chunk("sid", "c1", "done"), + lambda s: s.delete_by_source("sid"), + ], + ids=[ + "upsert_node", + "add_edge", + "link_node_chunk", + "apply_chunk", + "set_node_degrees", + "mark_chunk", + "delete_by_source", + ], + ) + def test_writes_ensure_tables_once(self, call): + store, _ = self._mock_store() + + call(store) + store._tables_ensured = True # what the real _ensure_tables_once sets + call(store) + + assert store._ensure_tables.call_count == 1 + + @pytest.mark.parametrize( + "call", + [ + lambda s: s.count_nodes("sid"), + lambda s: s.count_nodes_many(["sid"]), + lambda s: s.get_node_by_normalized("sid", "n"), + lambda s: s.search_nodes_by_embedding("sid", _embedding(1.0)), + lambda s: s.get_subgraph("sid", ["n"]), + lambda s: s.get_graph_overview("sid"), + lambda s: s.get_chunk_ids_for_nodes("sid", ["n"]), + lambda s: s.pending_chunks("sid", ["c1"]), + lambda s: s.get_progress("sid"), + ], + ids=[ + "count_nodes", + "count_nodes_many", + "get_node_by_normalized", + "search_nodes_by_embedding", + "get_subgraph", + "get_graph_overview", + "get_chunk_ids_for_nodes", + "pending_chunks", + "get_progress", + ], + ) + def test_reads_never_create_tables(self, call): + store, _ = self._mock_store() + + call(store) + + store._ensure_tables.assert_not_called() + + def test_create_schema_emits_the_ddl_without_committing(self): + conn, cursor = MagicMock(), MagicMock() + conn.cursor.return_value = cursor + + GraphStore.create_schema(conn, dimension=8) + + statements = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "CREATE EXTENSION IF NOT EXISTS vector" in statements + for table in ( + "graph_nodes", + "graph_edges", + "graph_node_chunks", + "graph_ingest_progress", + ): + assert f"CREATE TABLE IF NOT EXISTS {table}" in statements + assert "name_embedding vector(8)" in statements + assert statements.count("CREATE INDEX IF NOT EXISTS") == 5 + conn.commit.assert_not_called() + + def test_ensure_tables_locks_then_commits(self): + store, cursor = self._mock_store() + del store._ensure_tables # exercise the real method + + store._ensure_tables() + + statements = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "pg_advisory_xact_lock" in statements + assert "CREATE TABLE IF NOT EXISTS graph_nodes" in statements + store._connection.commit.assert_called_once() + + +@pytest.mark.unit +class TestGraphStorePooling: + """The graph store borrows from the same per-DSN pool as ``PGVectorStore``.""" + + def _store(self, dsn=POOL_DSN, pool_max_size=4): + store = GraphStore.__new__(GraphStore) + store._connection_string = dsn + store._connection = None + store._pooled = False + store._pool_max_size = pool_max_size + store._psycopg = MagicMock() + store._register_vector = MagicMock() + store._tables_ensured = True + return store + + def _fake_pool(self): + pooled_conn = MagicMock() + pooled_conn.closed = False + pool = MagicMock() + pool.getconn.return_value = pooled_conn + return pool, pooled_conn + + def test_get_connection_checks_out_of_the_pool(self, monkeypatch): + pool, pooled_conn = self._fake_pool() + monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool) + store = self._store() + + conn = store._get_connection() + + assert conn is pooled_conn + assert store._pooled is True + pool.getconn.assert_called_once() + store._psycopg.connect.assert_not_called() + # The pool's ``configure`` hook already registered the adapters. + store._register_vector.assert_not_called() + + def test_close_rolls_back_and_returns_the_connection(self, monkeypatch): + pool, pooled_conn = self._fake_pool() + monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool) + monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool) + store = self._store() + store._get_connection() + pooled_conn.info.transaction_status.name = "INTRANS" + + store.close() + + pooled_conn.rollback.assert_called_once() + pool.putconn.assert_called_once_with(pooled_conn) + pooled_conn.close.assert_not_called() + assert store._connection is None + + def test_close_does_not_roll_back_an_idle_connection(self, monkeypatch): + pool, pooled_conn = self._fake_pool() + monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool) + monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool) + store = self._store() + store._get_connection() + pooled_conn.info.transaction_status.name = "IDLE" + + store.close() + + pooled_conn.rollback.assert_not_called() + pool.putconn.assert_called_once_with(pooled_conn) + + def test_legacy_path_connects_directly_and_closes(self, monkeypatch): + def _never(dsn, n): + raise AssertionError("pooling is off; no pool must be built") + + monkeypatch.setattr(store_module.pgconn, "pool_for", _never) + store = self._store(pool_max_size=0) + direct = MagicMock() + direct.closed = False + store._psycopg.connect.return_value = direct + + conn = store._get_connection() + + assert conn is direct + assert store._pooled is False + store._register_vector.assert_called_once_with(direct) + + store.close() + direct.close.assert_called_once() + + def test_del_never_raises(self): + store = self._store() + broken = MagicMock() + broken.closed = False + broken.close.side_effect = RuntimeError("already gone") + store._connection = broken + + store.__del__() # must not propagate + + def test_the_graph_store_and_the_vector_store_share_one_pool(self): + """One DSN, one pool object — reached from either module.""" + pool, _ = self._fake_pool() + store = self._store() + + with patch("psycopg_pool.ConnectionPool", return_value=pool) as pool_cls: + store._get_connection() + # ``PGVectorStore``'s own entry point resolves to the same object. + assert pgvector_module._pool_for(POOL_DSN, 4) is pool + + assert pool_cls.call_count == 1 + assert pgconn._POOLS[POOL_DSN] is pool + assert pgvector_module._POOLS is pgconn._POOLS + + @pytest.mark.parametrize( + "value,expected", + [(0, 0), (2, 2), (None, 8), ("4", 8), (True, 8), (-1, 8)], + ) + def test_pool_size_is_resolved_defensively(self, monkeypatch, value, expected): + monkeypatch.setattr( + store_module.settings, "PGVECTOR_POOL_MAX_SIZE", value, raising=False + ) + assert store_module._resolve_pool_max_size() == expected + + +@pytest.mark.unit +class TestCountNodesMany: + """One ``ANY(%s)`` query replaces the retriever's per-source count fan-out.""" + + def _store_with_mock_conn(self, rows): + store = GraphStore.__new__(GraphStore) + cursor = MagicMock() + cursor.fetchall.return_value = rows + conn = MagicMock() + conn.cursor.return_value = cursor + store._connection = conn + store._get_connection = lambda: conn + store._tables_ensured = True + return store, cursor, conn + + def test_binds_the_ids_as_one_array_and_zero_fills(self): + store, cursor, _ = self._store_with_mock_conn([("a", 2), ("b", 1)]) + + counts = store.count_nodes_many(["a", "b", "c"]) + + assert counts == {"a": 2, "b": 1, "c": 0} + sql, params = ( + cursor.execute.call_args.args[0], + cursor.execute.call_args.args[1], + ) + assert "source_id = ANY(%s)" in sql + assert "GROUP BY source_id" in sql + assert cursor.execute.call_count == 1 + assert params == (["a", "b", "c"],) + + def test_empty_input_short_circuits(self): + store, cursor, _ = self._store_with_mock_conn([]) + + assert store.count_nodes_many([]) == {} + assert store.count_nodes_many([None, ""]) == {} + cursor.execute.assert_not_called() + + def test_a_failed_query_reports_every_source_as_graphless(self): + store, cursor, conn = self._store_with_mock_conn([]) + cursor.execute.side_effect = RuntimeError("no such table") + + assert store.count_nodes_many(["a", "b"]) == {"a": 0, "b": 0} + conn.rollback.assert_called_once() + + def test_the_callers_id_spelling_is_preserved(self): + """Postgres returns canonical lowercase UUID text; keys must still match.""" + source_id = str(uuid.uuid4()).upper() + store, _, _ = self._store_with_mock_conn([(source_id.lower(), 3)]) + + assert store.count_nodes_many([source_id]) == {source_id: 3} diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index dd97e715..04167ebf 100644 --- a/tests/llm/test_fallback.py +++ b/tests/llm/test_fallback.py @@ -5,12 +5,18 @@ per-agent backup model is used before the global FALLBACK_* settings. """ import copy +import logging +import types from unittest.mock import MagicMock import httpx import pytest +from application.llm.anthropic import AnthropicLLM from application.llm.base import BaseLLM +from application.llm.google_ai import GoogleLLM +from application.llm.groq import GroqLLM +from application.llm.openai import OpenAILLM # Concrete LLM stubs @@ -41,6 +47,7 @@ class FakeLLM(BaseLLM): self.gen_stream_called = False self.last_model_received = None # tracks the model kwarg passed to gen/gen_stream self.last_messages_received = None # tracks the messages kwarg at the raw level + self.last_kwargs_received = None # tracks the extra gen kwargs at the raw level # Track at the raw-method level. _execute_with_fallback applies # decorators to the fallback's raw method directly and @@ -50,6 +57,7 @@ class FakeLLM(BaseLLM): self.gen_called = True self.last_model_received = model self.last_messages_received = messages + self.last_kwargs_received = dict(kwargs) if self.fail_at is not None: raise RuntimeError("primary model unavailable") return self.responses[0] @@ -59,6 +67,7 @@ class FakeLLM(BaseLLM): self.stream_calls = getattr(self, "stream_calls", 0) + 1 self.last_model_received = model self.last_messages_received = messages + self.last_kwargs_received = dict(kwargs) yielded = 0 # Per-attempt failure schedule: `fail_schedule[n]` = fail_at value # for the n-th call (0-indexed). Falls back to the constant @@ -1296,3 +1305,489 @@ class TestFallbackMessageReshaping: assert isinstance(user_content, str) assert "EXTRACTED REPORT TEXT" in user_content assert "assistant-abc123" not in user_content + + +# Tests — cross-provider structured-output adaptation +# +# Structured output is provider-specific: OpenAI-wire classes take +# ``response_format``, Google takes ``response_schema``. Forwarding the +# primary's kwarg verbatim 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 turned the +# fallback into no fallback at all for every structured node. + + +SCHEMA = { + "type": "object", + "properties": { + "answer": {"type": "string"}, + "score": {"type": "integer"}, + }, + "required": ["answer"], +} + + +class _OpenAIWireFake(FakeLLM): + """OpenAI-wire double: real declaration + real preparer, fake transport.""" + + provider_name = "openai" + structured_output_kwarg = "response_format" + prepare_structured_output_format = OpenAILLM.prepare_structured_output_format + + def _supports_structured_output(self): + return True + + +class _GoogleFake(FakeLLM): + """Google double: real declaration + real preparer, fake transport.""" + + provider_name = "google" + structured_output_kwarg = "response_schema" + prepare_structured_output_format = GoogleLLM.prepare_structured_output_format + + def _supports_structured_output(self): + return True + + +class _AnthropicFake(FakeLLM): + """Provider with no structured-output kwarg at all.""" + + provider_name = "anthropic" + + +def _openai_envelope(schema=SCHEMA, strict=True): + """An OpenAI ``response_format`` built the way a provider would.""" + return OpenAILLM.prepare_structured_output_format( + _OpenAIWireFake(), schema, strict=strict + ) + + +def _google_schema(schema=SCHEMA): + """The Google ``response_schema`` conversion of ``schema``.""" + return GoogleLLM.prepare_structured_output_format(_GoogleFake(), schema) + + +@pytest.mark.unit +class TestStructuredOutputDeclarations: + """One source of truth: the kwarg name lives on the LLM class.""" + + def test_openai_declares_response_format(self): + assert OpenAILLM.structured_output_kwarg == "response_format" + + def test_openai_subclasses_inherit_the_declaration(self): + assert GroqLLM.structured_output_kwarg == "response_format" + + def test_google_declares_response_schema(self): + assert GoogleLLM.structured_output_kwarg == "response_schema" + + def test_anthropic_declares_nothing(self): + assert AnthropicLLM.structured_output_kwarg is None + + def test_base_declares_nothing(self): + assert BaseLLM.structured_output_kwarg is None + + def test_openai_prepare_records_source(self): + llm = _OpenAIWireFake() + llm.prepare_structured_output_format(SCHEMA, strict=False) + assert llm._structured_output_source == (SCHEMA, False) + + def test_google_prepare_records_source(self): + llm = _GoogleFake() + llm.prepare_structured_output_format(SCHEMA) + assert llm._structured_output_source == (SCHEMA, True) + + def test_empty_schema_clears_recorded_source(self): + llm = _OpenAIWireFake() + llm.prepare_structured_output_format(SCHEMA) + llm.prepare_structured_output_format(None) + assert llm._structured_output_source is None + + def test_base_records_nothing_by_default(self): + assert FakeLLM()._structured_output_source is None + + +@pytest.mark.unit +class TestAdaptStructuredOutputKwargs: + """Unit-level contract of the adapter itself.""" + + def test_no_structured_kwargs_returns_equal_copy(self): + primary = _OpenAIWireFake() + kwargs = {"model": "m", "messages": [], "temperature": 0.2} + + adapted = primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs) + + assert adapted == kwargs + assert adapted is not kwargs + + def test_does_not_mutate_the_callers_kwargs(self): + primary = _OpenAIWireFake() + primary.prepare_structured_output_format(SCHEMA) + kwargs = {"model": "m", "response_format": _openai_envelope()} + snapshot = copy.deepcopy(kwargs) + + primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs) + + assert kwargs == snapshot + + def test_recovers_schema_from_envelope_when_no_source_recorded(self): + """A hand-built ``response_format`` (research_agent) never went through + ``prepare_structured_output_format``, so nothing was recorded — the raw + schema is still readable out of the OpenAI envelope.""" + primary = _OpenAIWireFake() + assert primary._structured_output_source is None + + adapted = primary._adapt_structured_output_kwargs( + _GoogleFake(model_id="gemini-2.5-flash"), + {"model": "m", "response_format": _openai_envelope()}, + ) + + assert "response_format" not in adapted + assert adapted["response_schema"]["type"] == "OBJECT" + assert set(adapted["response_schema"]["properties"]) == {"answer", "score"} + assert adapted["model"] == "m" + + def test_google_schema_without_source_is_dropped(self, caplog): + """Google's conversion is lossy/type-mapped — not reversible.""" + primary = _GoogleFake() + assert primary._structured_output_source is None + + with caplog.at_level(logging.WARNING, logger="application.llm.base"): + adapted = primary._adapt_structured_output_kwargs( + _OpenAIWireFake(model_id="gpt-4o-mini"), + {"model": "m", "response_schema": _google_schema()}, + ) + + assert "response_schema" not in adapted + assert "response_format" not in adapted + assert "gpt-4o-mini" in caplog.text + + def test_fallback_without_structured_support_drops_and_warns(self, caplog): + primary = _OpenAIWireFake() + primary.prepare_structured_output_format(SCHEMA) + fallback = _GoogleFake(model_id="gemini-2.5-flash") + fallback._supports_structured_output = lambda: False + + with caplog.at_level(logging.WARNING, logger="application.llm.base"): + adapted = primary._adapt_structured_output_kwargs( + fallback, {"model": "m", "response_format": _openai_envelope()} + ) + + assert "response_format" not in adapted + assert "response_schema" not in adapted + assert "cannot enforce structured output" in caplog.text + + def test_non_callable_support_flag_is_honored(self): + """Test doubles sometimes set the capability as a plain bool.""" + primary = _OpenAIWireFake() + primary.prepare_structured_output_format(SCHEMA) + fallback = _GoogleFake(model_id="gemini-2.5-flash") + fallback._supports_structured_output = False + + adapted = primary._adapt_structured_output_kwargs( + fallback, {"response_format": _openai_envelope()} + ) + + assert adapted == {} + + def test_preparer_returning_none_drops_the_kwarg(self, caplog): + class _NullPreparer(_GoogleFake): + def prepare_structured_output_format(self, json_schema, strict=True): + return None + + primary = _OpenAIWireFake() + primary.prepare_structured_output_format(SCHEMA) + + with caplog.at_level(logging.WARNING, logger="application.llm.base"): + adapted = primary._adapt_structured_output_kwargs( + _NullPreparer(model_id="gemini-2.5-flash"), + {"response_format": _openai_envelope()}, + ) + + assert adapted == {} + assert "cannot enforce structured output" in caplog.text + + def test_preparer_raising_never_breaks_the_fallback(self, caplog): + class _ExplodingPreparer(_GoogleFake): + def prepare_structured_output_format(self, json_schema, strict=True): + raise ValueError("boom") + + primary = _OpenAIWireFake() + primary.prepare_structured_output_format(SCHEMA) + + with caplog.at_level(logging.WARNING, logger="application.llm.base"): + adapted = primary._adapt_structured_output_kwargs( + _ExplodingPreparer(model_id="gemini-2.5-flash"), + {"response_format": _openai_envelope()}, + ) + + assert adapted == {} + assert "Failed to prepare structured output" in caplog.text + + def test_strict_flag_survives_the_translation(self): + primary = _GoogleFake() + primary.prepare_structured_output_format(SCHEMA) + primary._structured_output_source = (SCHEMA, False) + + adapted = primary._adapt_structured_output_kwargs( + _OpenAIWireFake(model_id="gpt-4o-mini"), + {"response_schema": _google_schema()}, + ) + + assert adapted["response_format"]["json_schema"]["strict"] is False + # strict=False leaves the schema untouched (no additionalProperties). + assert "additionalProperties" not in ( + adapted["response_format"]["json_schema"]["schema"] + ) + + +@pytest.mark.integration +class TestCrossProviderStructuredOutputFallback: + """End-to-end through ``gen`` / ``gen_stream``: the backup must receive + the schema in *its own* provider's kwarg.""" + + def test_gen_openai_primary_google_fallback_gets_response_schema(self): + primary = _OpenAIWireFake(fail_at=0) + fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash") + primary._fallback_llm = fallback + response_format = primary.prepare_structured_output_format(SCHEMA) + + result = primary.gen(**CALL_ARGS, response_format=response_format) + + assert result == "fb" + assert fallback.last_kwargs_received["response_schema"] == _google_schema() + assert "response_format" not in fallback.last_kwargs_received + + def test_stream_openai_primary_google_fallback_gets_response_schema(self): + primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0) + fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash") + primary._fallback_llm = fallback + response_format = primary.prepare_structured_output_format(SCHEMA) + + chunks = list(primary.gen_stream(**CALL_ARGS, response_format=response_format)) + + assert chunks == ["fb"] + assert fallback.last_kwargs_received["response_schema"] == _google_schema() + assert "response_format" not in fallback.last_kwargs_received + + def test_gen_google_primary_openai_fallback_gets_response_format(self): + primary = _GoogleFake(fail_at=0) + fallback = _OpenAIWireFake(responses=["fb"], model_id="gpt-4o-mini") + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + result = primary.gen(**CALL_ARGS, response_schema=response_schema) + + assert result == "fb" + received = fallback.last_kwargs_received + assert "response_schema" not in received + assert received["response_format"]["type"] == "json_schema" + assert set(received["response_format"]["json_schema"]["schema"]["properties"]) == { + "answer", + "score", + } + + def test_stream_google_primary_openai_fallback_gets_response_format(self): + primary = _GoogleFake(stream_chunks=["x"], fail_at=0) + fallback = _OpenAIWireFake(stream_chunks=["fb"], model_id="gpt-4o-mini") + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema)) + + assert chunks == ["fb"] + received = fallback.last_kwargs_received + assert "response_schema" not in received + assert received["response_format"]["type"] == "json_schema" + + def test_same_wire_family_passes_response_format_verbatim(self): + """OpenAI -> openai_compatible: no re-preparation, byte-identical.""" + primary = _OpenAIWireFake(fail_at=0) + fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b") + primary._fallback_llm = fallback + response_format = primary.prepare_structured_output_format(SCHEMA) + + primary.gen(**CALL_ARGS, response_format=response_format) + + assert fallback.last_kwargs_received["response_format"] is response_format + assert "response_schema" not in fallback.last_kwargs_received + + def test_same_wire_family_passes_response_schema_verbatim(self): + primary = _GoogleFake(fail_at=0) + fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-pro") + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + primary.gen(**CALL_ARGS, response_schema=response_schema) + + assert fallback.last_kwargs_received["response_schema"] is response_schema + assert "response_format" not in fallback.last_kwargs_received + + def test_json_object_mode_kept_within_the_openai_family(self): + primary = _OpenAIWireFake(fail_at=0) + fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b") + primary._fallback_llm = fallback + + primary.gen(**CALL_ARGS, response_format={"type": "json_object"}) + + assert fallback.last_kwargs_received["response_format"] == { + "type": "json_object" + } + + def test_json_object_mode_dropped_for_google_fallback(self): + """Google has no json_object equivalent wired — drop, don't crash.""" + primary = _OpenAIWireFake(fail_at=0) + fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash") + primary._fallback_llm = fallback + + result = primary.gen(**CALL_ARGS, response_format={"type": "json_object"}) + + assert result == "fb" + assert "response_format" not in fallback.last_kwargs_received + assert "response_schema" not in fallback.last_kwargs_received + + def test_json_object_mode_dropped_for_google_fallback_streaming(self): + primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0) + fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash") + primary._fallback_llm = fallback + + chunks = list( + primary.gen_stream(**CALL_ARGS, response_format={"type": "json_object"}) + ) + + assert chunks == ["fb"] + assert "response_format" not in fallback.last_kwargs_received + assert "response_schema" not in fallback.last_kwargs_received + + def test_anthropic_fallback_gets_neither_kwarg(self): + """Anthropic has no structured-output kwarg: unstructured, not broken.""" + primary = _OpenAIWireFake(fail_at=0) + fallback = _AnthropicFake(responses=["fb"], model_id="claude-sonnet-4") + primary._fallback_llm = fallback + response_format = primary.prepare_structured_output_format(SCHEMA) + + result = primary.gen(**CALL_ARGS, response_format=response_format) + + assert result == "fb" + assert fallback.last_kwargs_received == {} + + def test_anthropic_fallback_gets_neither_kwarg_streaming(self): + primary = _GoogleFake(stream_chunks=["x"], fail_at=0) + fallback = _AnthropicFake(stream_chunks=["fb"], model_id="claude-sonnet-4") + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema)) + + assert chunks == ["fb"] + assert fallback.last_kwargs_received == {} + + def test_unrelated_gen_kwargs_are_forwarded_untouched(self): + primary = _OpenAIWireFake(fail_at=0) + fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash") + primary._fallback_llm = fallback + response_format = primary.prepare_structured_output_format(SCHEMA) + + primary.gen( + **CALL_ARGS, response_format=response_format, temperature=0.3 + ) + + assert fallback.last_kwargs_received["temperature"] == 0.3 + + +# The OpenAI SDK's ``chat.completions.create`` has an explicit keyword +# signature: a forwarded ``response_schema`` raises TypeError before the +# request is even built, so the Google -> OpenAI hop died with "Fallback LLM +# also failed". These tests drive the *real* ``OpenAILLM`` raw methods against +# a client double with the same strictness. + + +class _StrictChatCompletions: + """Chat-Completions double that rejects kwargs the real SDK rejects.""" + + _ACCEPTED = { + "model", + "messages", + "stream", + "stream_options", + "tools", + "tool_choice", + "parallel_tool_calls", + "response_format", + "temperature", + "top_p", + "max_completion_tokens", + "reasoning_effort", + "presence_penalty", + "frequency_penalty", + "seed", + "stop", + "n", + "user", + } + + def __init__(self): + self.last_kwargs = None + + def create(self, **kwargs): + unexpected = sorted(set(kwargs) - self._ACCEPTED) + if unexpected: + raise TypeError( + f"Completions.create() got an unexpected keyword argument " + f"'{unexpected[0]}'" + ) + self.last_kwargs = kwargs + if kwargs.get("stream"): + return [ + _stream_line(content="fb"), + _stream_line(finish_reason="stop"), + ] + message = types.SimpleNamespace(content="fb answer", tool_calls=None) + return types.SimpleNamespace( + choices=[types.SimpleNamespace(message=message)], usage=None + ) + + +def _stream_line(content=None, finish_reason=None): + delta = types.SimpleNamespace( + content=content, reasoning_content=None, tool_calls=None + ) + choice = types.SimpleNamespace(delta=delta, finish_reason=finish_reason) + return types.SimpleNamespace(choices=[choice], usage=None) + + +def _strict_openai_llm(): + llm = OpenAILLM(api_key="sk-test", user_api_key=None, model_id="gpt-4o-mini") + llm.client = types.SimpleNamespace( + chat=types.SimpleNamespace(completions=_StrictChatCompletions()) + ) + return llm + + +@pytest.mark.integration +class TestRealOpenAIFallbackRejectsForeignKwargs: + + def test_gen_google_primary_real_openai_fallback_does_not_typeerror(self): + primary = _GoogleFake(fail_at=0) + fallback = _strict_openai_llm() + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + result = primary.gen(**CALL_ARGS, response_schema=response_schema) + + assert result == "fb answer" + sent = fallback.client.chat.completions.last_kwargs + assert "response_schema" not in sent + assert sent["response_format"]["type"] == "json_schema" + + def test_stream_google_primary_real_openai_fallback_does_not_typeerror(self): + primary = _GoogleFake(stream_chunks=["x"], fail_at=0) + fallback = _strict_openai_llm() + primary._fallback_llm = fallback + response_schema = primary.prepare_structured_output_format(SCHEMA) + + chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema)) + + assert chunks == ["fb"] + sent = fallback.client.chat.completions.last_kwargs + assert "response_schema" not in sent + assert sent["response_format"]["type"] == "json_schema" diff --git a/tests/llm/test_openai.py b/tests/llm/test_openai.py index 65707845..d0d0653a 100644 --- a/tests/llm/test_openai.py +++ b/tests/llm/test_openai.py @@ -15,12 +15,19 @@ Extends coverage beyond test_openai_llm.py: """ import base64 +import logging import types from unittest.mock import MagicMock, patch +import httpx import pytest +from openai import BadRequestError -from application.llm.openai import OpenAILLM, _truncate_base64_for_logging +from application.llm.openai import ( + OpenAILLM, + _is_tools_unsupported_error, + _truncate_base64_for_logging, +) # Fake client helpers @@ -2080,3 +2087,235 @@ class TestKeylessConstruction: mock_settings.OPENAI_STT_MODEL = "whisper-1" stt = OpenAISTT(api_key=blank) assert stt.api_key == "sk-no-key" + + +# Tool-calling fallback for endpoints that can't serve tools + + +VLLM_TOOL_ERROR = ( + '"auto" tool choice requires --enable-auto-tool-choice and ' + "--tool-call-parser to be set" +) + + +def _bad_request(message): + """Build a real ``openai.BadRequestError`` carrying ``message``.""" + request = httpx.Request("POST", "http://localhost:8000/v1/chat/completions") + response = httpx.Response(400, request=request) + return BadRequestError(message, response=response, body={"message": message}) + + +class _RecordingCreate: + """``create`` stub that raises on the first call, then succeeds.""" + + def __init__(self, error, result): + self.error = error + self.result = result + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + if len(self.calls) == 1 and self.error is not None: + raise self.error + return self.result + + +@pytest.mark.unit +class TestIsToolsUnsupportedError: + + @pytest.mark.parametrize( + "message", + [ + VLLM_TOOL_ERROR, + "Tool calling is not supported by this model", + "This endpoint does not support tools", + "function calling is unavailable for this deployment", + "Unsupported parameter: 'tools'", + ], + ) + def test_matches_tool_messages(self, message): + assert _is_tools_unsupported_error(_bad_request(message)) is True + + @pytest.mark.parametrize( + "message", + [ + "This model's maximum context length is 8192 tokens", + "Invalid value for 'temperature': must be <= 2", + "Incorrect API key provided", + ], + ) + def test_ignores_unrelated_messages(self, message): + assert _is_tools_unsupported_error(_bad_request(message)) is False + + def test_matches_message_only_present_in_body(self): + request = httpx.Request("POST", "http://localhost:8000/v1/chat/completions") + error = BadRequestError( + "Error code: 400", + response=httpx.Response(400, request=request), + body={"message": VLLM_TOOL_ERROR}, + ) + assert _is_tools_unsupported_error(error) is True + + +@pytest.mark.unit +class TestToolCallingFallback: + """An OpenAI-compatible endpoint without tool support (vLLM started + without --enable-auto-tool-choice) 400s every chat, because agentless + chat always sends synthesized default tools. Degrade instead.""" + + TOOLS = [{"type": "function", "function": {"name": "t"}}] + MSGS = [{"role": "user", "content": "hi"}] + + def test_raw_gen_retries_once_without_tools(self, llm): + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), + _Response(choices=[_Choice(content="degraded answer")]), + ) + llm.client.chat.completions.create = create + + result = llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, tool_choice="auto", parallel_tool_calls=True, + ) + + assert len(create.calls) == 2 + assert create.calls[0]["tools"] == self.TOOLS + assert "tools" not in create.calls[1] + assert "tool_choice" not in create.calls[1] + assert "parallel_tool_calls" not in create.calls[1] + # The caller asked for tools, so the tool-shaped return contract holds. + assert result.message.content == "degraded answer" + + def test_raw_gen_stream_retries_once_without_tools(self, llm): + retry_response = _Response( + lines=[_StreamLine([_Choice(delta="degraded")])] + ) + create = _RecordingCreate(_bad_request(VLLM_TOOL_ERROR), retry_response) + llm.client.chat.completions.create = create + + chunks = list( + llm._raw_gen_stream( + llm, model="local-model", messages=self.MSGS, stream=True, + tools=self.TOOLS, tool_choice="auto", + ) + ) + + assert chunks == ["degraded"] + assert len(create.calls) == 2 + assert create.calls[0]["tools"] == self.TOOLS + assert "tools" not in create.calls[1] + assert "tool_choice" not in create.calls[1] + + def test_unrelated_bad_request_propagates(self, llm): + error = _bad_request("This model's maximum context length is 8192 tokens") + create = _RecordingCreate(error, _Response(choices=[_Choice(content="x")])) + llm.client.chat.completions.create = create + + with pytest.raises(BadRequestError): + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, + ) + + assert len(create.calls) == 1 + assert llm._supports_tools() is True + + def test_unrelated_bad_request_propagates_from_stream(self, llm): + error = _bad_request("This model's maximum context length is 8192 tokens") + create = _RecordingCreate(error, _Response(lines=[])) + llm.client.chat.completions.create = create + + with pytest.raises(BadRequestError): + list( + llm._raw_gen_stream( + llm, model="local-model", messages=self.MSGS, stream=True, + tools=self.TOOLS, + ) + ) + + assert len(create.calls) == 1 + + def test_tool_error_without_tools_in_request_propagates(self, llm): + """No tools were sent, so there is nothing to degrade to.""" + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), _Response(choices=[_Choice(content="x")]) + ) + llm.client.chat.completions.create = create + + with pytest.raises(BadRequestError): + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, tools=None + ) + + assert len(create.calls) == 1 + assert llm._supports_tools() is True + + def test_supports_tools_false_after_rejection(self, llm): + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), + _Response(choices=[_Choice(content="degraded")]), + ) + llm.client.chat.completions.create = create + + assert llm._supports_tools() is True + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, + ) + assert llm._supports_tools() is False + + def test_later_calls_skip_tools_without_another_round_trip(self, llm): + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), + _Response(choices=[_Choice(content="degraded")]), + ) + llm.client.chat.completions.create = create + + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, + ) + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, tool_choice="auto", + ) + + # 2 for the first call (reject + retry), 1 for the second. + assert len(create.calls) == 3 + assert "tools" not in create.calls[2] + assert "tool_choice" not in create.calls[2] + + def test_rejection_does_not_leak_to_other_instances(self, llm): + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), + _Response(choices=[_Choice(content="degraded")]), + ) + llm.client.chat.completions.create = create + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, + ) + + other = OpenAILLM(api_key="sk-test", user_api_key=None) + assert other._supports_tools() is True + + def test_warns_once_with_actionable_message(self, llm, caplog): + create = _RecordingCreate( + _bad_request(VLLM_TOOL_ERROR), + _Response(choices=[_Choice(content="degraded")]), + ) + llm.client.chat.completions.create = create + + with caplog.at_level(logging.WARNING): + llm._raw_gen( + llm, model="local-model", messages=self.MSGS, stream=False, + tools=self.TOOLS, + ) + llm._note_tools_rejected("local-model", "second time") + + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert len(warnings) == 1 + message = warnings[0].getMessage() + assert "supports_tools: false" in message + assert "local-model" in message + assert "--enable-auto-tool-choice" in message diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 1d2bb5d2..9ae2e645 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1471,3 +1471,82 @@ def test_to_responses_input_drops_out_of_order_pair_entirely(monkeypatch): items = llm._to_responses_input(messages) assert not [i for i in items if i.get("type") == "function_call"] assert not [i for i in items if i.get("type") == "function_call_output"] + + +# ── tool-calling fallback on the Responses path ────────────────────────────── + + +def _vllm_bad_request(): + """400 a tool-less OpenAI-compatible endpoint returns for a tools request.""" + import httpx + from openai import BadRequestError + + message = ( + '"auto" tool choice requires --enable-auto-tool-choice and ' + "--tool-call-parser to be set" + ) + request = httpx.Request("POST", "http://localhost:8000/v1/responses") + return BadRequestError( + message, response=httpx.Response(400, request=request), body={"message": message} + ) + + +class _RejectThenAccept: + def __init__(self, result): + self.result = result + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + if len(self.calls) == 1: + raise _vllm_bad_request() + return self.result + + +@pytest.mark.unit +def test_responses_gen_retries_without_tools(monkeypatch): + llm = _make_llm(monkeypatch, _responses_caps()) + response = _ns( + id="resp_degraded", + output=[_ns(type="message", content=[_ns(type="output_text", text="plain")])], + ) + create = _RejectThenAccept(response) + llm.client.responses.create = create + + choice = llm._responses_gen( + "local-model", + [{"role": "user", "content": "hi"}], + tools=[{"type": "function", "function": {"name": "t", "parameters": {}}}], + tool_choice="auto", + ) + + assert len(create.calls) == 2 + assert "tools" in create.calls[0] + assert "tools" not in create.calls[1] + assert "tool_choice" not in create.calls[1] + assert choice.message.content == "plain" + assert llm._supports_tools() is False + + +@pytest.mark.unit +def test_responses_gen_stream_retries_without_tools(monkeypatch): + llm = _make_llm(monkeypatch, _responses_caps()) + events = [ + _ns(type="response.output_text.delta", delta="plain"), + _ns(type="response.completed", response=_ns(id="resp_degraded")), + ] + create = _RejectThenAccept(events) + llm.client.responses.create = create + + chunks = list( + llm._responses_gen_stream( + "local-model", + [{"role": "user", "content": "hi"}], + tools=[{"type": "function", "function": {"name": "t", "parameters": {}}}], + ) + ) + + assert chunks == ["plain"] + assert len(create.calls) == 2 + assert "tools" in create.calls[0] + assert "tools" not in create.calls[1] diff --git a/tests/parser/file/test_docling_parser.py b/tests/parser/file/test_docling_parser.py index 83caaeef..28e612e1 100644 --- a/tests/parser/file/test_docling_parser.py +++ b/tests/parser/file/test_docling_parser.py @@ -259,7 +259,9 @@ class TestDoclingParserParseFile: mock_converter = MagicMock() mock_result = MagicMock() mock_doc = MagicMock() - mock_doc.export_to_markdown.return_value = "content" + # Long enough to clear the OCR chars-per-page floor, which a default + # (ocr_enabled=True) parser applies to PDFs. + mock_doc.export_to_markdown.return_value = "converted document content" mock_doc.texts = [] mock_result.document = mock_doc mock_converter.convert.return_value = mock_result @@ -268,7 +270,7 @@ class TestDoclingParserParseFile: parser._converter = mock_converter mock_init.return_value = {} result = parser.parse_file(Path("test.pdf")) - assert "content" in result + assert "converted document content" in result def test_parse_file_error_ignore_raises_instead_of_returning_error_text(self): """A failed conversion must never become the document's text. @@ -1011,3 +1013,352 @@ class TestApplyInferenceSettings: "torch.compile must be disabled before PdfPipelineOptions is " f"constructed, got {events}" ) + + +# ===================================================================== +# OCR dropout guard +# ===================================================================== + + +def _mock_conversion(markdown: str, pages: int = 1) -> MagicMock: + """Build a mock docling ConversionResult exporting ``markdown``.""" + document = MagicMock() + document.export_to_markdown.return_value = markdown + document.texts = [] + document.pages = {index: object() for index in range(pages)} + result = MagicMock() + result.document = document + return result + + +def _set_threshold(monkeypatch, value: int) -> None: + """Point the OCR dropout guard at a specific chars-per-page floor.""" + from application.core.settings import settings as real_settings + + class _Stub: + def __getattr__(self, name): + return getattr(real_settings, name) + + stub = _Stub() + stub.DOCLING_OCR_MIN_CHARS_PER_PAGE = value + monkeypatch.setattr("application.core.settings.settings", stub) + + +@pytest.mark.unit +class TestOCRDropoutGuard: + """A near-empty OCR parse must be loud, never stored as the document. + + In production, after a long scanned PDF, docling's pipeline degraded + inside the worker and returned zero characters for every later scanned + page — with no error. ``parse_file`` returned that empty string and + ingestion reported success, so the source was indexed as empty. + """ + + def test_near_empty_first_pass_recovers_on_fresh_converter(self, caplog): + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + degraded = MagicMock() + degraded.convert.return_value = _mock_conversion( + "\n\n\n\n", pages=3 + ) + parser._converter = degraded + + fresh = MagicMock() + fresh.convert.return_value = _mock_conversion( + "Recovered scanned text for page one, page two and page three.", + pages=3, + ) + forced_during_retry = [] + + def _create_converter(): + forced_during_retry.append(parser.force_full_page_ocr) + return fresh + + with patch.object( + parser, "_create_converter", side_effect=_create_converter + ) as create, caplog.at_level("WARNING"): + result = parser.parse_file(Path("scan.pdf")) + + assert "Recovered scanned text" in result + # Exactly one retry, on a converter that did not exist before. + assert create.call_count == 1 + assert fresh.convert.call_count == 1 + assert degraded.convert.call_count == 1 + # The retry forces full-page OCR... + assert forced_during_retry == [True] + # ...but that must not stick to the parser or its converter, or every + # later file in the worker pays for full-page OCR it does not need. + assert parser.force_full_page_ocr is False + assert parser._converter is None + messages = " ".join(caplog.messages) + assert "scan.pdf" in messages + assert "near-empty" in messages + assert "Recovered scan.pdf on retry" in messages + + def test_near_empty_both_passes_raises_document_parse_error(self): + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + degraded = MagicMock() + degraded.convert.return_value = _mock_conversion("", pages=100) + parser._converter = degraded + + fresh = MagicMock() + fresh.convert.return_value = _mock_conversion("", pages=100) + + with patch.object(parser, "_create_converter", return_value=fresh): + with pytest.raises(DocumentParseError) as excinfo: + parser.parse_file(Path("scan.pdf")) + + message = str(excinfo.value) + assert "scan.pdf" in message + assert "OCR produced 0 chars over 100 pages" in message + assert "OCR pipeline dropout" in message + assert "not indexed" in message + # The dropout diagnosis must survive: the generic handler must not + # re-wrap it into "Failed to parse ... with docling". + assert "Failed to parse" not in message + assert parser._converter is None + + def test_dropout_error_reports_an_unread_text_layer(self, monkeypatch): + """A text layer docling ignored points at the pipeline, not the scan.""" + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + monkeypatch.setattr( + "application.parser.file.docling_parser._pdf_text_layer_probe", + lambda file: (12, 48_000), + ) + parser = DoclingPDFParser(ocr_enabled=True) + parser._converter = MagicMock() + parser._converter.convert.return_value = _mock_conversion("", pages=12) + + with patch.object(parser, "_create_converter", return_value=parser._converter): + with pytest.raises(DocumentParseError) as excinfo: + parser.parse_file(Path("scan.pdf")) + + assert "The PDF carries a 48000-char text layer" in str(excinfo.value) + + def test_ocr_disabled_returns_near_empty_content_untouched(self): + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=False) + converter = MagicMock() + converter.convert.return_value = _mock_conversion("tiny", pages=40) + parser._converter = converter + + with patch.object(parser, "_create_converter") as create: + assert parser.parse_file(Path("text.pdf")) == "tiny" + + assert create.call_count == 0 + assert converter.convert.call_count == 1 + assert parser._converter is converter + + def test_threshold_zero_disables_the_guard(self, monkeypatch): + from application.parser.file.docling_parser import DoclingPDFParser + + _set_threshold(monkeypatch, 0) + parser = DoclingPDFParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion("", pages=100) + parser._converter = converter + + with patch.object(parser, "_create_converter") as create: + assert parser.parse_file(Path("scan.pdf")) == "" + + assert create.call_count == 0 + assert parser._converter is converter + + def test_healthy_ocr_parse_converts_once(self): + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion( + "A page of genuinely scanned text. " * 10, pages=2 + ) + parser._converter = converter + + with patch.object(parser, "_create_converter") as create: + result = parser.parse_file(Path("scan.pdf")) + + assert "genuinely scanned text" in result + assert converter.convert.call_count == 1 + assert create.call_count == 0 + assert parser._converter is converter + + def test_threshold_is_per_page_not_per_document(self, monkeypatch): + """200 chars is healthy for one page and a dropout for a hundred.""" + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + content = "x" * 200 + + parser = DoclingPDFParser(ocr_enabled=True) + parser._converter = MagicMock() + parser._converter.convert.return_value = _mock_conversion(content, pages=1) + with patch.object(parser, "_create_converter") as create: + assert parser.parse_file(Path("one-page.pdf")) == content + assert create.call_count == 0 + + long_doc = DoclingPDFParser(ocr_enabled=True) + long_doc._converter = MagicMock() + long_doc._converter.convert.return_value = _mock_conversion(content, pages=100) + with patch.object( + long_doc, "_create_converter", return_value=long_doc._converter + ): + with pytest.raises(DocumentParseError): + long_doc.parse_file(Path("hundred-page.pdf")) + + def test_guard_does_not_apply_to_non_ocr_formats(self): + """DOCX parsers inherit ocr_enabled=True but never OCR anything.""" + from application.parser.file.docling_parser import DoclingDocxParser + + parser = DoclingDocxParser() + assert parser.ocr_enabled is True + converter = MagicMock() + converter.convert.return_value = _mock_conversion("Hi", pages=1) + parser._converter = converter + + with patch.object(parser, "_create_converter") as create: + assert parser.parse_file(Path("memo.docx")) == "Hi" + + assert create.call_count == 0 + + def test_guard_applies_to_images(self): + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingImageParser + + parser = DoclingImageParser(ocr_enabled=True) + converter = MagicMock() + converter.convert.return_value = _mock_conversion("") + parser._converter = converter + + with patch.object(parser, "_create_converter", return_value=converter): + with pytest.raises(DocumentParseError, match="scan.png"): + parser.parse_file(Path("scan.png")) + + def test_retry_conversion_failure_is_still_a_parse_error(self): + from application.parser.file.base_parser import DocumentParseError + from application.parser.file.docling_parser import DoclingPDFParser + + parser = DoclingPDFParser(ocr_enabled=True) + parser._converter = MagicMock() + parser._converter.convert.return_value = _mock_conversion("", pages=2) + + fresh = MagicMock() + fresh.convert.side_effect = RuntimeError("OCR model unavailable") + + with patch.object(parser, "_create_converter", return_value=fresh): + with pytest.raises(DocumentParseError, match="OCR model unavailable"): + parser.parse_file(Path("scan.pdf")) + + assert parser.force_full_page_ocr is False + assert parser._converter is None + + +@pytest.mark.unit +class TestOCRDropoutHelpers: + + def test_image_placeholders_are_not_text(self): + from application.parser.file.docling_parser import _text_char_count + + assert _text_char_count(None) == 0 + assert _text_char_count("") == 0 + assert _text_char_count("\n\n") == 0 + assert _text_char_count(" abc ") == 3 + + def test_page_count_prefers_document_pages(self): + from application.parser.file.docling_parser import _result_page_count + + result = _mock_conversion("x", pages=7) + assert _result_page_count(result, Path("a.pdf")) == 7 + + def test_page_count_falls_back_to_num_pages(self): + from application.parser.file.docling_parser import _result_page_count + + result = _mock_conversion("x", pages=0) + result.document.num_pages.return_value = 5 + assert _result_page_count(result, Path("a.pdf")) == 5 + + def test_page_count_falls_back_to_result_pages(self): + from application.parser.file.docling_parser import _result_page_count + + result = _mock_conversion("x", pages=0) + result.document.num_pages.return_value = None + result.pages = [object(), object()] + assert _result_page_count(result, Path("a.pdf")) == 2 + + def test_page_count_defaults_to_one(self): + from application.parser.file.docling_parser import _result_page_count + + result = _mock_conversion("x", pages=0) + result.document.num_pages.return_value = None + result.pages = [] + assert _result_page_count(result, Path("a.pdf")) == 1 + + def test_images_are_always_one_page(self): + from application.parser.file.docling_parser import _result_page_count + + result = _mock_conversion("x", pages=9) + assert _result_page_count(result, Path("a.png")) == 1 + + def test_threshold_reads_settings_with_a_default(self, monkeypatch): + from application.parser.file.docling_parser import _ocr_min_chars_per_page + + assert _ocr_min_chars_per_page() == 20 + _set_threshold(monkeypatch, 5) + assert _ocr_min_chars_per_page() == 5 + _set_threshold(monkeypatch, "nonsense") + assert _ocr_min_chars_per_page() == 20 + + def test_text_layer_probe_survives_an_unreadable_pdf(self, tmp_path): + from application.parser.file.docling_parser import _pdf_text_layer_probe + + assert _pdf_text_layer_probe(tmp_path / "missing.pdf") == (0, 0) + broken = tmp_path / "broken.pdf" + broken.write_bytes(b"not a pdf") + assert _pdf_text_layer_probe(broken) == (0, 0) + + +@pytest.mark.unit +class TestForceFullPageOCRWiring: + """``force_full_page_ocr`` must reach the pipeline for any OCR engine.""" + + def _pipeline_options(self, monkeypatch, **kwargs): + import docling.datamodel.pipeline_options as dpo + + from application.parser.file.docling_parser import DoclingParser + + built = [] + real_options = dpo.PdfPipelineOptions + + def _tracking_options(*args, **opts): + options = real_options(*args, **opts) + built.append(options) + return options + + monkeypatch.setattr(dpo, "PdfPipelineOptions", _tracking_options) + monkeypatch.setattr("docling.document_converter.DocumentConverter", MagicMock()) + DoclingParser(**kwargs)._create_converter() + return built[0] + + def test_forced_without_rapidocr(self, monkeypatch): + options = self._pipeline_options( + monkeypatch, + ocr_enabled=True, + use_rapidocr=False, + force_full_page_ocr=True, + ) + assert options.ocr_options.force_full_page_ocr is True + + def test_not_forced_by_default(self, monkeypatch): + options = self._pipeline_options( + monkeypatch, + ocr_enabled=True, + use_rapidocr=False, + force_full_page_ocr=False, + ) + assert options.ocr_options.force_full_page_ocr is False diff --git a/tests/parser/test_chunking_strategies.py b/tests/parser/test_chunking_strategies.py index 62973a1d..a1847075 100644 --- a/tests/parser/test_chunking_strategies.py +++ b/tests/parser/test_chunking_strategies.py @@ -140,6 +140,13 @@ class _FakeEmbeddings: @pytest.mark.unit class TestSemantic: + @pytest.fixture(autouse=True) + def _no_remote_embeddings(self, monkeypatch): + """The resolver short-circuits to the remote API when this is set.""" + from application.core.settings import settings + + monkeypatch.setattr(settings, "EMBEDDINGS_BASE_URL", None) + def test_breakpoint_forces_split(self): # Two topics: sentences 0-1 vs 2-3, orthogonal embeddings between. text = "Alpha one. Alpha two. Beta one. Beta two." @@ -250,3 +257,25 @@ class TestClassicByteIdentical: assert [(c.doc_id, c.text, c.extra_info) for c in via] == [ (c.doc_id, c.text, c.extra_info) for c in direct ] + + +@pytest.mark.unit +class TestSemanticEmbeddingsResolution: + def test_uses_shared_resolver(self): + """Semantic chunking must resolve embeddings through ``get_embeddings``. + + Calling the singleton directly with the configured name skips the + bundled local model and downloads a second copy from the hub. + """ + vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]] + text = "Alpha one. Alpha two. Beta one. Beta two." + chunker = SemanticChunker(max_tokens=2000, min_tokens=0) + + with patch( + "application.vectorstore.base.get_embeddings", + return_value=_FakeEmbeddings(vectors), + ) as mock_resolver: + out = chunker.chunk([Document(text=text, doc_id="d")]) + + mock_resolver.assert_called_once_with() + assert len(out) == 2 diff --git a/tests/retriever/test_fanout.py b/tests/retriever/test_fanout.py new file mode 100644 index 00000000..b9b5d0be --- /dev/null +++ b/tests/retriever/test_fanout.py @@ -0,0 +1,176 @@ +"""Unit tests for the shared multi-source fan-out helpers. + +``application/retriever/fanout.py`` holds the pieces both ClassicRAG and the +search service reuse: the worker cap, the single query embedding, and the +order-preserving pool runner. These tests pin them directly, independent of +either caller. +""" + +import threading +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from application.retriever.fanout import ( + DEFAULT_MAX_PARALLEL_SOURCES, + embed_questions, + max_parallel_sources, + run_source_jobs, + store_embeddings, +) + + +@pytest.mark.unit +class TestMaxParallelSources: + def _cap(self, n_sources, configured=None): + stub = SimpleNamespace() + if configured is not None: + stub.RETRIEVAL_MAX_PARALLEL_SOURCES = configured + return max_parallel_sources(n_sources, stub) + + def test_default_when_setting_missing(self): + assert self._cap(10) == DEFAULT_MAX_PARALLEL_SOURCES + + def test_never_more_workers_than_sources(self): + assert self._cap(2) == 2 + + def test_setting_is_honoured(self): + assert self._cap(10, configured=8) == 8 + + def test_floor_is_one(self): + assert self._cap(10, configured=0) == 1 + assert self._cap(10, configured=-5) == 1 + + def test_garbage_setting_falls_back_to_default(self): + assert self._cap(10, configured="lots") == DEFAULT_MAX_PARALLEL_SOURCES + assert self._cap(10, configured=None) == DEFAULT_MAX_PARALLEL_SOURCES + + def test_zero_sources_still_returns_one(self): + assert self._cap(0) == 1 + + def test_falls_back_to_module_settings(self, monkeypatch): + import application.retriever.fanout as fanout + + monkeypatch.setattr( + fanout, "settings", SimpleNamespace(RETRIEVAL_MAX_PARALLEL_SOURCES=3) + ) + assert fanout.max_parallel_sources(10) == 3 + + +@pytest.mark.unit +class TestRunSourceJobs: + def test_empty_jobs_returns_empty(self): + assert run_source_jobs(Mock(side_effect=AssertionError), []) == [] + + def test_preserves_job_order(self): + import time + + delays = {"a": 0.06, "b": 0.0, "c": 0.0} + + def _fn(job): + time.sleep(delays[job]) + return job.upper() + + assert run_source_jobs(_fn, ["a", "b", "c"], workers=3) == ["A", "B", "C"] + + def test_jobs_run_concurrently(self): + """The barrier only clears if three jobs overlap; a serial loop hangs.""" + barrier = threading.Barrier(3, timeout=10) + + def _fn(job): + barrier.wait() + return job + + assert run_source_jobs(_fn, [1, 2, 3], workers=3) == [1, 2, 3] + + def test_single_job_runs_inline(self): + main_thread = threading.current_thread().name + seen = [] + + def _fn(job): + seen.append(threading.current_thread().name) + return job + + assert run_source_jobs(_fn, ["only"]) == ["only"] + assert seen == [main_thread] + + def test_one_worker_runs_inline(self): + main_thread = threading.current_thread().name + seen = [] + + def _fn(job): + seen.append(threading.current_thread().name) + return job + + run_source_jobs(_fn, ["a", "b"], workers=1) + assert seen == [main_thread, main_thread] + + def test_worker_count_defaults_to_the_cap(self): + with patch( + "application.retriever.fanout.max_parallel_sources", return_value=1 + ) as cap: + run_source_jobs(lambda job: job, ["a", "b", "c"]) + cap.assert_called_once_with(3) + + def test_accepts_any_iterable_of_jobs(self): + assert run_source_jobs(lambda job: job * 2, iter([1, 2])) == [2, 4] + + +@pytest.mark.unit +class TestStoreEmbeddings: + def test_finds_private_embedding_attribute(self): + embedder = Mock() + embedder.embed_query = Mock(return_value=[1.0]) + store = SimpleNamespace(_embedding=embedder) + assert store_embeddings(store) is embedder + + def test_probes_attributes_in_order(self): + first = Mock() + first.embed_query = Mock() + second = Mock() + second.embed_query = Mock() + store = SimpleNamespace(_embedding=first, _embeddings=second) + assert store_embeddings(store) is first + + def test_skips_attributes_without_embed_query(self): + embedder = Mock(spec=["embed_query"]) + store = SimpleNamespace(_embedding=object(), embeddings=embedder) + assert store_embeddings(store) is embedder + + def test_falls_back_to_get_embeddings(self): + embedder = Mock() + store = SimpleNamespace( + _embedding=None, _get_embeddings=Mock(return_value=embedder) + ) + assert store_embeddings(store) is embedder + + def test_returns_none_when_nothing_is_reachable(self): + store = SimpleNamespace(_embedding=None) + assert store_embeddings(store) is None + + def test_get_embeddings_failure_returns_none(self): + store = SimpleNamespace( + _embedding=None, _get_embeddings=Mock(side_effect=RuntimeError("boom")) + ) + assert store_embeddings(store) is None + + +@pytest.mark.unit +class TestEmbedQuestions: + def test_embeds_each_distinct_question_once(self): + embedder = Mock() + embedder.embed_query = Mock(side_effect=lambda q: [len(q)]) + store = SimpleNamespace(_embedding=embedder) + + assert embed_questions(store, ["ab", "cde"]) == {"ab": [2], "cde": [3]} + assert embedder.embed_query.call_count == 2 + + def test_no_embedder_returns_empty_map(self): + assert embed_questions(SimpleNamespace(_embedding=None), ["q"]) == {} + + def test_embedding_failure_returns_empty_map(self): + embedder = Mock() + embedder.embed_query = Mock(side_effect=RuntimeError("model down")) + store = SimpleNamespace(_embedding=embedder) + assert embed_questions(store, ["q"]) == {} diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py index 294eb524..9e585c79 100644 --- a/tests/retriever/test_graph_rag.py +++ b/tests/retriever/test_graph_rag.py @@ -56,7 +56,7 @@ class TestGraphRAGFallback: self, _avail, mock_store_cls, _patch_llm_creator ): store = MagicMock() - store.count_nodes.return_value = 0 + store.count_nodes_many.return_value = {"src1": 0} mock_store_cls.return_value = store rag = _make_retriever() @@ -68,6 +68,7 @@ class TestGraphRAGFallback: assert docs == classic_docs store.search_nodes_by_embedding.assert_not_called() store.get_subgraph.assert_not_called() + store.close.assert_called_once() @patch("application.retriever.graph_rag.GraphStore") @patch("application.retriever.graph_rag.graphrag_available", return_value=False) @@ -101,6 +102,9 @@ def _store_with_graph( ): store = MagicMock() store.count_nodes.return_value = len(nodes) + store.count_nodes_many.side_effect = lambda ids: { + source_id: len(nodes) for source_id in ids + } store.search_nodes_by_embedding.return_value = seed_rows store.get_subgraph.return_value = {"nodes": nodes, "edges": edges} store.get_chunk_ids_for_nodes.return_value = node_chunks @@ -166,7 +170,7 @@ class TestGraphRAGHappyPath: rag = _make_retriever(chunks=2) # Call the PPR path directly: _get_data would swallow a raise and fall # back to ClassicRAG, hiding the regression. - docs = rag._graph_docs_for_source(store, "src1") + docs = rag._graph_docs_for_source(store, "src1", [0.1, 0.2, 0.3]) assert len(docs) >= 1 @@ -278,7 +282,7 @@ class TestGraphRAGHappyPath: self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed ): store = _store_with_graph([], [], {}, {}, []) - store.count_nodes.return_value = 5 + store.count_nodes_many.side_effect = lambda ids: {s: 5 for s in ids} mock_store_cls.return_value = store rag = _make_retriever() @@ -449,3 +453,296 @@ class TestGraphRAGTopK: docs = rag._get_data() assert len(docs) == 2 + + +# ── Embeddings resolution ───────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestEmbedQueryResolution: + def test_embed_query_uses_shared_resolver(self): + """Query embedding must go through ``get_embeddings``. + + Only the resolver knows the bundled local-model path, so calling the + singleton directly loads a second copy of the model (or crashes on the + positional key). + """ + fake = Mock() + fake.embed_query.return_value = [0.1, 0.2, 0.3] + + with patch( + "application.retriever.graph_rag.get_embeddings", return_value=fake + ) as mock_resolver: + result = GraphRAGRetriever._embed_query(object(), "a question") + + mock_resolver.assert_called_once_with() + fake.embed_query.assert_called_once_with("a question") + assert result == [0.1, 0.2, 0.3] + + +# ── Batched retrieval across sources ───────────────────────────────────────── + + +def _multi_source_retriever(sources, **overrides): + """Retriever over several attached sources.""" + return _make_retriever( + source={"question": "q", "active_docs": list(sources)}, **overrides + ) + + +def _recording_classic(rag, docs): + """Stub ``ClassicRAG._get_data`` that records the sources it was handed.""" + seen = [] + + def _run(): + seen.append(list(rag._classic.vectorstores)) + return [dict(doc) for doc in docs] + + rag._classic._get_data = Mock(side_effect=_run) + return seen + + +def _single_node_store(counts): + """Graph store whose every source yields one chunk, with ``counts`` shape.""" + store = _store_with_graph( + [{"id": "n1", "doc_freq": 1}], + [], + {"n1": ["c1"]}, + {"c1": "graph text"}, + [{"id": "n1", "distance": 0.0}], + ) + store.count_nodes_many.side_effect = lambda ids: { + source_id: counts[source_id] for source_id in ids + } + return store + + +_CLASSIC_DOC = {"title": "cl", "text": "classic", "source": "a", "filename": "cl"} + + +class _SourceConfig: + """Minimal stand-in for the Dispatcher's per-source RetrievalConfig.""" + + def __init__(self, chunks: int): + self.chunks = chunks + + +@pytest.mark.unit +class TestGraphRAGBatching: + """N attached sources cost one count query and one classic run, not N of each.""" + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_node_counts_fetched_in_one_query( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = MagicMock() + store.count_nodes_many.return_value = {"a": 0, "b": 0, "c": 0} + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b", "c"]) + _recording_classic(rag, []) + + rag._get_data() + + store.count_nodes_many.assert_called_once_with(["a", "b", "c"]) + store.count_nodes.assert_not_called() + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_graphless_sources_share_one_classic_call( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + store = _single_node_store({"a": 0, "b": 3, "c": 0}) + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b", "c"], chunks=3) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + docs = rag._get_data() + + assert rag._classic._get_data.call_count == 1 + assert seen == [["a", "c"]] + # The classic batch occupies the slot of the first graphless source, so + # the graph source's docs still follow it in attachment order. + assert [doc["text"] for doc in docs] == ["classic", "graph text"] + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_only_the_batched_sources_keep_their_overrides( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + store = _single_node_store({"a": 0, "b": 3, "c": 0}) + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b", "c"], chunks=3) + configs = {sid: _SourceConfig(2) for sid in ("a", "b", "c")} + rag.per_source_retrieval = dict(configs) + captured = {} + + def _run(): + captured["overrides"] = dict(rag._classic.per_source_retrieval) + return [] + + rag._classic._get_data = Mock(side_effect=_run) + + rag._get_data() + + assert captured["overrides"] == {"a": configs["a"], "c": configs["c"]} + # Restored afterwards, exactly as the per-source path did. + assert rag._classic.per_source_retrieval == {} + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_query_is_embedded_once_for_several_graph_sources( + self, _avail, mock_store_cls, _tok, _patch_llm_creator + ): + store = _single_node_store({"a": 3, "b": 3}) + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b"], chunks=4) + rag._embed_query = Mock(return_value=[0.1, 0.2, 0.3]) + + docs = rag._get_data() + + assert rag._embed_query.call_count == 1 + assert store.search_nodes_by_embedding.call_count == 2 + # The one vector is what every source searches with. + for call in store.search_nodes_by_embedding.call_args_list: + assert call.args[1] == [0.1, 0.2, 0.3] + assert [doc["text"] for doc in docs] == ["graph text", "graph text"] + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_failed_graph_sources_land_in_one_batched_fallback( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + store = _single_node_store({"a": 3, "b": 3, "c": 3}) + + def _seed(source_id, embedding, k=10): + if source_id in ("b", "c"): + raise RuntimeError("graph exploded") + return [{"id": "n1", "distance": 0.0}] + + store.search_nodes_by_embedding.side_effect = _seed + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b", "c"], chunks=6) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + docs = rag._get_data() + + assert rag._classic._get_data.call_count == 1 + assert seen == [["b", "c"]] + # The retried batch is appended after the graph results. + assert [doc["text"] for doc in docs] == ["graph text", "classic"] + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_embedding_failure_falls_back_for_every_graph_source( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = _single_node_store({"a": 3, "b": 3}) + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b"]) + rag._embed_query = Mock(side_effect=RuntimeError("no embeddings")) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + docs = rag._get_data() + + assert seen == [["a", "b"]] + assert [doc["text"] for doc in docs] == ["classic"] + store.search_nodes_by_embedding.assert_not_called() + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_count_failure_falls_back_in_one_call( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = MagicMock() + store.count_nodes_many.side_effect = RuntimeError("no graph tables") + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a", "b"]) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + docs = rag._get_data() + + assert seen == [["a", "b"]] + assert [doc["text"] for doc in docs] == ["classic"] + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_unbuildable_store_falls_back_in_one_call( + self, _avail, mock_store_cls, _patch_llm_creator + ): + mock_store_cls.side_effect = RuntimeError("no connection string") + + rag = _multi_source_retriever(["a", "b"]) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + rag._get_data() + + assert seen == [["a", "b"]] + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=False) + def test_unavailable_graphrag_makes_one_batched_classic_call( + self, _avail, mock_store_cls, _patch_llm_creator + ): + rag = _multi_source_retriever(["a", "b", "c"]) + seen = _recording_classic(rag, [_CLASSIC_DOC]) + + rag._get_data() + + assert rag._classic._get_data.call_count == 1 + assert seen == [["a", "b", "c"]] + mock_store_cls.assert_not_called() + + @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10) + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_store_is_closed_on_the_success_path( + self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed + ): + store = _single_node_store({"a": 3}) + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a"], chunks=2) + rag._get_data() + + store.close.assert_called_once() + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_store_is_closed_when_retrieval_raises( + self, _avail, mock_store_cls, _patch_llm_creator + ): + store = MagicMock() + store.count_nodes_many.return_value = {"a": 0} + mock_store_cls.return_value = store + + rag = _multi_source_retriever(["a"]) + rag._classic._get_data = Mock(side_effect=RuntimeError("boom")) + + with pytest.raises(RuntimeError): + rag._get_data() + + store.close.assert_called_once() + + @patch("application.retriever.graph_rag.GraphStore") + @patch("application.retriever.graph_rag.graphrag_available", return_value=True) + def test_empty_source_list_never_builds_a_store( + self, _avail, mock_store_cls, _patch_llm_creator + ): + rag = _make_retriever(source={"question": "q", "active_docs": []}) + rag._classic._get_data = Mock(return_value=[]) + + assert rag._get_data() == [] + mock_store_cls.assert_not_called() + rag._classic._get_data.assert_not_called() diff --git a/tests/retriever/test_parallel_sources.py b/tests/retriever/test_parallel_sources.py new file mode 100644 index 00000000..86f3e8b1 --- /dev/null +++ b/tests/retriever/test_parallel_sources.py @@ -0,0 +1,374 @@ +"""Multi-source retrieval: one query embedding, concurrent per-source search. + +Six attached sources used to cost six query embeddings and six serial round +trips. These tests pin the two halves of the fix: the query is embedded once +and the vector is handed to every store, and the per-source searches run +concurrently while the merged output keeps the original source order. +""" + +import threading +import time +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from application.retriever.classic_rag import ClassicRAG +from application.retriever.hybrid_rag import HybridRetriever + + +@pytest.fixture +def _patch_llm_creator(mock_llm, monkeypatch): + monkeypatch.setattr( + "application.retriever.classic_rag.LLMCreator.create_llm", + Mock(return_value=mock_llm), + ) + return mock_llm + + +def _make_doc(page_content, source="s", title="t"): + doc = Mock() + doc.page_content = page_content + doc.metadata = {"title": title, "source": source} + return doc + + +def _make_embedder(vector=(0.1, 0.2, 0.3)): + embedder = Mock() + embedder.embed_query = Mock(return_value=list(vector)) + return embedder + + +def _make_store(embedder, hits=None, keyword_hits=None): + store = Mock() + store._embedding = embedder + store.score_kind = None + store.search.return_value = list(hits or []) + store.search_with_scores.return_value = [(doc, 0.5) for doc in (hits or [])] + store.keyword_search.return_value = list(keyword_hits or []) + return store + + +def _make_rag(cls=ClassicRAG, source=None, **overrides): + kwargs = dict( + source=source or {"question": "q", "active_docs": ["a", "b", "c"]}, + chat_history=None, + prompt="", + chunks=6, + doc_token_limit=50000, + model_id="test-model", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "user1"}, + ) + kwargs.update(overrides) + return cls(**kwargs) + + +def _run(rag, stores): + """Run ``_get_data`` with ``VectorCreator`` handing out ``stores`` by id.""" + with patch( + "application.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=lambda _type, source_id, _key: stores[source_id], + ), patch( + "application.retriever.classic_rag.num_tokens_from_string", return_value=10 + ): + return rag._get_data() + + +@pytest.mark.unit +class TestQueryEmbeddedOnce: + def test_three_sources_embed_the_query_once(self, _patch_llm_creator): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b", "c") + } + + docs = _run(_make_rag(), stores) + + assert len(docs) == 3 + assert embedder.embed_query.call_count == 1 + embedder.embed_query.assert_called_once_with("q") + + def test_every_store_receives_the_same_vector(self, _patch_llm_creator): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b", "c") + } + + _run(_make_rag(), stores) + + vectors = [ + store.search.call_args.kwargs["query_vector"] for store in stores.values() + ] + assert vectors == [[0.1, 0.2, 0.3]] * 3 + + def test_include_scores_path_passes_the_vector_too(self, _patch_llm_creator): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b") + } + + _run(_make_rag(include_scores=True, source={ + "question": "q", "active_docs": ["a", "b"] + }), stores) + + assert embedder.embed_query.call_count == 1 + for store in stores.values(): + kwargs = store.search_with_scores.call_args.kwargs + assert kwargs["query_vector"] == [0.1, 0.2, 0.3] + + def test_hybrid_reuses_the_shared_vector(self, _patch_llm_creator): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b") + } + + _run( + _make_rag( + cls=HybridRetriever, + source={"question": "q", "active_docs": ["a", "b"]}, + ), + stores, + ) + + assert embedder.embed_query.call_count == 1 + for store in stores.values(): + assert store.search.call_args.kwargs["query_vector"] == [0.1, 0.2, 0.3] + + def test_distinct_questions_are_embedded_once_each( + self, _patch_llm_creator, mock_llm + ): + """Per-source rephrase means two queries — one embedding each, no more.""" + from application.storage.db.source_config import RetrievalConfig + + mock_llm.gen = Mock(return_value="REPHRASED") + embedder = Mock() + embedder.embed_query = Mock(side_effect=lambda q: [len(q)]) + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b", "c") + } + + rag = _make_rag( + source={"question": "original", "active_docs": ["a", "b", "c"]}, + chat_history=[{"prompt": "hi", "response": "yo"}], + defer_rephrase=True, + ) + rag.per_source_retrieval = { + "a": RetrievalConfig(rephrase_query=False), + "b": RetrievalConfig(rephrase_query=True), + "c": RetrievalConfig(rephrase_query=True), + } + _run(rag, stores) + + assert embedder.embed_query.call_count == 2 + assert stores["a"].search.call_args.kwargs["query_vector"] == [len("original")] + assert stores["b"].search.call_args.kwargs["query_vector"] == [len("REPHRASED")] + assert stores["c"].search.call_args.kwargs["query_vector"] == [len("REPHRASED")] + + def test_embedding_failure_falls_back_to_per_store_embedding( + self, _patch_llm_creator + ): + """A broken embedder must not break retrieval — stores embed their own.""" + embedder = Mock() + embedder.embed_query = Mock(side_effect=RuntimeError("model down")) + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b") + } + + docs = _run( + _make_rag(source={"question": "q", "active_docs": ["a", "b"]}), stores + ) + + assert len(docs) == 2 + for store in stores.values(): + assert "query_vector" not in store.search.call_args.kwargs + + def test_store_without_embedder_still_searches(self, _patch_llm_creator): + """A store exposing no embeddings object keeps the old call shape.""" + store = _make_store(None, [_make_doc("hit")]) + store._embedding = None + store._embeddings = None + store.embeddings = None + store._get_embeddings = Mock(side_effect=AttributeError("nope")) + + docs = _run( + _make_rag(source={"question": "q", "active_docs": ["a"]}), {"a": store} + ) + + assert len(docs) == 1 + assert "query_vector" not in store.search.call_args.kwargs + + +@pytest.mark.unit +class TestConcurrentSourceSearch: + def test_sources_are_searched_concurrently(self, _patch_llm_creator): + """Three sources must be in flight at once, not one after another. + + The barrier only clears if three searches overlap; a serial loop + would trip its timeout and lose every source. + """ + barrier = threading.Barrier(3, timeout=10) + embedder = _make_embedder() + + def _search_factory(src): + def _search(*_args, **_kwargs): + barrier.wait() + return [_make_doc(f"hit {src}", source=src)] + + return _search + + stores = {} + for src in ("a", "b", "c"): + store = _make_store(embedder) + store.search.side_effect = _search_factory(src) + stores[src] = store + + docs = _run(_make_rag(), stores) + + assert [d["text"] for d in docs] == ["hit a", "hit b", "hit c"] + + def test_results_keep_source_order_regardless_of_completion_order( + self, _patch_llm_creator + ): + """The slowest source is first — it still leads the merged output.""" + embedder = _make_embedder() + delays = {"a": 0.06, "b": 0.0, "c": 0.0} + + def _search_factory(src): + def _search(*_args, **_kwargs): + time.sleep(delays[src]) + return [_make_doc(f"hit {src}", source=src)] + + return _search + + stores = {} + for src in ("a", "b", "c"): + store = _make_store(embedder) + store.search.side_effect = _search_factory(src) + stores[src] = store + + docs = _run(_make_rag(), stores) + + assert [d["text"] for d in docs] == ["hit a", "hit b", "hit c"] + assert [d["source"] for d in docs] == ["a", "b", "c"] + + def test_output_matches_the_serial_path_exactly(self, _patch_llm_creator): + """Parity: one worker (serial) and four workers agree, doc for doc.""" + + def _stores(): + embedder = _make_embedder() + return { + src: _make_store( + embedder, + [ + _make_doc(f"{src}-0", source=src), + _make_doc(f"{src}-1", source=src), + ], + ) + for src in ("a", "b", "c") + } + + with patch( + "application.retriever.classic_rag._max_parallel_sources", + side_effect=lambda n: 1, + ): + serial = _run(_make_rag(), _stores()) + with patch( + "application.retriever.classic_rag._max_parallel_sources", + side_effect=lambda n: 4, + ): + parallel = _run(_make_rag(), _stores()) + + assert serial == parallel + assert len(serial) == 6 + + def test_one_failing_source_does_not_kill_the_others(self, _patch_llm_creator): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b", "c") + } + stores["b"].search.side_effect = RuntimeError("store b is down") + + docs = _run(_make_rag(), stores) + + assert [d["text"] for d in docs] == ["hit a", "hit c"] + + def test_failing_store_construction_does_not_kill_the_others( + self, _patch_llm_creator + ): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)]) + for src in ("a", "b", "c") + } + + def _create(_type, source_id, _key): + if source_id == "a": + raise RuntimeError("connection failed") + return stores[source_id] + + rag = _make_rag() + with patch( + "application.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=_create, + ), patch( + "application.retriever.classic_rag.num_tokens_from_string", return_value=10 + ): + docs = rag._get_data() + + assert [d["text"] for d in docs] == ["hit b", "hit c"] + + def test_single_source_skips_the_pool(self, _patch_llm_creator): + """One source must not pay for a thread pool.""" + embedder = _make_embedder() + store = _make_store(embedder, [_make_doc("only hit")]) + main_thread = threading.current_thread().name + seen = [] + store.search.side_effect = lambda *_a, **_k: ( + seen.append(threading.current_thread().name), + [_make_doc("only hit")], + )[1] + + docs = _run( + _make_rag(source={"question": "q", "active_docs": ["a"]}), {"a": store} + ) + + assert len(docs) == 1 + assert seen == [main_thread] + + +@pytest.mark.unit +class TestWorkerCap: + """``RETRIEVAL_MAX_PARALLEL_SOURCES`` bounds the fan-out.""" + + def _cap(self, monkeypatch, n_sources, configured=None): + import application.retriever.classic_rag as classic_rag + + stub = SimpleNamespace() + if configured is not None: + stub.RETRIEVAL_MAX_PARALLEL_SOURCES = configured + monkeypatch.setattr(classic_rag, "settings", stub) + return classic_rag._max_parallel_sources(n_sources) + + def test_default_is_four(self, monkeypatch): + assert self._cap(monkeypatch, 10) == 4 + + def test_never_more_workers_than_sources(self, monkeypatch): + assert self._cap(monkeypatch, 2) == 2 + + def test_setting_is_honoured(self, monkeypatch): + assert self._cap(monkeypatch, 10, configured=8) == 8 + + def test_floor_is_one(self, monkeypatch): + assert self._cap(monkeypatch, 10, configured=0) == 1 + + def test_garbage_setting_falls_back_to_default(self, monkeypatch): + assert self._cap(monkeypatch, 10, configured="lots") == 4 diff --git a/tests/services/test_search_service.py b/tests/services/test_search_service.py index 25eee97f..5d866474 100644 --- a/tests/services/test_search_service.py +++ b/tests/services/test_search_service.py @@ -287,3 +287,340 @@ class TestSourceAuthorization: monkeypatch.setattr(ts, "can_access", _boom) assert _authorized_source_ids(None, self._agent(), ["s1"]) == [] + + +def _serial_search_sources(query, source_ids, chunks): + """The pre-fan-out ``_search_sources``, kept as a parity oracle.""" + from application.services.search_service import VectorCreator, settings + + if chunks <= 0 or not source_ids: + return [] + + results = [] + chunks_per_source = max(1, chunks // len(source_ids)) + seen_texts = set() + + for source_id in source_ids: + if not source_id or not source_id.strip(): + 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 + 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", {}) + text_hash = hash(page_content[:200]) + if text_hash in seen_texts: + continue + seen_texts.add(text_hash) + title = metadata.get("title", metadata.get("post_title", "")) + if not isinstance(title, str): + title = str(title) if title else "" + if title: + title = title.split("/")[-1] + else: + title = metadata.get("filename", page_content[:50] + "...") + source = metadata.get("source", source_id) + results.append( + {"text": page_content, "title": title, "source": source} + ) + if len(results) >= chunks: + break + except Exception: + continue + + return results[:chunks] + + +def _make_embedder(vector=(0.1, 0.2, 0.3)): + embedder = MagicMock() + embedder.embed_query = MagicMock(return_value=list(vector)) + return embedder + + +def _make_store(embedder, docs=None): + store = MagicMock() + store._embedding = embedder + store.search.return_value = list(docs or []) + return store + + +def _no_embedder_store(docs=None): + """A store exposing no reachable embeddings object.""" + store = MagicMock() + store._embedding = None + store._embeddings = None + store.embeddings = None + store._get_embeddings = MagicMock(side_effect=AttributeError("nope")) + store.search.return_value = list(docs or []) + return store + + +def _run_search_sources(query, source_ids, chunks, stores, serial=False): + """Run ``_search_sources`` (or its serial oracle) against ``stores``.""" + from application.services import search_service + + impl = _serial_search_sources if serial else search_service._search_sources + with patch( + "application.services.search_service.VectorCreator.create_vectorstore", + side_effect=lambda _type, source_id, _key: stores[source_id], + ): + return impl(query, source_ids, chunks) + + +@pytest.mark.unit +class TestSearchSourcesFanOut: + """Multi-source search: one query embedding, one bounded fan-out.""" + + def test_query_is_embedded_once_across_three_sources(self): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b", "c") + } + + results = _run_search_sources("q", ["a", "b", "c"], 5, stores) + + assert len(results) == 3 + embedder.embed_query.assert_called_once_with("q") + + def test_every_store_receives_the_same_vector(self): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b", "c") + } + + _run_search_sources("q", ["a", "b", "c"], 5, stores) + + vectors = [ + store.search.call_args.kwargs["query_vector"] for store in stores.values() + ] + assert vectors == [[0.1, 0.2, 0.3]] * 3 + + def test_k_per_source_is_unchanged(self): + embedder = _make_embedder() + stores = {src: _make_store(embedder) for src in ("a", "b", "c")} + + _run_search_sources("q", ["a", "b", "c"], 6, stores) + + for store in stores.values(): + assert store.search.call_args.kwargs["k"] == 4 + + def test_output_matches_the_serial_implementation(self): + """Parity: order, dedupe and the ``chunks`` cap are byte-identical.""" + duplicate = "shared chunk " * 30 + + def _stores(): + embedder = _make_embedder() + return { + "a": _make_store( + embedder, + [ + {"text": "a-0", "metadata": {"title": "A0", "source": "/a0"}}, + {"text": duplicate, "metadata": {"title": "dup"}}, + {"text": "a-2", "metadata": {"filename": "a2.pdf"}}, + ], + ), + "b": _make_store( + embedder, + [ + {"text": duplicate, "metadata": {"title": "dup again"}}, + {"text": "b-1", "metadata": {}}, + ], + ), + "c": _make_store( + embedder, + [ + {"text": "c-0", "metadata": {"post_title": "x/y/C0"}}, + {"text": "c-1", "metadata": {"title": 42}}, + ], + ), + } + + for chunks in (1, 3, 5, 6, 12): + serial = _run_search_sources("q", ["a", "b", "c"], chunks, _stores(), serial=True) + parallel = _run_search_sources("q", ["a", "b", "c"], chunks, _stores()) + assert parallel == serial, f"mismatch at chunks={chunks}" + + def test_langchain_documents_match_the_serial_implementation(self): + def _stores(): + embedder = _make_embedder() + stores = {} + for src in ("a", "b"): + doc = MagicMock() + doc.page_content = f"content {src}" + doc.metadata = {"title": f"path/to/{src}", "source": f"/{src}"} + stores[src] = _make_store(embedder, [doc]) + return stores + + serial = _run_search_sources("q", ["a", "b"], 4, _stores(), serial=True) + parallel = _run_search_sources("q", ["a", "b"], 4, _stores()) + assert parallel == serial + assert [r["title"] for r in parallel] == ["a", "b"] + + def test_results_keep_source_order_when_a_slow_source_is_first(self): + import time + + embedder = _make_embedder() + delays = {"a": 0.06, "b": 0.0, "c": 0.0} + + def _factory(src): + def _search(*_args, **_kwargs): + time.sleep(delays[src]) + return [{"text": f"hit {src}", "metadata": {}}] + + return _search + + stores = {} + for src in ("a", "b", "c"): + store = _make_store(embedder) + store.search.side_effect = _factory(src) + stores[src] = store + + results = _run_search_sources("q", ["a", "b", "c"], 6, stores) + + assert [r["text"] for r in results] == ["hit a", "hit b", "hit c"] + + def test_sources_are_searched_concurrently(self): + """The barrier only clears if three searches overlap.""" + import threading + + barrier = threading.Barrier(3, timeout=10) + embedder = _make_embedder() + + def _factory(src): + def _search(*_args, **_kwargs): + barrier.wait() + return [{"text": f"hit {src}", "metadata": {}}] + + return _search + + stores = {} + for src in ("a", "b", "c"): + store = _make_store(embedder) + store.search.side_effect = _factory(src) + stores[src] = store + + results = _run_search_sources("q", ["a", "b", "c"], 6, stores) + + assert [r["text"] for r in results] == ["hit a", "hit b", "hit c"] + + def test_single_source_skips_the_pool(self): + import threading + + main_thread = threading.current_thread().name + seen = [] + embedder = _make_embedder() + store = _make_store(embedder) + store.search.side_effect = lambda *_a, **_k: ( + seen.append(threading.current_thread().name), + [{"text": "only", "metadata": {}}], + )[1] + + results = _run_search_sources("q", ["a"], 5, {"a": store}) + + assert [r["text"] for r in results] == ["only"] + assert seen == [main_thread] + + def test_one_failing_search_does_not_kill_the_others(self): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b", "c") + } + stores["b"].search.side_effect = RuntimeError("index missing") + + results = _run_search_sources("q", ["a", "b", "c"], 6, stores) + + assert [r["text"] for r in results] == ["hit a", "hit c"] + + def test_failing_first_store_does_not_kill_the_others(self): + embedder = _make_embedder() + stores = { + src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b", "c") + } + + def _create(_type, source_id, _key): + if source_id == "a": + raise RuntimeError("connection failed") + return stores[source_id] + + from application.services import search_service + + with patch( + "application.services.search_service.VectorCreator.create_vectorstore", + side_effect=_create, + ): + results = search_service._search_sources("q", ["a", "b", "c"], 6) + + assert [r["text"] for r in results] == ["hit b", "hit c"] + + def test_no_embedder_means_no_query_vector_kwarg(self): + stores = { + src: _no_embedder_store([{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b") + } + + results = _run_search_sources("q", ["a", "b"], 6, stores) + + assert [r["text"] for r in results] == ["hit a", "hit b"] + for store in stores.values(): + assert "query_vector" not in store.search.call_args.kwargs + + def test_embedding_failure_falls_back_to_per_store_embedding(self): + embedder = MagicMock() + embedder.embed_query = MagicMock(side_effect=RuntimeError("model down")) + stores = { + src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}]) + for src in ("a", "b") + } + + results = _run_search_sources("q", ["a", "b"], 6, stores) + + assert [r["text"] for r in results] == ["hit a", "hit b"] + for store in stores.values(): + assert "query_vector" not in store.search.call_args.kwargs + + def test_blank_source_ids_are_skipped_but_still_split_the_budget(self): + """Blank ids build no store, yet ``chunks_per_source`` counts them.""" + embedder = _make_embedder() + stores = {"a": _make_store(embedder, [{"text": "hit a", "metadata": {}}])} + + from application.services import search_service + + with patch( + "application.services.search_service.VectorCreator.create_vectorstore", + side_effect=lambda _t, source_id, _k: stores[source_id], + ): + results = search_service._search_sources("q", [" ", "a", ""], 6) + + assert [r["text"] for r in results] == ["hit a"] + assert stores["a"].search.call_args.kwargs["k"] == 4 + + def test_all_blank_source_ids_build_no_store(self): + from application.services import search_service + + with patch( + "application.services.search_service.VectorCreator.create_vectorstore" + ) as mock_create: + assert search_service._search_sources("q", [" ", ""], 5) == [] + mock_create.assert_not_called() + + def test_zero_chunks_short_circuits(self): + from application.services import search_service + + with patch( + "application.services.search_service.VectorCreator.create_vectorstore" + ) as mock_create: + assert search_service._search_sources("q", ["a"], 0) == [] + mock_create.assert_not_called() diff --git a/tests/storage/db/test_bootstrap_vector_schema.py b/tests/storage/db/test_bootstrap_vector_schema.py new file mode 100644 index 00000000..a4bf9fc9 --- /dev/null +++ b/tests/storage/db/test_bootstrap_vector_schema.py @@ -0,0 +1,200 @@ +"""Boot-time ownership of the pgvector schema. + +``PGVectorStore`` used to run its DDL on every instantiation — once per source +per request. The schema is created here instead, once per process, and the boot +hook is also the only place that can catch an embedding-dimension mismatch: +a ``documents`` table built for a different model silently retrieves garbage. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from application.storage.db.bootstrap import ensure_vector_schema + +_APP_PY = Path(__file__).resolve().parents[3] / "application" / "app.py" + + +@pytest.fixture +def vector_settings(monkeypatch): + """Settings configured for a pgvector deployment.""" + from application.core import settings as settings_module + + settings = settings_module.settings + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False) + monkeypatch.setattr( + settings, + "PGVECTOR_CONNECTION_STRING", + "postgresql://user:pass@localhost/db", + raising=False, + ) + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False) + monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "test-model", raising=False) + return settings + + +def _embeddings(dimension): + stub = MagicMock() + stub.dimension = dimension + return stub + + +@pytest.mark.unit +class TestEnsureVectorSchemaSkips: + def test_skips_when_the_store_is_not_pgvector(self, vector_settings, monkeypatch): + monkeypatch.setattr(vector_settings, "VECTOR_STORE", "faiss", raising=False) + + with patch("psycopg.connect") as connect: + ensure_vector_schema() + + connect.assert_not_called() + + def test_skips_when_no_connection_string_is_configured( + self, vector_settings, monkeypatch + ): + monkeypatch.setattr( + vector_settings, "PGVECTOR_CONNECTION_STRING", None, raising=False + ) + monkeypatch.setattr(vector_settings, "POSTGRES_URI", None, raising=False) + + with patch("psycopg.connect") as connect: + ensure_vector_schema() + + connect.assert_not_called() + + +@pytest.mark.unit +class TestEnsureVectorSchemaCreates: + def _run(self, dimension=768, table_dimension=768): + conn = MagicMock() + cursor = MagicMock() + conn.cursor.return_value = cursor + with patch("psycopg.connect", return_value=conn) as connect, patch( + "application.vectorstore.base.get_embeddings", + return_value=_embeddings(dimension), + ), patch( + "application.vectorstore.pgvector.PGVectorStore.create_schema" + ) as vector_schema, patch( + "application.vectorstore.pgvector.PGVectorStore.table_dimension", + return_value=table_dimension, + ), patch( + "application.graphrag.store.GraphStore.create_schema" + ) as graph_schema, patch( + "application.vectorstore.pgvector._pool_for" + ) as pool_for: + ensure_vector_schema() + return connect, conn, cursor, vector_schema, graph_schema, pool_for + + def test_creates_the_vector_schema_under_the_advisory_lock(self, vector_settings): + connect, conn, cursor, vector_schema, graph_schema, _ = self._run() + + connect.assert_called_once_with("postgresql://user:pass@localhost/db") + statements = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "pg_advisory_xact_lock" in statements + vector_schema.assert_called_once() + assert vector_schema.call_args.kwargs["dimension"] == 768 + graph_schema.assert_not_called() + conn.commit.assert_called_once() + conn.close.assert_called_once() + + def test_creates_the_graph_schema_only_when_graphrag_is_enabled( + self, vector_settings, monkeypatch + ): + monkeypatch.setattr(vector_settings, "GRAPHRAG_ENABLED", True, raising=False) + + _, _, _, _, graph_schema, _ = self._run() + + graph_schema.assert_called_once() + assert graph_schema.call_args.kwargs["dimension"] == 768 + + def test_uses_a_direct_connection_never_the_store_pool(self, vector_settings): + # The hook can run pre-fork under ``gunicorn --preload``; a pooled + # socket inherited by a worker is a corrupted connection. + _, _, _, _, _, pool_for = self._run() + + pool_for.assert_not_called() + + +@pytest.mark.unit +class TestEnsureVectorSchemaDimensionCheck: + def test_raises_when_the_table_width_disagrees_with_the_model( + self, vector_settings + ): + conn = MagicMock() + with patch("psycopg.connect", return_value=conn), patch( + "application.vectorstore.base.get_embeddings", + return_value=_embeddings(1536), + ), patch("application.vectorstore.pgvector.PGVectorStore.create_schema"), patch( + "application.vectorstore.pgvector.PGVectorStore.table_dimension", + return_value=768, + ): + with pytest.raises(RuntimeError) as excinfo: + ensure_vector_schema() + + message = str(excinfo.value) + assert "768" in message and "1536" in message + assert "test-model" in message + conn.close.assert_called_once() + + def test_skips_the_check_when_the_embeddings_expose_no_dimension( + self, vector_settings + ): + conn = MagicMock() + stub = MagicMock() + del stub.dimension + with patch("psycopg.connect", return_value=conn), patch( + "application.vectorstore.base.get_embeddings", return_value=stub + ), patch( + "application.vectorstore.pgvector.PGVectorStore.create_schema" + ) as vector_schema, patch( + "application.vectorstore.pgvector.PGVectorStore.table_dimension", + return_value=1536, + ): + ensure_vector_schema() # must not raise + + assert vector_schema.call_args.kwargs["dimension"] == 768 + + def test_skips_the_check_when_the_model_cannot_be_loaded(self, vector_settings): + conn = MagicMock() + with patch("psycopg.connect", return_value=conn), patch( + "application.vectorstore.base.get_embeddings", + side_effect=RuntimeError("no model"), + ), patch( + "application.vectorstore.pgvector.PGVectorStore.create_schema" + ) as vector_schema, patch( + "application.vectorstore.pgvector.PGVectorStore.table_dimension", + return_value=1536, + ): + ensure_vector_schema() # must not raise + + assert vector_schema.call_args.kwargs["dimension"] == 768 + + +@pytest.mark.unit +class TestBootGating: + """The boot hook must stay behind AUTO_VECTOR_SCHEMA. + + Asserted on the source because importing ``application.app`` runs the hook, + and a test that imports it cannot observe the gate it just executed. + """ + + def test_app_calls_the_hook_behind_the_setting(self): + source = _APP_PY.read_text() + + assert "if settings.AUTO_VECTOR_SCHEMA:" in source + assert "ensure_vector_schema(" in source + + def test_setting_defaults_on(self): + from application.core.settings import Settings + + assert Settings.model_fields["AUTO_VECTOR_SCHEMA"].default is True + + def test_test_suite_opts_out_by_default(self): + conftest = Path(__file__).resolve().parents[2] / "conftest.py" + + assert 'os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false")' in ( + conftest.read_text() + ) diff --git a/tests/test_cache.py b/tests/test_cache.py index 0f61eda1..9667e0c2 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -658,3 +658,254 @@ def test_gen_cache_key_handles_bytearray_and_memoryview(): # All three should hash the same content to the same key. assert gen_cache_key(msgs_ba, "m") == gen_cache_key(msgs_b, "m") assert gen_cache_key(msgs_mv, "m") == gen_cache_key(msgs_b, "m") + + +# ===================================================================== +# Generation kwargs are part of the cache key +# +# The decorators wrap ``_raw_gen``/``_raw_gen_stream``, whose extra kwargs +# carry ``response_format`` (OpenAI structured output) and +# ``response_schema`` (Google). Before these were hashed, changing a +# workflow node's JSON schema replayed the previous schema's answer for +# the whole 30-minute TTL. +# ===================================================================== + + +class _FakeRedis: + """Dict-backed stand-in for the redis client used by the decorators.""" + + def __init__(self): + self.store = {} + + def get(self, key): + return self.store.get(key) + + def set(self, key, value, ex=None): + self.store[key] = value.encode("utf-8") if isinstance(value, str) else value + + def delete(self, key): + self.store.pop(key, None) + + +_SCHEMA_A = { + "type": "json_schema", + "json_schema": {"name": "r", "schema": {"properties": {"a": {"type": "string"}}}}, +} +_SCHEMA_B = { + "type": "json_schema", + "json_schema": {"name": "r", "schema": {"properties": {"b": {"type": "number"}}}}, +} + + +@pytest.mark.unit +def test_gen_cache_key_differs_for_different_response_format(): + messages = [{"role": "user", "content": "test"}] + key_a = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_A}) + key_b = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_B}) + assert key_a != key_b + + +@pytest.mark.unit +def test_gen_cache_key_stable_for_same_response_format(): + messages = [{"role": "user", "content": "test"}] + key_a = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_A}) + key_b = gen_cache_key(messages, "m", None, extra={"response_format": dict(_SCHEMA_A)}) + assert key_a == key_b + + +@pytest.mark.unit +def test_gen_cache_key_response_format_differs_from_no_format(): + messages = [{"role": "user", "content": "test"}] + assert gen_cache_key(messages, "m") != gen_cache_key( + messages, "m", None, extra={"response_format": _SCHEMA_A} + ) + + +@pytest.mark.unit +def test_gen_cache_key_differs_for_different_response_schema(): + """Google's structured-output kwarg is keyed just like OpenAI's.""" + messages = [{"role": "user", "content": "test"}] + key_a = gen_cache_key(messages, "m", None, extra={"response_schema": _SCHEMA_A}) + key_b = gen_cache_key(messages, "m", None, extra={"response_schema": _SCHEMA_B}) + assert key_a != key_b + + +@pytest.mark.unit +def test_gen_cache_key_covers_other_generation_kwargs(): + messages = [{"role": "user", "content": "test"}] + assert gen_cache_key( + messages, "m", None, extra={"temperature": 0.1} + ) != gen_cache_key(messages, "m", None, extra={"temperature": 0.9}) + assert gen_cache_key( + messages, "m", None, extra={"reasoning_effort": "low"} + ) != gen_cache_key(messages, "m", None, extra={"reasoning_effort": "high"}) + + +@pytest.mark.unit +def test_gen_cache_key_ignores_usage_attachments(): + """``_usage_attachments`` is a token-accounting side channel that never + reaches the provider — and the gen/stream decorator stacks disagree on + whether it is still in kwargs — so it must not move the key.""" + messages = [{"role": "user", "content": "test"}] + plain = gen_cache_key(messages, "m") + with_attachments = gen_cache_key( + messages, "m", None, extra={"_usage_attachments": [{"id": "att1"}]} + ) + assert plain == with_attachments + + +@pytest.mark.unit +def test_gen_cache_key_ignores_none_valued_kwargs(): + """``response_format=None`` is the default, not a distinct request.""" + messages = [{"role": "user", "content": "test"}] + assert gen_cache_key(messages, "m") == gen_cache_key( + messages, "m", None, extra={"response_format": None, "response_schema": None} + ) + + +@pytest.mark.unit +def test_gen_cache_key_extra_ordering_is_irrelevant(): + messages = [{"role": "user", "content": "test"}] + first = gen_cache_key(messages, "m", None, extra={"a": 1, "b": 2}) + second = gen_cache_key(messages, "m", None, extra={"b": 2, "a": 1}) + assert first == second + + +@pytest.mark.unit +def test_gen_cache_key_unserializable_extra_raises_value_error(): + """A key we cannot compute must raise so the decorators bypass the + cache rather than reuse a wrong entry.""" + messages = [{"role": "user", "content": "test"}] + circular = {} + circular["self"] = circular + with pytest.raises(ValueError): + gen_cache_key(messages, "m", None, extra={"response_format": circular}) + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_gen_cache_does_not_serve_entry_from_other_response_format(mock_make_redis): + fake = _FakeRedis() + mock_make_redis.return_value = fake + calls = [] + + @gen_cache + def mock_function(self, model, messages, stream, tools, **kwargs): + calls.append(kwargs.get("response_format")) + return f"answer-{len(calls)}" + + messages = [{"role": "user", "content": "test"}] + first = mock_function( + None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_A + ) + second = mock_function( + None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_B + ) + cached = mock_function( + None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_A + ) + + assert first == "answer-1" + assert second == "answer-2" + assert cached == "answer-1" + assert calls == [_SCHEMA_A, _SCHEMA_B] + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_stream_cache_does_not_replay_entry_from_other_response_format(mock_make_redis): + """The reported bug: a workflow node whose schema changed replayed the + old schema's cached stream for the rest of the TTL.""" + fake = _FakeRedis() + mock_make_redis.return_value = fake + calls = [] + + @stream_cache + def mock_function(self, model, messages, stream, tools, **kwargs): + calls.append(kwargs.get("response_format")) + yield f"chunk-{len(calls)}" + + messages = [{"role": "user", "content": "test"}] + first = list( + mock_function( + None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_A + ) + ) + second = list( + mock_function( + None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_B + ) + ) + replay = list( + mock_function( + None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_A + ) + ) + + assert first == ["chunk-1"] + assert second == ["chunk-2"] + assert replay == ["chunk-1"] + assert calls == [_SCHEMA_A, _SCHEMA_B] + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_gen_cache_bypassed_for_previous_response_id(mock_make_redis): + """A Responses API turn chained to a server-held id depends on state no + key can capture, so it must not read or write the cache.""" + fake = MagicMock() + mock_make_redis.return_value = fake + + @gen_cache + def mock_function(self, model, messages, stream, tools, **kwargs): + return "fresh" + + messages = [{"role": "user", "content": "test"}] + result = mock_function( + None, "m", messages, stream=False, tools=None, previous_response_id="resp_1" + ) + + assert result == "fresh" + fake.get.assert_not_called() + fake.set.assert_not_called() + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_stream_cache_bypassed_for_previous_response_id(mock_make_redis): + fake = MagicMock() + mock_make_redis.return_value = fake + + @stream_cache + def mock_function(self, model, messages, stream, tools, **kwargs): + yield "fresh" + + messages = [{"role": "user", "content": "test"}] + result = list( + mock_function( + None, "m", messages, stream=True, tools=None, previous_response_id="resp_1" + ) + ) + + assert result == ["fresh"] + fake.get.assert_not_called() + fake.set.assert_not_called() + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_gen_cache_ignores_stream_payload_stored_under_same_key(mock_make_redis): + """``gen_cache`` and ``stream_cache`` share a key space; a stream + envelope must never be handed back as a non-streaming answer.""" + fake = _FakeRedis() + mock_make_redis.return_value = fake + messages = [{"role": "user", "content": "test"}] + key = gen_cache_key(messages, "m", None) + fake.set(key, json.dumps({"version": 1, "chunks": ["a", "b"]})) + + @gen_cache + def mock_function(self, model, messages, stream, tools, **kwargs): + return "real answer" + + result = mock_function(None, "m", messages, stream=False, tools=None) + assert result == "real answer" diff --git a/tests/vectorstore/test_base.py b/tests/vectorstore/test_base.py index 12d3ef79..a727d07a 100644 --- a/tests/vectorstore/test_base.py +++ b/tests/vectorstore/test_base.py @@ -6,8 +6,12 @@ from application.vectorstore.base import ( BaseVectorStore, EmbeddingsSingleton, RemoteEmbeddings, + get_embeddings, ) +HF_MPNET = "huggingface_sentence-transformers/all-mpnet-base-v2" +LOCAL_MPNET = "/app/models/all-mpnet-base-v2" + # --- RemoteEmbeddings --- @@ -254,6 +258,46 @@ class TestEmbeddingsSingleton: assert result.headers["Authorization"] == "Bearer sk-from-settings" + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base._get_embeddings_wrapper") + def test_get_instance_hf_ignores_positional_key( + self, mock_get_wrapper, mock_settings + ): + """A stray key must not reach the zero-arg HuggingFace factory. + + The factories are ``lambda: EmbeddingsWrapper(...)``, so a caller that + passed ``settings.EMBEDDINGS_KEY`` positionally used to blow up with + ``TypeError: () takes 0 positional arguments``. + """ + mock_settings.EMBEDDINGS_BASE_URL = None + mock_wrapper_cls = Mock() + mock_instance = Mock() + mock_wrapper_cls.return_value = mock_instance + mock_get_wrapper.return_value = mock_wrapper_cls + + result = EmbeddingsSingleton.get_instance(HF_MPNET, None) + + assert result is mock_instance + mock_wrapper_cls.assert_called_once_with( + "sentence-transformers/all-mpnet-base-v2" + ) + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base._get_embeddings_wrapper") + def test_get_instance_hf_ignores_keyword_args( + self, mock_get_wrapper, mock_settings + ): + mock_settings.EMBEDDINGS_BASE_URL = None + mock_wrapper_cls = Mock() + mock_get_wrapper.return_value = mock_wrapper_cls + + EmbeddingsSingleton.get_instance(HF_MPNET, openai_api_key="sk-nope") + + mock_wrapper_cls.assert_called_once_with( + "sentence-transformers/all-mpnet-base-v2" + ) + + # --- BaseVectorStore --- @@ -415,3 +459,159 @@ class TestSearchWithScoresDefault: return [] assert _Store().search_with_scores("q") == [] + + +# --- get_embeddings (the single resolver) --- + + +@pytest.mark.unit +class TestGetEmbeddingsResolver: + """``get_embeddings`` is the one entry point every caller must use. + + Calling ``EmbeddingsSingleton.get_instance`` directly reproduces neither the + bundled local-model path nor the OpenAI/Azure key handling. + """ + + def setup_method(self): + EmbeddingsSingleton._instances = {} + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base._get_embeddings_wrapper") + @patch("os.path.exists", return_value=False) + def test_defaults_from_settings_do_not_raise( + self, _mock_exists, mock_get_wrapper, mock_settings + ): + """The default config (HF mpnet name, no key) must resolve, not crash.""" + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.EMBEDDINGS_NAME = HF_MPNET + mock_settings.EMBEDDINGS_KEY = None + mock_wrapper_cls = Mock() + mock_instance = Mock() + mock_wrapper_cls.return_value = mock_instance + mock_get_wrapper.return_value = mock_wrapper_cls + + result = get_embeddings() + + assert result is mock_instance + assert set(EmbeddingsSingleton._instances) == {HF_MPNET} + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base._get_embeddings_wrapper") + @patch("os.path.exists", return_value=False) + def test_shares_cache_entry_with_vectorstore_helper( + self, _mock_exists, mock_get_wrapper, mock_settings + ): + """Same object, same cache key as the vector stores get — one model.""" + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.EMBEDDINGS_NAME = HF_MPNET + mock_settings.EMBEDDINGS_KEY = None + mock_wrapper_cls = Mock() + mock_wrapper_cls.return_value = Mock() + mock_get_wrapper.return_value = mock_wrapper_cls + + store_result = ConcreteVectorStore()._get_embeddings(HF_MPNET, None) + resolver_result = get_embeddings() + + assert resolver_result is store_result + assert set(EmbeddingsSingleton._instances) == {HF_MPNET} + mock_wrapper_cls.assert_called_once() + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base._get_embeddings_wrapper") + @patch("os.path.exists") + def test_uses_local_model_path_and_caches_it( + self, mock_exists, mock_get_wrapper, mock_settings + ): + """With the bundled model present, the instance is keyed by its path. + + A second call must not load a second copy of the model. + """ + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.EMBEDDINGS_NAME = HF_MPNET + mock_settings.EMBEDDINGS_KEY = None + mock_exists.side_effect = lambda path: path == LOCAL_MPNET + mock_wrapper_cls = Mock() + mock_wrapper_cls.return_value = Mock() + mock_get_wrapper.return_value = mock_wrapper_cls + + first = get_embeddings() + second = get_embeddings() + + assert first is second + assert set(EmbeddingsSingleton._instances) == {LOCAL_MPNET} + mock_wrapper_cls.assert_called_once_with(LOCAL_MPNET) + + @patch("application.vectorstore.base.settings") + def test_remote_when_base_url_configured(self, mock_settings): + mock_settings.EMBEDDINGS_BASE_URL = "http://remote:8080" + mock_settings.EMBEDDINGS_NAME = HF_MPNET + mock_settings.EMBEDDINGS_KEY = "sk-remote" + + result = get_embeddings() + + assert isinstance(result, RemoteEmbeddings) + assert result.api_url == "http://remote:8080" + assert result.model_name == HF_MPNET + assert result.headers["Authorization"] == "Bearer sk-remote" + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance") + def test_openai_passes_key(self, mock_get_instance, mock_settings): + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.OPENAI_API_BASE = None + mock_settings.OPENAI_API_VERSION = None + mock_settings.AZURE_DEPLOYMENT_NAME = None + mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002" + mock_settings.EMBEDDINGS_KEY = "sk-from-settings" + + get_embeddings() + + mock_get_instance.assert_called_once_with( + "openai_text-embedding-ada-002", openai_api_key="sk-from-settings" + ) + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance") + def test_openai_azure_uses_deployment_name( + self, mock_get_instance, mock_settings + ): + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.OPENAI_API_BASE = "https://azure.openai.com" + mock_settings.OPENAI_API_VERSION = "2023-05-15" + mock_settings.AZURE_DEPLOYMENT_NAME = "deploy" + mock_settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME = "embed-deploy" + mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002" + mock_settings.EMBEDDINGS_KEY = "sk-key" + + get_embeddings() + + mock_get_instance.assert_called_once_with( + "openai_text-embedding-ada-002", model="embed-deploy" + ) + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance") + def test_explicit_arguments_win_over_settings( + self, mock_get_instance, mock_settings + ): + mock_settings.EMBEDDINGS_BASE_URL = None + mock_settings.EMBEDDINGS_NAME = HF_MPNET + mock_settings.EMBEDDINGS_KEY = "sk-from-settings" + + get_embeddings("some_custom_embedding", "sk-explicit") + + mock_get_instance.assert_called_once_with("some_custom_embedding") + + @patch("application.vectorstore.base.settings") + @patch("application.vectorstore.base.get_embeddings") + def test_vectorstore_helper_delegates_to_resolver( + self, mock_resolver, _mock_settings + ): + """``BaseVectorStore._get_embeddings`` is a thin delegate now.""" + sentinel = Mock() + mock_resolver.return_value = sentinel + + result = ConcreteVectorStore()._get_embeddings("a-name", "a-key") + + assert result is sentinel + mock_resolver.assert_called_once_with("a-name", "a-key") diff --git a/tests/vectorstore/test_pgvector.py b/tests/vectorstore/test_pgvector.py index 0611653c..fa59621b 100644 --- a/tests/vectorstore/test_pgvector.py +++ b/tests/vectorstore/test_pgvector.py @@ -31,13 +31,17 @@ def _make_store( from application.vectorstore.pgvector import PGVectorStore - # Patch _ensure_table_exists to avoid DB calls during init - with patch.object(PGVectorStore, "_ensure_table_exists"): - store = PGVectorStore( - source_id=source_id, - embeddings_key=embeddings_key, - connection_string=connection_string, - ) + store = PGVectorStore( + source_id=source_id, + embeddings_key=embeddings_key, + connection_string=connection_string, + ) + # Boot owns the schema, so these unit stores start already ensured; + # the write-path safety net has its own tests. + store._schema_ensured = True + # Legacy direct-connection mode: pooling would build a real + # psycopg_pool against the fake DSN. The pooled path has its own tests. + store._pool_max_size = 0 # Provide a mock connection mock_conn = MagicMock() mock_cursor = MagicMock() @@ -438,3 +442,47 @@ class TestPGVectorSearchWithScores: mock_cursor.execute.side_effect = Exception("connection lost") assert store.search_with_scores("query") == [] + + +@pytest.mark.unit +class TestPGVectorPrecomputedQueryVector: + """A caller that already embedded the query must not pay for it again. + + Multi-source retrieval embeds the query once and hands the same vector to + every source's store; without this the query is embedded once per source. + """ + + def test_search_with_scores_skips_embedding(self): + store, _, mock_cursor, mock_emb = _make_store() + mock_cursor.fetchall.return_value = [("hello world", {}, 0.1)] + + results = store.search_with_scores("query", k=2, query_vector=[0.9, 0.8, 0.7]) + + mock_emb.embed_query.assert_not_called() + assert [doc.page_content for doc, _ in results] == ["hello world"] + + def test_search_skips_embedding(self): + store, _, mock_cursor, mock_emb = _make_store() + mock_cursor.fetchall.return_value = [("hello world", {}, 0.1)] + + results = store.search("query", k=2, query_vector=[0.9, 0.8, 0.7]) + + mock_emb.embed_query.assert_not_called() + assert [doc.page_content for doc in results] == ["hello world"] + + def test_supplied_vector_is_what_the_query_binds(self): + store, _, mock_cursor, _ = _make_store(source_id="src1") + mock_cursor.fetchall.return_value = [("a", {}, 0.1)] + + store.search_with_scores("query", k=1, query_vector=[0.9, 0.8, 0.7]) + + _sql, params = mock_cursor.execute.call_args_list[0].args + assert params == ([0.9, 0.8, 0.7], "src1", [0.9, 0.8, 0.7], 1) + + def test_none_vector_still_embeds(self): + store, _, mock_cursor, mock_emb = _make_store() + mock_cursor.fetchall.return_value = [] + + store.search_with_scores("query", k=2, query_vector=None) + + mock_emb.embed_query.assert_called_once_with("query") diff --git a/tests/vectorstore/test_pgvector_live_schema.py b/tests/vectorstore/test_pgvector_live_schema.py new file mode 100644 index 00000000..273cea3e --- /dev/null +++ b/tests/vectorstore/test_pgvector_live_schema.py @@ -0,0 +1,261 @@ +"""Live pgvector end-to-end: boot-owned schema, pure-read queries, pooling. + +Runs against the ephemeral pytest-postgresql cluster (never the operator's dev +database) with an 8-dimension stub embeddings model, so a real ``CREATE +EXTENSION vector`` / ``vector(8)`` round trip is exercised in a few seconds. +Skips when the cluster's Postgres has no pgvector build installed. +""" + +from __future__ import annotations + +import threading +from unittest.mock import patch + +import pytest + +from application.storage.db.bootstrap import ensure_vector_schema +from application.vectorstore import pgvector as pgvector_module +from application.vectorstore.pgvector import PGVectorStore + +pytestmark = pytest.mark.integration + +STUB_DIM = 8 + + +class _StubEmbeddings: + """Deterministic tiny embeddings; no model download, no network.""" + + dimension = STUB_DIM + + def embed_documents(self, texts): + return [self._vector(float(i + 1)) for i, _ in enumerate(texts)] + + def embed_query(self, query): + return self._vector(1.0) + + @staticmethod + def _vector(seed: float): + vector = [0.0] * STUB_DIM + vector[0] = seed + return vector + + +class _WideStubEmbeddings(_StubEmbeddings): + dimension = 16 + + +def _dsn(info) -> str: + password = f":{info.password}" if info.password else "" + return ( + f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}" + ) + + +def _regclass(conn, name: str): + """Look a relation up on ``conn``, in a transaction that starts fresh. + + ``to_regclass`` takes no lock, so a negative catalog-cache entry from an + earlier lookup in the same transaction survives another backend's DDL. A + new transaction accepts the invalidation. + """ + conn.rollback() + with conn.cursor() as cursor: + cursor.execute("SELECT to_regclass(%s);", (name,)) + return cursor.fetchone()[0] + + +@pytest.fixture(autouse=True) +def _close_pools(): + """Never leak a pool into another test; the DSN dies with the test DB.""" + yield + for dsn, pool in list(pgvector_module._POOLS.items()): + try: + pool.close() + except Exception: + pass + pgvector_module._POOLS.pop(dsn, None) + + +@pytest.fixture +def live_dsn(postgresql, monkeypatch): + """DSN of an empty ephemeral DB, with pgvector proven installable.""" + try: + with postgresql.cursor() as cursor: + cursor.execute("CREATE EXTENSION vector;") + postgresql.rollback() # leave the database pristine for the hook + except Exception as exc: + postgresql.rollback() + pytest.skip(f"pgvector extension unavailable: {exc}") + + dsn = _dsn(postgresql.info) + from application.core import settings as settings_module + + settings = settings_module.settings + monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False) + monkeypatch.setattr( + settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False + ) + monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False) + monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False) + monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False) + return dsn + + +@pytest.fixture +def stub_embeddings(): + stub = _StubEmbeddings() + with patch( + "application.vectorstore.base.get_embeddings", return_value=stub + ), patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings", + return_value=stub, + ): + yield stub + + +def _store(dsn, source_id="live-source"): + return PGVectorStore(source_id=source_id, connection_string=dsn) + + +class TestBootHookCreatesTheSchema: + def test_creates_the_table_and_indexes_and_is_idempotent( + self, live_dsn, postgresql, stub_embeddings + ): + ensure_vector_schema() + ensure_vector_schema() # second boot must be a no-op + + assert _regclass(postgresql, "documents") is not None + with postgresql.cursor() as cursor: + cursor.execute( + "SELECT indexname FROM pg_indexes WHERE tablename = 'documents';" + ) + indexes = {row[0] for row in cursor.fetchall()} + assert "documents_source_id_idx" in indexes + assert "documents_text_fts_idx" in indexes + assert PGVectorStore.table_dimension(postgresql) == STUB_DIM + + def test_mismatched_model_width_fails_loudly(self, live_dsn, stub_embeddings): + ensure_vector_schema() + + wide = _WideStubEmbeddings() + with patch( + "application.vectorstore.base.get_embeddings", return_value=wide + ): + with pytest.raises(RuntimeError) as excinfo: + ensure_vector_schema() + + message = str(excinfo.value) + assert f"vector({STUB_DIM})" in message + assert "16-dim" in message + + def test_creates_graph_tables_when_graphrag_is_enabled( + self, live_dsn, postgresql, stub_embeddings, monkeypatch + ): + from application.core import settings as settings_module + + monkeypatch.setattr( + settings_module.settings, "GRAPHRAG_ENABLED", True, raising=False + ) + + ensure_vector_schema() + + for table in ( + "graph_nodes", + "graph_edges", + "graph_node_chunks", + "graph_ingest_progress", + ): + assert _regclass(postgresql, table) is not None + + +class TestStoreRoundTrip: + def test_construction_is_free_and_writes_then_reads_survive( + self, live_dsn, stub_embeddings + ): + ensure_vector_schema() + + store = _store(live_dsn) + # Construction must not open a connection or build a pool: the + # retriever makes one of these per source, per request. + assert store._connection is None + assert pgvector_module._POOLS == {} + + ids = store.add_texts( + ["alpha chunk", "beta chunk"], + [{"source": "a.txt"}, {"source": "b.txt"}], + ) + assert len(ids) == 2 + + results = store.search_with_scores("alpha", k=2) + assert [doc.page_content for doc, _ in results] == [ + "alpha chunk", + "beta chunk", + ] + assert results[0][1] == pytest.approx(1.0) + + keyword_hits = store.keyword_search("alpha", k=5) + assert [doc.page_content for doc in keyword_hits] == ["alpha chunk"] + + # The read paths must not leave the pooled connection in a transaction. + assert store._connection.info.transaction_status.name == "IDLE" + store.close() + assert store._connection is None + + +class TestWritePathSafetyNet: + """A process that never ran the boot hook must still be able to ingest. + + ``register_vector`` raises when the ``vector`` type is absent, which is the + state of a brand-new database — so the store has to tolerate that, create + the extension, and re-register before it inserts a vector. + """ + + def test_first_write_bootstraps_a_brand_new_database( + self, live_dsn, postgresql, stub_embeddings + ): + assert _regclass(postgresql, "documents") is None + + store = _store(live_dsn) + store.add_texts(["alpha chunk"], [{"source": "a.txt"}]) + + + assert _regclass(postgresql, "documents") is not None + assert [doc.page_content for doc in store.search("alpha", k=2)] == [ + "alpha chunk" + ] + store.close() + + +class TestPooling: + def test_concurrent_stores_share_one_pool_and_return_connections( + self, live_dsn, stub_embeddings + ): + ensure_vector_schema() + seed = _store(live_dsn) + seed.add_texts(["alpha chunk"], [{"source": "a.txt"}]) + seed.close() + + errors: list[Exception] = [] + hits: list[int] = [] + + def _search(): + store = _store(live_dsn) + try: + hits.append(len(store.search("alpha", k=2))) + except Exception as exc: # surfaced below + errors.append(exc) + finally: + store.close() + + threads = [threading.Thread(target=_search) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=20) + + assert errors == [] + assert hits == [1, 1] + assert list(pgvector_module._POOLS) == [live_dsn] + stats = pgvector_module._POOLS[live_dsn].get_stats() + assert stats["pool_size"] <= 4 + assert stats["pool_available"] == stats["pool_size"] diff --git a/tests/vectorstore/test_pgvector_probes.py b/tests/vectorstore/test_pgvector_probes.py index 013a5173..6ceacc23 100644 --- a/tests/vectorstore/test_pgvector_probes.py +++ b/tests/vectorstore/test_pgvector_probes.py @@ -109,10 +109,10 @@ class TestNoVectorIndexOnEmptyTable: sources with hundreds of chunks, silently. """ - def test_ensure_table_exists_creates_no_vector_index(self): + def test_create_schema_creates_no_vector_index(self): import inspect - source = inspect.getsource(PGVectorStore._ensure_table_exists) + source = inspect.getsource(PGVectorStore.create_schema) assert "USING ivfflat" not in source assert "USING hnsw" not in source # the non-vector indexes are still expected diff --git a/tests/vectorstore/test_pgvector_schema_and_pool.py b/tests/vectorstore/test_pgvector_schema_and_pool.py new file mode 100644 index 00000000..3b0cae95 --- /dev/null +++ b/tests/vectorstore/test_pgvector_schema_and_pool.py @@ -0,0 +1,312 @@ +"""Schema ownership and connection pooling for ``PGVectorStore``. + +Constructing a store used to run six DDL statements plus a fresh unpooled +connect, on every instantiation — and the retriever builds one store per source +per request. Schema is now owned at boot (``ensure_vector_schema``); the query +path is pure reads over a pooled connection, and the write path keeps a +one-shot safety net for processes that never ran the boot hook. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from application.vectorstore import pgvector as pgvector_module + +CONNECTION_STRING = "postgresql://user:pass@localhost/db" + + +def _make_store(source_id="test-source", connection_string=CONNECTION_STRING): + """Build a store with every external dependency mocked, no DDL patching.""" + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings" + ) as mock_get_emb, patch( + "application.vectorstore.pgvector.settings" + ) as mock_settings, patch.dict( + "sys.modules", + { + "psycopg": MagicMock(), + "pgvector": MagicMock(), + "pgvector.psycopg": MagicMock(), + }, + ): + mock_emb = Mock() + mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3]) + mock_emb.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]]) + mock_emb.dimension = 768 + mock_get_emb.return_value = mock_emb + mock_settings.EMBEDDINGS_NAME = "test_model" + mock_settings.PGVECTOR_CONNECTION_STRING = connection_string + + from application.vectorstore.pgvector import PGVectorStore + + store = PGVectorStore( + source_id=source_id, + embeddings_key="key", + connection_string=connection_string, + ) + + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_conn.closed = False + store._connection = mock_conn + store._pooled = False + store._ensure_table_exists = Mock() + return store, mock_conn, mock_cursor, mock_emb + + +@pytest.mark.unit +class TestConstructionTouchesNothing: + def test_init_opens_no_connection_and_runs_no_ddl(self): + with patch( + "application.vectorstore.base.BaseVectorStore._get_embeddings" + ) as mock_get_emb, patch( + "application.vectorstore.pgvector.settings" + ) as mock_settings, patch.dict( + "sys.modules", + { + "psycopg": MagicMock(), + "pgvector": MagicMock(), + "pgvector.psycopg": MagicMock(), + }, + ): + mock_get_emb.return_value = Mock(dimension=768) + mock_settings.EMBEDDINGS_NAME = "test_model" + mock_settings.PGVECTOR_CONNECTION_STRING = CONNECTION_STRING + + from application.vectorstore.pgvector import PGVectorStore + + with patch.object( + PGVectorStore, "_get_connection" + ) as mock_get_conn, patch.object( + PGVectorStore, "_ensure_table_exists" + ) as mock_ensure: + store = PGVectorStore( + source_id="s1", connection_string=CONNECTION_STRING + ) + + mock_get_conn.assert_not_called() + mock_ensure.assert_not_called() + assert store._schema_ensured is False + + def test_pool_size_falls_back_when_setting_is_not_an_int(self): + # The unit suite patches ``settings`` with a MagicMock; a MagicMock + # attribute must not become the pool size. + store, _, _, _ = _make_store() + assert store._pool_max_size == 8 + + +@pytest.mark.unit +class TestWritePathEnsuresSchemaOnce: + def test_add_texts_ensures_once_across_two_calls(self): + store, _, mock_cursor, _ = _make_store() + mock_cursor.fetchone.return_value = (1,) + + store.add_texts(["a"]) + store.add_texts(["b"]) + + assert store._ensure_table_exists.call_count == 1 + + def test_add_chunk_ensures_once_across_two_calls(self): + store, _, mock_cursor, _ = _make_store() + mock_cursor.fetchone.return_value = (1,) + + store.add_chunk("a") + store.add_chunk("b") + + assert store._ensure_table_exists.call_count == 1 + + def test_reads_and_deletes_never_ensure_schema(self): + store, _, mock_cursor, _ = _make_store() + mock_cursor.fetchall.return_value = [] + mock_cursor.fetchone.return_value = (0,) + mock_cursor.rowcount = 0 + + store.search("q") + store.keyword_search("q") + store.get_chunks() + store.delete_index() + store.delete_chunk("1") + + store._ensure_table_exists.assert_not_called() + + +@pytest.mark.unit +class TestCreateSchema: + def test_emits_the_ddl_and_leaves_the_commit_to_the_caller(self): + from application.vectorstore.pgvector import PGVectorStore + + conn, cursor = MagicMock(), MagicMock() + conn.cursor.return_value = cursor + + PGVectorStore.create_schema(conn, dimension=8) + + statements = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "CREATE EXTENSION IF NOT EXISTS vector" in statements + assert "CREATE TABLE IF NOT EXISTS documents" in statements + assert "vector(8)" in statements + assert "documents_source_id_idx" in statements + assert "documents_text_fts_idx" in statements + assert cursor.execute.call_count == 4 + conn.commit.assert_not_called() + + def test_ensure_table_exists_locks_then_commits(self): + store, mock_conn, mock_cursor, _ = _make_store() + del store._ensure_table_exists # exercise the real method + + store._ensure_table_exists() + + statements = " ".join(str(c) for c in mock_cursor.execute.call_args_list) + assert "pg_advisory_xact_lock" in statements + assert "CREATE TABLE IF NOT EXISTS documents" in statements + mock_conn.commit.assert_called_once() + + +@pytest.mark.unit +class TestTableDimension: + def _conn(self, rows): + conn, cursor = MagicMock(), MagicMock() + conn.cursor.return_value = cursor + cursor.fetchone.side_effect = rows + return conn + + def test_parses_the_declared_vector_width(self): + from application.vectorstore.pgvector import PGVectorStore + + conn = self._conn([("documents",), ("vector(768)",)]) + assert PGVectorStore.table_dimension(conn) == 768 + + def test_returns_none_when_the_table_is_absent(self): + from application.vectorstore.pgvector import PGVectorStore + + conn = self._conn([(None,)]) + assert PGVectorStore.table_dimension(conn) is None + + def test_returns_none_when_the_column_type_is_not_a_vector(self): + from application.vectorstore.pgvector import PGVectorStore + + conn = self._conn([("documents",), ("text",)]) + assert PGVectorStore.table_dimension(conn) is None + + +@pytest.mark.unit +class TestReadPathsCloseTheirTransaction: + """A read that never commits leaves the connection idle-in-transaction. + + On a persistent (now pooled) connection that holds a snapshot open for the + life of the process and makes psycopg_pool roll back every returned + connection. + """ + + def test_search_commits(self): + store, mock_conn, mock_cursor, _ = _make_store() + mock_cursor.fetchall.return_value = [("hello", {}, 0.1), ("bye", {}, 0.2)] + + store.search_with_scores("q", k=2) + + mock_conn.commit.assert_called_once() + + def test_keyword_search_commits(self): + store, mock_conn, mock_cursor, _ = _make_store() + mock_cursor.fetchall.return_value = [] + + store.keyword_search("q") + + mock_conn.commit.assert_called_once() + + def test_get_chunks_commits(self): + store, mock_conn, mock_cursor, _ = _make_store() + mock_cursor.fetchall.return_value = [] + + store.get_chunks() + + mock_conn.commit.assert_called_once() + + +@pytest.mark.unit +class TestConnectionPooling: + def _pooled_store(self, monkeypatch): + store, _, _, _ = _make_store() + store._connection = None + store._pool_max_size = 4 + pooled_conn = MagicMock() + pooled_conn.closed = False + fake_pool = MagicMock() + fake_pool.getconn.return_value = pooled_conn + monkeypatch.setattr( + pgvector_module, "_pool_for", lambda dsn, size: fake_pool + ) + return store, fake_pool, pooled_conn + + def test_get_connection_checks_out_of_the_pool(self, monkeypatch): + store, fake_pool, pooled_conn = self._pooled_store(monkeypatch) + + conn = store._get_connection() + + assert conn is pooled_conn + assert store._pooled is True + fake_pool.getconn.assert_called_once() + store._psycopg.connect.assert_not_called() + + def test_probes_are_applied_once_per_physical_connection(self, monkeypatch): + store, _, pooled_conn = self._pooled_store(monkeypatch) + applied = Mock() + monkeypatch.setattr(store, "_apply_ivfflat_probes", applied) + + store._get_connection() + store._connection = None # simulate a second checkout of the same conn + store._get_connection() + + assert applied.call_count == 1 + + def test_close_rolls_back_and_returns_the_connection(self, monkeypatch): + store, fake_pool, pooled_conn = self._pooled_store(monkeypatch) + store._get_connection() + pooled_conn.info.transaction_status.name = "INTRANS" + monkeypatch.setitem( + pgvector_module._POOLS, store._connection_string, fake_pool + ) + + store.close() + + pooled_conn.rollback.assert_called_once() + fake_pool.putconn.assert_called_once_with(pooled_conn) + pooled_conn.close.assert_not_called() + assert store._connection is None + + def test_close_does_not_roll_back_an_idle_connection(self, monkeypatch): + store, fake_pool, pooled_conn = self._pooled_store(monkeypatch) + store._get_connection() + pooled_conn.info.transaction_status.name = "IDLE" + monkeypatch.setitem( + pgvector_module._POOLS, store._connection_string, fake_pool + ) + + store.close() + + pooled_conn.rollback.assert_not_called() + fake_pool.putconn.assert_called_once_with(pooled_conn) + + def test_legacy_path_connects_directly_and_closes(self): + store, mock_conn, _, _ = _make_store() + store._pool_max_size = 0 + mock_conn.closed = True + new_conn = MagicMock() + new_conn.closed = False + store._psycopg.connect.return_value = new_conn + + conn = store._get_connection() + assert conn is new_conn + assert store._pooled is False + + store.close() + new_conn.close.assert_called_once() + + def test_del_never_raises(self): + store, mock_conn, _, _ = _make_store() + mock_conn.close.side_effect = RuntimeError("already gone") + + store.__del__() # must not propagate