mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-11-01 04:25:44 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a6c55b6c47 |
No files matched your search
+34
-31
@@ -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)
|
||||
@@ -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",
|
||||
]
|
||||
Reference in new issue
Block a user