From a6c55b6c47d33d5b3303d45e58e43321636fd278 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:55:22 -0700 Subject: [PATCH] Fix: release the LLM index read lock before streaming the chat answer Document chat held the shared index read lock for the whole request, including the LLM stream, which can take minutes. A compaction or migration swap needs exclusive access, so it had to wait for every in-flight chat to finish. Now that the answer is synthesized from the nodes already retrieved, nothing after the vector search touches the store. Hold the lock only while loading the index and searching, and release it before synthesis. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01SRX1rCtSx62vSDaiFuBFCM --- src/paperless_ai/chat.py | 65 ++++++++------- src/paperless_ai/tests/test_chat_read_lock.py | 81 +++++++++++++++++++ 2 files changed, 115 insertions(+), 31 deletions(-) create mode 100644 src/paperless_ai/tests/test_chat_read_lock.py diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index 56d1d14ff3..df6ea7b638 100644 --- a/src/paperless_ai/chat.py +++ b/src/paperless_ai/chat.py @@ -141,8 +141,10 @@ def _stream_chat_with_documents( str(pk) for pk in documents.values_list("pk", flat=True) ) - # Hold the shared read lock for the whole operation so the compaction swap - # cannot replace the database file while this connection is open. + # Hold the shared read lock only while reading the index, so the compaction + # swap cannot replace the database file while this connection is open. + # Everything after retrieval works on the nodes already in memory, so a slow + # LLM stream never holds the lock. with read_store() as store: index = load_or_build_index(config, store) retriever = VectorIndexRetriever( @@ -156,37 +158,38 @@ def _stream_chat_with_documents( # #12976. with db_connection_released(): top_nodes = retriever.retrieve(query_str) - if not top_nodes: - logger.warning("No nodes found for the given documents.") - yield CHAT_NO_CONTENT_MESSAGE - return - client = AIClient() + if not top_nodes: + logger.warning("No nodes found for the given documents.") + yield CHAT_NO_CONTENT_MESSAGE + return - references = _get_document_references(documents, top_nodes) + client = AIClient() - prompt_template = PromptTemplate(template=_build_chat_prompt(output_language)) - refine_template = PromptTemplate(template=_build_refine_prompt(output_language)) - response_synthesizer = get_response_synthesizer( - llm=client.llm, - prompt_helper=get_rag_prompt_helper( - chunk_size=config.llm_embedding_chunk_size, - context_size=config.llm_context_size, - ), - text_qa_template=prompt_template, - refine_template=refine_template, - streaming=True, - ) + references = _get_document_references(documents, top_nodes) - logger.debug("Document chat query: %s", query_str) - # Release the pooled DB connection for the slow streaming LLM response - # so it is not pinned for the whole stream; see paperless_ai.db and - # #12976. - with db_connection_released(): - response_stream = response_synthesizer.synthesize(query_str, top_nodes) - for chunk in response_stream.response_gen: - yield chunk - sys.stdout.flush() + prompt_template = PromptTemplate(template=_build_chat_prompt(output_language)) + refine_template = PromptTemplate(template=_build_refine_prompt(output_language)) + response_synthesizer = get_response_synthesizer( + llm=client.llm, + prompt_helper=get_rag_prompt_helper( + chunk_size=config.llm_embedding_chunk_size, + context_size=config.llm_context_size, + ), + text_qa_template=prompt_template, + refine_template=refine_template, + streaming=True, + ) - if references: - yield _format_chat_metadata_trailer(references) + logger.debug("Document chat query: %s", query_str) + # Release the pooled DB connection for the slow streaming LLM response + # so it is not pinned for the whole stream; see paperless_ai.db and + # #12976. + with db_connection_released(): + response_stream = response_synthesizer.synthesize(query_str, top_nodes) + for chunk in response_stream.response_gen: + yield chunk + sys.stdout.flush() + + if references: + yield _format_chat_metadata_trailer(references) diff --git a/src/paperless_ai/tests/test_chat_read_lock.py b/src/paperless_ai/tests/test_chat_read_lock.py new file mode 100644 index 0000000000..ec2c08c775 --- /dev/null +++ b/src/paperless_ai/tests/test_chat_read_lock.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import pytest + +from documents.models import Document +from paperless_ai.chat import stream_chat_with_documents +from paperless_testing.factories import DocumentFactory + +if TYPE_CHECKING: + from collections.abc import Iterator + + import pytest_mock + + +@pytest.mark.django_db +class TestChatReadLock: + def test_read_lock_is_released_before_the_answer_streams( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN a document with indexed content + WHEN a chat question is streamed + THEN the index read lock is held for the vector search only + AND is released before the answer is synthesized and streamed + """ + events: list[str] = [] + document = DocumentFactory.create(title="Test Document", content="ignored") + + @contextmanager + def fake_read_store() -> Iterator[object]: + events.append("lock acquired") + try: + yield mocker.MagicMock() + finally: + events.append("lock released") + + retriever = mocker.MagicMock() + retriever.retrieve.side_effect = lambda _query: ( + events.append("retrieve") + or [ + mocker.MagicMock( + metadata={ + "document_id": str(document.pk), + "title": "Test Document", + }, + ), + ] + ) + + def synthesize(*_args: object) -> object: + events.append("synthesize") + return mocker.MagicMock(response_gen=iter(["chunk"])) + + mocker.patch("paperless_ai.chat.read_store", fake_read_store) + mocker.patch("paperless_ai.chat.AIClient") + mocker.patch("paperless_ai.chat.load_or_build_index") + mocker.patch( + "llama_index.core.retrievers.VectorIndexRetriever", + return_value=retriever, + ) + mocker.patch( + "llama_index.core.response_synthesizers.get_response_synthesizer", + ).return_value.synthesize.side_effect = synthesize + + list( + stream_chat_with_documents( + "What is this?", + Document.objects.filter(pk=document.pk), + ), + ) + + assert events == [ + "lock acquired", + "retrieve", + "lock released", + "synthesize", + ]