fix: refactors and optimisations

application/agents/tools/read_document.py

@@ -15,8 +15,6 @@ from __future__ import annotations
 import logging
 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
@@ -318,14 +316,22 @@ class ReadDocumentTool(Tool):
     @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."""
+        completed = False
 
         def _on_alarm(signum: int, frame: Any) -> None:
+            # A straggler racing the disarm below: ``fn`` already returned, and
+            # raising here would escape the ``finally`` and replace its value
+            # with a timeout for a document that parsed fine.
+            if completed:
+                return
             raise _InlineParseTimeout()
 
         previous = signal.signal(signal.SIGALRM, _on_alarm)
         signal.setitimer(signal.ITIMER_REAL, timeout)
         try:
-            return fn()
+            result = fn()
+            completed = True
+            return result
         finally:
             # Nested so the handler is restored even if the alarm lands in this very
             # window (fn returned just as the timer expired).
@@ -336,24 +342,37 @@ class ReadDocumentTool(Tool):
 
     @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)
+        """Run ``fn`` in a daemon helper thread, giving up at ``timeout`` (the parse is NOT killed)."""
+        # A raw daemon thread rather than a ThreadPoolExecutor: executor threads
+        # are non-daemon and registered with ``concurrent.futures``' atexit hook,
+        # which joins them -- so an abandoned parse would hold up worker
+        # shutdown for the rest of its (size-scaled) window. Same reasoning as
+        # application/guardrails/engine.py. ``shutdown(cancel_futures=True)``
+        # is not an alternative: it only drops queued work items, never the one
+        # already running.
+        slot: Dict[str, Any] = {}
+
+        def _fill() -> None:
             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)
+                slot["value"] = fn()
+            except BaseException as exc:  # noqa: BLE001 - re-raised on the caller's thread
+                slot["error"] = exc
+
+        thread = threading.Thread(target=_fill, daemon=True, name="read-document-inline")
+        thread.start()
+        thread.join(timeout)
+        if thread.is_alive():
+            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()
+        # ``fn`` raising TimeoutError is a parse failure, not our deadline; with
+        # join()+is_alive() the two are distinguishable without a done() probe.
+        if "error" in slot:
+            raise slot["error"]
+        return slot.get("value")
 
     # ------------------------------------------------------------------
     # Input resolution (run-scoped gate, before enqueue)

application/cache.py

@@ -77,12 +77,20 @@ def _bypasses_cache(extra: dict | None) -> bool:
     return any(extra.get(key) for key in _CACHE_BYPASS_KWARGS)
 
 
+# ``gen_cache`` and ``stream_cache`` hash the same (messages, model, kwargs)
+# tuple, so they would collide on any call served both ways: whichever wrote
+# last replaced the other's value. Namespacing the key is what keeps the two
+# payload shapes apart -- sniffing the shape on read cannot, because the write
+# side has already destroyed the other entry by then.
+_GEN_KEY_PREFIX = "gen:"
+_STREAM_KEY_PREFIX = "stream:"
+
+
 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.
+    Belt-and-braces behind the key prefixes above: a stream envelope must
+    never be handed back as a non-streaming answer, whatever put it there.
     """
     if not raw.startswith("{"):
         return False
@@ -220,7 +228,7 @@ def gen_cache(func):
             return func(self, model, messages, stream, tools, *args, **kwargs)
 
         try:
-            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
+            cache_key = f"{_GEN_KEY_PREFIX}{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)
@@ -255,7 +263,7 @@ def stream_cache(func):
             return
 
         try:
-            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
+            cache_key = f"{_STREAM_KEY_PREFIX}{gen_cache_key(messages, model, tools, extra=kwargs)}"
         except ValueError as e:
             logger.error(f"Cache key generation failed: {e}")
             yield from func(self, model, messages, stream, tools, *args, **kwargs)

application/graphrag/store.py

@@ -24,10 +24,12 @@ 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
+# Bound here (same objects, not copies) from the shared pool module, which both
+# stores already import. Reaching through ``pgvector`` instead would drag the
+# embeddings stack in at import time; ``pgconn`` imports nothing heavier than
+# ``logging`` and ``threading``. 0 disables pooling.
+DEFAULT_POOL_MAX_SIZE = pgconn.DEFAULT_POOL_MAX_SIZE
+_resolve_pool_max_size = pgconn.resolve_pool_max_size
 
 MAX_SUBGRAPH_NODES = 500
 MAX_SUBGRAPH_EDGES = 2000
@@ -38,14 +40,6 @@ 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.
 

application/llm/openai.py

@@ -4,6 +4,7 @@ import io
 import json
 import logging
 import os.path
