Fix document query index reuse
Reuse the DocumentQueryStore per chat context so repeated document_query calls can see an existing vector DB instead of re-parsing and re-embedding the same document every time. Add a regression test that proves a second store lookup in the same context reuses the indexed document while a separate context stays isolated.
Alessandro committed
Jun 3, 2026 at 01:55 UTC
94065dbe8602d00218cc5839334dc8ce69bf8bf5
2 files changed
+105
-6
plugins/_document_query/helpers/document_query.py
+18
-3
@@ -72,15 +72,30 @@ def _load_config(agent: Agent) -> dict:
72
class DocumentQueryStore:
73
"""FAISS Store for document query results."""
74
75
+ CONTEXT_DATA_KEY = "_document_query_store"
76
DEFAULT_CHUNK_SIZE = 1000
77
DEFAULT_CHUNK_OVERLAP = 100
78
DEFAULT_MAX_INDEX_CHUNKS = 1200
79
+ _GET_LOCK = threading.RLock()
80
79
- @staticmethod
80
- def get(agent: Agent):
81
+ @classmethod
82
+ def get(cls, agent: Agent):
83
if not agent or not agent.config:
84
raise ValueError("Agent and agent config must be provided")
83
- return DocumentQueryStore(agent)
85
+
86
+ context = getattr(agent, "context", None)
87
+ if context is None:
88
+ return cls(agent)
89
+
90
+ with cls._GET_LOCK:
91
+ store = context.get_data(cls.CONTEXT_DATA_KEY, recursive=False)
92
+ if not isinstance(store, cls):
93
+ store = cls(agent)
94
+ context.set_data(cls.CONTEXT_DATA_KEY, store, recursive=False)
95
+ else:
96
+ store.agent = agent
97
+ store.config = _load_config(agent)
98
+ return store
99
100
def __init__(self, agent: Agent):
101
self.agent = agent
tests/test_document_query_plugin.py
+87
-3
@@ -1,13 +1,19 @@
1
from __future__ import annotations
2
3
import asyncio
4
+import sys
5
from pathlib import Path
6
7
import pytest
8
from PIL import Image
9
10
+ROOT = Path(__file__).resolve().parents[1]
11
+if str(ROOT) not in sys.path:
12
+ sys.path.insert(0, str(ROOT))
13
+
14
from plugins._document_query import hooks as document_query_hooks
15
from plugins._document_query.helpers.fetch import FetchedDocument, fetch_public_resource
16
+import plugins._document_query.helpers.document_query as document_query_module
17
from plugins._document_query.helpers.document_query import (
18
DocumentQueryHelper,
19
DocumentQueryStore,
@@ -19,9 +25,6 @@ from plugins._document_query.helpers.parsers.liteparse import LiteParseParser
25
from plugins._document_query.helpers.parsers.text import TextParser
26
27
22
-ROOT = Path(__file__).resolve().parents[1]
23
-
24
-
28
def run_async(coro):
29
with asyncio.Runner() as runner:
30
return runner.run(coro)
@@ -52,6 +55,52 @@ class CountingAsyncParser(BaseParser):
55
return document.uri
56
57
58
+class _StoreContext:
59
+ def __init__(self, context_id: str):
60
+ self.id = context_id
61
+ self.data = {}
62
+
63
+ def get_data(self, key: str, recursive: bool = True):
64
+ return self.data.get(key)
65
+
66
+ def set_data(self, key: str, value, recursive: bool = True):
67
+ self.data[key] = value
68
+
69
+
70
+class _StoreAgent:
71
+ def __init__(self, context_id: str):
72
+ self.config = object()
73
+ self.context = _StoreContext(context_id)
74
+
75
+
76
+class _FakeVectorDB:
77
+ def __init__(self):
78
+ self.docs = []
79
+
80
+ async def insert_documents(self, docs):
81
+ ids = []
82
+ for doc in docs:
83
+ doc_id = f"doc-{len(self.docs)}"
84
+ doc.metadata["id"] = doc_id
85
+ ids.append(doc_id)
86
+ self.docs.append(doc)
87
+ return ids
88
+
89
+ async def search_by_metadata(self, filter: str, limit: int = 0):
90
+ document_uri = filter.split("'", 2)[1]
91
+ docs = [
92
+ doc
93
+ for doc in self.docs
94
+ if doc.metadata.get("document_uri") == document_uri
95
+ ]
96
+ return docs[:limit] if limit > 0 else docs
97
+
98
+ async def delete_documents_by_ids(self, ids: list[str]):
99
+ removed = [doc for doc in self.docs if doc.metadata.get("id") in ids]
100
+ self.docs = [doc for doc in self.docs if doc.metadata.get("id") not in ids]
101
+ return removed
102
+
103
+
104
def test_fetch_file_detects_mimetype_and_reads_once(tmp_path):
105
document = tmp_path / "notes.txt"
106
document.write_text("hello\nworld\n", encoding="utf-8")
@@ -196,6 +245,41 @@ def test_document_query_allows_uncapped_index_chunks():
245
assert len(chunks) > 10
246
247
248
+def test_document_query_store_reuses_vector_db_per_context(monkeypatch):
249
+ monkeypatch.setattr(
250
+ document_query_module,
251
+ "_load_config",
252
+ lambda _agent: {
253
+ "chunk_size": 100,
254
+ "chunk_overlap": 10,
255
+ "max_index_chunks": 20,
256
+ },
257
+ )
258
+ monkeypatch.setattr(
259
+ DocumentQueryStore,
260
+ "init_vector_db",
261
+ lambda _self: _FakeVectorDB(),
262
+ )
263
+
264
+ agent = _StoreAgent("ctx-one")
265
+ store = DocumentQueryStore.get(agent)
266
+
267
+ success, ids = run_async(
268
+ store.add_document("alpha beta gamma " * 20, "/tmp/book.txt")
269
+ )
270
+ second_store = DocumentQueryStore.get(agent)
271
+
272
+ assert success is True
273
+ assert ids
274
+ assert second_store is store
275
+ assert second_store.vector_db is store.vector_db
276
+ assert run_async(second_store.document_exists("/tmp/book.txt")) is True
277
+
278
+ isolated_store = DocumentQueryStore.get(_StoreAgent("ctx-two"))
279
+ assert isolated_store is not store
280
+ assert run_async(isolated_store.document_exists("/tmp/book.txt")) is False
281
+
282
+
283
def test_document_query_thumbnail_matches_plugin_hub_limits():
284
thumbnail = ROOT / "plugins" / "_document_query" / "webui" / "thumbnail.jpg"
285