From 82ac2dcee98cc16d1aa98e0cf6d9647f1b74dff3 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 20 Aug 2026 13:56:37 +0100 Subject: [PATCH] fix: refactors and optimisations --- application/agents/tools/read_document.py | 59 +++++--- application/cache.py | 18 ++- application/graphrag/store.py | 18 +-- application/llm/openai.py | 7 +- application/retriever/classic_rag.py | 68 ++-------- application/retriever/fanout.py | 63 +++++++++ application/retriever/hybrid_rag.py | 13 +- application/services/search_service.py | 48 ++----- application/vectorstore/pgconn.py | 23 ++++ application/vectorstore/pgvector.py | 17 ++- tests/agents/tools/test_read_document_tool.py | 73 ++++++++++ tests/retriever/test_fanout.py | 127 ++++++++++++++++++ tests/test_cache.py | 70 +++++++++- .../test_pgvector_schema_and_pool.py | 37 +++++ 14 files changed, 497 insertions(+), 144 deletions(-) diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index bff32d45..52ce7183 100644 --- a/application/agents/tools/read_document.py +++ b/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) diff --git a/application/cache.py b/application/cache.py index 73efb53b..c7d73cb9 100644 --- a/application/cache.py +++ b/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) diff --git a/application/graphrag/store.py b/application/graphrag/store.py index 1ee642c0..a7b0facf 100644 --- a/application/graphrag/store.py +++ b/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. diff --git a/application/llm/openai.py b/application/llm/openai.py index 439c6766..1f60e568 100644 --- a/application/llm/openai.py +++ b/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: diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py index 21c855a3..c034d232 100644 --- a/application/retriever/classic_rag.py +++ b/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): diff --git a/application/retriever/fanout.py b/application/retriever/fanout.py index 81b0e1e5..bb9f0605 100644 --- a/application/retriever/fanout.py +++ b/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) + ] + ) diff --git a/application/retriever/hybrid_rag.py b/application/retriever/hybrid_rag.py index f37fca1a..d63752ca 100644 --- a/application/retriever/hybrid_rag.py +++ b/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 diff --git a/application/services/search_service.py b/application/services/search_service.py index c739bf1c..af8141f0 100644 --- a/application/services/search_service.py +++ b/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, ) diff --git a/application/vectorstore/pgconn.py b/application/vectorstore/pgconn.py index 4538972a..0f825a4b 100644 --- a/application/vectorstore/pgconn.py +++ b/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. diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index c3fe9f6d..f00ee73c 100644 --- a/application/vectorstore/pgvector.py +++ b/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: diff --git a/tests/agents/tools/test_read_document_tool.py b/tests/agents/tools/test_read_document_tool.py index ed9808dc..02493d05 100644 --- a/tests/agents/tools/test_read_document_tool.py +++ b/tests/agents/tools/test_read_document_tool.py @@ -481,6 +481,79 @@ def test_inline_parse_times_out_on_the_signal_path(monkeypatch, caplog): assert any("timed out" in rec.getMessage() for rec in caplog.records if rec.levelname == "WARNING") +@pytest.mark.unit +def test_a_straggler_alarm_does_not_discard_a_finished_parse(monkeypatch): + """The alarm can land between ``fn()`` returning and the disarm call. + + ``_InlineParseTimeout`` is a BaseException, so nothing catches it locally: + raised from the ``finally`` it replaces the return value, and the caller + reports a timeout for a document that parsed correctly. + """ + real_setitimer = signal.setitimer + + def _fire_on_disarm(which, seconds, *rest): + # Reproduce the race deterministically: deliver SIGALRM at the exact + # instant the timer is being cancelled. + if seconds == 0: + handler = signal.getsignal(signal.SIGALRM) + if callable(handler): + handler(signal.SIGALRM, None) + return real_setitimer(which, seconds, *rest) + + monkeypatch.setattr(signal, "setitimer", _fire_on_disarm) + + parsed = {"status": "ok", "content": "the document really did parse"} + assert rd.ReadDocumentTool._run_with_sigalrm(lambda: parsed, 30.0) == parsed + + +@pytest.mark.unit +def test_a_genuine_sigalrm_timeout_still_raises(monkeypatch): + """The completion flag must not disarm the guard itself.""" + with pytest.raises(rd._InlineParseTimeout): + rd.ReadDocumentTool._run_with_sigalrm(lambda: time.sleep(1.0), 0.15) + + assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0) + + +@pytest.mark.unit +def test_the_inline_helper_thread_is_a_daemon(): + """A timed-out parse is abandoned, so its thread must not outlive the process. + + ``concurrent.futures`` registers its (non-daemon) workers with an atexit + hook that joins them, which is exactly why application/guardrails/engine.py + uses raw daemon threads for the same abandon-on-timeout shape. + """ + seen: Dict[str, Any] = {} + + def _capture(): + seen["thread"] = threading.current_thread() + return {"status": "ok"} + + assert rd.ReadDocumentTool._run_in_thread(_capture, 5.0) == {"status": "ok"} + assert seen["thread"].name.startswith("read-document-inline") + assert seen["thread"].daemon is True + + +@pytest.mark.unit +def test_the_thread_path_propagates_a_parse_failure(): + """An error raised by ``fn`` surfaces as itself, not as a timeout.""" + def _boom(): + raise RuntimeError("parser exploded") + + with pytest.raises(RuntimeError, match="parser exploded"): + rd.ReadDocumentTool._run_in_thread(_boom, 5.0) + + +@pytest.mark.unit +def test_the_thread_path_does_not_mistake_a_parse_timeout_for_its_own(): + """``fn`` raising TimeoutError is a failure, not the deadline elapsing.""" + def _raises_timeout(): + raise TimeoutError("the parser's own network timeout") + + with pytest.raises(TimeoutError, match="own network timeout"): + rd.ReadDocumentTool._run_in_thread(_raises_timeout, 5.0) + + @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.""" diff --git a/tests/retriever/test_fanout.py b/tests/retriever/test_fanout.py index b9b5d0be..64a4020a 100644 --- a/tests/retriever/test_fanout.py +++ b/tests/retriever/test_fanout.py @@ -15,6 +15,7 @@ import pytest from application.retriever.fanout import ( DEFAULT_MAX_PARALLEL_SOURCES, embed_questions, + fetch_per_source, max_parallel_sources, run_source_jobs, store_embeddings, @@ -174,3 +175,129 @@ class TestEmbedQuestions: embedder.embed_query = Mock(side_effect=RuntimeError("model down")) store = SimpleNamespace(_embedding=embedder) assert embed_questions(store, ["q"]) == {} + + +def _fanout_store(vector=(9.0,)): + """A store whose embedder returns a fixed, recognisable vector.""" + embedder = Mock() + embedder.embed_query = Mock(return_value=list(vector)) + return SimpleNamespace(_embedding=embedder) + + +@pytest.mark.unit +class TestFetchPerSource: + """The whole per-source fan-out both ClassicRAG and the search service use.""" + + def test_only_the_first_store_is_built_on_the_calling_thread(self): + built = [] + store = _fanout_store() + + def _build(item): + built.append(item) + return store + + seen = [] + results = fetch_per_source( + ["a", "b", "c"], + _build, + lambda item, docsearch, vector: seen.append((item, docsearch, vector)) + or f"hit-{item}", + lambda item: "q", + ) + + assert results == ["hit-a", "hit-b", "hit-c"] + # One construction up front; the workers build their own. + assert built == ["a"] + assert seen[0][1] is store + assert [s[1] for s in seen[1:]] == [None, None] + + def test_the_query_is_embedded_once_for_every_source(self): + store = _fanout_store(vector=(7.0,)) + + vectors = [] + fetch_per_source( + ["a", "b", "c"], + lambda item: store, + lambda item, docsearch, vector: vectors.append(vector), + lambda item: "same question", + ) + + assert vectors == [[7.0], [7.0], [7.0]] + assert store._embedding.embed_query.call_count == 1 + + def test_per_item_questions_get_their_own_vectors(self): + embedder = Mock() + embedder.embed_query = Mock(side_effect=lambda q: [len(q)]) + store = SimpleNamespace(_embedding=embedder) + + vectors = [] + fetch_per_source( + ["ab", "cde"], + lambda item: store, + lambda item, docsearch, vector: vectors.append(vector), + lambda item: item, + ) + + assert vectors == [[2], [3]] + + def test_results_keep_source_order_regardless_of_finish_order(self): + store = _fanout_store() + started = threading.Barrier(3, timeout=5) + + def _search(item, docsearch, vector): + started.wait() # force genuine interleaving + return item + + results = fetch_per_source( + ["a", "b", "c"], lambda item: store, _search, lambda item: "q", + ) + + assert results == ["a", "b", "c"] + + def test_a_failed_first_store_still_runs_the_rest(self): + def _build(item): + raise RuntimeError("index is gone") + + results = fetch_per_source( + ["a", "b", "c"], + _build, + lambda item, docsearch, vector: (item, docsearch, vector), + lambda item: "q", + ) + + # The first slot is a logged failure; the rest run store-less and + # vector-less, so each embeds its own query exactly as before. + assert results[0] is None + assert results[1:] == [("b", None, None), ("c", None, None)] + + def test_no_items_is_not_a_fan_out(self): + build = Mock() + + assert fetch_per_source([], build, Mock(), Mock()) == [] + build.assert_not_called() + + def test_worker_cap_is_delegated_to_the_caller(self): + store = _fanout_store() + widths = [] + + fetch_per_source( + ["a", "b", "c"], + lambda item: store, + lambda item, docsearch, vector: item, + lambda item: "q", + workers_for=lambda n: widths.append(n) or 1, + ) + + assert widths == [3] + + +@pytest.mark.unit +class TestBothCallersShareTheFanOut: + """The duplication this helper replaced must not creep back in.""" + + def test_classic_rag_and_the_search_service_both_delegate(self): + import application.retriever.classic_rag as classic_rag + import application.services.search_service as search_service + + assert classic_rag.fetch_per_source is fetch_per_source + assert search_service.fetch_per_source is fetch_per_source diff --git a/tests/test_cache.py b/tests/test_cache.py index 9667e0c2..a3f11969 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -895,12 +895,12 @@ def test_stream_cache_bypassed_for_previous_response_id(mock_make_redis): @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 + """Belt-and-braces: even planted directly under the gen key, 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) + key = f"gen:{gen_cache_key(messages, 'm', None)}" fake.set(key, json.dumps({"version": 1, "chunks": ["a", "b"]})) @gen_cache @@ -909,3 +909,69 @@ def test_gen_cache_ignores_stream_payload_stored_under_same_key(mock_make_redis) result = mock_function(None, "m", messages, stream=False, tools=None) assert result == "real answer" + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_gen_and_stream_caches_do_not_share_a_key_space(mock_make_redis): + """A gen write must not overwrite the stream envelope for the same call. + + Both wrappers hash the same (messages, model, kwargs) tuple, so before the + namespacing they collided: the gen write replaced the envelope with a bare + string and the next stream read failed to decode it. + """ + fake = _FakeRedis() + mock_make_redis.return_value = fake + messages = [{"role": "user", "content": "test"}] + + @stream_cache + def streamer(self, model, messages, stream, tools, **kwargs): + yield "alpha" + yield "beta" + + @gen_cache + def generator(self, model, messages, stream, tools, **kwargs): + return "a plain answer" + + assert list(streamer(None, "m", messages, stream=True, tools=None)) == [ + "alpha", "beta", + ] + assert generator(None, "m", messages, stream=False, tools=None) == "a plain answer" + + upstream_calls = [] + + @stream_cache + def streamer_again(self, model, messages, stream, tools, **kwargs): + upstream_calls.append(1) + yield "SHOULD NOT REACH UPSTREAM" + + replayed = list(streamer_again(None, "m", messages, stream=True, tools=None)) + + assert replayed == ["alpha", "beta"] + assert not upstream_calls + + +@pytest.mark.unit +@patch("application.cache.get_redis_instance") +def test_a_json_array_answer_is_never_replayed_as_stream_chunks(mock_make_redis): + """A gen answer that happens to be a JSON array is not a chunk list. + + The stream reader's pre-v1 compatibility branch accepts a bare JSON array, + so a colliding gen entry used to be replayed verbatim as the streamed + answer -- silent corruption rather than a cache miss. + """ + fake = _FakeRedis() + mock_make_redis.return_value = fake + messages = [{"role": "user", "content": "test"}] + + @gen_cache + def generator(self, model, messages, stream, tools, **kwargs): + return '["alpha", "beta"]' + + generator(None, "m", messages, stream=False, tools=None) + + @stream_cache + def streamer(self, model, messages, stream, tools, **kwargs): + yield "fresh" + + assert list(streamer(None, "m", messages, stream=True, tools=None)) == ["fresh"] diff --git a/tests/vectorstore/test_pgvector_schema_and_pool.py b/tests/vectorstore/test_pgvector_schema_and_pool.py index 9f3c6816..ba4d499c 100644 --- a/tests/vectorstore/test_pgvector_schema_and_pool.py +++ b/tests/vectorstore/test_pgvector_schema_and_pool.py @@ -99,6 +99,43 @@ class TestConstructionTouchesNothing: assert store._pool_max_size == 8 +@pytest.mark.unit +class TestPoolSizingHasOneHome: + """``pgconn`` owns the pool defaults; both stores bind to it, not copies. + + ``pool_for`` keys its registry by DSN and honours ``max_size`` only when it + first builds the pool, so if the vector store and the graph store carried + their own defaults the effective size would depend on which one touched + the DSN first. + """ + + def test_both_stores_share_one_implementation(self): + from application.graphrag import store as store_module + from application.vectorstore import pgconn + + assert pgvector_module.DEFAULT_POOL_MAX_SIZE is pgconn.DEFAULT_POOL_MAX_SIZE + assert store_module.DEFAULT_POOL_MAX_SIZE is pgconn.DEFAULT_POOL_MAX_SIZE + assert store_module._resolve_pool_max_size is pgconn.resolve_pool_max_size + assert ( + pgvector_module.PGVectorStore._resolve_pool_max_size + is pgconn.resolve_pool_max_size + ) + + @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): + from application.core import settings as settings_module + from application.vectorstore import pgconn + + monkeypatch.setattr( + settings_module.settings, "PGVECTOR_POOL_MAX_SIZE", value, raising=False + ) + + assert pgconn.resolve_pool_max_size() == expected + + @pytest.mark.unit class TestWritePathEnsuresSchemaOnce: def test_add_texts_ensures_once_across_two_calls(self):