mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-11-01 12:35:43 +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)
|
str(pk) for pk in documents.values_list("pk", flat=True)
|
||||||
)
|
)
|
||||||
|
|
||||||
# Hold the shared read lock for the whole operation so the compaction swap
|
# Hold the shared read lock only while reading the index, so the compaction
|
||||||
# cannot replace the database file while this connection is open.
|
# 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:
|
with read_store() as store:
|
||||||
index = load_or_build_index(config, store)
|
index = load_or_build_index(config, store)
|
||||||
retriever = VectorIndexRetriever(
|
retriever = VectorIndexRetriever(
|
||||||
@@ -156,37 +158,38 @@ def _stream_chat_with_documents(
|
|||||||
# #12976.
|
# #12976.
|
||||||
with db_connection_released():
|
with db_connection_released():
|
||||||
top_nodes = retriever.retrieve(query_str)
|
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))
|
references = _get_document_references(documents, top_nodes)
|
||||||
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,
|
|
||||||
)
|
|
||||||
|
|
||||||
logger.debug("Document chat query: %s", query_str)
|
prompt_template = PromptTemplate(template=_build_chat_prompt(output_language))
|
||||||
# Release the pooled DB connection for the slow streaming LLM response
|
refine_template = PromptTemplate(template=_build_refine_prompt(output_language))
|
||||||
# so it is not pinned for the whole stream; see paperless_ai.db and
|
response_synthesizer = get_response_synthesizer(
|
||||||
# #12976.
|
llm=client.llm,
|
||||||
with db_connection_released():
|
prompt_helper=get_rag_prompt_helper(
|
||||||
response_stream = response_synthesizer.synthesize(query_str, top_nodes)
|
chunk_size=config.llm_embedding_chunk_size,
|
||||||
for chunk in response_stream.response_gen:
|
context_size=config.llm_context_size,
|
||||||
yield chunk
|
),
|
||||||
sys.stdout.flush()
|
text_qa_template=prompt_template,
|
||||||
|
refine_template=refine_template,
|
||||||
|
streaming=True,
|
||||||
|
)
|
||||||
|
|
||||||
if references:
|
logger.debug("Document chat query: %s", query_str)
|
||||||
yield _format_chat_metadata_trailer(references)
|
# 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