| 1 | """Document query helper with thread-safe parsers and configurable timeouts. |
| 2 | |
| 3 | Extracted from the monolithic helpers/document_query.py into a plugin |
| 4 | with parser strategy pattern where every document parser is offloaded to |
| 5 | a thread pool and bounded by configurable timeouts. |
| 6 | """ |
| 7 | |
| 8 | import asyncio |
| 9 | import json |
| 10 | import threading |
| 11 | from datetime import datetime |
| 12 | from typing import Any, Callable, List, Optional, Sequence, Tuple |
| 13 | from urllib.parse import urlparse |
| 14 | |
| 15 | from langchain.schema import SystemMessage, HumanMessage |
| 16 | from langchain.text_splitter import RecursiveCharacterTextSplitter |
| 17 | from langchain_core.documents import Document |
| 18 | |
| 19 | from helpers import files, errors |
| 20 | from helpers.print_style import PrintStyle |
| 21 | from helpers.vector_db import VectorDB |
| 22 | from agent import Agent |
| 23 | |
| 24 | from plugins._document_query.helpers.fetch import FetchedDocument, fetch_public_resource |
| 25 | from plugins._document_query.helpers.parsers import BaseParser, get_parsers_for_mimetype |
| 26 | |
| 27 | |
| 28 | DEFAULT_SEARCH_THRESHOLD = 0.5 |
| 29 | DEFAULT_PARSER_CONCURRENCY = 1 |
| 30 | SMALL_DOCUMENT_FALLBACK_MAX_CHARS = 12000 |
| 31 | _PARSER_SEMAPHORES: dict[tuple[int, int], asyncio.Semaphore] = {} |
| 32 | _PARSER_SEMAPHORES_LOCK = threading.Lock() |
| 33 | |
| 34 | |
| 35 | def _positive_int(value: Any, default: int) -> int: |
| 36 | try: |
| 37 | parsed = int(value) |
| 38 | except (TypeError, ValueError): |
| 39 | return default |
| 40 | return parsed if parsed > 0 else default |
| 41 | |
| 42 | |
| 43 | def _nonnegative_int(value: Any, default: int) -> int: |
| 44 | try: |
| 45 | parsed = int(value) |
| 46 | except (TypeError, ValueError): |
| 47 | return default |
| 48 | return parsed if parsed >= 0 else default |
| 49 | |
| 50 | |
| 51 | def _parser_semaphore(config: dict) -> asyncio.Semaphore: |
| 52 | concurrency = _positive_int( |
| 53 | config.get("parser_concurrency"), |
| 54 | DEFAULT_PARSER_CONCURRENCY, |
| 55 | ) |
| 56 | loop = asyncio.get_running_loop() |
| 57 | key = (id(loop), concurrency) |
| 58 | with _PARSER_SEMAPHORES_LOCK: |
| 59 | semaphore = _PARSER_SEMAPHORES.get(key) |
| 60 | if semaphore is None: |
| 61 | semaphore = asyncio.Semaphore(concurrency) |
| 62 | _PARSER_SEMAPHORES[key] = semaphore |
| 63 | return semaphore |
| 64 | |
| 65 | |
| 66 | def _load_config(agent: Agent) -> dict: |
| 67 | """Load plugin config with fallback to defaults.""" |
| 68 | from helpers.plugins import get_plugin_config |
| 69 | return get_plugin_config("_document_query", agent=agent) or {} |
| 70 | |
| 71 | |
| 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 | |
| 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") |
| 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 |
| 102 | self.vector_db: VectorDB | None = None |
| 103 | self.config = _load_config(agent) |
| 104 | |
| 105 | @staticmethod |
| 106 | def normalize_uri(uri: str) -> str: |
| 107 | normalized = uri.strip() |
| 108 | parsed = urlparse(normalized) |
| 109 | scheme = parsed.scheme or "file" |
| 110 | if scheme == "file": |
| 111 | path = files.fix_dev_path( |
| 112 | normalized.removeprefix("file://").removeprefix("file:") |
| 113 | ) |
| 114 | normalized = f"file://{path}" |
| 115 | elif scheme in ["http", "https"]: |
| 116 | normalized = normalized.replace("http://", "https://") |
| 117 | return normalized |
| 118 | |
| 119 | def init_vector_db(self): |
| 120 | return VectorDB(self.agent, cache=True) |
| 121 | |
| 122 | async def add_document( |
| 123 | self, text: str, document_uri: str, metadata: dict | None = None |
| 124 | ) -> tuple[bool, list[str]]: |
| 125 | document_uri = self.normalize_uri(document_uri) |
| 126 | await self.delete_document(document_uri) |
| 127 | doc_metadata = metadata or {} |
| 128 | doc_metadata["document_uri"] = document_uri |
| 129 | doc_metadata["timestamp"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") |
| 130 | chunks = self._split_text_for_index(text) |
| 131 | docs = [] |
| 132 | for i, chunk in enumerate(chunks): |
| 133 | chunk_metadata = doc_metadata.copy() |
| 134 | chunk_metadata["chunk_index"] = i |
| 135 | chunk_metadata["total_chunks"] = len(chunks) |
| 136 | docs.append(Document(page_content=chunk, metadata=chunk_metadata)) |
| 137 | if not docs: |
| 138 | PrintStyle.error(f"No chunks created for document: {document_uri}") |
| 139 | return False, [] |
| 140 | try: |
| 141 | if not self.vector_db: |
| 142 | self.vector_db = self.init_vector_db() |
| 143 | ids = await self.vector_db.insert_documents(docs) |
| 144 | PrintStyle.standard(f"Added document '{document_uri}' with {len(docs)} chunks") |
| 145 | return True, ids |
| 146 | except Exception as e: |
| 147 | err_text = errors.format_error(e) |
| 148 | PrintStyle.error(f"Error adding document '{document_uri}': {err_text}") |
| 149 | return False, [] |
| 150 | |
| 151 | def _split_text_for_index(self, text: str) -> list[str]: |
| 152 | chunk_size = _positive_int( |
| 153 | self.config.get("chunk_size"), |
| 154 | self.DEFAULT_CHUNK_SIZE, |
| 155 | ) |
| 156 | chunk_overlap = min( |
| 157 | _nonnegative_int( |
| 158 | self.config.get("chunk_overlap"), |
| 159 | self.DEFAULT_CHUNK_OVERLAP, |
| 160 | ), |
| 161 | max(0, chunk_size - 1), |
| 162 | ) |
| 163 | chunks = self._split_text(text, chunk_size, chunk_overlap) |
| 164 | |
| 165 | max_chunks = _nonnegative_int( |
| 166 | self.config.get("max_index_chunks"), |
| 167 | self.DEFAULT_MAX_INDEX_CHUNKS, |
| 168 | ) |
| 169 | if not max_chunks or len(chunks) <= max_chunks: |
| 170 | return chunks |
| 171 | |
| 172 | overlap_ratio = chunk_overlap / chunk_size if chunk_size else 0 |
| 173 | overlap_ratio = max(0, min(overlap_ratio, 0.5)) |
| 174 | target_size = max( |
| 175 | chunk_size + 1, |
| 176 | int(len(text) / max(1, max_chunks * (1 - overlap_ratio))) + 1, |
| 177 | ) |
| 178 | |
| 179 | for _ in range(8): |
| 180 | target_overlap = min(int(target_size * overlap_ratio), target_size - 1) |
| 181 | chunks = self._split_text(text, target_size, target_overlap) |
| 182 | if len(chunks) <= max_chunks or target_size >= len(text): |
| 183 | return chunks |
| 184 | target_size = min(len(text), int(target_size * 1.25) + 1) |
| 185 | |
| 186 | return chunks |
| 187 | |
| 188 | @staticmethod |
| 189 | def _split_text(text: str, chunk_size: int, chunk_overlap: int) -> list[str]: |
| 190 | text_splitter = RecursiveCharacterTextSplitter( |
| 191 | chunk_size=chunk_size, |
| 192 | chunk_overlap=chunk_overlap, |
| 193 | ) |
| 194 | return text_splitter.split_text(text) |
| 195 | |
| 196 | async def get_document(self, document_uri: str) -> Optional[Document]: |
| 197 | if not self.vector_db: |
| 198 | return None |
| 199 | document_uri = self.normalize_uri(document_uri) |
| 200 | docs = await self._get_document_chunks(document_uri) |
| 201 | if not docs: |
| 202 | return None |
| 203 | chunks = sorted(docs, key=lambda x: x.metadata.get("chunk_index", 0)) |
| 204 | full_content = "\n".join(chunk.page_content for chunk in chunks) |
| 205 | metadata = chunks[0].metadata.copy() |
| 206 | metadata.pop("chunk_index", None) |
| 207 | metadata.pop("total_chunks", None) |
| 208 | return Document(page_content=full_content, metadata=metadata) |
| 209 | |
| 210 | async def _get_document_chunks(self, document_uri: str) -> List[Document]: |
| 211 | if not self.vector_db: |
| 212 | return [] |
| 213 | document_uri = self.normalize_uri(document_uri) |
| 214 | chunks = await self.vector_db.search_by_metadata( |
| 215 | filter=f"document_uri == '{document_uri}'", |
| 216 | ) |
| 217 | PrintStyle.standard(f"Found {len(chunks)} chunks for document: {document_uri}") |
| 218 | return chunks |
| 219 | |
| 220 | async def document_exists(self, document_uri: str) -> bool: |
| 221 | if not self.vector_db: |
| 222 | return False |
| 223 | document_uri = self.normalize_uri(document_uri) |
| 224 | chunks = await self._get_document_chunks(document_uri) |
| 225 | return len(chunks) > 0 |
| 226 | |
| 227 | async def delete_document(self, document_uri: str) -> bool: |
| 228 | if not self.vector_db: |
| 229 | return False |
| 230 | document_uri = self.normalize_uri(document_uri) |
| 231 | chunks = await self.vector_db.search_by_metadata( |
| 232 | filter=f"document_uri == '{document_uri}'", |
| 233 | ) |
| 234 | if not chunks: |
| 235 | return False |
| 236 | ids_to_delete = [chunk.metadata["id"] for chunk in chunks] |
| 237 | if ids_to_delete: |
| 238 | dels = await self.vector_db.delete_documents_by_ids(ids_to_delete) |
| 239 | PrintStyle.standard(f"Deleted document '{document_uri}' with {len(dels)} chunks") |
| 240 | return True |
| 241 | return False |
| 242 | |
| 243 | async def search_documents( |
| 244 | self, query: str, limit: int = 10, threshold: float = 0.5, filter: str = "" |
| 245 | ) -> List[Document]: |
| 246 | if not self.vector_db: |
| 247 | return [] |
| 248 | if not query: |
| 249 | return [] |
| 250 | try: |
| 251 | results = await self.vector_db.search_by_similarity_threshold( |
| 252 | query=query, limit=limit, threshold=threshold, filter=filter |
| 253 | ) |
| 254 | PrintStyle.standard(f"Search '{query}' returned {len(results)} results") |
| 255 | return results |
| 256 | except Exception as e: |
| 257 | PrintStyle.error(f"Error searching documents: {str(e)}") |
| 258 | return [] |
| 259 | |
| 260 | async def search_document( |
| 261 | self, document_uri: str, query: str, limit: int = 10, threshold: float = 0.5 |
| 262 | ) -> List[Document]: |
| 263 | return await self.search_documents( |
| 264 | query, limit, threshold, f"document_uri == '{document_uri}'" |
| 265 | ) |
| 266 | |
| 267 | async def list_documents(self) -> List[str]: |
| 268 | if not self.vector_db: |
| 269 | return [] |
| 270 | uris = set() |
| 271 | for doc in self.vector_db.db.get_all_docs().values(): |
| 272 | if isinstance(doc.metadata, dict): |
| 273 | uri = doc.metadata.get("document_uri") |
| 274 | if uri: |
| 275 | uris.add(uri) |
| 276 | return sorted(list(uris)) |
| 277 | |
| 278 | |
| 279 | class DocumentQueryHelper: |
| 280 | |
| 281 | def __init__( |
| 282 | self, agent: Agent, progress_callback: Callable[[str], None] | None = None |
| 283 | ): |
| 284 | self.agent = agent |
| 285 | self.store = DocumentQueryStore.get(agent) |
| 286 | self.progress_callback = progress_callback or (lambda x: None) |
| 287 | self.store_lock = asyncio.Lock() |
| 288 | self.config = _load_config(agent) |
| 289 | |
| 290 | async def document_qa( |
| 291 | self, document_uris: List[str] | str, questions: Sequence[str] | str |
| 292 | ) -> Tuple[bool, str]: |
| 293 | if isinstance(document_uris, str): |
| 294 | document_uris = [document_uris] |
| 295 | if isinstance(questions, str): |
| 296 | questions = [questions] |
| 297 | self.progress_callback(f"Starting Q&A process for {len(document_uris)} documents") |
| 298 | await self.agent.handle_intervention() |
| 299 | |
| 300 | gather_timeout = self.config.get("gather_timeout", 120) |
| 301 | try: |
| 302 | document_contents = await asyncio.wait_for( |
| 303 | asyncio.gather( |
| 304 | *[self.document_get_content(uri, True) for uri in document_uris] |
| 305 | ), |
| 306 | timeout=gather_timeout, |
| 307 | ) |
| 308 | except asyncio.TimeoutError: |
| 309 | raise ValueError(f"Document indexing timed out after {gather_timeout}s") |
| 310 | |
| 311 | await self.agent.handle_intervention() |
| 312 | search_threshold = self.config.get("search_threshold", DEFAULT_SEARCH_THRESHOLD) |
| 313 | search_limit = self.config.get("search_limit", 100) |
| 314 | selected_chunks = {} |
| 315 | normalized_uris = [self.store.normalize_uri(uri) for uri in document_uris] |
| 316 | intro_chunk_count = _positive_int( |
| 317 | self.config.get("context_intro_chunks"), |
| 318 | 2, |
| 319 | ) |
| 320 | for uri in normalized_uris: |
| 321 | for chunk in await self._get_document_intro_chunks(uri, intro_chunk_count): |
| 322 | selected_chunks[chunk.metadata["id"]] = chunk |
| 323 | |
| 324 | for question in questions: |
| 325 | self.progress_callback(f"Optimizing query: {question}") |
| 326 | await self.agent.handle_intervention() |
| 327 | system_content = self.agent.parse_prompt("fw.document_query.optimize_query.md") |
| 328 | optimized_query = ( |
| 329 | await self.agent.call_utility_model( |
| 330 | system=system_content, message=f'Search Query: "{question}"', |
| 331 | ) |
| 332 | ).strip() |
| 333 | |
| 334 | await self.agent.handle_intervention() |
| 335 | self.progress_callback(f"Searching documents with query: {optimized_query}") |
| 336 | doc_filter = " or ".join( |
| 337 | [f"document_uri == '{uri}'" for uri in normalized_uris] |
| 338 | ) |
| 339 | chunks = await self.store.search_documents( |
| 340 | query=optimized_query, limit=search_limit, |
| 341 | threshold=search_threshold, filter=doc_filter, |
| 342 | ) |
| 343 | self.progress_callback(f"Found {len(chunks)} chunks") |
| 344 | for chunk in chunks: |
| 345 | selected_chunks[chunk.metadata["id"]] = chunk |
| 346 | |
| 347 | if not selected_chunks: |
| 348 | fallback_content = self._small_document_fallback_content( |
| 349 | document_uris, |
| 350 | document_contents, |
| 351 | ) |
| 352 | if fallback_content: |
| 353 | self.progress_callback( |
| 354 | "No matching chunks found; using extracted document content" |
| 355 | ) |
| 356 | ai_response = await self._answer_questions_from_content( |
| 357 | fallback_content, |
| 358 | questions, |
| 359 | "extracted document content", |
| 360 | ) |
| 361 | self.progress_callback(f"Q&A process completed") |
| 362 | return True, ai_response |
| 363 | |
| 364 | self.progress_callback("No relevant content found in the documents") |
| 365 | content = f"!!! No content found for documents: {json.dumps(document_uris)} matching queries: {json.dumps(questions)}" |
| 366 | return False, content |
| 367 | |
| 368 | content = "\n\n----\n\n".join( |
| 369 | [chunk.page_content for chunk in selected_chunks.values()] |
| 370 | ) |
| 371 | ai_response = await self._answer_questions_from_content( |
| 372 | content, |
| 373 | questions, |
| 374 | f"{len(selected_chunks)} chunks", |
| 375 | ) |
| 376 | self.progress_callback(f"Q&A process completed") |
| 377 | return True, ai_response |
| 378 | |
| 379 | async def _answer_questions_from_content( |
| 380 | self, |
| 381 | content: str, |
| 382 | questions: Sequence[str], |
| 383 | context_label: str, |
| 384 | ) -> str: |
| 385 | self.progress_callback( |
| 386 | f"Processing {len(questions)} questions in context of {context_label}" |
| 387 | ) |
| 388 | await self.agent.handle_intervention() |
| 389 | questions_str = "\n".join([f" * {question}" for question in questions]) |
| 390 | qa_system_message = self.agent.parse_prompt("fw.document_query.system_prompt.md") |
| 391 | qa_user_message = f"# Document:\n{content}\n\n# Queries:\n{questions_str}" |
| 392 | ai_response, _reasoning = await self.agent.call_chat_model( |
| 393 | messages=[ |
| 394 | SystemMessage(content=qa_system_message), |
| 395 | HumanMessage(content=qa_user_message), |
| 396 | ], |
| 397 | explicit_caching=False, |
| 398 | ) |
| 399 | return str(ai_response) |
| 400 | |
| 401 | @staticmethod |
| 402 | def _small_document_fallback_content( |
| 403 | document_uris: Sequence[str], |
| 404 | document_contents: Sequence[str], |
| 405 | max_chars: int = SMALL_DOCUMENT_FALLBACK_MAX_CHARS, |
| 406 | ) -> str: |
| 407 | blocks = [] |
| 408 | for document_uri, document_content in zip(document_uris, document_contents): |
| 409 | text = (document_content or "").strip() |
| 410 | if text: |
| 411 | blocks.append(f"# Source: {document_uri}\n\n{text}") |
| 412 | |
| 413 | if not blocks: |
| 414 | return "" |
| 415 | |
| 416 | content = "\n\n----\n\n".join(blocks) |
| 417 | if len(content) > max_chars: |
| 418 | return "" |
| 419 | return content |
| 420 | |
| 421 | async def _get_document_intro_chunks( |
| 422 | self, |
| 423 | document_uri: str, |
| 424 | limit: int, |
| 425 | ) -> list[Document]: |
| 426 | if limit <= 0: |
| 427 | return [] |
| 428 | if not hasattr(self.store, "_get_document_chunks"): |
| 429 | return [] |
| 430 | chunks = await self.store._get_document_chunks(document_uri) |
| 431 | return sorted(chunks, key=lambda chunk: chunk.metadata.get("chunk_index", 0))[ |
| 432 | :limit |
| 433 | ] |
| 434 | |
| 435 | async def document_get_content( |
| 436 | self, document_uri: str, add_to_db: bool = False |
| 437 | ) -> str: |
| 438 | self.progress_callback(f"Fetching document content") |
| 439 | await self.agent.handle_intervention() |
| 440 | document = await fetch_public_resource( |
| 441 | document_uri, |
| 442 | self.config, |
| 443 | self.agent.handle_intervention, |
| 444 | ) |
| 445 | document_uri_norm = self.store.normalize_uri(document.uri) |
| 446 | await self.agent.handle_intervention() |
| 447 | exists = await self.store.document_exists(document_uri_norm) |
| 448 | document_content = "" |
| 449 | |
| 450 | if not exists: |
| 451 | await self.agent.handle_intervention() |
| 452 | parsers = get_parsers_for_mimetype(document.mimetype, self.config) |
| 453 | if not parsers: |
| 454 | raise ValueError( |
| 455 | f"No parser found for mimetype '{document.mimetype}' ({document.uri})" |
| 456 | ) |
| 457 | per_doc_timeout = self.config.get("per_document_timeout", 60) |
| 458 | thread_offload = self.config.get("thread_offload", True) |
| 459 | document_content = await self._parse_document( |
| 460 | document=document, |
| 461 | parsers=parsers, |
| 462 | timeout=per_doc_timeout, |
| 463 | thread_offload=thread_offload, |
| 464 | ) |
| 465 | if add_to_db: |
| 466 | self.progress_callback(f"Indexing document") |
| 467 | await self.agent.handle_intervention() |
| 468 | async with self.store_lock: |
| 469 | success, ids = await self.store.add_document( |
| 470 | document_content, document_uri_norm |
| 471 | ) |
| 472 | if not success: |
| 473 | self.progress_callback(f"Failed to index document") |
| 474 | raise ValueError(f"Failed to index document: {document_uri_norm}") |
| 475 | self.progress_callback(f"Indexed {len(ids)} chunks") |
| 476 | else: |
| 477 | await self.agent.handle_intervention() |
| 478 | doc = await self.store.get_document(document_uri_norm) |
| 479 | if doc: |
| 480 | document_content = doc.page_content |
| 481 | else: |
| 482 | raise ValueError(f"Document not found: {document_uri_norm}") |
| 483 | return document_content |
| 484 | |
| 485 | async def _parse_document( |
| 486 | self, |
| 487 | document: FetchedDocument, |
| 488 | parsers: list[BaseParser], |
| 489 | timeout: float, |
| 490 | thread_offload: bool, |
| 491 | ) -> str: |
| 492 | errors_seen = [] |
| 493 | semaphore = _parser_semaphore(self.config) |
| 494 | for parser in parsers: |
| 495 | try: |
| 496 | async with semaphore: |
| 497 | self.progress_callback("Parsing document content") |
| 498 | content = await parser.parse( |
| 499 | document=document, |
| 500 | config=self.config, |
| 501 | timeout=timeout, |
| 502 | thread_offload=thread_offload, |
| 503 | ) |
| 504 | if content: |
| 505 | return content |
| 506 | errors_seen.append(f"{parser.__class__.__name__}: no content") |
| 507 | except Exception as e: |
| 508 | errors_seen.append(f"{parser.__class__.__name__}: {e}") |
| 509 | PrintStyle.error(f"Document parser failed: {errors_seen[-1]}") |
| 510 | |
| 511 | raise ValueError( |
| 512 | f"No parser succeeded for mimetype '{document.mimetype}' ({document.uri}): " |
| 513 | + "; ".join(errors_seen) |
| 514 | ) |