From b2cc4e10be8d5c603a46dd69262a9e6d55b8c68e Mon Sep 17 00:00:00 2001 From: Giuseppe La Rocca <52716342+JustBeGiusee@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:48:17 +0200 Subject: [PATCH 1/8] feat: implement long document protocol across profiles and enhance document ingestion - Added `long_document_protocol` to multiple profiles including `aion_std.yaml`, `document_extractor.yaml`, `generic_assistant.yaml`, and `research_docs.yaml` to support efficient processing of lengthy documents. - Updated the `ocr_file` and `doc_ingest` functions to improve handling of multi-page PDFs, ensuring better performance and error management during document extraction. - Introduced a timeout message for document tools to provide clearer guidance on handling timeouts during extraction processes. --- config_std/profiles/aion_std.yaml | 1 + config_std/profiles/document_extractor.yaml | 4 +- config_std/profiles/generic_assistant.yaml | 1 + config_std/profiles/research_docs.yaml | 1 + config_std/skills/long_document_protocol.md | 141 ++++++++ .../long_document/cases/altomonte_rumore.yaml | 97 +++++ .../long_document/cases/synthetic_smoke.yaml | 39 ++ mcp_servers_std/ocr_mcp/server.py | 130 ++++++- src/agent_pipeline.py | 55 ++- src/api/session_uploads.py | 2 + src/api/v1/files.py | 2 + src/benchmarks/cli.py | 26 ++ src/benchmarks/long_document/__init__.py | 5 + src/benchmarks/long_document/runner.py | 215 +++++++++++ src/benchmarks/long_document/scoring.py | 144 ++++++++ src/benchmarks/long_document/synthetic.py | 46 +++ src/main.py | 12 +- src/runtime/mcp_tool_result.py | 42 +++ src/test/test_config_std_profile_skills.py | 34 ++ src/test/test_doc_auto_ingest.py | 90 +++++ src/test/test_doc_ingest.py | 333 ++++++++++++++++++ src/test/test_long_document_eval.py | 148 ++++++++ src/test/test_mcp_tool_result.py | 31 ++ src/tools/doc_auto_ingest.py | 151 ++++++++ src/tools/doc_ingest.py | 273 ++++++++++++++ 25 files changed, 1996 insertions(+), 27 deletions(-) create mode 100644 config_std/skills/long_document_protocol.md create mode 100644 evals/long_document/cases/altomonte_rumore.yaml create mode 100644 evals/long_document/cases/synthetic_smoke.yaml create mode 100644 src/benchmarks/long_document/__init__.py create mode 100644 src/benchmarks/long_document/runner.py create mode 100644 src/benchmarks/long_document/scoring.py create mode 100644 src/benchmarks/long_document/synthetic.py create mode 100644 src/test/test_doc_auto_ingest.py create mode 100644 src/test/test_doc_ingest.py create mode 100644 src/test/test_long_document_eval.py create mode 100644 src/tools/doc_auto_ingest.py create mode 100644 src/tools/doc_ingest.py diff --git a/config_std/profiles/aion_std.yaml b/config_std/profiles/aion_std.yaml index f9038e91..96481c6c 100644 --- a/config_std/profiles/aion_std.yaml +++ b/config_std/profiles/aion_std.yaml @@ -33,6 +33,7 @@ skills: - promotional_graphics - web_research_protocol - incremental_execution_protocol +- long_document_protocol - docx - pdf - xlsx diff --git a/config_std/profiles/document_extractor.yaml b/config_std/profiles/document_extractor.yaml index 746b246b..93ed37df 100644 --- a/config_std/profiles/document_extractor.yaml +++ b/config_std/profiles/document_extractor.yaml @@ -6,7 +6,7 @@ instructions: | ## Workflow 1. **Clarify the goal**: what must be extracted (entities, fields, constraints). If the user does not provide a schema, propose a field list consistent with the document type and ask for confirmation, or apply a default and state it clearly. 2. **Identify the source**: files in the session workspace (`session_sandbox`), pasted text. - 3. **Obtain text**: for scanned PDFs/images or sources that cannot be read as plain text, use **OCR**; when text is already readable (e.g. text files or text-based PDFs), prefer direct reads from the workspace and avoid unnecessary OCR. + 3. **Obtain text**: for any PDF longer than a few pages use **`doc_ingest`** and follow `long_document_protocol` — it extracts the text layer page by page and falls back to OCR only where needed. Reserve `ocr_file` for single images or a small page range. When text is already readable (text files), read it directly and skip OCR entirely. 4. **Extract**: populate only what is **supported by the text**; for missing values use `null` or empty arrays—do not invent. 5. **Validate**: consistent formats (dates `YYYY-MM-DD`, amounts as numbers where possible), internal consistency across derived fields, and notes on ambiguity. @@ -28,9 +28,11 @@ skills: - filesystem_tools_protocol - web_research_protocol - incremental_execution_protocol +- long_document_protocol critical_skills: - core_protocol - artifact_protocol +- long_document_protocol native_tool_groups: - web_research mcp_servers: diff --git a/config_std/profiles/generic_assistant.yaml b/config_std/profiles/generic_assistant.yaml index 837c10ba..4db20490 100644 --- a/config_std/profiles/generic_assistant.yaml +++ b/config_std/profiles/generic_assistant.yaml @@ -24,6 +24,7 @@ skills: - web_research_protocol - incremental_execution_protocol - email_imap_mcp +- long_document_protocol - docx - pdf - xlsx diff --git a/config_std/profiles/research_docs.yaml b/config_std/profiles/research_docs.yaml index 7f50a0e5..db30d2a7 100644 --- a/config_std/profiles/research_docs.yaml +++ b/config_std/profiles/research_docs.yaml @@ -15,6 +15,7 @@ skills: - presentation_design - web_research_protocol - incremental_execution_protocol +- long_document_protocol critical_skills: - core_protocol - artifact_protocol diff --git a/config_std/skills/long_document_protocol.md b/config_std/skills/long_document_protocol.md new file mode 100644 index 00000000..9944acae --- /dev/null +++ b/config_std/skills/long_document_protocol.md @@ -0,0 +1,141 @@ +--- +name: long_document_protocol +description: "Mandatory protocol for finding and extracting facts from PDFs longer than a few pages: doc_ingest into one file per page, document identity check, exhaustive grep, verbatim citations with page numbers." +tags: [core, protocol, documents, pdf, ocr, extraction] +status: verified +source: curated +version: 2 +--- + +# Long Document Protocol + +## When this applies + +Any request of the form "find / extract / summarize X from this document" where the +source is a PDF or scan of **more than ~10 pages**. Typical cases: decrees, +authorizations, contracts, technical reports, regulations. + +If the document is short, or you only need to read it front to back, this protocol is +optional. Everything else below is mandatory. + +## Never do these + +| Shortcut | Why it fails | +|----------|--------------| +| `ocr_file` on a whole multi-page PDF | One vision-model call per page. On a 200-page file it always exceeds the MCP timeout, and the run is lost. | +| A custom `pypdf`/`pdfplumber` script that accumulates text in a variable | Killed by the OOM killer (exit code `-9`) on large files. | +| Concatenating everything into one big `.txt` and grepping it | `sandbox_grep_content` **silently skips** files above `AION_GREP_MAX_FILE_BYTES` (500 KB). You get zero hits and conclude, wrongly, that the term is absent. | +| Answering after the first few grep hits | Prescriptions and clauses repeat across chapters. Partial retrieval reads as a complete answer and is worse than no answer. | + +## Step 1 — Verify extraction (auto-ingest on upload) + +PDF uploads are **automatically extracted** to `derived/docs//pages/pNNNN.txt` +(text layer only). Check the attachments block for `PDF extraction ready` or read +`derived/docs//manifest.json`. + +Call `doc_ingest(relative_path="uploads/.pdf")` **only if**: +- manifest is missing (extraction still running), +- `"partial": true` (resume with `first_page=`), +- `empty_pages` need OCR (`ocr_mode="auto"`), +- or the user re-uploaded with `force=True`. + +Never call `ocr_file` on the whole PDF. Never write custom extraction scripts. + +Leave `write_full` off. A `full.txt` is exactly the file grep would skip. + +## Step 2 — Identity gate (do not skip) + +Read `title_guess` and `first_page_excerpt` from the manifest and confirm the document +is the one the user is asking about — the right plant, company, year, protocol number. + +**If it does not match, stop and tell the user.** Do not proceed, do not name the +deliverable after what the user asked for. Producing a well-formatted report from the +wrong source document is the single most damaging failure mode in this workflow, +because nothing in the output looks wrong. + +Also check `ocr_failed_pages` and `empty_pages`: those pages contain no usable text and +any conclusion of the form "the document does not mention X" is unsound while they +remain unread. + +## Step 3 — Search + +``` +sandbox_grep_content( + pattern="", + relative_root="derived", + glob_filter="docs//pages/*.txt", + max_matches=200, +) +``` + +The file name of each grep hit **is** the page number: `pages/p0101.txt` → page 101. + +Widen the pattern before you run it. One spelling is never enough: + +* case and accents: `[Rr]umor|[Aa]custic|acùstic` +* morphological variants: `rumor(e|osità)|acustic(o|a|he|i)|sonor` +* the domain synonyms a drafter would use: `dB\(A\)|Leq|fonometr|immission|emission` +* numbered markers, if the document uses them: `\[[0-9]{1,3}\]` + +If a search returns `truncated: true`, narrow the glob to a page range and repeat — +do not accept a truncated result set as complete. + +## Step 4 — Read every hit, plus one page either side + +For each distinct page in the hit list: + +``` +sandbox_read_file_chunk(relative_path="derived/docs//pages/p0101.txt") +``` + +Always read `page-1` and `page+1` as well. Clauses run across page breaks: a +prescription that starts at the bottom of one page and ends on the next is silently +truncated if you only read the page that matched. + +## Step 5 — Record findings as you go + +Append to `workspace/_findings.json`, one object per extracted item, never in +your reasoning only: + +```json +{ + "id": "53", + "source_doc": "", + "page": 101, + "section": "8.9 Rumore", + "verbatim_quote": "Il Gestore è tenuto al rispetto dei valori limite...", + "summary": "..." +} +``` + +`source_doc` is mandatory on every record. When more than one document is in the +session — for example a reference example supplied by the user alongside the document +to analyse — it is the only thing that stops content from one leaking into the other. + +Never build a deliverable from what you remember of an earlier document in the +conversation. Re-read the page. + +## Step 6 — Coverage gate before answering + +Before writing the answer or generating a file, verify that **every** page in the hit +list has been either read or explicitly discarded with a reason. If any remain, go back +to step 4. + +State the coverage in the answer: how many pages matched, how many were retained, and +which pages could not be extracted. + +## Step 7 — Deliverable + +Every claim carries its `source_doc` and `page`. For a Word/Excel deliverable, load the +`docx` / `xlsx` skill and build it **from `workspace/_findings.json`**, not from +the conversation, so the citations cannot drift. + +## Troubleshooting + +| Symptom | Action | +|---------|--------| +| `doc_ingest` returns `partial: true` | Call again with `first_page=`. | +| `ocr_failed_pages` is not empty | Retry those pages with `ocr_file(relative_path, first_page=N, last_page=N)`; if OCR is unavailable, say so in the answer. | +| Grep returns nothing | Your pattern is too narrow, or you grepped the wrong root. Confirm with a pattern you know is present, e.g. a word from `first_page_excerpt`. | +| Grep returns `truncated: true` | Split by page range; do not treat it as the full result. | +| Tool timeout | Reduce the page range. Do not repeat the same call unchanged. | diff --git a/evals/long_document/cases/altomonte_rumore.yaml b/evals/long_document/cases/altomonte_rumore.yaml new file mode 100644 index 00000000..a8f59091 --- /dev/null +++ b/evals/long_document/cases/altomonte_rumore.yaml @@ -0,0 +1,97 @@ +title: Altomonte — prescrizioni rumore PIC + PMC +description: > + Golden set from the production failure case (decreto IPPC ~200 pagine). + Requires the real PDF via AION_EVAL_ALTOMONTE_PDF; skipped in CI when unset. + Validate page numbers against your copy before tightening min_recall. +tier: manual +cases: + - id: altomonte_rumore_pic_pmc + skip_if_pdf_missing: true + min_recall: 0.85 + pdf: + source: env + path_env: AION_EVAL_ALTOMONTE_PDF + ingest: + ocr_mode: auto + budget_sec: 90 + identity: + must_contain: + - altomonte + must_not_contain: [] + required_hits: + - label: PIC prescrizione 53 + grep_pattern: "\\[53\\]" + page: 101 + must_contain: + - "[53]" + - label: PIC prescrizione 54 + grep_pattern: "\\[54\\]" + page: 101 + must_contain: + - "[54]" + - label: PIC prescrizione 55 + grep_pattern: "\\[55\\]" + page: 101 + must_contain: + - "[55]" + - label: PIC prescrizione 56 + grep_pattern: "\\[56\\]" + page: 102 + must_contain: + - "[56]" + - label: PIC prescrizione 57 + grep_pattern: "\\[57\\]" + page: 102 + must_contain: + - "[57]" + - label: PIC prescrizione 58 + grep_pattern: "\\[58\\]" + page: 102 + must_contain: + - "[58]" + - label: PIC prescrizione 59 + grep_pattern: "\\[59\\]" + page: 102 + must_contain: + - "[59]" + - label: PMC sezione Rumore + grep_pattern: "(?i)rumor" + page: null + must_contain: + - "monitoraggio" + agent_expectations: + user_query: > + ESTRAI DA DECRETO - PIC E PMC LE PRESCRIZIONI AUTORIZZATIVE RELATIVE AL RUMORE. + Genera un elenco con documento, riferimento e pagina. + output_must_contain: + - "[53]" + - "PIC" + - "PMC" + - "101" + output_must_not_contain: + - "presenzano" + + - id: altomonte_wrong_plant_guard + skip_if_pdf_missing: true + min_recall: 1.0 + pdf: + source: env + path_env: AION_EVAL_ALTOMONTE_PDF + ingest: + ocr_mode: auto + budget_sec: 90 + identity: + must_contain: + - altomonte + must_not_contain: + - presenzano + required_hits: [] + agent_expectations: + user_query: > + Estrai le prescrizioni rumore per l'impianto PRESENZANO da questo decreto. + output_must_not_contain: + - "[53]" + - "altomonte" + notes: > + If the PDF is Altomonte, a correct agent must refuse or ask for the right + document instead of fabricating Presenzano prescriptions. diff --git a/evals/long_document/cases/synthetic_smoke.yaml b/evals/long_document/cases/synthetic_smoke.yaml new file mode 100644 index 00000000..cc1c3a99 --- /dev/null +++ b/evals/long_document/cases/synthetic_smoke.yaml @@ -0,0 +1,39 @@ +title: Long document — synthetic smoke (CI) +description: > + Pipeline eval without LLM: doc_ingest on a 200-page synthetic decree, then + grep for PIC [53] and PMC rumore markers on known pages. +tier: ci +cases: + - id: synthetic_rumore_smoke + min_recall: 1.0 + pdf: + source: synthetic + profile: rumore_decreto + pages: 200 + prescription_page: 101 + pmc_page: 150 + ingest: + ocr_mode: never + budget_sec: 600 + identity: + must_contain: [] + must_not_contain: [] + required_hits: + - label: PIC prescrizione 53 + grep_pattern: "\\[53\\]" + page: 101 + must_contain: + - "[53]" + - "rumore" + - label: PIC sezione 8.9 Rumore + grep_pattern: "8\\.9 Rumore" + page: 101 + must_contain: + - "Parere Istruttorio Conclusivo" + - label: PMC tabella rumore + grep_pattern: "Piano di Monitoraggio" + page: 150 + must_contain: + - "Rumore" + - "monitoraggio acustico" + nl_query: prescrizioni autorizzative relative al rumore PIC e PMC monitoraggio acustico diff --git a/mcp_servers_std/ocr_mcp/server.py b/mcp_servers_std/ocr_mcp/server.py index edad4b1c..bfd985d1 100644 --- a/mcp_servers_std/ocr_mcp/server.py +++ b/mcp_servers_std/ocr_mcp/server.py @@ -145,16 +145,115 @@ async def _ocr_via_api_async( return content if isinstance(content, str) else str(content) +def _clamp_ingest_budget(requested: float) -> float: + """Keep the internal deadline safely below the MCP bridge timeout.""" + bridge = _env_float("AION_MCP_TOOL_RESULT_TIMEOUT", 120.0) + ceiling = max(10.0, bridge - 25.0) + try: + value = float(requested) + except (TypeError, ValueError): + value = 90.0 + if value <= 0: + value = ceiling + return min(max(value, 10.0), ceiling) + + +@mcp.tool() +async def doc_ingest( + relative_path: str, + first_page: int = 1, + last_page: int = 0, + ocr_mode: str = "auto", + budget_sec: float = 90.0, + force: bool = False, + write_full: bool = False, +) -> str: + """ + Extract a PDF into one text file per page under derived/docs//pages/. + + This is the entry point for ANY multi-page document. Prefer it over ocr_file and + over custom extraction scripts: it never loads the whole document in memory, it + skips pages already extracted, and if it runs out of time it returns + ``partial: true`` with ``resume_from`` so the next call continues where it stopped. + + ``ocr_mode``: ``auto`` runs OCR only on pages with no usable text layer (cheap on + born-digital PDFs), ``never`` disables it, ``always`` forces OCR on every page. + ``write_full`` additionally concatenates everything into ``full.txt``; leave it off + unless you need sequential reading, because a single large file is skipped by + ``sandbox_grep_content`` above AION_GREP_MAX_FILE_BYTES. + + Returns a small JSON manifest: page counts, empty pages, the grep pattern to use, + and an excerpt of the first page to confirm the document is the one requested. + """ + import json + + from src.session_workspace import ensure_session_dirs, safe_resolve, session_root + from src.tools.doc_ingest import ingest_document + + sid = _require_session() + ensure_session_dirs(sid) + try: + path = safe_resolve(sid, relative_path, must_exist=True) + except Exception as e: + return json.dumps( + {"ok": False, "error": "path_error", "message": str(e)}, + ensure_ascii=False, + ) + + async def _ocr_page(page_no: int, image_bytes: bytes, mime: str) -> str: + return await _ocr_via_api_async( + image_bytes, + mime, + f"Page {page_no}: Extract all visible text. Preserve reading order.", + ) + + use_ocr = ocr_mode != "never" and _is_advanced_ocr_enabled() + + try: + manifest = await ingest_document( + path, + session_root(sid), + first_page=first_page, + last_page=last_page, + ocr_mode=ocr_mode, + budget_sec=_clamp_ingest_budget(budget_sec), + force=force, + write_full=write_full, + ocr_page=_ocr_page if use_ocr else None, + ) + except Exception as e: + logger.exception("doc_ingest failed for %s", relative_path) + return json.dumps( + {"ok": False, "error": "ingest_failed", "message": str(e)}, + ensure_ascii=False, + ) + + if manifest.get("ok") and not use_ocr and manifest.get("empty_pages_count"): + manifest["warning"] = ( + f"{manifest['empty_pages_count']} page(s) have no text layer and OCR is " + "unavailable (ocr_mode=never or OCR service not configured). Those pages " + "are empty in the extraction." + ) + return json.dumps(manifest, ensure_ascii=False) + + @mcp.tool() async def ocr_file( relative_path: str, instruction: str = "Extract all visible text. Preserve reading order.", max_pages: int = 20, + first_page: int = 1, + last_page: int = 0, ) -> str: """ - Extract text from a session file (uploads/, derived/, workspace/). - ALWAYS use the vision-based OCR model (vLLM/OpenAI) per la massima precisione. - Per i PDF, elabora le pagine in parallelo. + Extract text from a session file (uploads/, derived/, workspace/) via vision OCR. + + For multi-page PDFs prefer ``doc_ingest``: it is far cheaper, writes one file per + page and resumes after a timeout. Use ``ocr_file`` for single images or a small + page range of a scanned PDF. + + Page range (PDF only): ``first_page``/``last_page`` are 1-based and inclusive. + Leave ``last_page=0`` to read ``max_pages`` pages starting at ``first_page``. """ from src.session_workspace import ensure_session_dirs, safe_resolve import asyncio @@ -207,8 +306,15 @@ async def ocr_file( from pdf2image import convert_from_path # Carichiamo le impostazioni o usiamo il parametro - limit = _env_int("AION_OCR_PDF_MAX_PAGES", max_pages) - images = convert_from_path(str(path), first_page=1, last_page=limit) + span = _env_int("AION_OCR_PDF_MAX_PAGES", max_pages) + start = max(1, int(first_page or 1)) + if last_page and int(last_page) >= start: + end = int(last_page) + else: + end = start + max(1, span) - 1 + # Never let an explicit range exceed the configured per-call page budget. + end = min(end, start + max(1, span) - 1) + images = convert_from_path(str(path), first_page=start, last_page=end) # Limit parallel calls to avoid overloading the OCR server sem = asyncio.Semaphore(5) @@ -223,20 +329,26 @@ async def limited_ocr(img_data, mime, page_instr): img.save(buf, format="JPEG", quality=85) img_data = buf.getvalue() tasks.append( - limited_ocr(img_data, "image/jpeg", f"Page {i + 1}: {instruction}") + limited_ocr( + img_data, "image/jpeg", f"Page {start + i}: {instruction}" + ) ) logger.info( - f"Starting parallel OCR (limit 5) for {len(tasks)} pages of {path.name}" + "Starting parallel OCR (limit 5) for pages %d-%d of %s", + start, + start + len(tasks) - 1, + path.name, ) results = await asyncio.gather(*tasks, return_exceptions=True) all_text = [] for i, res in enumerate(results): + page_no = start + i if isinstance(res, Exception): - all_text.append(f"--- PAGE {i + 1} ERROR ---\n{res}") + all_text.append(f"--- PAGE {page_no} ERROR ---\n{res}") else: - all_text.append(f"--- PAGE {i + 1} ---\n{res}") + all_text.append(f"--- PAGE {page_no} ---\n{res}") return "\n\n".join(all_text) except Exception as e: diff --git a/src/agent_pipeline.py b/src/agent_pipeline.py index 49a26082..5b0029e9 100644 --- a/src/agent_pipeline.py +++ b/src/agent_pipeline.py @@ -645,21 +645,60 @@ def _format_attachments_block( "You MUST read, analyze, and consider ALL of these newly uploaded documents to answer the current request. " "CRITICAL RULES:\n" "- NEVER process documents in parallel\n" - "- ALWAYS process documents sequentially by using OCR tool.\n" - "- Call ocr_file on one document at a time\n" - "- Wait for the ocr_file call to complete before processing the next document\n" - "- Do not ignore any of these newly uploaded documents\n" - "Ensure you use your tools (like read_file, ocr, or custom scripts) to inspect all of them." + "- For PDFs, check the extraction manifest below before calling doc_ingest\n" + "- Process documents sequentially\n" + "- Do not ignore any of these newly uploaded documents" ) else: lines.append( "IMPORTANT: A new document has been uploaded in this prompt. " "You MUST read, analyze, and consider this document to answer the current request." - "Always use OCR tools to read the document." - "If OCR tools are not available, try to read the document anyway using any " - "available tool." ) + from src.tools.doc_auto_ingest import load_manifest + from src.tools.doc_ingest import slugify_document_name + + for a in new_files: + rp = a.get("relative_path", "") + mime = (a.get("mime") or "").lower() + if not mime.startswith("application/pdf") and not rp.lower().endswith( + ".pdf" + ): + continue + slug = slugify_document_name( + a.get("original_name") or Path(rp).name + ) + manifest = load_manifest(self.session_id, rp) + if manifest and manifest.get("ok"): + lines.append(f"\n### PDF extraction ready: `{slug}`") + lines.append( + f"- pages_total: {manifest.get('pages_total')}, " + f"written: {manifest.get('pages_written')}, " + f"partial: {manifest.get('partial')}" + ) + if manifest.get("partial") and manifest.get("resume_from"): + lines.append( + f"- INCOMPLETE: call doc_ingest(first_page={manifest['resume_from']}) to resume" + ) + if manifest.get("empty_pages_count"): + lines.append( + f"- empty_pages (no text layer): {manifest.get('empty_pages_count')} " + "— use doc_ingest(ocr_mode='auto') or ocr_file on those pages" + ) + excerpt = (manifest.get("first_page_excerpt") or "").strip() + if excerpt: + lines.append(f"- first_page_excerpt: {excerpt[:300]}") + hint = manifest.get("grep_hint") or ( + f"sandbox_grep_content(pattern=..., relative_root='derived', " + f"glob_filter='docs/{slug}/pages/*.txt')" + ) + lines.append(f"- search: {hint}") + elif mime.startswith("application/pdf") or rp.lower().endswith(".pdf"): + lines.append( + f"\n### PDF `{slug}`: extraction in progress or not started. " + "If pages are missing after a few seconds, call doc_ingest(relative_path=...)." + ) + if old_files: lines.append( "NOTE: The historical files listed above are available in your sandbox but were uploaded in previous turns. " diff --git a/src/api/session_uploads.py b/src/api/session_uploads.py index 72b84b2e..4cd9d5bc 100644 --- a/src/api/session_uploads.py +++ b/src/api/session_uploads.py @@ -13,6 +13,7 @@ from sse_starlette.sse import EventSourceResponse from src.session_workspace import list_dir, save_upload +from src.tools.doc_auto_ingest import schedule_auto_ingest from .auth_login import ChatAuthIdentity, require_chat_auth logger = logging.getLogger(__name__) @@ -69,6 +70,7 @@ async def upload_session_files( for f in files: data = await f.read() meta = save_upload(session_id, f.filename or "upload", data) + schedule_auto_ingest(session_id, meta) out.append(meta) logger.info( diff --git a/src/api/v1/files.py b/src/api/v1/files.py index bedfd7e6..89b59c59 100644 --- a/src/api/v1/files.py +++ b/src/api/v1/files.py @@ -11,6 +11,7 @@ from src.api.auth import AuthContext, Scope, require_scope from src.session_workspace import save_upload from src.storage import get_storage_backend +from src.tools.doc_auto_ingest import schedule_auto_ingest router = APIRouter() @@ -29,6 +30,7 @@ async def upload_files( for f in files: data = await f.read() meta = save_upload(conversation_id, f.filename or "upload", data) + schedule_auto_ingest(conversation_id, meta) key = f"{tenant}/conversations/{conversation_id}/uploads/{uuid.uuid4().hex[:12]}_{meta.get('name', 'file')}" try: backend.put_bytes(key, data, meta.get("mime") or "application/octet-stream") diff --git a/src/benchmarks/cli.py b/src/benchmarks/cli.py index 16709cbb..7fa3c252 100644 --- a/src/benchmarks/cli.py +++ b/src/benchmarks/cli.py @@ -12,6 +12,7 @@ from .general_agent import run_general_agent_benchmark from .longmemeval_v2.runner import run_longmemeval_v2_small +from .long_document.runner import run_long_document_pipeline_eval from .mnemos_bench.runner import run_mnemos_bench from .registry import register_benchmark, BenchmarkSpec, catalog_entries from .longmemeval_v2.prepare import is_dataset_ready @@ -45,6 +46,15 @@ def _register_defaults() -> None: ), _run_mnemos_wrapper, ) + register_benchmark( + BenchmarkSpec( + id="long_document_pipeline", + title="Long document pipeline eval", + description="doc_ingest + grep golden cases (no LLM); see evals/long_document/cases/", + tier="ci", + ), + _run_long_document_wrapper, + ) async def _run_general_wrapper( @@ -95,6 +105,19 @@ async def _run_mnemos_wrapper( ) +async def _run_long_document_wrapper( + run_id: str, + profile_name: str, + config: dict | None = None, + dataset_path: str | None = None, + **_: object, +) -> dict: + del run_id, profile_name, config # pipeline eval is profile-agnostic + if not dataset_path: + dataset_path = "evals/long_document/cases/synthetic_smoke.yaml" + return await run_long_document_pipeline_eval(dataset_path) + + async def _async_main(args: argparse.Namespace) -> int: set_event_loop(asyncio.get_running_loop()) _register_defaults() @@ -139,6 +162,9 @@ async def _async_main(args: argparse.Namespace) -> int: dataset_path=dataset, config=config, ) + elif args.benchmark == "long_document_pipeline": + dataset = args.dataset or "evals/long_document/cases/synthetic_smoke.yaml" + metrics = await run_long_document_pipeline_eval(dataset) else: raise ValueError(f"unknown benchmark: {args.benchmark}") if ( diff --git a/src/benchmarks/long_document/__init__.py b/src/benchmarks/long_document/__init__.py new file mode 100644 index 00000000..ea91672b --- /dev/null +++ b/src/benchmarks/long_document/__init__.py @@ -0,0 +1,5 @@ +"""Long-document benchmark package.""" + +from .runner import run_long_document_pipeline_eval + +__all__ = ["run_long_document_pipeline_eval"] diff --git a/src/benchmarks/long_document/runner.py b/src/benchmarks/long_document/runner.py new file mode 100644 index 00000000..0efe3792 --- /dev/null +++ b/src/benchmarks/long_document/runner.py @@ -0,0 +1,215 @@ +"""Long-document eval runner: doc_ingest + grep against YAML golden cases.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import yaml + +from src.tools.doc_ingest import ingest_document + +from .scoring import score_identity_gate, score_required_hits +from .synthetic import build_rumore_decreto_pdf + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_CASES_DIR = REPO_ROOT / "evals" / "long_document" / "cases" + + +def load_dataset(path: str | Path) -> dict[str, Any]: + p = Path(path) + raw = p.read_text(encoding="utf-8") + if p.suffix in {".yaml", ".yml"}: + data = yaml.safe_load(raw) or {} + else: + data = json.loads(raw) + if not isinstance(data, dict): + raise ValueError(f"dataset must be a mapping: {p}") + return data + + +def _resolve_pdf(case: dict[str, Any], work_dir: Path) -> Path | None: + pdf = case.get("pdf") or {} + source = str(pdf.get("source") or "synthetic") + + if source == "synthetic": + profile = str(pdf.get("profile") or "rumore_decreto") + if profile != "rumore_decreto": + raise ValueError(f"unknown synthetic pdf profile: {profile}") + out = work_dir / str(pdf.get("filename") or "synthetic_decreto.pdf") + return build_rumore_decreto_pdf( + out, + pages=int(pdf.get("pages") or 200), + prescription_page=int(pdf.get("prescription_page") or 101), + pmc_page=int(pdf.get("pmc_page") or 150), + ) + + if source == "env": + env_name = str(pdf.get("path_env") or "AION_EVAL_LONG_DOC_PDF") + raw = os.environ.get(env_name, "").strip() + if not raw: + return None + path = Path(raw).expanduser() + return path if path.is_file() else None + + if source == "path": + path = Path(str(pdf.get("path") or "")).expanduser() + return path if path.is_file() else None + + raise ValueError(f"unsupported pdf.source: {source}") + + +async def _ingest_until_complete( + pdf_path: Path, + session_root: Path, + ingest_cfg: dict[str, Any], +) -> dict[str, Any]: + """Call doc_ingest until ``partial`` is false (resume-friendly).""" + first = int(ingest_cfg.get("first_page") or 1) + last = int(ingest_cfg.get("last_page") or 0) + budget = float(ingest_cfg.get("budget_sec") or 600) + ocr_mode = str(ingest_cfg.get("ocr_mode") or "never") + force = bool(ingest_cfg.get("force") or False) + write_full = bool(ingest_cfg.get("write_full") or False) + + manifest: dict[str, Any] = {"ok": False} + resume_from = first + for _ in range(50): + manifest = await ingest_document( + pdf_path, + session_root, + first_page=resume_from, + last_page=last, + ocr_mode=ocr_mode, + budget_sec=budget, + force=force, + write_full=write_full, + ) + if not manifest.get("ok"): + return manifest + if not manifest.get("partial"): + return manifest + nxt = manifest.get("resume_from") + if not nxt or int(nxt) <= resume_from: + return manifest + resume_from = int(nxt) + manifest["warning"] = "ingest stopped after 50 resume iterations" + return manifest + + +async def run_pipeline_case( + case: dict[str, Any], + *, + work_dir: Path, +) -> dict[str, Any]: + """Run ingest + scoring for a single YAML case.""" + case_id = str(case.get("id") or "case") + pdf_path = _resolve_pdf(case, work_dir) + if pdf_path is None: + skipped = bool(case.get("skip_if_pdf_missing", True)) + return { + "case_id": case_id, + "skipped": skipped, + "reason": "pdf not available (set env or path)", + "recall": None, + "passed": None if skipped else False, + } + + session_root = work_dir / "session" + session_root.mkdir(parents=True, exist_ok=True) + uploads = session_root / "uploads" + uploads.mkdir(parents=True, exist_ok=True) + linked = uploads / pdf_path.name + if not linked.exists(): + try: + linked.hardlink_to(pdf_path) + except OSError: + import shutil + + shutil.copy2(pdf_path, linked) + + ingest_cfg = case.get("ingest") or {} + manifest = await _ingest_until_complete(linked, session_root, ingest_cfg) + if not manifest.get("ok"): + return { + "case_id": case_id, + "skipped": False, + "passed": False, + "recall": 0.0, + "error": manifest, + } + + slug = str(manifest["slug"]) + hit_score = score_required_hits( + session_root, + slug, + list(case.get("required_hits") or []), + ) + identity_score = score_identity_gate(manifest, case.get("identity") or {}) + + passed = ( + hit_score["recall"] >= float(case.get("min_recall") or 1.0) + and identity_score["passed"] + ) + return { + "case_id": case_id, + "skipped": False, + "passed": passed, + "recall": hit_score["recall"], + "slug": slug, + "manifest": { + "pages_total": manifest.get("pages_total"), + "pages_written": manifest.get("pages_written"), + "partial": manifest.get("partial"), + "title_guess": manifest.get("title_guess"), + }, + "hits": hit_score, + "identity": identity_score, + } + + +async def run_long_document_pipeline_eval( + dataset_path: str | Path, + *, + work_dir: Path | None = None, +) -> dict[str, Any]: + """Evaluate all cases in a dataset directory or single YAML file.""" + path = Path(dataset_path) + if path.is_dir(): + case_files = sorted(path.glob("*.yaml")) + sorted(path.glob("*.yml")) + else: + case_files = [path] + + root_work = work_dir or (REPO_ROOT / "data" / "eval_runs" / "long_document") + root_work.mkdir(parents=True, exist_ok=True) + + results: list[dict[str, Any]] = [] + for case_file in case_files: + dataset = load_dataset(case_file) + cases = dataset.get("cases") + if cases is None: + cases = [dataset] + for case in cases: + case_work = root_work / str(case.get("id") or case_file.stem) + case_work.mkdir(parents=True, exist_ok=True) + row = await run_pipeline_case(case, work_dir=case_work) + row["dataset"] = str(case_file) + results.append(row) + + scored = [r for r in results if not r.get("skipped")] + skipped = [r for r in results if r.get("skipped")] + passed = [r for r in scored if r.get("passed")] + recall_vals = [float(r["recall"]) for r in scored if r.get("recall") is not None] + + metrics = { + "case_count": len(results), + "scored_count": len(scored), + "skipped_count": len(skipped), + "passed_count": len(passed), + "accuracy_overall": (len(passed) / len(scored)) if scored else 1.0, + "mean_recall": (sum(recall_vals) / len(recall_vals)) if recall_vals else 1.0, + "results": results, + } + return metrics diff --git a/src/benchmarks/long_document/scoring.py b/src/benchmarks/long_document/scoring.py new file mode 100644 index 00000000..f447afaf --- /dev/null +++ b/src/benchmarks/long_document/scoring.py @@ -0,0 +1,144 @@ +"""Scoring helpers for long-document pipeline evals (ingest + grep, no LLM).""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any + +from src.tools.session_fs_tools import grep_content + + +def page_no_from_hit_file(file_path: str) -> int | None: + """Extract page number from ``derived/docs//pages/p0101.txt``.""" + name = Path(file_path.replace("\\", "/")).name + m = re.match(r"^p(\d{4})\.txt$", name) + return int(m.group(1)) if m else None + + +def grep_pages( + session_root: Path, + slug: str, + pattern: str, + *, + max_matches: int = 500, +) -> list[dict]: + return grep_content( + session_root, + session_root / "derived", + pattern, + glob_filter=f"docs/{slug}/pages/*.txt", + max_matches=max_matches, + ) + + +def score_required_hits( + session_root: Path, + slug: str, + required_hits: list[dict[str, Any]], +) -> dict[str, Any]: + """Score PIC/PMC-style expectations after ``doc_ingest``.""" + details: list[dict[str, Any]] = [] + passed = 0 + + for item in required_hits: + label = str(item.get("label") or item.get("id") or "hit") + pattern = str(item["grep_pattern"]) + expected_page = item.get("page") + must_contain = [str(s) for s in (item.get("must_contain") or [])] + + hits = grep_pages(session_root, slug, pattern) + page_hits = [h for h in hits if page_no_from_hit_file(h["file"]) is not None] + on_page = [ + h + for h in page_hits + if page_no_from_hit_file(h["file"]) == int(expected_page) + ] if expected_page is not None else page_hits + + row: dict[str, Any] = { + "label": label, + "grep_pattern": pattern, + "expected_page": expected_page, + "hit_count": len(hits), + "hits_on_expected_page": len(on_page), + "passed": False, + "reason": "", + } + + if not hits: + row["reason"] = "no grep hits" + details.append(row) + continue + + if expected_page is not None and not on_page: + found_pages = sorted( + {page_no_from_hit_file(h["file"]) for h in page_hits} + - {None} + ) + row["reason"] = f"hits on pages {found_pages}, expected {expected_page}" + details.append(row) + continue + + target_hit = on_page[0] if on_page else hits[0] + page_no = page_no_from_hit_file(target_hit["file"]) + page_path = session_root / target_hit["file"] + body = page_path.read_text(encoding="utf-8", errors="replace") + missing = [s for s in must_contain if s not in body] + if missing: + row["reason"] = f"missing substrings on page {page_no}: {missing}" + row["page"] = page_no + details.append(row) + continue + + row["passed"] = True + row["page"] = page_no + row["file"] = target_hit["file"] + passed += 1 + details.append(row) + + total = len(required_hits) + return { + "required_total": total, + "required_passed": passed, + "recall": (passed / total) if total else 1.0, + "details": details, + } + + +def score_identity_gate(manifest: dict[str, Any], identity: dict[str, Any]) -> dict[str, Any]: + """Check title/excerpt against must / must-not lists.""" + haystack = " ".join( + [ + str(manifest.get("title_guess") or ""), + str(manifest.get("first_page_excerpt") or ""), + str(manifest.get("source") or ""), + ] + ).lower() + + must = [str(s).lower() for s in (identity.get("must_contain") or [])] + must_not = [str(s).lower() for s in (identity.get("must_not_contain") or [])] + + missing = [s for s in must if s not in haystack] + forbidden = [s for s in must_not if s in haystack] + passed = not missing and not forbidden + return { + "passed": passed, + "missing": missing, + "forbidden": forbidden, + } + + +def score_agent_output(text: str, expectations: dict[str, Any]) -> dict[str, Any]: + """Lightweight substring checks on a final agent answer (opt-in LLM eval).""" + low = (text or "").lower() + must = [str(s) for s in (expectations.get("output_must_contain") or [])] + must_not = [str(s) for s in (expectations.get("output_must_not_contain") or [])] + + missing = [s for s in must if s.lower() not in low] + forbidden = [s for s in must_not if s.lower() in low] + passed = not missing and not forbidden + return { + "passed": passed, + "missing": missing, + "forbidden": forbidden, + } diff --git a/src/benchmarks/long_document/synthetic.py b/src/benchmarks/long_document/synthetic.py new file mode 100644 index 00000000..f793ef03 --- /dev/null +++ b/src/benchmarks/long_document/synthetic.py @@ -0,0 +1,46 @@ +"""Synthetic PDF fixtures for long-document eval cases (no external assets).""" + +from __future__ import annotations + +from pathlib import Path + +_FILLER = ( + "Commissione Istruttoria IPPC - Centrale termoelettrica - testo di riempimento " + "per riprodurre la densita tipica di una pagina di decreto autorizzativo." +) + + +def build_rumore_decreto_pdf( + path: Path, + *, + pages: int = 200, + prescription_page: int = 101, + pmc_page: int = 150, + blank_pages: frozenset[int] | None = None, +) -> Path: + """Born-digital decree-like PDF with PIC + PMC noise markers on known pages.""" + from reportlab.lib.pagesizes import A4 + from reportlab.pdfgen import canvas + + blanks = blank_pages if blank_pages is not None else frozenset({7, 42}) + path.parent.mkdir(parents=True, exist_ok=True) + + c = canvas.Canvas(str(path), pagesize=A4) + for page_no in range(1, pages + 1): + if page_no not in blanks: + c.drawString(72, 780, f"PAGINA {page_no} marcatore acustico") + for row, offset in enumerate(range(750, 690, -15)): + c.drawString(72, offset, f"{_FILLER} riga {row}") + if page_no == prescription_page: + c.drawString( + 72, + 660, + "[53] Il Gestore e tenuto al rispetto dei valori limite di rumore", + ) + c.drawString(72, 645, "8.9 Rumore - Parere Istruttorio Conclusivo") + if page_no == pmc_page: + c.drawString(72, 660, "Piano di Monitoraggio e Controllo (PMC)") + c.drawString(72, 645, "Tabella parametri Rumore e monitoraggio acustico") + c.showPage() + c.save() + return path diff --git a/src/main.py b/src/main.py index 5c8aeec0..f532e6df 100644 --- a/src/main.py +++ b/src/main.py @@ -621,17 +621,11 @@ def _emit_tool_outcome(*, is_error: bool, body: str) -> str: ) err_text = format_exception_for_tool(tool_name, e) if isinstance(e, TimeoutError): - pg_cap = os.getenv("AION_PG_QUERY_TIMEOUT_SEC", "60") + from src.runtime.mcp_tool_result import build_timeout_message + err_text = format_exception_for_tool( tool_name, - TimeoutError( - f"Query timed out ({server_name}/{tool_name}). " - f"PostgreSQL cap AION_PG_QUERY_TIMEOUT_SEC={pg_cap}s; " - f"MCP bridge cap AION_MCP_TOOL_RESULT_TIMEOUT=" - f"{os.getenv('AION_MCP_TOOL_RESULT_TIMEOUT', '120')}s. " - "Heavy JOINs may need indexes or a narrower filter (e.g. codice_ditta). " - "The MCP worker was recycled; retry with a simpler query." - ), + TimeoutError(build_timeout_message(server_name, tool_name)), ) mark_sql_exec_tool_failed(session_id, tool_name) diff --git a/src/runtime/mcp_tool_result.py b/src/runtime/mcp_tool_result.py index 5be4d026..476a7e12 100644 --- a/src/runtime/mcp_tool_result.py +++ b/src/runtime/mcp_tool_result.py @@ -180,3 +180,45 @@ def format_exception_for_tool(tool_name: str, exc: BaseException) -> str: "tool": tool_name, } return json.dumps(payload, ensure_ascii=False) + + +_DOC_TOOL_NAMES = frozenset({"ocr_file", "doc_ingest"}) + + +def build_timeout_message(server_name: str, tool_name: str) -> str: + """Timeout guidance tailored to the tool family. + + A generic message that mentions PostgreSQL derails the model when the tool + that timed out is an OCR or document call, so each family gets the remedy + that actually applies to it. + """ + import os + + from src.runtime.pg_query_guard import is_postgres_query_tool + + mcp_cap = os.getenv("AION_MCP_TOOL_RESULT_TIMEOUT", "120") + head = f"Tool timed out ({server_name}/{tool_name}). MCP bridge cap AION_MCP_TOOL_RESULT_TIMEOUT={mcp_cap}s." + + if is_postgres_query_tool(server_name, tool_name): + pg_cap = os.getenv("AION_PG_QUERY_TIMEOUT_SEC", "60") + return ( + f"{head} PostgreSQL cap AION_PG_QUERY_TIMEOUT_SEC={pg_cap}s. " + "Heavy JOINs may need indexes or a narrower filter (e.g. codice_ditta). " + "The MCP worker was recycled; retry with a simpler query." + ) + + base_tool = (tool_name or "").split("-")[-1] + if base_tool in _DOC_TOOL_NAMES: + return ( + f"{head} Document extraction is too slow for the whole file in one call. " + "Use doc_ingest(relative_path, first_page=..., last_page=...) to process a page " + "range: it writes one file per page under derived/docs//pages/, skips pages " + "already done, and can be called again to resume. Do NOT retry the same call " + "unchanged, and do NOT fall back to a custom extraction script." + ) + + return ( + f"{head} The call exceeded the budget. Retry with a narrower scope " + "(fewer items, a smaller range, or a more specific filter) rather than repeating " + "the same arguments." + ) diff --git a/src/test/test_config_std_profile_skills.py b/src/test/test_config_std_profile_skills.py index 36658b1a..57591020 100644 --- a/src/test/test_config_std_profile_skills.py +++ b/src/test/test_config_std_profile_skills.py @@ -17,6 +17,40 @@ def _artifact_skill_loaded(reg: SkillRegistry) -> bool: return bool(reg.get_skill_full("artifact_protocol")) +def _profiles() -> list[tuple[Path, dict]]: + profiles_dir = _repo_root() / "config_std" / "profiles" + out = [] + for path in sorted(profiles_dir.glob("*.yaml")): + out.append((path, yaml.safe_load(path.read_text(encoding="utf-8")) or {})) + return out + + +def test_long_document_protocol_skill_exists(): + reg = SkillRegistry() + reg.reload() + assert reg.get_skill_full("long_document_protocol"), ( + "config_std/skills/long_document_protocol.md must be loadable" + ) + + +def test_profiles_with_ocr_declare_long_document_protocol(): + """A profile that can ingest documents must carry the protocol for long ones. + + Without it the model falls back to ocr_file on the whole PDF, which times out. + """ + # Image-only workflows do not need the multi-page protocol. + exempt = {"graphic_designer.yaml"} + for path, data in _profiles(): + if path.name in exempt: + continue + if "ocr" not in (data.get("mcp_servers") or []): + continue + assert "long_document_protocol" in (data.get("skills") or []), ( + f"{path.name} mounts the ocr MCP server but does not list " + "long_document_protocol" + ) + + def test_config_std_profile_skills_resolve(): reg = SkillRegistry() reg.reload() diff --git a/src/test/test_doc_auto_ingest.py b/src/test/test_doc_auto_ingest.py new file mode 100644 index 00000000..32346c7f --- /dev/null +++ b/src/test/test_doc_auto_ingest.py @@ -0,0 +1,90 @@ +"""Tests for automatic PDF ingest on upload.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from fastapi import UploadFile +from io import BytesIO + +from src.benchmarks.long_document.synthetic import build_rumore_decreto_pdf +from src.session_workspace import save_upload, session_root +from src.tools.doc_auto_ingest import ( + auto_ingest_enabled, + load_manifest, + run_auto_ingest_background, + schedule_auto_ingest, +) +from src.tools.doc_ingest import slugify_document_name + + +@pytest.fixture() +def data_dir(tmp_path, monkeypatch): + monkeypatch.setenv("AION_DATA_DIR", str(tmp_path)) + return tmp_path + + +@pytest.mark.asyncio +async def test_run_auto_ingest_writes_pages(data_dir): + sid = "sess_auto_ingest" + pdf_bytes = BytesIO() + path = build_rumore_decreto_pdf(data_dir / "src.pdf", pages=30) + meta = save_upload(sid, "decreto.pdf", path.read_bytes()) + + await run_auto_ingest_background(sid, meta["relative_path"], meta["mime"]) + + slug = slugify_document_name("decreto.pdf") + pages = session_root(sid) / "derived" / "docs" / slug / "pages" + assert len(list(pages.glob("p*.txt"))) == 30 + + manifest = load_manifest(sid, meta["relative_path"]) + assert manifest is not None + assert manifest.get("ok") is True + assert manifest.get("pages_total") == 30 + + +@pytest.mark.asyncio +async def test_schedule_auto_ingest_skips_non_pdf(data_dir): + sid = "sess_txt" + meta = save_upload(sid, "note.txt", b"hello world") + + async def _noop(): + return None + + # Should not raise; schedules only PDFs + schedule_auto_ingest(sid, meta) + await asyncio.sleep(0.05) + slug_dir = session_root(sid) / "derived" / "docs" + assert not slug_dir.exists() or not any(slug_dir.rglob("pages")) + + +@pytest.mark.asyncio +async def test_upload_endpoint_schedules_ingest(data_dir, monkeypatch): + monkeypatch.setenv("AION_DOC_AUTO_INGEST", "1") + from src.api.session_uploads import upload_session_files + from src.api.auth_login import ChatAuthIdentity + + sid = "sess_upload_ep" + pdf = build_rumore_decreto_pdf(data_dir / "u.pdf", pages=15) + upload = UploadFile(filename="decreto.pdf", file=BytesIO(pdf.read_bytes())) + + class _Auth: + identifier = "tester" + + result = await upload_session_files(sid, files=[upload], _auth=_Auth()) + assert len(result["files"]) == 1 + + # Background task needs a tick + await asyncio.sleep(0.5) + + rel = result["files"][0]["relative_path"] + manifest = load_manifest(sid, rel) + assert manifest is not None + assert manifest.get("pages_total") == 15 + + +def test_auto_ingest_disabled_by_env(monkeypatch): + monkeypatch.setenv("AION_DOC_AUTO_INGEST", "0") + assert auto_ingest_enabled() is False diff --git a/src/test/test_doc_ingest.py b/src/test/test_doc_ingest.py new file mode 100644 index 00000000..0fe24e82 --- /dev/null +++ b/src/test/test_doc_ingest.py @@ -0,0 +1,333 @@ +"""Deterministic coverage for src.tools.doc_ingest (no LLM, no MCP, no OCR service). + +The regression these tests protect against is the one seen in production: a +200-page decree that could not be extracted, and a concatenated text file that +grep silently skips once it grows past AION_GREP_MAX_FILE_BYTES. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from src.tools.doc_ingest import ingest_document, slugify_document_name + +PAGE_COUNT = 200 +# Sentinel reproducing the real task: a numbered prescription on a known page. +PRESCRIPTION_PAGE = 101 +PRESCRIPTION_TEXT = "[53] Il Gestore e tenuto al rispetto dei valori limite di rumore" +BLANK_PAGES = {7, 42} + + +_FILLER = ( + "Commissione Istruttoria IPPC - Centrale termoelettrica - testo di riempimento " + "per riprodurre la densita tipica di una pagina di decreto autorizzativo." +) + + +def _build_pdf(path: Path, pages: int = PAGE_COUNT) -> Path: + from reportlab.lib.pagesizes import A4 + from reportlab.pdfgen import canvas + + c = canvas.Canvas(str(path), pagesize=A4) + for page_no in range(1, pages + 1): + if page_no not in BLANK_PAGES: + c.drawString(72, 780, f"PAGINA {page_no} marcatore acustico") + # Real pages carry hundreds of characters; a sparse fixture would + # trip the missing-text-layer heuristic for the wrong reason. + for row, offset in enumerate(range(750, 690, -15)): + c.drawString(72, offset, f"{_FILLER} riga {row}") + if page_no == PRESCRIPTION_PAGE: + c.drawString(72, 660, PRESCRIPTION_TEXT) + c.showPage() + c.save() + return path + + +@pytest.fixture(scope="module") +def sample_pdf(tmp_path_factory) -> Path: + root = tmp_path_factory.mktemp("docsrc") + return _build_pdf(root / "08b3392800_DECRETOCOMPLETO.pdf") + + +@pytest.fixture() +def session_root(tmp_path) -> Path: + return tmp_path + + +def _pages_dir(session_root: Path, slug: str) -> Path: + return session_root / "derived" / "docs" / slug / "pages" + + +def test_slugify_strips_upload_prefix(): + assert slugify_document_name("08b3392800_DECRETOCOMPLETO.pdf") == "decretocompleto" + assert slugify_document_name("esempio per ai.pdf") == "esempio_per_ai" + assert slugify_document_name("___.pdf") == "document" + + +@pytest.mark.asyncio +async def test_writes_one_file_per_page(sample_pdf, session_root): + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + + assert manifest["ok"] is True + assert manifest["partial"] is False + assert manifest["pages_total"] == PAGE_COUNT + assert manifest["pages_written"] == PAGE_COUNT + + pages = sorted(_pages_dir(session_root, manifest["slug"]).glob("p*.txt")) + assert len(pages) == PAGE_COUNT + assert pages[0].name == "p0001.txt" + assert pages[-1].name == "p0200.txt" + + target = _pages_dir(session_root, manifest["slug"]) / "p0101.txt" + assert "[53]" in target.read_text(encoding="utf-8") + + +@pytest.mark.asyncio +async def test_manifest_is_persisted_and_small(sample_pdf, session_root): + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + + on_disk = ( + session_root / "derived" / "docs" / manifest["slug"] / "manifest.json" + ).read_text(encoding="utf-8") + assert json.loads(on_disk)["slug"] == manifest["slug"] + # Must stay well under the tool-offload threshold (8000 chars) so the model + # sees the manifest inline instead of a pointer to a file. + assert len(json.dumps(manifest, ensure_ascii=False)) < 4000 + + +@pytest.mark.asyncio +async def test_page_files_stay_greppable(sample_pdf, session_root): + """Every page must sit below the grep size cap that silently skips files.""" + from src.tools.session_fs_tools import _grep_max_file_bytes, grep_content + + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + pages_dir = _pages_dir(session_root, manifest["slug"]) + + cap = _grep_max_file_bytes() + assert all(p.stat().st_size < cap for p in pages_dir.glob("p*.txt")) + + hits = grep_content( + session_root, + session_root / "derived", + "marcatore acustico", + glob_filter=f"docs/{manifest['slug']}/pages/*.txt", + max_matches=PAGE_COUNT + 10, + ) + assert len(hits) == PAGE_COUNT - len(BLANK_PAGES) + + +@pytest.mark.asyncio +async def test_grep_hit_filename_carries_page_number(sample_pdf, session_root): + from src.tools.session_fs_tools import grep_content + + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + hits = grep_content( + session_root, + session_root / "derived", + r"\[53\]", + glob_filter=f"docs/{manifest['slug']}/pages/*.txt", + ) + assert len(hits) == 1 + assert hits[0]["file"].endswith(f"p{PRESCRIPTION_PAGE:04d}.txt") + + +@pytest.mark.asyncio +async def test_partial_run_reports_resume_point(sample_pdf, session_root): + ticks = iter(range(0, 10_000)) + manifest = await ingest_document( + sample_pdf, session_root, budget_sec=5, clock=lambda: float(next(ticks)) + ) + + assert manifest["partial"] is True + assert manifest["resume_from"] == 5 + assert manifest["pages_written"] == 4 + assert "doc_ingest again with first_page=5" in manifest["next_step"] + + +@pytest.mark.asyncio +async def test_second_call_resumes_and_skips_existing(sample_pdf, session_root): + ticks = iter(range(0, 10_000)) + first = await ingest_document( + sample_pdf, session_root, budget_sec=5, clock=lambda: float(next(ticks)) + ) + assert first["partial"] is True + + second = await ingest_document(sample_pdf, session_root, budget_sec=600) + + assert second["partial"] is False + assert second["pages_skipped_existing"] == first["pages_written"] + assert second["pages_written"] == PAGE_COUNT - first["pages_written"] + assert len(list(_pages_dir(session_root, second["slug"]).glob("p*.txt"))) == PAGE_COUNT + + +@pytest.mark.asyncio +async def test_ocr_runs_only_on_pages_without_text_layer(sample_pdf, session_root): + seen: list[int] = [] + + async def fake_ocr(page_no: int, image_bytes: bytes, mime: str) -> str: + seen.append(page_no) + assert image_bytes[:4] == b"\x89PNG" + assert mime == "image/png" + return f"OCR PAGINA {page_no}" + + manifest = await ingest_document( + sample_pdf, session_root, budget_sec=600, ocr_page=fake_ocr + ) + + assert sorted(seen) == sorted(BLANK_PAGES) + assert manifest["ocr_pages"] == len(BLANK_PAGES) + assert manifest["text_layer_pages"] == PAGE_COUNT - len(BLANK_PAGES) + assert manifest["empty_pages"] == [] + + recovered = _pages_dir(session_root, manifest["slug"]) / "p0007.txt" + assert recovered.read_text(encoding="utf-8") == "OCR PAGINA 7" + + +@pytest.mark.asyncio +async def test_min_text_chars_controls_the_ocr_trigger(sample_pdf, session_root): + """A page is sent to OCR only when its text layer is below the threshold.""" + seen: list[int] = [] + + async def fake_ocr(page_no: int, image_bytes: bytes, mime: str) -> str: + seen.append(page_no) + return f"OCR PAGINA {page_no}" + + await ingest_document( + sample_pdf, + session_root, + budget_sec=600, + last_page=5, + ocr_page=fake_ocr, + min_text_chars=10_000, + ) + assert seen == [1, 2, 3, 4, 5] + + +@pytest.mark.asyncio +async def test_ocr_mode_always_bypasses_the_text_layer(sample_pdf, session_root): + seen: list[int] = [] + + async def fake_ocr(page_no: int, image_bytes: bytes, mime: str) -> str: + seen.append(page_no) + return f"OCR PAGINA {page_no}" + + manifest = await ingest_document( + sample_pdf, + session_root, + budget_sec=600, + last_page=3, + ocr_mode="always", + ocr_page=fake_ocr, + ) + + assert seen == [1, 2, 3] + assert manifest["ocr_pages"] == 3 + assert manifest["text_layer_pages"] == 0 + + +@pytest.mark.asyncio +async def test_ocr_failure_keeps_the_text_layer(sample_pdf, session_root): + async def failing_ocr(page_no, image_bytes, mime): + raise RuntimeError("OCR service unreachable") + + manifest = await ingest_document( + sample_pdf, session_root, budget_sec=600, ocr_page=failing_ocr + ) + + # Pages with text are untouched; the blank ones are reported as failed rather + # than silently counted as successfully extracted. + assert manifest["ocr_pages"] == 0 + assert manifest["text_layer_pages"] == PAGE_COUNT - len(BLANK_PAGES) + assert manifest["ocr_failed_pages"] == sorted(BLANK_PAGES) + assert manifest["ocr_failed_pages_count"] == len(BLANK_PAGES) + blank = _pages_dir(session_root, manifest["slug"]) / "p0007.txt" + assert "OCR failed for page 7" in blank.read_text(encoding="utf-8") + + +@pytest.mark.asyncio +async def test_ocr_mode_never_reports_empty_pages(sample_pdf, session_root): + async def unexpected_ocr(page_no, image_bytes, mime): # pragma: no cover + raise AssertionError("OCR must not run with ocr_mode='never'") + + manifest = await ingest_document( + sample_pdf, + session_root, + budget_sec=600, + ocr_mode="never", + ocr_page=unexpected_ocr, + ) + + assert manifest["ocr_pages"] == 0 + assert manifest["empty_pages"] == sorted(BLANK_PAGES) + assert manifest["empty_pages_count"] == len(BLANK_PAGES) + + +@pytest.mark.asyncio +async def test_full_text_is_opt_in(sample_pdf, session_root): + default = await ingest_document(sample_pdf, session_root, budget_sec=600) + root = session_root / "derived" / "docs" / default["slug"] + assert not (root / "full.txt").exists() + assert default["full_text"] is None + + with_full = await ingest_document( + sample_pdf, session_root, budget_sec=600, write_full=True + ) + full = root / "full.txt" + assert full.exists() + assert with_full["full_text"].endswith("full.txt") + body = full.read_text(encoding="utf-8") + assert "=== PAGE 101 ===" in body + assert "[53]" in body + + +@pytest.mark.asyncio +async def test_force_reextracts_pages(sample_pdf, session_root): + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + stale = _pages_dir(session_root, manifest["slug"]) / "p0101.txt" + stale.write_text("contenuto obsoleto", encoding="utf-8") + + refreshed = await ingest_document( + sample_pdf, session_root, budget_sec=600, last_page=PRESCRIPTION_PAGE, force=True + ) + + assert refreshed["pages_skipped_existing"] == 0 + assert "[53]" in stale.read_text(encoding="utf-8") + + +@pytest.mark.asyncio +async def test_page_range_limits_extraction(sample_pdf, session_root): + manifest = await ingest_document( + sample_pdf, session_root, budget_sec=600, first_page=10, last_page=12 + ) + + assert manifest["range"] == [10, 12] + assert manifest["pages_written"] == 3 + names = sorted(p.name for p in _pages_dir(session_root, manifest["slug"]).glob("*.txt")) + assert names == ["p0010.txt", "p0011.txt", "p0012.txt"] + + +@pytest.mark.asyncio +async def test_first_page_excerpt_supports_identity_check(sample_pdf, session_root): + manifest = await ingest_document(sample_pdf, session_root, budget_sec=600) + assert "PAGINA 1" in manifest["first_page_excerpt"] + + +@pytest.mark.asyncio +async def test_invalid_inputs_return_structured_errors(sample_pdf, session_root, tmp_path): + bad_mode = await ingest_document(sample_pdf, session_root, ocr_mode="sometimes") + assert bad_mode == { + "ok": False, + "error": "invalid_ocr_mode", + "message": "ocr_mode must be auto|never|always, got 'sometimes'", + } + + missing = await ingest_document(tmp_path / "nope.pdf", session_root) + assert missing["error"] == "not_a_file" + + out_of_range = await ingest_document( + sample_pdf, session_root, first_page=PAGE_COUNT + 5, budget_sec=600 + ) + assert out_of_range["error"] == "empty_range" diff --git a/src/test/test_long_document_eval.py b/src/test/test_long_document_eval.py new file mode 100644 index 00000000..c1ee76a8 --- /dev/null +++ b/src/test/test_long_document_eval.py @@ -0,0 +1,148 @@ +"""Eval harness tests for long-document extraction (pipeline + dataset schema).""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import yaml + +from src.benchmarks.long_document.runner import ( + DEFAULT_CASES_DIR, + load_dataset, + run_long_document_pipeline_eval, + run_pipeline_case, +) +from src.benchmarks.long_document.scoring import ( + page_no_from_hit_file, + score_identity_gate, + score_required_hits, +) +from src.benchmarks.long_document.synthetic import build_rumore_decreto_pdf +from src.tools.doc_ingest import ingest_document + +SYNTHETIC = DEFAULT_CASES_DIR / "synthetic_smoke.yaml" +ALTOMONTE = DEFAULT_CASES_DIR / "altomonte_rumore.yaml" + + +def _validate_case(case: dict, *, path: Path) -> None: + assert case.get("id"), f"{path}: case missing id" + pdf = case.get("pdf") or {} + assert pdf.get("source") in {"synthetic", "env", "path"}, ( + f"{case['id']}: invalid pdf.source" + ) + for hit in case.get("required_hits") or []: + assert hit.get("grep_pattern"), f"{case['id']}: hit missing grep_pattern" + + +@pytest.mark.parametrize("path", [SYNTHETIC, ALTOMONTE]) +def test_long_document_dataset_schema(path: Path): + assert path.is_file(), f"missing {path}" + data = load_dataset(path) + cases = data.get("cases") or [data] + assert cases, f"{path}: no cases" + ids = [c["id"] for c in cases] + assert len(ids) == len(set(ids)), f"{path}: duplicate case ids" + for case in cases: + _validate_case(case, path=path) + + +def test_page_no_from_hit_file(): + assert page_no_from_hit_file("derived/docs/foo/pages/p0101.txt") == 101 + assert page_no_from_hit_file("p0042.txt") == 42 + assert page_no_from_hit_file("readme.txt") is None + + +@pytest.mark.asyncio +async def test_synthetic_smoke_pipeline_eval_passes(tmp_path): + metrics = await run_long_document_pipeline_eval(SYNTHETIC, work_dir=tmp_path) + assert metrics["scored_count"] == 1 + assert metrics["passed_count"] == 1 + assert metrics["accuracy_overall"] == 1.0 + assert metrics["mean_recall"] == 1.0 + + +@pytest.mark.asyncio +async def test_altomonte_cases_skip_without_pdf(tmp_path, monkeypatch): + monkeypatch.delenv("AION_EVAL_ALTOMONTE_PDF", raising=False) + metrics = await run_long_document_pipeline_eval(ALTOMONTE, work_dir=tmp_path) + assert metrics["case_count"] == 2 + assert metrics["skipped_count"] == 2 + assert metrics["scored_count"] == 0 + + +@pytest.mark.asyncio +async def test_scoring_required_hits_on_ingested_pages(tmp_path): + pdf = build_rumore_decreto_pdf(tmp_path / "decreto.pdf", pages=20, prescription_page=10, pmc_page=15) + session = tmp_path / "session" + manifest = await ingest_document(pdf, session, ocr_mode="never", budget_sec=60) + assert manifest["ok"] + + score = score_required_hits( + session, + manifest["slug"], + [ + { + "label": "pic53", + "grep_pattern": "\\[53\\]", + "page": 10, + "must_contain": ["[53]", "rumore"], + } + ], + ) + assert score["recall"] == 1.0 + assert score["details"][0]["passed"] + + +def test_identity_gate_detects_wrong_plant(): + manifest = { + "title_guess": "Decreto impianto Altomonte", + "first_page_excerpt": "Centrale termoelettrica Altomonte", + "source": "decreto.pdf", + } + ok = score_identity_gate( + manifest, + {"must_contain": ["altomonte"], "must_not_contain": ["presenzano"]}, + ) + assert ok["passed"] is True + + bad = score_identity_gate( + manifest, + {"must_contain": ["presenzano"], "must_not_contain": []}, + ) + assert bad["passed"] is False + assert "presenzano" in bad["missing"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not os.getenv("AION_EVAL_LLM"), + reason="Set AION_EVAL_LLM=1 to run agentic long-document eval", +) +async def test_agentic_altomonte_optional(tmp_path): + """End-to-end agent eval — requires LLM, OCR service, and AION_EVAL_ALTOMONTE_PDF.""" + from src.agent_pipeline import AgentPipeline + from src.benchmarks.long_document.scoring import score_agent_output + from src.main import get_agent, set_event_loop + + pdf_env = os.getenv("AION_EVAL_ALTOMONTE_PDF", "").strip() + if not pdf_env or not Path(pdf_env).is_file(): + pytest.skip("AION_EVAL_ALTOMONTE_PDF not set") + + data = yaml.safe_load(ALTOMONTE.read_text(encoding="utf-8")) + case = data["cases"][0] + session_id = f"eval_{case['id']}" + set_event_loop(__import__("asyncio").get_event_loop()) + + # Seed session with ingested PDF via pipeline eval helper first. + await run_pipeline_case(case, work_dir=tmp_path / "seed") + + agent, profile = await get_agent("document_extractor", session_id=session_id, user_id="eval") + pipeline = AgentPipeline( + agent, session_id=session_id, profile_name=profile, user_id="eval" + ) + agent_cfg = case.get("agent_expectations") or {} + res = await pipeline.run(str(agent_cfg.get("user_query") or "")) + scored = score_agent_output(res.get("text", ""), agent_cfg) + assert scored["passed"], scored diff --git a/src/test/test_mcp_tool_result.py b/src/test/test_mcp_tool_result.py index 3141608e..84e457a9 100644 --- a/src/test/test_mcp_tool_result.py +++ b/src/test/test_mcp_tool_result.py @@ -5,6 +5,7 @@ import json from src.runtime.mcp_tool_result import ( + build_timeout_message, classify_tool_result_text, format_exception_for_tool, ) @@ -54,6 +55,36 @@ def test_format_exception_for_tool(): assert "connection reset" in data["message"] +def test_timeout_message_for_postgres_keeps_sql_guidance(): + msg = build_timeout_message("toolbox-postgres", "query") + assert "PostgreSQL cap" in msg + assert "Heavy JOINs" in msg + + +def test_timeout_message_for_document_tools_is_not_about_sql(): + """A PostgreSQL hint on an OCR timeout actively derails the model.""" + for tool in ("ocr_file", "doc_ingest"): + msg = build_timeout_message("ocr", tool) + assert "PostgreSQL" not in msg + assert "JOIN" not in msg + assert "doc_ingest" in msg + assert "first_page" in msg + + +def test_timeout_message_generic_tool_has_no_sql_or_false_recycle_claim(): + msg = build_timeout_message("session_sandbox", "sandbox_run_python_file") + assert "PostgreSQL" not in msg + # Only the Postgres path actually restarts the worker. + assert "recycled" not in msg + assert "AION_MCP_TOOL_RESULT_TIMEOUT" in msg + + +def test_timeout_message_handles_prefixed_tool_names(): + msg = build_timeout_message("ocr", "ocr-doc_ingest") + assert "doc_ingest" in msg + assert "PostgreSQL" not in msg + + def test_classify_skill_view_not_error_despite_keywords(): raw = "Plane Project Management. Error: this tool failed sometimes due to Pydantic Validation exception." is_err, body = classify_tool_result_text(raw, "skill_view") diff --git a/src/tools/doc_auto_ingest.py b/src/tools/doc_auto_ingest.py new file mode 100644 index 00000000..a5f74e34 --- /dev/null +++ b/src/tools/doc_auto_ingest.py @@ -0,0 +1,151 @@ +"""Fire-and-forget PDF text-layer extraction triggered on upload. + +Runs ``ingest_document`` in a worker thread so PyMuPDF does not block the FastAPI +event loop (mandatory with ``--workers 1``). +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from pathlib import Path +from typing import Any + +from src.session_workspace import safe_resolve, session_root +from src.tools.doc_ingest import DOCS_SUBDIR, ingest_document, slugify_document_name + +logger = logging.getLogger("aion.doc_auto_ingest") + +_PDF_MIME = "application/pdf" + + +def auto_ingest_enabled() -> bool: + return os.getenv("AION_DOC_AUTO_INGEST", "1").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _auto_ingest_max_pages() -> int: + try: + return max(1, int(os.getenv("AION_DOC_AUTO_INGEST_MAX_PAGES", "500"))) + except ValueError: + return 500 + + +def _auto_ingest_budget_sec() -> float: + try: + return float(os.getenv("AION_DOC_AUTO_INGEST_BUDGET_SEC", "120")) + except ValueError: + return 120.0 + + +def manifest_path(session_id: str, slug: str) -> Path: + return session_root(session_id) / DOCS_SUBDIR / slug / "manifest.json" + + +def load_manifest(session_id: str, relative_path: str) -> dict[str, Any] | None: + """Return ingest manifest for an uploaded PDF if extraction has run.""" + slug = slugify_document_name(Path(relative_path).name) + path = manifest_path(session_id, slug) + if not path.is_file(): + return None + try: + import json + + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def _ingest_sync(session_id: str, relative_path: str) -> dict[str, Any]: + """Blocking ingest loop with resume until complete or page cap.""" + import asyncio + + path = safe_resolve(session_id, relative_path, must_exist=True) + root = session_root(session_id) + max_pages = _auto_ingest_max_pages() + budget = _auto_ingest_budget_sec() + resume_from = 1 + manifest: dict[str, Any] = {"ok": False} + + async def _run_once(start: int) -> dict[str, Any]: + return await ingest_document( + path, + root, + first_page=start, + last_page=max_pages, + ocr_mode="never", + budget_sec=budget, + force=False, + write_full=False, + ) + + for _ in range(50): + manifest = asyncio.run(_run_once(resume_from)) + if not manifest.get("ok"): + return manifest + if not manifest.get("partial"): + break + nxt = manifest.get("resume_from") + if not nxt or int(nxt) <= resume_from: + break + resume_from = int(nxt) + + return manifest + + +async def run_auto_ingest_background(session_id: str, relative_path: str, mime: str) -> None: + """Schedule-safe background ingest for one uploaded PDF.""" + if not auto_ingest_enabled(): + return + if (mime or "").split(";")[0].strip().lower() != _PDF_MIME: + return + + try: + manifest = await asyncio.to_thread(_ingest_sync, session_id, relative_path) + except Exception as exc: + logger.exception( + "auto_ingest failed session=%s path=%s: %s", + session_id[:8], + relative_path, + exc, + ) + return + + if manifest.get("ok"): + slug = str(manifest.get("slug") or "") + logger.info( + "auto_ingest complete session=%s slug=%s pages=%s partial=%s", + session_id[:8], + slug, + manifest.get("pages_written"), + manifest.get("partial"), + ) + else: + logger.warning( + "auto_ingest error session=%s path=%s: %s", + session_id[:8], + relative_path, + manifest, + ) + + +def schedule_auto_ingest( + session_id: str, + upload_meta: dict[str, Any], +) -> None: + """Fire-and-forget task for a single ``save_upload`` result.""" + mime = str(upload_meta.get("mime") or "") + rel = str(upload_meta.get("relative_path") or "") + if not rel: + return + try: + asyncio.get_running_loop().create_task( + run_auto_ingest_background(session_id, rel, mime) + ) + except RuntimeError: + logger.debug("no running loop; skip auto_ingest for %s", rel) diff --git a/src/tools/doc_ingest.py b/src/tools/doc_ingest.py new file mode 100644 index 00000000..204928ff --- /dev/null +++ b/src/tools/doc_ingest.py @@ -0,0 +1,273 @@ +"""Page-by-page ingestion of documents into the session ``derived/docs`` tree. + +The artefact layout is one text file per page. That single choice is what makes +large documents workable: + +* ``sandbox_grep_content`` silently skips files above ``AION_GREP_MAX_FILE_BYTES`` + (500 KB); single pages stay orders of magnitude below it, so a long decree can + never be searched into a false negative. +* The page number travels in the file name, so a grep hit already carries the + citation the caller needs. +* ``read_file_chunk`` loads a whole file before slicing it, so reading one page + costs kilobytes instead of the entire document. + +Extraction is idempotent and deadline-aware: pages already on disk are skipped and +a run that hits ``budget_sec`` returns ``partial`` plus ``resume_from`` instead of +being killed by the MCP bridge timeout. +""" + +from __future__ import annotations + +import json +import re +import time +from pathlib import Path +from typing import Awaitable, Callable, Optional + +__all__ = ["ingest_document", "slugify_document_name", "DOCS_SUBDIR"] + +DOCS_SUBDIR = "derived/docs" + +# Uploads are stored with a random hex prefix; it carries no meaning for the slug. +_UPLOAD_PREFIX_RE = re.compile(r"^[0-9a-f]{6,}_") +_NON_SLUG_RE = re.compile(r"[^a-z0-9]+") + +# Below this many characters a page is treated as having no usable text layer. +DEFAULT_MIN_TEXT_CHARS = 60 +DEFAULT_OCR_DPI = 200 +_EXCERPT_CHARS = 400 +_MAX_LISTED_EMPTY_PAGES = 50 + +OcrCallback = Callable[[int, bytes, str], Awaitable[str]] + + +def slugify_document_name(name: str) -> str: + """Stable, filesystem-safe slug for a document file name.""" + stem = Path(name).stem + stem = _UPLOAD_PREFIX_RE.sub("", stem) + slug = _NON_SLUG_RE.sub("_", stem.lower()).strip("_") + return slug or "document" + + +def _page_filename(page_no: int) -> str: + return f"p{page_no:04d}.txt" + + +def _open_pdf(path: Path): + try: + import pymupdf # type: ignore + except ImportError: # pragma: no cover - older PyMuPDF only exposes `fitz` + import fitz as pymupdf # type: ignore + return pymupdf.open(str(path)) + + +def _rebuild_full_text(pages_dir: Path, full_path: Path, page_numbers: list[int]) -> int: + """Concatenate page files with explicit markers, streaming to avoid buffering.""" + written = 0 + with full_path.open("w", encoding="utf-8") as out: + for page_no in page_numbers: + page_file = pages_dir / _page_filename(page_no) + if not page_file.is_file(): + continue + header = f"\n=== PAGE {page_no} ===\n" + out.write(header) + body = page_file.read_text(encoding="utf-8", errors="replace") + out.write(body) + written += len(body) + return written + + +async def ingest_document( + src_path: Path, + session_root: Path, + *, + first_page: int = 1, + last_page: int = 0, + ocr_mode: str = "auto", + budget_sec: float = 90.0, + force: bool = False, + write_full: bool = False, + ocr_page: Optional[OcrCallback] = None, + min_text_chars: int = DEFAULT_MIN_TEXT_CHARS, + ocr_dpi: int = DEFAULT_OCR_DPI, + clock: Callable[[], float] = time.monotonic, +) -> dict: + """Extract ``src_path`` into ``/derived/docs//pages``. + + ``ocr_mode`` is one of ``auto`` (OCR only pages whose text layer is shorter + than ``min_text_chars``), ``never`` or ``always``. OCR is performed through + ``ocr_page``; when it is not supplied, pages without a text layer are written + empty and reported in ``empty_pages``. + """ + started = clock() + + if ocr_mode not in ("auto", "never", "always"): + return { + "ok": False, + "error": "invalid_ocr_mode", + "message": f"ocr_mode must be auto|never|always, got {ocr_mode!r}", + } + if not src_path.is_file(): + return { + "ok": False, + "error": "not_a_file", + "message": f"{src_path.name} is not a readable file", + } + + slug = slugify_document_name(src_path.name) + root = session_root / DOCS_SUBDIR / slug + pages_dir = root / "pages" + pages_dir.mkdir(parents=True, exist_ok=True) + + try: + doc = _open_pdf(src_path) + except Exception as exc: # noqa: BLE001 - surfaced to the model as JSON + return { + "ok": False, + "error": "open_failed", + "message": f"Cannot open {src_path.name}: {exc}", + } + + try: + pages_total = doc.page_count + start = max(1, int(first_page or 1)) + end = int(last_page) if last_page and int(last_page) > 0 else pages_total + end = min(end, pages_total) + if start > end: + return { + "ok": False, + "error": "empty_range", + "message": f"first_page={start} is past the last page ({pages_total})", + } + + title_guess = "" + try: + title_guess = (doc.metadata or {}).get("title") or "" + except Exception: # noqa: BLE001 - metadata is best effort + title_guess = "" + + written = 0 + skipped = 0 + text_layer_pages = 0 + ocr_pages = 0 + chars_total = 0 + empty_pages: list[int] = [] + ocr_failed_pages: list[int] = [] + first_excerpt = "" + partial = False + resume_from: Optional[int] = None + + for page_no in range(start, end + 1): + if clock() - started >= budget_sec: + partial = True + resume_from = page_no + break + + target = pages_dir / _page_filename(page_no) + if target.is_file() and not force: + skipped += 1 + body = target.read_text(encoding="utf-8", errors="replace") + chars_total += len(body) + if not body.strip(): + empty_pages.append(page_no) + if page_no == start and not first_excerpt: + first_excerpt = body.strip()[:_EXCERPT_CHARS] + continue + + page = doc.load_page(page_no - 1) + text = "" + if ocr_mode != "always": + try: + text = page.get_text("text") or "" + except Exception: # noqa: BLE001 - fall through to OCR / empty + text = "" + + used_ocr = False + ocr_failed = False + needs_ocr = ocr_mode == "always" or ( + ocr_mode == "auto" and len(text.strip()) < min_text_chars + ) + if needs_ocr and ocr_page is not None: + try: + pixmap = page.get_pixmap(dpi=ocr_dpi) + ocr_text = await ocr_page(page_no, pixmap.tobytes("png"), "image/png") + if ocr_text and ocr_text.strip(): + text = ocr_text + used_ocr = True + except Exception as exc: # noqa: BLE001 - keep the text layer we have + if not text.strip(): + # A placeholder must never be mistaken for extracted content: + # the page is reported as failed, not as a text-layer page. + text = f"[OCR failed for page {page_no}: {exc}]" + ocr_failed = True + + target.write_text(text, encoding="utf-8") + written += 1 + chars_total += len(text) + if ocr_failed: + ocr_failed_pages.append(page_no) + elif used_ocr: + ocr_pages += 1 + elif text.strip(): + text_layer_pages += 1 + else: + empty_pages.append(page_no) + if page_no == start and not first_excerpt: + first_excerpt = text.strip()[:_EXCERPT_CHARS] + + rel_root = f"{DOCS_SUBDIR}/{slug}" + full_rel = None + if write_full and not partial: + full_path = root / "full.txt" + _rebuild_full_text(pages_dir, full_path, list(range(start, end + 1))) + full_rel = f"{rel_root}/full.txt" + + manifest = { + "ok": True, + "slug": slug, + "source": src_path.name, + "title_guess": title_guess, + "root": rel_root, + "pages_total": pages_total, + "range": [start, end], + "pages_written": written, + "pages_skipped_existing": skipped, + "text_layer_pages": text_layer_pages, + "ocr_pages": ocr_pages, + "empty_pages": empty_pages[:_MAX_LISTED_EMPTY_PAGES], + "empty_pages_count": len(empty_pages), + "ocr_failed_pages": ocr_failed_pages[:_MAX_LISTED_EMPTY_PAGES], + "ocr_failed_pages_count": len(ocr_failed_pages), + "chars_total": chars_total, + "partial": partial, + "resume_from": resume_from, + "page_file_pattern": f"{rel_root}/pages/pNNNN.txt", + "full_text": full_rel, + "grep_hint": ( + "sandbox_grep_content(pattern=, relative_root='derived', " + f"glob_filter='docs/{slug}/pages/*.txt', max_matches=200) — the file " + "name of each hit is the page number." + ), + "first_page_excerpt": first_excerpt, + } + if partial: + manifest["next_step"] = ( + f"Budget of {budget_sec:.0f}s reached at page {resume_from}. Call " + f"doc_ingest again with first_page={resume_from} to resume; pages " + "already written are skipped." + ) + else: + manifest["next_step"] = ( + "Verify the document identity against the user request using " + "title_guess/first_page_excerpt, then grep the pages as per grep_hint." + ) + + (root / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8" + ) + return manifest + finally: + try: + doc.close() + except Exception: # noqa: BLE001 - best effort + pass From e04abca424ed86648c0e7c0b266280bf38120018 Mon Sep 17 00:00:00 2001 From: Giuseppe La Rocca <52716342+JustBeGiusee@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:09:49 +0200 Subject: [PATCH 2/8] feat: enhance file upload management in chat workspace - Integrated a new upload session file handling mechanism with progress tracking in the ChatWorkspace component. - Replaced the previous pending files state management with a more robust system using `usePendingSessionUploads` hook. - Updated UI to reflect the status of file uploads, including progress indicators and error handling options. - Added localized messages for upload status updates in English and Italian. --- chat-ui/components/chat/ChatWorkspace.tsx | 134 +++++++++----- .../chat/CircularUploadProgress.tsx | 84 +++++++++ chat-ui/hooks/use-pending-session-uploads.ts | 164 ++++++++++++++++++ chat-ui/lib/api/aion.ts | 84 +++++++++ chat-ui/lib/i18n/locales/en.json | 7 + chat-ui/lib/i18n/locales/it.json | 7 + 6 files changed, 437 insertions(+), 43 deletions(-) create mode 100644 chat-ui/components/chat/CircularUploadProgress.tsx create mode 100644 chat-ui/hooks/use-pending-session-uploads.ts diff --git a/chat-ui/components/chat/ChatWorkspace.tsx b/chat-ui/components/chat/ChatWorkspace.tsx index e467fd9e..30de4d57 100644 --- a/chat-ui/components/chat/ChatWorkspace.tsx +++ b/chat-ui/components/chat/ChatWorkspace.tsx @@ -17,6 +17,7 @@ import { ShimmerText } from "@/components/chat/ShimmerText"; import Link from "next/link"; import { AgentModeSelectChip } from "@/components/chat/AgentModeSelectChip"; import { ChatEmptyState } from "@/components/chat/ChatEmptyState"; +import { CircularUploadProgress } from "@/components/chat/CircularUploadProgress"; import { ComposerOptionRow } from "@/components/chat/ComposerOptionRow"; import { ChatDragDrop } from "@/components/chat/ChatDragDrop"; import { mergeAttachmentRefs } from "@/lib/attachments"; @@ -38,7 +39,6 @@ import { waitForChatPrepare, type ChatPrepareMcpError, sessionDownloadUrl, - uploadSessionFiles, listSessionUploads, fetchConversationHistory, fetchStreamStatus, @@ -60,6 +60,7 @@ import { type SessionChart, } from "@/lib/api/aion"; import { useStoredToken, useStoredUserId } from "@/lib/auth/use-stored-auth"; +import { usePendingSessionUploads } from "@/hooks/use-pending-session-uploads"; import { useT } from "@/lib/i18n/use-t"; import { extractStreamingPlanMarkdown, @@ -413,6 +414,19 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve const sidebarOpen = useSidebarOpen(); const [conversationId, setConversationId] = useState(initialConversationId); + const { + items: pendingUploadItems, + queueFiles: queuePendingUploads, + removeItem: removePendingUpload, + retryItem: retryPendingUpload, + clearAll: clearPendingUploads, + isUploading: pendingUploadsInProgress, + hasUploadErrors: pendingUploadsFailed, + completedAttachments: pendingUploadedAttachments, + } = usePendingSessionUploads(conversationId, userId, token); + + const sendBlockedByUploads = pendingUploadsInProgress || pendingUploadsFailed; + const [dockTab, setDockTab] = useState("none"); const [lastActiveTab, setLastActiveTab] = useState("plan"); @@ -1278,16 +1292,9 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve minHeight: COMPOSER_TEXTAREA_MIN, maxHeight: composerTextMax, }); - const [pendingFiles, setPendingFiles] = useState([]); - const handleFilesDropped = useCallback((files: File[]) => { - setPendingFiles((prev) => { - const filtered = files.filter( - (sf) => !prev.some((pf) => pf.name === sf.name && pf.size === sf.size) - ); - return [...prev, ...filtered]; - }); - }, []); + queuePendingUploads(files); + }, [queuePendingUploads]); const [streamEpoch, setStreamEpoch] = useState(0); const abortRef = useRef(null); @@ -1831,18 +1838,20 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve opts?.deepResearchModeOverride !== undefined ? opts.deepResearchModeOverride : effectiveAgentMode === "deep_research"; + if (pendingUploadsInProgress || pendingUploadsFailed) { + return; + } + let uidMsg = crypto.randomUUID(); let aid = crypto.randomUUID(); setActiveMessageId(aid); - const hasPendingFiles = pendingFiles.length > 0; - const uploads = await uploadSessionFiles(conversationId, userId, pendingFiles, token); - setPendingFiles([]); + const uploads = [...pendingUploadedAttachments]; + const hasNewUploads = uploads.length > 0; + clearPendingUploads(); void fetchSessionFiles(); - // Fetch existing session uploads only when new files were uploaded, - // so previous uploads aren't incorrectly attached to the current message. - const existing = hasPendingFiles ? await listSessionUploads(conversationId, userId, token) : []; - const attachments = hasPendingFiles ? mergeAttachmentRefs(uploads, existing) : []; + const existing = hasNewUploads ? await listSessionUploads(conversationId, userId, token) : []; + const attachments = hasNewUploads ? mergeAttachmentRefs(uploads, existing) : []; const userArtifacts: ChatHistoryArtifact[] = uploads.map((a, i) => ({ id: `att-${i}-${Date.now()}`, @@ -2245,7 +2254,10 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve token, activeProfileSlug, effectiveEffort, - pendingFiles, + pendingUploadsInProgress, + pendingUploadsFailed, + pendingUploadedAttachments, + clearPendingUploads, thinkingEnabled, markStreamConversation, transcriptStreaming, @@ -3834,24 +3846,58 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve ) ) : null} - {pendingFiles.length > 0 && ( -
- {pendingFiles.map((f) => ( - - {f.name} - - - ))} + + {item.file.name} + {item.status === "error" ? ( + + ) : null} + + + ); + })}
)} {isProjectRequiredButMissing && ( @@ -3914,7 +3960,7 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve handleAgentModeChange(agentMode === "plan" ? "normal" : "plan"); } else if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); - void send(); + if (!sendBlockedByUploads) void send(); } }} placeholder={isProjectRequiredButMissing ? t("chat.project_required.textarea_placeholder") : t("chat.composer_placeholder")} @@ -3930,12 +3976,7 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve className="hidden" onChange={(e) => { const selected = Array.from(e.target.files || []); - setPendingFiles((prev) => { - const filtered = selected.filter( - (sf) => !prev.some((pf) => pf.name === sf.name && pf.size === sf.size) - ); - return [...prev, ...filtered]; - }); + queuePendingUploads(selected); e.target.value = ""; }} /> @@ -4492,7 +4533,14 @@ export function ChatWorkspace({ conversationId: initialConversationId }: { conve