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", + ]