@samitouri / QOSami_Agents_Zero / commits / 94065dbe

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