rag tool progress and optimization
frdel committed
Jun 17, 2025 at 23:05 UTC
11f7c6009fb6bfbc38e2d2af71b38afcbea4ea7a
2 files changed
+43
-14
python/helpers/document_query.py
+34
-13
@@ -10,7 +10,7 @@ os.environ["USER_AGENT"] = "@mixedbread-ai/unstructured" # noqa E402
10
from langchain_unstructured import UnstructuredLoader # noqa E402
11
12
from urllib.parse import urlparse
13
-from typing import Sequence, List, Optional, Tuple
13
+from typing import Callable, Sequence, List, Optional, Tuple
14
from datetime import datetime
15
16
from langchain_community.document_loaders import AsyncHtmlLoader
@@ -109,7 +109,7 @@ class DocumentQueryStore:
109
110
async def add_document(
111
self, text: str, document_uri: str, metadata: dict | None = None
112
- ) -> bool:
112
+ ) -> tuple[bool, list[str]]:
113
"""
114
Add a document to the store with the given URI.
115
@@ -148,7 +148,7 @@ class DocumentQueryStore:
148
149
if not docs:
150
PrintStyle.error(f"No chunks created for document: {document_uri}")
151
- return False
151
+ return False, []
152
153
# Apply rate limiter
154
try:
@@ -165,10 +165,10 @@ class DocumentQueryStore:
165
PrintStyle.standard(
166
f"Added document '{document_uri}' with {len(docs)} chunks"
167
)
168
- return True
168
+ return True, ids
169
except Exception as e:
170
PrintStyle.error(f"Error adding document '{document_uri}': {str(e)}")
171
- return False
171
+ return False, []
172
173
async def get_document(self, document_uri: str) -> Optional[Document]:
174
"""
@@ -368,17 +368,21 @@ class DocumentQueryStore:
368
369
class DocumentQueryHelper:
370
371
- def __init__(self, agent: Agent):
371
+ def __init__(self, agent: Agent, progress_callback: Callable[[str], None] | None = None):
372
self.agent = agent
373
self.store = DocumentQueryStore.get(agent)
374
+ self.progress_callback = progress_callback or (lambda x: None)
375
376
async def document_qa(
377
self, document_uri: str, questions: Sequence[str]
378
) -> Tuple[bool, str]:
379
+ self.progress_callback(f"Starting Q&A process")
380
+
381
# index document
382
_ = await self.document_get_content(document_uri, True)
380
- content = ""
383
+ selected_chunks = {}
384
for question in questions:
385
+ self.progress_callback(f"Optimizing query: {question}")
386
human_content = f'Search Query: "{question}"'
387
system_content = self.agent.parse_prompt(
388
"fw.document_query.optmimize_query.md"
@@ -390,6 +394,8 @@ class DocumentQueryHelper:
394
)
395
).strip()
396
397
+ self.progress_callback(f"Searching document with query: {optimized_query}")
398
+
399
normalized_uri = self.store.normalize_uri(document_uri)
400
chunks = await self.store.search_document(
401
document_uri=normalized_uri,
@@ -397,16 +403,21 @@ class DocumentQueryHelper:
403
limit=100,
404
threshold=DEFAULT_SEARCH_THRESHOLD,
405
)
400
- content += (
401
- "\n\n----\n\n".join([chunk.page_content for chunk in chunks])
402
- + "\n\n----\n\n"
403
- )
406
405
- if not content:
407
+ self.progress_callback(f"Found {len(chunks)} chunks")
408
+
409
+ for chunk in chunks:
410
+ selected_chunks[chunk.metadata["id"]] = chunk
411
+
412
+ if not selected_chunks:
413
+ self.progress_callback(f"No relevant content found in the document")
414
content = f"!!! No content found for document: {document_uri} matching queries: {json.dumps(questions)}"
415
return False, content
416
417
+ self.progress_callback(f"Processing {len(questions)} questions in context of {len(selected_chunks)} chunks")
418
+
419
questions_str = "\n".join([f" * {question}" for question in questions])
420
+ content = "\n\n----\n\n".join([chunk.page_content for chunk in selected_chunks.values()])
421
422
qa_system_message = self.agent.parse_prompt(
423
"fw.document_query.system_prompt.md"
@@ -422,9 +433,12 @@ class DocumentQueryHelper:
433
)
434
)
435
436
+ self.progress_callback(f"Q&A process completed")
437
+
438
return True, str(ai_response)
439
440
async def document_get_content(self, document_uri: str, add_to_db: bool = False) -> str:
441
+ self.progress_callback(f"Fetching document content")
442
url = urlparse(document_uri)
443
scheme = url.scheme or "file"
444
mimetype, encoding = mimetypes.guess_type(document_uri)
@@ -503,7 +517,14 @@ class DocumentQueryHelper:
517
document_uri, scheme
518
)
519
if add_to_db:
506
- await self.store.add_document(document_content, document_uri_norm)
520
+ self.progress_callback(f"Indexing document")
521
+ success, ids = await self.store.add_document(document_content, document_uri_norm)
522
+ if not success:
523
+ self.progress_callback(f"Failed to index document")
524
+ raise ValueError(
525
+ f"DocumentQueryHelper::document_get_content: Failed to index document: {document_uri_norm}"
526
+ )
527
+ self.progress_callback(f"Indexed {len(ids)} chunks")
528
else:
529
doc = await self.store.get_document(document_uri_norm)
530
if doc:
python/tools/document_query.py
+9
-1
@@ -10,7 +10,15 @@ class DocumentQueryTool(Tool):
10
if not isinstance(document_uri, str) or not document_uri:
11
return Response(message="Error: no document provided", break_loop=False)
12
try:
13
- helper = DocumentQueryHelper(self.agent)
13
+
14
+ progress = []
15
+
16
+ # logging callback
17
+ def progress_callback(msg):
18
+ progress.append(msg)
19
+ self.log.update(progress="\n".join(progress))
20
+
21
+ helper = DocumentQueryHelper(self.agent, progress_callback)
22
if not queries:
23
content = await helper.document_get_content(document_uri)
24
else: