| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import ipaddress |
| 5 | import sys |
| 6 | from pathlib import Path |
| 7 | |
| 8 | import pytest |
| 9 | from PIL import Image |
| 10 | |
| 11 | ROOT = Path(__file__).resolve().parents[1] |
| 12 | if str(ROOT) not in sys.path: |
| 13 | sys.path.insert(0, str(ROOT)) |
| 14 | |
| 15 | from helpers import network as network_helper |
| 16 | from plugins._document_query import hooks as document_query_hooks |
| 17 | from plugins._document_query.helpers.fetch import FetchedDocument, fetch_public_resource |
| 18 | import plugins._document_query.helpers.document_query as document_query_module |
| 19 | from plugins._document_query.helpers.document_query import ( |
| 20 | DocumentQueryHelper, |
| 21 | DocumentQueryStore, |
| 22 | ) |
| 23 | from plugins._document_query.helpers.parsers.base import BaseParser |
| 24 | from plugins._document_query.helpers.parsers import get_parsers_for_mimetype |
| 25 | from plugins._document_query.helpers.parsers import liteparse as liteparse_module |
| 26 | from plugins._document_query.helpers.parsers.liteparse import LiteParseParser |
| 27 | from plugins._document_query.helpers.parsers.text import TextParser |
| 28 | |
| 29 | |
| 30 | def run_async(coro): |
| 31 | with asyncio.Runner() as runner: |
| 32 | return runner.run(coro) |
| 33 | |
| 34 | |
| 35 | class ParserNameShouldNotLeak(BaseParser): |
| 36 | mimetypes = ["text/plain"] |
| 37 | |
| 38 | def _parse_sync(self, document: FetchedDocument, config: dict) -> str: |
| 39 | return "parsed" |
| 40 | |
| 41 | |
| 42 | class CountingAsyncParser(BaseParser): |
| 43 | mimetypes = ["text/plain"] |
| 44 | active = 0 |
| 45 | max_active = 0 |
| 46 | |
| 47 | async def _parse_async(self, document: FetchedDocument, config: dict) -> str: |
| 48 | type(self).active += 1 |
| 49 | type(self).max_active = max(type(self).max_active, type(self).active) |
| 50 | try: |
| 51 | await asyncio.sleep(0.02) |
| 52 | return document.uri |
| 53 | finally: |
| 54 | type(self).active -= 1 |
| 55 | |
| 56 | def _parse_sync(self, document: FetchedDocument, config: dict) -> str: |
| 57 | return document.uri |
| 58 | |
| 59 | |
| 60 | class _StoreContext: |
| 61 | def __init__(self, context_id: str): |
| 62 | self.id = context_id |
| 63 | self.data = {} |
| 64 | |
| 65 | def get_data(self, key: str, recursive: bool = True): |
| 66 | return self.data.get(key) |
| 67 | |
| 68 | def set_data(self, key: str, value, recursive: bool = True): |
| 69 | self.data[key] = value |
| 70 | |
| 71 | |
| 72 | class _StoreAgent: |
| 73 | def __init__(self, context_id: str): |
| 74 | self.config = object() |
| 75 | self.context = _StoreContext(context_id) |
| 76 | |
| 77 | |
| 78 | class _FakeVectorDB: |
| 79 | def __init__(self): |
| 80 | self.docs = [] |
| 81 | |
| 82 | async def insert_documents(self, docs): |
| 83 | ids = [] |
| 84 | for doc in docs: |
| 85 | doc_id = f"doc-{len(self.docs)}" |
| 86 | doc.metadata["id"] = doc_id |
| 87 | ids.append(doc_id) |
| 88 | self.docs.append(doc) |
| 89 | return ids |
| 90 | |
| 91 | async def search_by_metadata(self, filter: str, limit: int = 0): |
| 92 | document_uri = filter.split("'", 2)[1] |
| 93 | docs = [ |
| 94 | doc |
| 95 | for doc in self.docs |
| 96 | if doc.metadata.get("document_uri") == document_uri |
| 97 | ] |
| 98 | return docs[:limit] if limit > 0 else docs |
| 99 | |
| 100 | async def delete_documents_by_ids(self, ids: list[str]): |
| 101 | removed = [doc for doc in self.docs if doc.metadata.get("id") in ids] |
| 102 | self.docs = [doc for doc in self.docs if doc.metadata.get("id") not in ids] |
| 103 | return removed |
| 104 | |
| 105 | |
| 106 | def test_fetch_file_detects_mimetype_and_reads_once(tmp_path): |
| 107 | document = tmp_path / "notes.txt" |
| 108 | document.write_text("hello\nworld\n", encoding="utf-8") |
| 109 | |
| 110 | fetched = run_async(fetch_public_resource(str(document), {})) |
| 111 | |
| 112 | assert fetched.scheme == "file" |
| 113 | assert fetched.mimetype == "text/plain" |
| 114 | assert fetched.local_path == str(document) |
| 115 | assert fetched.text() == "hello\nworld\n" |
| 116 | |
| 117 | |
| 118 | def test_fetch_http_blocks_non_public_destinations(): |
| 119 | with pytest.raises(ValueError, match="Blocked non-public address"): |
| 120 | run_async( |
| 121 | fetch_public_resource( |
| 122 | "http://127.0.0.1/internal.txt", |
| 123 | {"fetch_retries": 1}, |
| 124 | ) |
| 125 | ) |
| 126 | |
| 127 | |
| 128 | def test_fetch_http_blocks_redirects_to_non_public_destinations(monkeypatch): |
| 129 | source = "https://public.example/start" |
| 130 | calls = [] |
| 131 | |
| 132 | class RedirectResponse: |
| 133 | status_code = 302 |
| 134 | headers = {"Location": "http://127.0.0.1/internal.txt"} |
| 135 | encoding = None |
| 136 | |
| 137 | def __enter__(self): |
| 138 | return self |
| 139 | |
| 140 | def __exit__(self, *_args): |
| 141 | return False |
| 142 | |
| 143 | class FakeSession: |
| 144 | trust_env = True |
| 145 | |
| 146 | def get(self, url, **kwargs): |
| 147 | calls.append((url, kwargs, self.trust_env)) |
| 148 | return RedirectResponse() |
| 149 | |
| 150 | def resolve_host_ips(hostname): |
| 151 | address = "127.0.0.1" if hostname == "127.0.0.1" else "93.184.216.34" |
| 152 | return (ipaddress.ip_address(address),) |
| 153 | |
| 154 | monkeypatch.setattr(network_helper, "resolve_host_ips", resolve_host_ips) |
| 155 | monkeypatch.setattr(network_helper.requests, "Session", FakeSession) |
| 156 | |
| 157 | with pytest.raises(ValueError, match="Blocked non-public address"): |
| 158 | run_async(fetch_public_resource(source, {"fetch_retries": 1})) |
| 159 | |
| 160 | assert [url for url, _kwargs, _trust_env in calls] == [source] |
| 161 | |
| 162 | |
| 163 | def test_fetch_http_preserves_public_redirects_and_request_compatibility(monkeypatch): |
| 164 | source = "https://public.example/start" |
| 165 | destination = "https://public.example/report.txt" |
| 166 | calls = [] |
| 167 | |
| 168 | class FakeResponse: |
| 169 | def __init__(self, status_code, headers, body=b""): |
| 170 | self.status_code = status_code |
| 171 | self.headers = headers |
| 172 | self.encoding = "utf-8" |
| 173 | self.body = body |
| 174 | |
| 175 | def __enter__(self): |
| 176 | return self |
| 177 | |
| 178 | def __exit__(self, *_args): |
| 179 | return False |
| 180 | |
| 181 | def iter_content(self, chunk_size): |
| 182 | assert chunk_size == 64 * 1024 |
| 183 | yield self.body |
| 184 | |
| 185 | class FakeSession: |
| 186 | trust_env = True |
| 187 | |
| 188 | def get(self, url, **kwargs): |
| 189 | calls.append((url, kwargs, self.trust_env)) |
| 190 | if url == source: |
| 191 | return FakeResponse(302, {"Location": "/report.txt"}) |
| 192 | return FakeResponse( |
| 193 | 200, |
| 194 | { |
| 195 | "Content-Length": "9", |
| 196 | "Content-Type": "text/plain; charset=utf-8", |
| 197 | }, |
| 198 | b"public ok", |
| 199 | ) |
| 200 | |
| 201 | monkeypatch.setattr( |
| 202 | network_helper, |
| 203 | "resolve_host_ips", |
| 204 | lambda _hostname: (ipaddress.ip_address("93.184.216.34"),), |
| 205 | ) |
| 206 | monkeypatch.setattr(network_helper.requests, "Session", FakeSession) |
| 207 | monkeypatch.setenv("USER_AGENT", "AgentZeroTest") |
| 208 | |
| 209 | fetched = run_async( |
| 210 | fetch_public_resource( |
| 211 | source, |
| 212 | { |
| 213 | "fetch_retries": 1, |
| 214 | "fetch_timeout": 2, |
| 215 | "max_remote_bytes": 1024, |
| 216 | }, |
| 217 | ) |
| 218 | ) |
| 219 | |
| 220 | assert fetched.uri == destination |
| 221 | assert fetched.mimetype == "text/plain" |
| 222 | assert fetched.text() == "public ok" |
| 223 | assert [url for url, _kwargs, _trust_env in calls] == [source, destination] |
| 224 | assert all(not trust_env for _url, _kwargs, trust_env in calls) |
| 225 | assert all( |
| 226 | kwargs["allow_redirects"] is False |
| 227 | for _url, kwargs, _trust_env in calls |
| 228 | ) |
| 229 | assert all(kwargs["timeout"] == (2.0, 2.0) for _url, kwargs, _trust_env in calls) |
| 230 | assert all( |
| 231 | kwargs["headers"] == {"User-Agent": "AgentZeroTest"} |
| 232 | for _url, kwargs, _trust_env in calls |
| 233 | ) |
| 234 | |
| 235 | |
| 236 | def test_parser_registry_prefers_liteparse_for_pdf(): |
| 237 | parsers = get_parsers_for_mimetype("application/pdf", {"liteparse_enabled": True}) |
| 238 | |
| 239 | assert [parser.__class__.__name__ for parser in parsers[:2]] == [ |
| 240 | "LiteParseParser", |
| 241 | "PdfParser", |
| 242 | ] |
| 243 | |
| 244 | |
| 245 | def test_parser_registry_can_disable_liteparse(): |
| 246 | parsers = get_parsers_for_mimetype("application/pdf", {"liteparse_enabled": False}) |
| 247 | |
| 248 | assert parsers |
| 249 | assert parsers[0].__class__.__name__ == "PdfParser" |
| 250 | |
| 251 | |
| 252 | def test_text_parser_uses_prefetched_content(): |
| 253 | fetched = FetchedDocument( |
| 254 | uri="/tmp/example.json", |
| 255 | source_uri="/tmp/example.json", |
| 256 | scheme="file", |
| 257 | mimetype="application/json", |
| 258 | content=b'{"ok": true}', |
| 259 | local_path=None, |
| 260 | ) |
| 261 | |
| 262 | text = run_async(TextParser().parse(fetched, {}, timeout=1)) |
| 263 | |
| 264 | assert text == '{"ok": true}' |
| 265 | |
| 266 | |
| 267 | def test_compatibility_imports_point_to_plugin_classes(): |
| 268 | pytest.importorskip("langchain_core") |
| 269 | |
| 270 | from helpers.document_query import DocumentQueryHelper as CompatHelper |
| 271 | from plugins._document_query.helpers.document_query import DocumentQueryHelper |
| 272 | from plugins._document_query.tools.document_query import DocumentQueryTool |
| 273 | from tools.document_query import DocumentQueryTool as CompatTool |
| 274 | |
| 275 | assert CompatHelper is DocumentQueryHelper |
| 276 | assert CompatTool is DocumentQueryTool |
| 277 | |
| 278 | |
| 279 | def test_liteparse_is_installed_by_docker_and_plugin_hook_requirements(): |
| 280 | root_requirements = (ROOT / "requirements.txt").read_text(encoding="utf-8") |
| 281 | hooks_source = ( |
| 282 | ROOT / "plugins" / "_document_query" / "hooks.py" |
| 283 | ).read_text(encoding="utf-8") |
| 284 | |
| 285 | assert "liteparse==2.0.3" in root_requirements |
| 286 | assert "_ROOT_REQUIREMENTS_FILE" in hooks_source |
| 287 | assert document_query_hooks._liteparse_requirement() == "liteparse==2.0.3" |
| 288 | assert not (ROOT / "plugins" / "_document_query" / "requirements.txt").exists() |
| 289 | |
| 290 | |
| 291 | def test_default_config_bounds_liteparse_runtime_concurrency(): |
| 292 | default_config = ( |
| 293 | ROOT / "plugins" / "_document_query" / "default_config.yaml" |
| 294 | ).read_text(encoding="utf-8") |
| 295 | |
| 296 | assert "parser_concurrency: 1" in default_config |
| 297 | assert "context_intro_chunks: 2" in default_config |
| 298 | assert "max_index_chunks: 1200" in default_config |
| 299 | assert "liteparse_num_workers: 2" in default_config |
| 300 | assert "liteparse_ocr_auto_disable_pages: 30" in default_config |
| 301 | assert "liteparse_subprocess" not in default_config |
| 302 | |
| 303 | |
| 304 | def test_config_panel_exposes_document_query_settings(): |
| 305 | config_html = ( |
| 306 | ROOT / "plugins" / "_document_query" / "webui" / "config.html" |
| 307 | ).read_text(encoding="utf-8") |
| 308 | |
| 309 | assert "Max parser concurrency" in config_html |
| 310 | for setting in [ |
| 311 | "parser_concurrency", |
| 312 | "per_document_timeout", |
| 313 | "gather_timeout", |
| 314 | "chunk_size", |
| 315 | "chunk_overlap", |
| 316 | "max_index_chunks", |
| 317 | "search_threshold", |
| 318 | "search_limit", |
| 319 | "context_intro_chunks", |
| 320 | "fetch_timeout", |
| 321 | "fetch_retries", |
| 322 | "fetch_retry_backoff", |
| 323 | "max_remote_bytes", |
| 324 | "liteparse_enabled", |
| 325 | "liteparse_ocr_enabled", |
| 326 | "liteparse_ocr_language", |
| 327 | "liteparse_ocr_server_url", |
| 328 | "liteparse_tessdata_path", |
| 329 | "liteparse_max_pages", |
| 330 | "liteparse_target_pages", |
| 331 | "liteparse_dpi", |
| 332 | "liteparse_preserve_very_small_text", |
| 333 | "liteparse_output_format", |
| 334 | "liteparse_num_workers", |
| 335 | "pdf_ocr_fallback", |
| 336 | "thread_offload", |
| 337 | ]: |
| 338 | assert f"config.{setting}" in config_html |
| 339 | assert "liteparse_subprocess" not in config_html |
| 340 | |
| 341 | |
| 342 | def test_document_query_adapts_chunk_size_for_large_documents(): |
| 343 | store = object.__new__(DocumentQueryStore) |
| 344 | store.config = { |
| 345 | "chunk_size": 100, |
| 346 | "chunk_overlap": 10, |
| 347 | "max_index_chunks": 10, |
| 348 | } |
| 349 | |
| 350 | chunks = store._split_text_for_index(("alpha beta gamma delta " * 400).strip()) |
| 351 | |
| 352 | assert 1 < len(chunks) <= 10 |
| 353 | |
| 354 | |
| 355 | def test_document_query_allows_uncapped_index_chunks(): |
| 356 | store = object.__new__(DocumentQueryStore) |
| 357 | store.config = { |
| 358 | "chunk_size": 100, |
| 359 | "chunk_overlap": 10, |
| 360 | "max_index_chunks": 0, |
| 361 | } |
| 362 | |
| 363 | chunks = store._split_text_for_index(("alpha beta gamma delta " * 400).strip()) |
| 364 | |
| 365 | assert len(chunks) > 10 |
| 366 | |
| 367 | |
| 368 | def test_document_query_store_reuses_vector_db_per_context(monkeypatch): |
| 369 | monkeypatch.setattr( |
| 370 | document_query_module, |
| 371 | "_load_config", |
| 372 | lambda _agent: { |
| 373 | "chunk_size": 100, |
| 374 | "chunk_overlap": 10, |
| 375 | "max_index_chunks": 20, |
| 376 | }, |
| 377 | ) |
| 378 | monkeypatch.setattr( |
| 379 | DocumentQueryStore, |
| 380 | "init_vector_db", |
| 381 | lambda _self: _FakeVectorDB(), |
| 382 | ) |
| 383 | |
| 384 | agent = _StoreAgent("ctx-one") |
| 385 | store = DocumentQueryStore.get(agent) |
| 386 | |
| 387 | success, ids = run_async( |
| 388 | store.add_document("alpha beta gamma " * 20, "/tmp/book.txt") |
| 389 | ) |
| 390 | second_store = DocumentQueryStore.get(agent) |
| 391 | |
| 392 | assert success is True |
| 393 | assert ids |
| 394 | assert second_store is store |
| 395 | assert second_store.vector_db is store.vector_db |
| 396 | assert run_async(second_store.document_exists("/tmp/book.txt")) is True |
| 397 | |
| 398 | isolated_store = DocumentQueryStore.get(_StoreAgent("ctx-two")) |
| 399 | assert isolated_store is not store |
| 400 | assert run_async(isolated_store.document_exists("/tmp/book.txt")) is False |
| 401 | |
| 402 | |
| 403 | def test_document_query_thumbnail_matches_plugin_hub_limits(): |
| 404 | thumbnail = ROOT / "plugins" / "_document_query" / "webui" / "thumbnail.jpg" |
| 405 | |
| 406 | assert thumbnail.exists() |
| 407 | assert thumbnail.stat().st_size <= 20 * 1024 |
| 408 | with Image.open(thumbnail) as image: |
| 409 | assert image.format == "JPEG" |
| 410 | assert image.size == (256, 256) |
| 411 | |
| 412 | |
| 413 | def test_liteparse_parser_caps_workers_by_default(): |
| 414 | parser = LiteParseParser() |
| 415 | |
| 416 | assert parser._liteparse_kwargs({})["num_workers"] == 2 |
| 417 | assert parser._liteparse_kwargs({"liteparse_num_workers": "3"})["num_workers"] == 3 |
| 418 | assert parser._liteparse_kwargs({"liteparse_num_workers": ""})["num_workers"] == 2 |
| 419 | |
| 420 | |
| 421 | def test_liteparse_parser_always_uses_subprocess(monkeypatch): |
| 422 | fetched = FetchedDocument( |
| 423 | uri="/tmp/report.pdf", |
| 424 | source_uri="/tmp/report.pdf", |
| 425 | scheme="file", |
| 426 | mimetype="application/pdf", |
| 427 | content=b"", |
| 428 | local_path="/tmp/report.pdf", |
| 429 | ) |
| 430 | parser = LiteParseParser() |
| 431 | |
| 432 | monkeypatch.setattr(parser, "_parse_subprocess", lambda _document, _config: "ok") |
| 433 | |
| 434 | def fail_in_process(_document, _config): |
| 435 | raise AssertionError("LiteParse must stay isolated from the Web UI process") |
| 436 | |
| 437 | monkeypatch.setattr(parser, "_parse_in_process", fail_in_process) |
| 438 | |
| 439 | assert parser._parse_sync(fetched, {"liteparse_subprocess": False}) == "ok" |
| 440 | |
| 441 | |
| 442 | def test_liteparse_auto_disables_ocr_for_large_text_pdf(monkeypatch): |
| 443 | parser = LiteParseParser() |
| 444 | fetched = FetchedDocument( |
| 445 | uri="/tmp/report.pdf", |
| 446 | source_uri="/tmp/report.pdf", |
| 447 | scheme="file", |
| 448 | mimetype="application/pdf", |
| 449 | content=b"", |
| 450 | local_path="/tmp/report.pdf", |
| 451 | ) |
| 452 | monkeypatch.setattr( |
| 453 | liteparse_module, |
| 454 | "_pdf_text_profile", |
| 455 | lambda _file_path, _config: liteparse_module._PdfTextProfile( |
| 456 | page_count=277, |
| 457 | sampled_pages=5, |
| 458 | text_chars=2500, |
| 459 | ), |
| 460 | ) |
| 461 | |
| 462 | kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/report.pdf") |
| 463 | |
| 464 | assert kwargs["ocr_enabled"] is False |
| 465 | |
| 466 | |
| 467 | def test_liteparse_keeps_ocr_for_small_pdf(monkeypatch): |
| 468 | parser = LiteParseParser() |
| 469 | fetched = FetchedDocument( |
| 470 | uri="/tmp/bill.pdf", |
| 471 | source_uri="/tmp/bill.pdf", |
| 472 | scheme="file", |
| 473 | mimetype="application/pdf", |
| 474 | content=b"", |
| 475 | local_path="/tmp/bill.pdf", |
| 476 | ) |
| 477 | monkeypatch.setattr( |
| 478 | liteparse_module, |
| 479 | "_pdf_text_profile", |
| 480 | lambda _file_path, _config: liteparse_module._PdfTextProfile( |
| 481 | page_count=10, |
| 482 | sampled_pages=5, |
| 483 | text_chars=2500, |
| 484 | ), |
| 485 | ) |
| 486 | |
| 487 | kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/bill.pdf") |
| 488 | |
| 489 | assert kwargs["ocr_enabled"] is True |
| 490 | |
| 491 | |
| 492 | def test_liteparse_disables_ocr_for_large_text_sparse_pdf(monkeypatch): |
| 493 | parser = LiteParseParser() |
| 494 | fetched = FetchedDocument( |
| 495 | uri="/tmp/scan.pdf", |
| 496 | source_uri="/tmp/scan.pdf", |
| 497 | scheme="file", |
| 498 | mimetype="application/pdf", |
| 499 | content=b"", |
| 500 | local_path="/tmp/scan.pdf", |
| 501 | ) |
| 502 | monkeypatch.setattr( |
| 503 | liteparse_module, |
| 504 | "_pdf_text_profile", |
| 505 | lambda _file_path, _config: liteparse_module._PdfTextProfile( |
| 506 | page_count=277, |
| 507 | sampled_pages=5, |
| 508 | text_chars=20, |
| 509 | ), |
| 510 | ) |
| 511 | |
| 512 | kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/scan.pdf") |
| 513 | |
| 514 | assert kwargs["ocr_enabled"] is False |
| 515 | |
| 516 | |
| 517 | def test_liteparse_respects_explicit_ocr_disabled(monkeypatch): |
| 518 | parser = LiteParseParser() |
| 519 | fetched = FetchedDocument( |
| 520 | uri="/tmp/bill.pdf", |
| 521 | source_uri="/tmp/bill.pdf", |
| 522 | scheme="file", |
| 523 | mimetype="application/pdf", |
| 524 | content=b"", |
| 525 | local_path="/tmp/bill.pdf", |
| 526 | ) |
| 527 | monkeypatch.setattr( |
| 528 | liteparse_module, |
| 529 | "_pdf_text_profile", |
| 530 | lambda _file_path, _config: liteparse_module._PdfTextProfile( |
| 531 | page_count=10, |
| 532 | sampled_pages=5, |
| 533 | text_chars=0, |
| 534 | ), |
| 535 | ) |
| 536 | |
| 537 | kwargs = parser._liteparse_kwargs( |
| 538 | {"liteparse_ocr_enabled": False}, |
| 539 | fetched, |
| 540 | "/tmp/bill.pdf", |
| 541 | ) |
| 542 | |
| 543 | assert kwargs["ocr_enabled"] is False |
| 544 | |
| 545 | |
| 546 | def test_liteparse_target_pages_can_keep_ocr_enabled_for_large_pdf(monkeypatch): |
| 547 | parser = LiteParseParser() |
| 548 | fetched = FetchedDocument( |
| 549 | uri="/tmp/report.pdf", |
| 550 | source_uri="/tmp/report.pdf", |
| 551 | scheme="file", |
| 552 | mimetype="application/pdf", |
| 553 | content=b"", |
| 554 | local_path="/tmp/report.pdf", |
| 555 | ) |
| 556 | monkeypatch.setattr( |
| 557 | liteparse_module, |
| 558 | "_pdf_text_profile", |
| 559 | lambda _file_path, _config: liteparse_module._PdfTextProfile( |
| 560 | page_count=277, |
| 561 | sampled_pages=5, |
| 562 | text_chars=2500, |
| 563 | ), |
| 564 | ) |
| 565 | |
| 566 | small_range = parser._liteparse_kwargs( |
| 567 | {"liteparse_target_pages": "1-10"}, |
| 568 | fetched, |
| 569 | "/tmp/report.pdf", |
| 570 | ) |
| 571 | large_range = parser._liteparse_kwargs( |
| 572 | {"liteparse_target_pages": "1-40"}, |
| 573 | fetched, |
| 574 | "/tmp/report.pdf", |
| 575 | ) |
| 576 | |
| 577 | assert small_range["ocr_enabled"] is True |
| 578 | assert large_range["ocr_enabled"] is False |
| 579 | |
| 580 | |
| 581 | def test_query_optimize_prompt_filename_is_spelled_correctly(): |
| 582 | prompt_dir = ROOT / "plugins" / "_document_query" / "prompts" |
| 583 | helper_source = ( |
| 584 | ROOT / "plugins" / "_document_query" / "helpers" / "document_query.py" |
| 585 | ).read_text(encoding="utf-8") |
| 586 | |
| 587 | assert (prompt_dir / "fw.document_query.optimize_query.md").exists() |
| 588 | assert "fw.document_query.optimize_query.md" in helper_source |
| 589 | |
| 590 | |
| 591 | def test_parser_progress_is_user_facing_and_generic(): |
| 592 | fetched = FetchedDocument( |
| 593 | uri="/tmp/example.txt", |
| 594 | source_uri="/tmp/example.txt", |
| 595 | scheme="file", |
| 596 | mimetype="text/plain", |
| 597 | content=b"content", |
| 598 | local_path=None, |
| 599 | ) |
| 600 | progress = [] |
| 601 | helper = object.__new__(DocumentQueryHelper) |
| 602 | helper.config = {} |
| 603 | helper.progress_callback = progress.append |
| 604 | |
| 605 | content = run_async( |
| 606 | helper._parse_document( |
| 607 | document=fetched, |
| 608 | parsers=[ParserNameShouldNotLeak()], |
| 609 | timeout=1, |
| 610 | thread_offload=False, |
| 611 | ) |
| 612 | ) |
| 613 | |
| 614 | assert content == "parsed" |
| 615 | assert progress == ["Parsing document content"] |
| 616 | |
| 617 | |
| 618 | def test_parse_document_limits_parser_concurrency_across_helpers(): |
| 619 | CountingAsyncParser.active = 0 |
| 620 | CountingAsyncParser.max_active = 0 |
| 621 | fetched_a = FetchedDocument( |
| 622 | uri="/tmp/a.txt", |
| 623 | source_uri="/tmp/a.txt", |
| 624 | scheme="file", |
| 625 | mimetype="text/plain", |
| 626 | content=b"a", |
| 627 | local_path=None, |
| 628 | ) |
| 629 | fetched_b = FetchedDocument( |
| 630 | uri="/tmp/b.txt", |
| 631 | source_uri="/tmp/b.txt", |
| 632 | scheme="file", |
| 633 | mimetype="text/plain", |
| 634 | content=b"b", |
| 635 | local_path=None, |
| 636 | ) |
| 637 | helper_a = object.__new__(DocumentQueryHelper) |
| 638 | helper_a.config = {"parser_concurrency": 1} |
| 639 | helper_a.progress_callback = lambda _msg: None |
| 640 | helper_b = object.__new__(DocumentQueryHelper) |
| 641 | helper_b.config = {"parser_concurrency": 1} |
| 642 | helper_b.progress_callback = lambda _msg: None |
| 643 | |
| 644 | async def parse_both(): |
| 645 | return await asyncio.gather( |
| 646 | helper_a._parse_document( |
| 647 | document=fetched_a, |
| 648 | parsers=[CountingAsyncParser()], |
| 649 | timeout=1, |
| 650 | thread_offload=False, |
| 651 | ), |
| 652 | helper_b._parse_document( |
| 653 | document=fetched_b, |
| 654 | parsers=[CountingAsyncParser()], |
| 655 | timeout=1, |
| 656 | thread_offload=False, |
| 657 | ), |
| 658 | ) |
| 659 | |
| 660 | assert sorted(run_async(parse_both())) == ["/tmp/a.txt", "/tmp/b.txt"] |
| 661 | assert CountingAsyncParser.max_active == 1 |
| 662 | |
| 663 | |
| 664 | def test_document_query_prompt_uses_progressive_skill_disclosure(): |
| 665 | from helpers.skills import find_skill |
| 666 | |
| 667 | prompt = ( |
| 668 | ROOT |
| 669 | / "plugins" |
| 670 | / "_document_query" |
| 671 | / "prompts" |
| 672 | / "agent.system.tool.document_query.md" |
| 673 | ).read_text(encoding="utf-8") |
| 674 | main_prompt = (ROOT / "prompts" / "agent.system.main.tips.md").read_text( |
| 675 | encoding="utf-8" |
| 676 | ) |
| 677 | skill = find_skill("document-query", include_content=True) |
| 678 | |
| 679 | assert skill is not None |
| 680 | assert "document_query for Q&A" in main_prompt |
| 681 | assert "specific code files" in main_prompt |
| 682 | assert "use vision_load first for image files" in main_prompt |
| 683 | assert "document_query for image OCR only when vision tools cannot read" in main_prompt |
| 684 | assert "skills_tool:load" in prompt |
| 685 | assert "document-query" in prompt |
| 686 | assert "document_query" in prompt |
| 687 | assert "Use vision tools first" in prompt |
| 688 | assert "fallback OCR" in prompt |
| 689 | assert "answering questions over local or remote documents" in skill.description |
| 690 | assert "fallback OCR" in skill.description |
| 691 | assert "### Answer Questions Over A Document" in skill.content |
| 692 | assert "Use vision tools first" in skill.content |
| 693 | assert "### Fallback OCR After Vision Cannot Read A Document Image" in skill.content |