+from typing import Any, Callable
 
 from openai import BadRequestError, OpenAI
 
@@ -638,7 +639,7 @@ class OpenAILLM(BaseLLM):
                 return normalized
         return ""
 
-    def _note_tools_rejected(self, model, error) -> None:
+    def _note_tools_rejected(self, model: str, error: Exception) -> None:
         """Record that this endpoint refuses tool calling, warning once."""
         if not self._tools_rejected:
             logging.warning(
@@ -661,7 +662,9 @@ class OpenAILLM(BaseLLM):
             if key not in _TOOL_REQUEST_KEYS
         }
 
-    def _create_with_tool_fallback(self, create, params: dict, model):
+    def _create_with_tool_fallback(
+        self, create: Callable[..., Any], params: dict, model: str
+    ) -> Any:
         """Send a request, retrying once without tools on a tools-unsupported 400.
 
         Args:

application/retriever/classic_rag.py

@@ -4,14 +4,7 @@ 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.fanout import fetch_per_source, max_parallel_sources
 from application.retriever.labels import labels_from_metadata
 from application.utils import num_tokens_from_string
 from application.vectorstore.vector_creator import VectorCreator
@@ -260,17 +253,6 @@ class ClassicRAG(BaseRetriever):
                 )
         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],
@@ -309,44 +291,18 @@ class ClassicRAG(BaseRetriever):
     ) -> 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.
+        Shares :func:`~application.retriever.fanout.fetch_per_source` with the
+        search service so both paths order, embed and degrade identically.
         """
-        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)),
+        return fetch_per_source(
+            plans,
+            lambda plan: VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, plan["id"], settings.EMBEDDINGS_KEY
+            ),
+            self._search_source,
+            lambda plan: plan["question"],
+            label_of=lambda plan: plan["id"],
+            workers_for=_max_parallel_sources,
         )
 
     def _get_data(self):

application/retriever/fanout.py

@@ -125,3 +125,66 @@ def run_source_jobs(
         max_workers=workers, thread_name_prefix="rag-source"
     ) as pool:
         return list(pool.map(fn, jobs))
