Enhance document_query to support multiple documents

linuztx committed Oct 18, 2025 at 16:04 UTC 63b98451821e28d2789fcd521ae7b1c8f5dfc1eb
3 files changed +85 -56
prompts/agent.system.tool.document_query.md
+43 -41
@@ -1,60 +1,62 @@
1 -### document_query:
2 -This tool can be used to read or analyze remote and local documents.
3 -It can be used to:
4 - * Get webpage or remote document text content
5 - * Get local document text content
6 - * Answer queries about a webpage, remote or local document
7 -By default, when the "queries" argument is empty, this tool returns the text content of the document retrieved using OCR.
8 -Additionally, you can pass a list of "queries" - in this case, the tool returns the answers to all the passed queries about the document.
9 -!!! This is a universal document reader qnd query tool
10 -!!! Supported document formats: HTML, PDF, Office Documents (word,excel, powerpoint), Textfiles and many more.
1 +### document_query
2 +read and analyze remote/local documents get text content or answer questions
3 +pass a single url/path or a list for multiple documents in "document"
4 +for web documents use "http://" or "https://"" prefix
5 +for local files "file://" prefix is optional but full path is required
6 +if "queries" is empty tool returns document content
7 +if "queries" is a list of strings tool returns answers
8 +supports various formats HTML PDF Office Text etc
9 +usage:
10
12 -#### Arguments:
13 - * "document" (string) : The web address or local path to the document in question. Webdocuments need "http://" or "https://" protocol prefix. For local files the "file:" protocol prefix is optional. Local files MUST be passed with full filesystem path.
14 - * "queries" (Optional, list[str]) : Optionally, here you can pass one or more queries to be answered (using and/or about) the document
15 -
16 -#### Usage example 1:
17 -##### Request:
18 -```json
11 +1 get content
12 +~~~json
13 {
14 "thoughts": [
21 - "...",
15 + "I need to read..."
16 ],
23 - "headline": "Reading web document content",
17 + "headline": "...",
18 "tool_name": "document_query",
19 "tool_args": {
26 - "document": "https://...somexample",
20 + "document": "https://.../document"
21 }
22 }
29 -```
30 -##### Response:
31 -```plaintext
32 -... Here is the entire content of the web document requested ...
33 -```
23 +~~~
24
35 -#### Usage example 2:
36 -##### Request:
37 -```json
25 +2 query document
26 +~~~json
27 {
28 "thoughts": [
40 - "...",
29 + "I need to answer..."
30 ],
42 - "headline": "Analyzing document to answer specific questions",
31 + "headline": "...",
32 "tool_name": "document_query",
33 "tool_args": {
45 - "document": "https://...somexample",
34 + "document": "https://.../document",
35 "queries": [
47 - "What is the topic?",
48 - "Who is the audience?"
36 + "What is...",
37 + "Who is..."
38 ]
39 }
40 }
52 -```
53 -##### Response:
54 -```plaintext
55 -# What is the topic?
56 -... Description of the document topic ...
41 +~~~
42
58 -# Who is the audience?
59 -... The intended document audience list with short descriptions ...
60 -```
43 +3 query multiple documents
44 +~~~json
45 +{
46 + "thoughts": [
47 + "I need to compare..."
48 + ],
49 + "headline": "...",
50 + "tool_name": "document_query",
51 + "tool_args": {
52 + "document": [
53 + "https://.../document-one",
54 + "file:///path/to/document-two"
55 + ],
56 + "queries": [
57 + "Compare the main conclusions...",
58 + "What are the key differences..."
59 + ]
60 + }
61 +}
62 +~~~
python/helpers/document_query.py
+18 -10
@@ -361,12 +361,16 @@ class DocumentQueryHelper:
361 self.progress_callback = progress_callback or (lambda x: None)
362
363 async def document_qa(
364 - self, document_uri: str, questions: Sequence[str]
364 + self, document_uris: List[str], questions: Sequence[str]
365 ) -> Tuple[bool, str]:
366 - self.progress_callback(f"Starting Q&A process")
366 + self.progress_callback(
367 + f"Starting Q&A process for {len(document_uris)} documents"
368 + )
369
368 - # index document
369 - _ = await self.document_get_content(document_uri, True)
370 + # index documents
371 + await asyncio.gather(
372 + *[self.document_get_content(uri, True) for uri in document_uris]
373 + )
374 selected_chunks = {}
375 for question in questions:
376 self.progress_callback(f"Optimizing query: {question}")
@@ -381,14 +385,18 @@ class DocumentQueryHelper:
385 )
386 ).strip()
387
384 - self.progress_callback(f"Searching document with query: {optimized_query}")
388 + self.progress_callback(f"Searching documents with query: {optimized_query}")
389 +
390 + normalized_uris = [self.store.normalize_uri(uri) for uri in document_uris]
391 + doc_filter = " or ".join(
392 + [f"document_uri == '{uri}'" for uri in normalized_uris]
393 + )
394
386 - normalized_uri = self.store.normalize_uri(document_uri)
387 - chunks = await self.store.search_document(
388 - document_uri=normalized_uri,
395 + chunks = await self.store.search_documents(
396 query=optimized_query,
397 limit=100,
398 threshold=DEFAULT_SEARCH_THRESHOLD,
399 + filter=doc_filter,
400 )
401
402 self.progress_callback(f"Found {len(chunks)} chunks")
@@ -397,8 +405,8 @@ class DocumentQueryHelper:
405 selected_chunks[chunk.metadata["id"]] = chunk
406
407 if not selected_chunks:
400 - self.progress_callback(f"No relevant content found in the document")
401 - content = f"!!! No content found for document: {document_uri} matching queries: {json.dumps(questions)}"
408 + self.progress_callback("No relevant content found in the documents")
409 + content = f"!!! No content found for documents: {json.dumps(document_uris)} matching queries: {json.dumps(questions)}"
410 return False, content
411
412 self.progress_callback(
python/tools/document_query.py
+24 -5
@@ -1,3 +1,5 @@
1 +import asyncio
2 +
3 from python.helpers.tool import Tool, Response
4 from python.helpers.document_query import DocumentQueryHelper
5
@@ -5,10 +7,24 @@ from python.helpers.document_query import DocumentQueryHelper
7 class DocumentQueryTool(Tool):
8
9 async def execute(self, **kwargs):
8 - document_uri = kwargs["document"] or None
9 - queries = kwargs["queries"] if "queries" in kwargs else [kwargs["query"]] if ("query" in kwargs and kwargs["query"]) else []
10 - if not isinstance(document_uri, str) or not document_uri:
10 + document_uri = kwargs.get("document")
11 + document_uris = []
12 +
13 + if isinstance(document_uri, list):
14 + document_uris = document_uri
15 + elif isinstance(document_uri, str):
16 + document_uris = [document_uri]
17 +
18 + if not document_uris:
19 return Response(message="Error: no document provided", break_loop=False)
20 +
21 + queries = (
22 + kwargs["queries"]
23 + if "queries" in kwargs
24 + else [kwargs["query"]]
25 + if ("query" in kwargs and kwargs["query"])
26 + else []
27 + )
28 try:
29
30 progress = []
@@ -20,9 +36,12 @@ class DocumentQueryTool(Tool):
36
37 helper = DocumentQueryHelper(self.agent, progress_callback)
38 if not queries:
23 - content = await helper.document_get_content(document_uri)
39 + contents = await asyncio.gather(
40 + *[helper.document_get_content(uri) for uri in document_uris]
41 + )
42 + content = "\n\n---\n\n".join(contents)
43 else:
25 - _, content = await helper.document_qa(document_uri, queries)
44 + _, content = await helper.document_qa(document_uris, queries)
45 return Response(message=content, break_loop=False)
46 except Exception as e: # pylint: disable=broad-exception-caught
47 return Response(message=f"Error processing document: {e}", break_loop=False)