vrr/agent-zero
2
0
Fork 0
mirror of https://github.com/agent0ai/agent-zero.git synced 2026-05-04 22:50:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

document query tool pausable

Browse source
This commit is contained in:
linuztx 2025-10-11 18:14:17 +08:00
parent 4283f5a937
commit 587b06ed33
1 changed files with 11 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

11
python/helpers/document_query.py
View file

@ -364,12 +364,15 @@ class DocumentQueryHelper:
self, document_uri: str, questions: Sequence[str]
) -> Tuple[bool, str]:
self.progress_callback(f"Starting Q&A process")
await self.agent.handle_intervention()
# index document
_ = await self.document_get_content(document_uri, True)
await self.agent.handle_intervention()
selected_chunks = {}
for question in questions:
self.progress_callback(f"Optimizing query: {question}")
await self.agent.handle_intervention()
human_content = f'Search Query: "{question}"'
system_content = self.agent.parse_prompt(
"fw.document_query.optmimize_query.md"
@ -381,6 +384,7 @@ class DocumentQueryHelper:
)
).strip()
await self.agent.handle_intervention()
self.progress_callback(f"Searching document with query: {optimized_query}")
normalized_uri = self.store.normalize_uri(document_uri)
@ -404,6 +408,7 @@ class DocumentQueryHelper:
self.progress_callback(
f"Processing {len(questions)} questions in context of {len(selected_chunks)} chunks"
)
await self.agent.handle_intervention()
questions_str = "\n".join([f" * {question}" for question in questions])
content = "\n\n----\n\n".join(
@ -430,6 +435,7 @@ class DocumentQueryHelper:
self, document_uri: str, add_to_db: bool = False
) -> str:
self.progress_callback(f"Fetching document content")
await self.agent.handle_intervention()
url = urlparse(document_uri)
scheme = url.scheme or "file"
mimetype, encoding = mimetypes.guess_type(document_uri)
@ -455,6 +461,7 @@ class DocumentQueryHelper:
await asyncio.sleep(1)
last_error = str(e)
retries += 1
await self.agent.handle_intervention()
if not response:
raise ValueError(
@ -492,9 +499,11 @@ class DocumentQueryHelper:
# Use the store's normalization method
document_uri_norm = self.store.normalize_uri(document_uri)
await self.agent.handle_intervention()
exists = await self.store.document_exists(document_uri_norm)
document_content = ""
if not exists:
await self.agent.handle_intervention()
if mimetype.startswith("image/"):
document_content = self.handle_image_document(document_uri, scheme)
elif mimetype == "text/html":
@ -509,6 +518,7 @@ class DocumentQueryHelper:
)
if add_to_db:
self.progress_callback(f"Indexing document")
await self.agent.handle_intervention()
success, ids = await self.store.add_document(
document_content, document_uri_norm
)
@ -519,6 +529,7 @@ class DocumentQueryHelper:
)
self.progress_callback(f"Indexed {len(ids)} chunks")
else:
await self.agent.handle_intervention()
doc = await self.store.get_document(document_uri_norm)
if doc:
document_content = doc.page_content