+
+
+def fetch_per_source(
+    items: List[Job],
+    build_store: Callable[[Job], Any],
+    search: Callable[[Job, Optional[Any], Optional[List[float]]], Result],
+    question_of: Callable[[Job], str],
+    label_of: Callable[[Job], Any] = str,
+    workers_for: Optional[Callable[[int], int]] = None,
+) -> List[Optional[Result]]:
+    """Search every source: one store built up front, one embedding, 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 caller's merge stays deterministic.
+
+    Args:
+        items: One entry per source, in the order results must come back.
+        build_store: Builds the vector store for an item; may raise.
+        search: ``(item, store_or_None, query_vector_or_None) -> result``. It
+            must swallow its own errors and report failure as ``None``, since a
+            raise here aborts the whole fan-out.
+        question_of: The query for an item, used to look up its shared vector.
+        label_of: Item identifier for the store-construction error log.
+        workers_for: Pool size from the job count; defaults to
+            :func:`max_parallel_sources`.
+
+    Returns:
+        list: One entry per item, in ``items`` order. ``None`` marks a source
+        whose store could not be built or whose search reported failure.
+    """
+    items = list(items)
+    if not items:
+        return []
+
+    def _dispatch(jobs: List[Any]) -> List[Optional[Result]]:
+        if not jobs:
+            return []
+        workers = workers_for(len(jobs)) if workers_for is not None else None
+        return run_source_jobs(lambda job: search(*job), jobs, workers=workers)
+
+    first_store = None
+    try:
+        first_store = build_store(items[0])
+    except Exception as e:
+        logger.error(
+            "Error searching vectorstore %s: %s", label_of(items[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] + _dispatch([(item, None, None) for item in items[1:]])
+
+    vectors = embed_questions(first_store, [question_of(item) for item in items])
+    return _dispatch(
+        [
+            (item, first_store if idx == 0 else None, vectors.get(question_of(item)))
+            for idx, item in enumerate(items)
+        ]
+    )

application/retriever/hybrid_rag.py

@@ -8,6 +8,8 @@ support (``keyword_search`` returns ``[]``) reduce to exact vector-only
 behaviour.
 """
 
+from typing import Any, List, Optional
+
 from application.retriever.classic_rag import ClassicRAG
 
 RRF_K = 60
@@ -49,15 +51,20 @@ def fuse_with_scores(vector_hits, keyword_hits, k=RRF_K):
 class HybridRetriever(ClassicRAG):
     """ClassicRAG variant that fuses vector + keyword search with RRF."""
 
-    def _score_kind(self, docsearch):
+    def _score_kind(self, docsearch: Any) -> str:
         """RRF scores rank hits against each other, not against a similarity
         cutoff — they are not comparable to the store's cosine scores, so they
         carry their own label."""
         return "rrf"
 
     def _fetch_candidates(
-        self, docsearch, question, src_k, score_threshold, query_vector=None
-    ):
+        self,
+        docsearch: Any,
+        question: str,
+        src_k: int,
+        score_threshold: Optional[float],
+        query_vector: Optional[List[float]] = None,
+    ) -> List[Any]:
         """Return RRF-fused vector+keyword hits for one vector store.
 
         Inherits the per-source resolution and budgeting from

application/services/search_service.py

@@ -11,7 +11,7 @@ import logging
 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.retriever.fanout import fetch_per_source
 from application.storage.db.repositories.agents import AgentsRepository
 from application.storage.db.session import db_readonly
 from application.vectorstore.vector_creator import VectorCreator
@@ -125,41 +125,19 @@ def _fetch_sources(
 ) -> 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.
+    Shares :func:`~application.retriever.fanout.fetch_per_source` with
+    ``ClassicRAG`` so the search route and the answer path order, embed and
+    degrade identically.
     """
-
-    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)
-        ],
+    return fetch_per_source(
+        source_ids,
+        lambda source_id: VectorCreator.create_vectorstore(
+            settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+        ),
+        lambda source_id, docsearch, query_vector: _search_one(
+            source_id, docsearch, query, k, query_vector
+        ),
+        lambda _source_id: query,
     )
 
 

application/vectorstore/pgconn.py

@@ -21,6 +21,12 @@ from typing import Any, Dict
 # Seconds a caller waits for a free pooled connection before giving up.
 POOL_TIMEOUT_SECONDS = 30.0
 
+# Physical connections one process may hold per DSN. Lives here, not in either
+# store: ``pool_for`` keys its registry by DSN and applies ``max_size`` only on
+# creation, so two stores disagreeing about the default would make the
+# effective pool size depend on which of them touched the DSN first.
+DEFAULT_POOL_MAX_SIZE = 8
+
 # dsn -> psycopg_pool.ConnectionPool, one per process.
 _POOLS: Dict[str, Any] = {}
 _POOLS_LOCK = threading.Lock()
@@ -42,6 +48,23 @@ def configure_pooled_connection(conn) -> None:
         logging.debug("pgvector types not registered yet: %s", e)
 
 
+def resolve_pool_max_size() -> int:
+    """Pool size from settings, defensively — 0 means one direct connection.
+
+    Returns:
+        int: ``PGVECTOR_POOL_MAX_SIZE`` when it is a non-negative, non-bool
+        int, else :data:`DEFAULT_POOL_MAX_SIZE`. The unit suite replaces
+        ``settings`` with a MagicMock, whose attributes must never become
+        a pool size.
+    """
+    from application.core.settings import settings
+
+    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 pool_for(dsn: str, max_size: int):
     """Return this process's connection pool for ``dsn``, creating it lazily.
 

application/vectorstore/pgvector.py

@@ -14,7 +14,6 @@ from application.vectorstore.document_class import Document
 _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"
 
@@ -23,6 +22,7 @@ SCHEMA_LOCK_KEY = "docsgpt:vectors:ddl"
 # here — same objects, not copies — because callers and tests reach for
 # ``pgvector._POOLS`` / patch ``pgvector._pool_for``.
 POOL_TIMEOUT_SECONDS = pgconn.POOL_TIMEOUT_SECONDS
+DEFAULT_POOL_MAX_SIZE = pgconn.DEFAULT_POOL_MAX_SIZE
 _POOLS = pgconn._POOLS
 _POOLS_LOCK = pgconn._POOLS_LOCK
 _configure_pooled_connection = pgconn.configure_pooled_connection
@@ -87,13 +87,8 @@ class PGVectorStore(BaseVectorStore):
         # 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
+    # Shared with ``GraphStore`` via pgconn so the two can never disagree.
+    _resolve_pool_max_size = staticmethod(pgconn.resolve_pool_max_size)
 
     def _get_connection(self):
         """Get or create this store's connection, pooled unless pooling is off."""
@@ -117,7 +112,11 @@ class PGVectorStore(BaseVectorStore):
         return self._connection
 
     def _register_pgvector_types(self, conn) -> None:
-        """Register pgvector's adapters, tolerating a not-yet-created extension."""
+        """Register pgvector's adapters, tolerating a not-yet-created extension.
+
+        Kept as a method rather than calling :func:`pgconn.configure_pooled_connection`
+        directly because ``self._register_vector`` is the seam the tests patch.
+        """
         try:
             self._register_vector(conn)
         except Exception as e: