diff --git a/.gitignore b/.gitignore index b6d0978..4d2ca98 100644 --- a/.gitignore +++ b/.gitignore @@ -31,8 +31,12 @@ scripts/ # * Project instructions (not for public repo) CLAUDE.md +# * Internal improvement plan (contains security PoC — never publish) +docs/improvement-plan.md + # * MAC OS .DS_Store # * MCP -.serena/.claude/ +.serena/ +.claude/ diff --git a/.serena/.gitignore b/.serena/.gitignore deleted file mode 100644 index 14d86ad..0000000 --- a/.serena/.gitignore +++ /dev/null @@ -1 +0,0 @@ -/cache diff --git a/.serena/memories/xlstruct/release-readiness-audit.md b/.serena/memories/xlstruct/release-readiness-audit.md deleted file mode 100644 index 50bd455..0000000 --- a/.serena/memories/xlstruct/release-readiness-audit.md +++ /dev/null @@ -1,34 +0,0 @@ -# XLStruct Release Readiness Audit - -## Overview -Comprehensive documentation and packaging review for public release on PyPI. - -## Checklist Results - -### PASS - Already Complete -1. ✓ README.md - Comprehensive with badges (implicit in structure), installation, quick start, examples, contributing reference -2. ✓ LICENSE - MIT license complete and correct -3. ✓ pyproject.toml - Well-structured with metadata, dependencies, extras, build system, CLI entry point -4. ✓ examples/ directory - 3 quality examples (basic_extraction, cloud_storage, custom_instructions) -5. ✓ CI/CD - GitHub Actions with lint, test, integration test jobs -6. ✓ Code quality tools - Ruff + mypy configured - -### CRITICAL GAPS - Block Release -1. ✗ No py.typed marker (PEP 561 compliance for type hints) -2. ✗ No CHANGELOG/HISTORY file -3. ✗ No CONTRIBUTING.md -4. ✗ No __version__ in __init__.py (single source of truth for versioning) -5. ✗ No badges in README (GitHub Actions, PyPI, Python version, etc.) -6. ✗ docs/ directory lacks user-facing documentation (only has PRD) - -### MEDIUM PRIORITY -1. No comprehensive API reference documentation -2. No troubleshooting guide -3. No migration guide for future versions -4. README links to docs/ missing -5. API reference link missing from README - -### LOW PRIORITY -1. No MANIFEST.in (though hatchling auto-includes common files) -2. No security policy (SECURITY.md) -3. No code of conduct (CODE_OF_CONDUCT.md) - optional but nice to have diff --git a/.serena/project.yml b/.serena/project.yml deleted file mode 100644 index 73ecf99..0000000 --- a/.serena/project.yml +++ /dev/null @@ -1,149 +0,0 @@ -# the name by which the project can be referenced within Serena -project_name: "LLMForge" - - -# list of languages for which language servers are started; choose from: -# al bash clojure cpp csharp -# csharp_omnisharp dart elixir elm erlang -# fortran fsharp go groovy haskell -# java julia kotlin lua markdown -# matlab nix pascal perl php -# php_phpactor powershell python python_jedi r -# rego ruby ruby_solargraph rust scala -# swift terraform toml typescript typescript_vts -# vue yaml zig -# (This list may be outdated. For the current list, see values of Language enum here: -# https://github.com/oraios/serena/blob/main/src/solidlsp/ls_config.py -# For some languages, there are alternative language servers, e.g. csharp_omnisharp, ruby_solargraph.) -# Note: -# - For C, use cpp -# - For JavaScript, use typescript -# - For Free Pascal/Lazarus, use pascal -# Special requirements: -# Some languages require additional setup/installations. -# See here for details: https://oraios.github.io/serena/01-about/020_programming-languages.html#language-servers -# When using multiple languages, the first language server that supports a given file will be used for that file. -# The first language is the default language and the respective language server will be used as a fallback. -# Note that when using the JetBrains backend, language servers are not used and this list is correspondingly ignored. -languages: -- python - -# the encoding used by text files in the project -# For a list of possible encodings, see https://docs.python.org/3.11/library/codecs.html#standard-encodings -encoding: "utf-8" - -# The language backend to use for this project. -# If not set, the global setting from serena_config.yml is used. -# Valid values: LSP, JetBrains -# Note: the backend is fixed at startup. If a project with a different backend -# is activated post-init, an error will be returned. -language_backend: - -# whether to use project's .gitignore files to ignore files -ignore_all_files_in_gitignore: true - -# list of additional paths to ignore in this project. -# Same syntax as gitignore, so you can use * and **. -# Note: global ignored_paths from serena_config.yml are also applied additively. -ignored_paths: [] - -# whether the project is in read-only mode -# If set to true, all editing tools will be disabled and attempts to use them will result in an error -# Added on 2025-04-18 -read_only: false - -# list of tool names to exclude. We recommend not excluding any tools, see the readme for more details. -# Below is the complete list of tools for convenience. -# To make sure you have the latest list of tools, and to view their descriptions, -# execute `uv run scripts/print_tool_overview.py`. -# -# * `activate_project`: Activates a project by name. -# * `check_onboarding_performed`: Checks whether project onboarding was already performed. -# * `create_text_file`: Creates/overwrites a file in the project directory. -# * `delete_lines`: Deletes a range of lines within a file. -# * `delete_memory`: Deletes a memory from Serena's project-specific memory store. -# * `execute_shell_command`: Executes a shell command. -# * `find_referencing_code_snippets`: Finds code snippets in which the symbol at the given location is referenced. -# * `find_referencing_symbols`: Finds symbols that reference the symbol at the given location (optionally filtered by type). -# * `find_symbol`: Performs a global (or local) search for symbols with/containing a given name/substring (optionally filtered by type). -# * `get_current_config`: Prints the current configuration of the agent, including the active and available projects, tools, contexts, and modes. -# * `get_symbols_overview`: Gets an overview of the top-level symbols defined in a given file. -# * `initial_instructions`: Gets the initial instructions for the current project. -# Should only be used in settings where the system prompt cannot be set, -# e.g. in clients you have no control over, like Claude Desktop. -# * `insert_after_symbol`: Inserts content after the end of the definition of a given symbol. -# * `insert_at_line`: Inserts content at a given line in a file. -# * `insert_before_symbol`: Inserts content before the beginning of the definition of a given symbol. -# * `list_dir`: Lists files and directories in the given directory (optionally with recursion). -# * `list_memories`: Lists memories in Serena's project-specific memory store. -# * `onboarding`: Performs onboarding (identifying the project structure and essential tasks, e.g. for testing or building). -# * `prepare_for_new_conversation`: Provides instructions for preparing for a new conversation (in order to continue with the necessary context). -# * `read_file`: Reads a file within the project directory. -# * `read_memory`: Reads the memory with the given name from Serena's project-specific memory store. -# * `remove_project`: Removes a project from the Serena configuration. -# * `replace_lines`: Replaces a range of lines within a file with new content. -# * `replace_symbol_body`: Replaces the full definition of a symbol. -# * `restart_language_server`: Restarts the language server, may be necessary when edits not through Serena happen. -# * `search_for_pattern`: Performs a search for a pattern in the project. -# * `summarize_changes`: Provides instructions for summarizing the changes made to the codebase. -# * `switch_modes`: Activates modes by providing a list of their names -# * `think_about_collected_information`: Thinking tool for pondering the completeness of collected information. -# * `think_about_task_adherence`: Thinking tool for determining whether the agent is still on track with the current task. -# * `think_about_whether_you_are_done`: Thinking tool for determining whether the task is truly completed. -# * `write_memory`: Writes a named memory (for future reference) to Serena's project-specific memory store. -excluded_tools: [] - -# list of tools to include that would otherwise be disabled (particularly optional tools that are disabled by default) -included_optional_tools: [] - -# fixed set of tools to use as the base tool set (if non-empty), replacing Serena's default set of tools. -# This cannot be combined with non-empty excluded_tools or included_optional_tools. -fixed_tools: [] - -# list of mode names to that are always to be included in the set of active modes -# The full set of modes to be activated is base_modes + default_modes. -# If the setting is undefined, the base_modes from the global configuration (serena_config.yml) apply. -# Otherwise, this setting overrides the global configuration. -# Set this to [] to disable base modes for this project. -# Set this to a list of mode names to always include the respective modes for this project. -base_modes: - -# list of mode names that are to be activated by default. -# The full set of modes to be activated is base_modes + default_modes. -# If the setting is undefined, the default_modes from the global configuration (serena_config.yml) apply. -# Otherwise, this overrides the setting from the global configuration (serena_config.yml). -# This setting can, in turn, be overridden by CLI parameters (--mode). -default_modes: - -# initial prompt for the project. It will always be given to the LLM upon activating the project -# (contrary to the memories, which are loaded on demand). -initial_prompt: "" - -# time budget (seconds) per tool call for the retrieval of additional symbol information -# such as docstrings or parameter information. -# This overrides the corresponding setting in the global configuration; see the documentation there. -# If null or missing, use the setting from the global configuration. -symbol_info_budget: - -# list of regex patterns which, when matched, mark a memory entry as read‑only. -# Extends the list from the global configuration, merging the two lists. -read_only_memory_patterns: [] - -# line ending convention to use when writing source files. -# Possible values: unset (use global setting), "lf", "crlf", or "native" (platform default) -# This does not affect Serena's own files (e.g. memories and configuration files), which always use native line endings. -line_ending: - -# list of regex patterns for memories to completely ignore. -# Matching memories will not appear in list_memories or activate_project output -# and cannot be accessed via read_memory or write_memory. -# To access ignored memory files, use the read_file tool on the raw file path. -# Extends the list from the global configuration, merging the two lists. -# Example: ["_archive/.*", "_episodes/.*"] -ignored_memory_patterns: [] - -# advanced configuration option allowing to configure language server-specific options. -# Maps the language key to the options. -# Have a look at the docstring of the constructors of the LS implementations within solidlsp (e.g., for C# or PHP) to see which options are available. -# No documentation on options means no options are available. -ls_specific_settings: {} diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f3887b..e8f3085 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,48 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.7.0] - 2026-06-02 + +### Security + +- **OS-level sandbox is now the default for codegen execution (S1)** — untrusted, LLM-generated scripts run in `DockerBackend` when `xlstruct[docker]` is installed. The pre-execution AST scan is a best-effort filter, not a boundary (it has known bypasses), so it is no longer relied upon as the sole defense. +- **Hardened `SubprocessBackend` and demoted it to trusted/dev-only (S1)** — runs the script via `python -I` with a reduced builtins namespace (no `eval`/`exec`/`compile`/`breakpoint`/`input`), adds `RLIMIT_CPU`/`RLIMIT_NPROC`/`RLIMIT_NOFILE`/`RLIMIT_FSIZE`, kills the whole process group on timeout with a bounded drain, and fails closed when limits cannot be applied. It is explicitly NOT a security boundary. +- **Hardened `DockerBackend` (S2)** — the execution container now runs as a non-root user (`65534:65534`), drops all Linux capabilities (`CapDrop=["ALL"]`), mounts a read-only root filesystem with a writable `/tmp` tmpfs, and keeps `NetworkDisabled` + `no-new-privileges` + PID/memory/CPU limits. Optional `runtime` (e.g. gVisor `runsc`) and custom `seccomp_profile` knobs added to `DockerConfig`. The one-time package-install step stays permissive. +- **Integrity-protected codegen script cache (S3)** — the cache directory is `0700`, entries are `0600`, and every cached script carries an HMAC keyed by a per-user secret. Entries that fail verification are refused (never executed) and regenerated. The structure signature was widened from 64-bit to the full 256-bit SHA-256 digest. + +### Added + +- **`ExtractorConfig.codegen_sandbox`** (`"auto"` | `"docker"` | `"subprocess"`, default `"auto"`) — selects the codegen execution sandbox. An explicit `execution_backend` always overrides it. +- **`CodegenSecurityError`** + `ErrorCode.CODEGEN_NO_SANDBOX` — raised when codegen would run untrusted code with no sandbox available. +- **`SubprocessBackend(trusted=...)`** — acknowledge trusted/dev-only use and silence the "not a security boundary" warning. +- **`ExtractorConfig.csv_encoding`** — text encoding for CSV files (e.g. `euc-kr`, `cp1252`, `shift-jis`); `utf-8` strips a BOM if present. Ignored for Excel formats. +- **`ExtractorConfig.max_concurrent_chunks`** (default `5`) — bounds how many chunk LLM calls run concurrently for a single chunked sheet, to respect provider rate limits. +- **`xlstruct.__version__`** — the package version is now exposed in code, mirroring `pyproject.toml`. + +### Changed + +- **Extended thinking now applies to direct extraction and `suggest_schema`** (C1) — previously honored only on the codegen path. LLM client construction is unified so `thinking=True` reaches every call. +- **Chunked single-sheet extraction runs concurrently** (P1) — independent chunks are dispatched together under a bounded semaphore instead of serially, cutting wall-clock latency on large sheets. +- **Anthropic prompt caching marks only the static system prompt** (P2) — the variable per-call sheet data is no longer tagged, so a cache breakpoint is not wasted on content that never repeats. +- **Chunk sizing uses measured per-row token cost** (P4) — `ChunkSplitter` greedy bin-packs rows by real cost (reserving the header) so every chunk fits `token_budget` regardless of where heavy rows sit; the previous sampled, header-blind estimate could over-budget chunks. +- **Internal: `Extractor` decomposed into a thin facade over `ExtractionPipeline`** (A1, A2, A3) — orchestration, report assembly, and concurrency moved into `extraction/`; reader dispatch centralized into `READER_REGISTRY` shared by the extractor and MCP server; the unused `ExcelReader` protocol replaced by `WorkbookReader` + `ReaderOptions`. Public API and behavior unchanged. + +### Changed (breaking) + +- **Codegen with the default config now requires a sandbox.** When `xlstruct[docker]` is not installed, codegen fails closed with `CodegenSecurityError` instead of silently executing in a non-isolating subprocess. To restore the previous (unsandboxed) behavior in a trusted environment, set `ExtractorConfig(codegen_sandbox="subprocess")` or pass `execution_backend=SubprocessBackend(trusted=True)`. Direct (non-codegen) extraction is unaffected. +- **Existing codegen script caches are invalidated.** The widened structure signature changes cache keys, so cached scripts regenerate once on first use. +- `ScriptValidator` now requires an explicit `backend` argument (no implicit subprocess default). + +### Fixed + +- **Codegen no longer silently drops invalid records** (C2) — records that fail schema validation now raise instead of being skipped, so a partially mis-extracted result is surfaced rather than returned truncated. +- **Non-UTF-8 CSV files raise `ReaderError`** (C3) — a decode failure is wrapped with guidance and honors the configurable `csv_encoding`, instead of an unhandled `UnicodeDecodeError`. +- **Extraction error handling no longer mislabels internal bugs** (C4) — only the provider call is wrapped, so a post-processing bug is not reported as an LLM failure. +- **Empty (0-row) sheets fail fast** (B3) — both the configured and legacy `extract()` paths raise before spending an LLM call. +- **Confidence scoring no longer silently defaults** (B2) — a missing confidence level surfaces as an error instead of a fabricated `moderate` (0.5). +- **`find_empty_rows` no longer allocates a full row-range set** (B1) — wide/tall sheets skip the `set(range(...))` allocation; row-skipping behavior is unchanged. +- **Token counting tolerates special-token literals** — cell content like `<|endoftext|>` no longer raises during token estimation (now uses `encode_ordinary`). + ## [0.6.0] - 2026-04-03 ### Added diff --git a/README.md b/README.md index 51e4bb9..5bf306a 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ Excel File + Pydantic Schema → LLM → Validated Structured Data - **Fast hybrid reader** — calamine (Rust) for speed + openpyxl for formula extraction. Both passes in one call. - **Token-aware encoding** — Compressed markdown encoding with head+tail sampling. Auto-chunks large sheets to fit within token budget. - **Prompt caching** — Anthropic cache_control markers applied automatically; OpenAI cached_tokens tracked -- **Sandboxed execution** — Generated scripts run in a subprocess with blocked imports (network, subprocess) and stripped credentials. Optional Docker backend for OS-level isolation. +- **Sandboxed execution** — Untrusted, generated scripts run in an OS-level Docker sandbox by default (non-root, no network, dropped capabilities, read-only rootfs). Without Docker, codegen fails closed; the non-isolating subprocess backend is an explicit trusted/dev-only opt-in. - **Multi-provider LLM** — OpenAI, Anthropic, Gemini via [Instructor](https://github.com/jxnl/instructor) - **Cloud storage** — Read from S3, Azure Blob, GCS via fsspec - **Async-first** — Async API with sync convenience wrappers. Jupyter-compatible via nest_asyncio. @@ -289,7 +289,7 @@ config = ExtractionConfig( 1. **Header Detection** — Auto-detect header rows via lightweight LLM call 2. **Phase 0 (Analyzer)** — LLM analyzes spreadsheet structure → `MappingPlan` -3. **Phase 1 (Parser Agent)** — LLM generates openpyxl-based parsing script → validated via subprocess +3. **Phase 1 (Parser Agent)** — LLM generates openpyxl-based parsing script → validated by executing it in the sandbox (Docker by default) Each phase includes self-correction — errors are fed back to the LLM (up to `max_codegen_retries`). @@ -408,9 +408,18 @@ config = ExtractorConfig( extractor = Extractor(config=config) ``` -### Docker Backend +### Codegen Execution Sandbox -For OS-level sandboxing, pass a `DockerBackend` via `execution_backend`: +Codegen executes untrusted, LLM-generated Python. The execution backend is chosen by +`ExtractorConfig.codegen_sandbox`: + +| `codegen_sandbox` | Behavior | +|-------------------|----------| +| `"auto"` (default) | Use Docker when `xlstruct[docker]` is installed; otherwise **fail closed** with `CodegenSecurityError` | +| `"docker"` | Always use Docker (errors at run time if the daemon/package is missing) | +| `"subprocess"` | Use the non-isolating subprocess backend — **NOT a security boundary; trusted/dev only** | + +An explicit `execution_backend` always overrides `codegen_sandbox`. ```bash pip install "xlstruct[docker]" @@ -420,13 +429,28 @@ pip install "xlstruct[docker]" from xlstruct import Extractor from xlstruct.codegen.backends.docker import DockerBackend, DockerConfig +# Customize the Docker sandbox (gVisor runtime, custom seccomp, image, limits) extractor = Extractor( execution_backend=DockerBackend( - config=DockerConfig(image="python:3.12-slim", mem_limit="1g"), + config=DockerConfig(image="python:3.12-slim", mem_limit="1g", runtime="runsc"), ), ) ``` +If you cannot run Docker and accept the risk in a trusted environment, opt into the +subprocess backend explicitly: + +```python +from xlstruct import Extractor, ExtractorConfig + +# Either via config… +extractor = Extractor(config=ExtractorConfig(codegen_sandbox="subprocess")) + +# …or by injecting the backend directly (acknowledging it is not a boundary) +from xlstruct.codegen.backends.subprocess import SubprocessBackend +extractor = Extractor(execution_backend=SubprocessBackend(trusted=True)) +``` + ## Architecture ``` @@ -464,8 +488,9 @@ src/xlstruct/ │ ├── schema_utils.py # Pydantic schema → source code utilities │ └── backends/ │ ├── base.py # ExecutionBackend protocol -│ ├── subprocess.py # SubprocessBackend (default sandbox) -│ └── docker.py # DockerBackend (OS-level isolation) +│ ├── resolver.py # Backend resolution (sandbox-by-default, fail-closed) +│ ├── subprocess.py # SubprocessBackend (trusted/dev-only — NOT a boundary) +│ └── docker.py # DockerBackend (OS-level isolation — default sandbox) ├── schemas/ │ ├── core.py # SheetData, WorkbookData, CellData │ ├── codegen.py # GeneratedScript, MappingPlan, CodegenAttempt @@ -497,11 +522,26 @@ src/xlstruct/ ### Sandboxed Execution -Generated scripts run in `SubprocessBackend` (or optionally `DockerBackend`) with security layers: -- **Allowlist imports** — only safe modules (openpyxl, pydantic, stdlib math/data) permitted via AST scanning -- **Blocked builtins** — `exec`, `eval`, `__import__`, `open`, etc. rejected before execution -- **Stripped credentials** — API keys and cloud credentials removed from subprocess environment -- **Timeout** — enforced via `codegen_timeout` config +Codegen runs untrusted, LLM-generated Python, so the real boundary is an OS-level +sandbox. By default (`codegen_sandbox="auto"`) scripts run in **`DockerBackend`** when +`xlstruct[docker]` is installed; otherwise codegen **fails closed**. + +**`DockerBackend` (default sandbox)** — full OS-level isolation: +- **Non-root** — runs as `65534:65534` (nobody) +- **No capabilities** — `CapDrop=["ALL"]` + `no-new-privileges` +- **Read-only root filesystem** — with a small writable `/tmp` tmpfs +- **No network** — `NetworkDisabled` +- **Resource limits** — memory, CPU, and `PidsLimit` +- **seccomp** — Docker's default profile (or a custom one); optional gVisor (`runsc`) runtime + +**`SubprocessBackend` (trusted/dev-only — NOT a security boundary)** — defense-in-depth only: +- **Allowlist imports / blocked builtins** — best-effort AST scan (has known bypasses) +- **Stripped credentials** — API keys and cloud credentials removed from the environment +- **Isolated interpreter** — `python -I` with a reduced builtins namespace +- **Resource limits** — CPU/process/file-descriptor/file-size limits; process-group kill on timeout + +The pre-execution AST scan applies in both backends but is a filter, not a boundary — do +not rely on it alone for untrusted output. ### Exceptions diff --git a/pyproject.toml b/pyproject.toml index 0d75cca..11658cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "xlstruct" -version = "0.6.0" +version = "0.7.0" description = "LLM-powered Excel parser — define a Pydantic schema, get structured data from any Excel file" license = "MIT" readme = "README.md" diff --git a/src/xlstruct/__init__.py b/src/xlstruct/__init__.py index cd5972f..4e123da 100644 --- a/src/xlstruct/__init__.py +++ b/src/xlstruct/__init__.py @@ -1,7 +1,10 @@ """XLStruct — LLM-powered Excel parser.""" +__version__ = "0.7.0" + from xlstruct.config import ExtractionConfig, ExtractionMode, ExtractorConfig from xlstruct.exceptions import ( + CodegenSecurityError, CodegenValidationError, ErrorCode, ExtractionError, @@ -15,6 +18,7 @@ from xlstruct.schemas.usage import TokenUsage __all__ = [ + "__version__", "Extractor", "ExtractionResult", "ExtractorConfig", @@ -30,4 +34,5 @@ "ReaderError", "ExtractionError", "CodegenValidationError", + "CodegenSecurityError", ] diff --git a/src/xlstruct/_tokens.py b/src/xlstruct/_tokens.py index 25dc460..72ef68b 100644 --- a/src/xlstruct/_tokens.py +++ b/src/xlstruct/_tokens.py @@ -9,7 +9,7 @@ import tiktoken if TYPE_CHECKING: - from xlstruct.schemas.core import SheetData + from xlstruct.schemas.core import CellData, SheetData # * Module-level singleton _encoding: tiktoken.Encoding | None = None @@ -24,31 +24,71 @@ def _get_encoding() -> tiktoken.Encoding: def count_tokens(text: str) -> int: """Count tokens in a text string.""" - return len(_get_encoding().encode(text)) + # ^ encode_ordinary, not encode: cell content may contain literal special-token + # ^ text (e.g. "<|endoftext|>"), which the strict encode() raises on. + return len(_get_encoding().encode_ordinary(text)) -def estimate_sheet_tokens(sheet: "SheetData") -> int: - """Fast token estimation for a sheet without full encoding. +# ^ Formatting overhead applied to raw cell tokens (markdown pipes, separators, +# ^ row numbers, metadata). Shared by the sheet estimate and per-row costing. +_OVERHEAD_FACTOR = 1.3 - Approximation: avg tokens per cell value * cell count + overhead. - Used for encoder strategy selection (exact count not needed). + +def estimate_cells_tokens(cells: list["CellData"]) -> int: + """Fast, rough token estimation for an arbitrary list of cells. + + Samples up to 50 cells spread evenly across the list (not just the first 50) + for an average per-cell token count, scales to the full list, and adds a + formatting-overhead factor. Even spacing keeps the estimate representative + when token density varies by position. This is the chunking *gate*; chunk + sizing itself uses exact per-row costs (see estimate_row_token_costs). """ - if not sheet.cells: + n = len(cells) + if n == 0: return 0 - # ^ Sample up to 50 cells for average token count - sample_size = min(50, len(sheet.cells)) - sample = sheet.cells[:sample_size] - - total_sample_tokens = 0 + # ^ Stratified sample: evenly spaced indices across [0, n) so a dense head + # ^ or dense tail does not skew the average (a plain cells[:50] would). + sample_size = min(50, n) enc = _get_encoding() - for cell in sample: + total_sample_tokens = 0 + for i in range(sample_size): + cell = cells[i * n // sample_size] val = str(cell.display_value) if cell.display_value is not None else "" - total_sample_tokens += len(enc.encode(val)) + # ^ encode_ordinary: never raise on cell content resembling a special token + total_sample_tokens += len(enc.encode_ordinary(val)) avg_tokens_per_cell = total_sample_tokens / sample_size - estimated = int(avg_tokens_per_cell * len(sheet.cells)) + estimated = int(avg_tokens_per_cell * n) + return int(estimated * _OVERHEAD_FACTOR) + + +def estimate_row_token_costs(rows_cells: list[list["CellData"]]) -> list[int]: + """Per-row token cost (row tokens x overhead) for chunk sizing. - # ^ Add overhead for formatting (headers, separators, metadata) - overhead_factor = 1.3 - return int(estimated * overhead_factor) + Joins each row's cell values and encodes one string per row in a single + batch. Unlike estimate_cells_tokens this samples nothing: chunk sizing must + reflect each row's real weight, not a sheet-wide average, so a chunk packed + to the budget fits regardless of where the heavy rows sit. Encoding per + joined row (not per cell) is ~17x faster on large sheets — far fewer, longer + strings — and the join mirrors how cells render with separators, so the + count tracks the per-cell sum closely. + """ + if not rows_cells: + return [] + + enc = _get_encoding() + joined = [ + " ".join(str(c.display_value) if c.display_value is not None else "" for c in cells) + for cells in rows_cells + ] + return [int(len(tokens) * _OVERHEAD_FACTOR) for tokens in enc.encode_ordinary_batch(joined)] + + +def estimate_sheet_tokens(sheet: "SheetData") -> int: + """Fast token estimation for a sheet without full encoding. + + Approximation: avg tokens per cell value * cell count + overhead. + Used for encoder strategy selection (exact count not needed). + """ + return estimate_cells_tokens(sheet.cells) diff --git a/src/xlstruct/codegen/backends/_subprocess_bootstrap.py b/src/xlstruct/codegen/backends/_subprocess_bootstrap.py new file mode 100644 index 0000000..2842e01 --- /dev/null +++ b/src/xlstruct/codegen/backends/_subprocess_bootstrap.py @@ -0,0 +1,52 @@ +"""Bootstrap that runs an untrusted codegen script with reduced builtins. + +Invoked as ``python -I ``. It executes +the user script under a builtins namespace stripped of dynamic-execution +primitives (``eval``/``exec``/``compile``/``breakpoint``/``input``) that a +legitimate spreadsheet parser never needs. + +This is defense-in-depth only, NOT a security boundary: ``__import__`` must stay +available for the script's ``import`` statements, so a determined escape is still +possible. Real isolation requires DockerBackend. +""" + +import builtins +import sys + +# ^ Primitives a parser never needs; removing them raises the bar for casual escapes. +_STRIPPED_BUILTINS = ("eval", "exec", "compile", "breakpoint", "input") + + +def _main() -> int: + # ^ argv: [bootstrap, user_script_path, source_path] + if len(sys.argv) < 3: + print("bootstrap: expected ", file=sys.stderr) + return 2 + + user_script = sys.argv[1] + source_path = sys.argv[2] + + with open(user_script, encoding="utf-8") as f: + code = f.read() + + safe_builtins: dict[str, object] = {name: getattr(builtins, name) for name in dir(builtins)} + for name in _STRIPPED_BUILTINS: + safe_builtins.pop(name, None) + + script_globals: dict[str, object] = { + "__name__": "__main__", + "__file__": user_script, + "__builtins__": safe_builtins, + } + + # ^ Present the user script with the argv it expects (source path as argv[1]), + # and compile with the real filename so tracebacks keep correct line numbers + # for the self-correction loop. + sys.argv = [user_script, source_path] + compiled = builtins.compile(code, user_script, "exec") + exec(compiled, script_globals) # noqa: S102 + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/src/xlstruct/codegen/backends/base.py b/src/xlstruct/codegen/backends/base.py index 43d8b50..1b2c3b5 100644 --- a/src/xlstruct/codegen/backends/base.py +++ b/src/xlstruct/codegen/backends/base.py @@ -7,8 +7,8 @@ class ExecutionBackend(Protocol): """Protocol for script execution backends. Implementations: - - SubprocessBackend: Hardened subprocess (default). - - DockerBackend: Full OS-level isolation via Docker. + - SubprocessBackend: Hardened subprocess (trusted/dev-only — not a security boundary). + - DockerBackend: Full OS-level isolation via Docker (the default for untrusted codegen). """ async def execute( diff --git a/src/xlstruct/codegen/backends/docker.py b/src/xlstruct/codegen/backends/docker.py index 8331d9d..0e6515d 100644 --- a/src/xlstruct/codegen/backends/docker.py +++ b/src/xlstruct/codegen/backends/docker.py @@ -16,6 +16,9 @@ # ^ Suffix appended to base image name for the prepared image _PREPARED_IMAGE_TAG = "xlstruct-ready" +# ^ Default UID:GID for the execution container (nobody:nogroup) +_DEFAULT_NONROOT_USER = "65534:65534" + class DockerConfig(BaseModel): """Docker backend configuration for codegen execution.""" @@ -40,6 +43,34 @@ class DockerConfig(BaseModel): default=True, description="Pull the Docker image if not found locally.", ) + # * Hardening — applied to the execution container, never to the one-time install step. + user: str = Field( + default=_DEFAULT_NONROOT_USER, + description="UID:GID for the execution container. Defaults to nobody. " + "Empty string uses the image default (root) — not recommended.", + ) + read_only_rootfs: bool = Field( + default=True, + description="Mount the container root filesystem read-only (a /tmp tmpfs stays writable).", + ) + cap_drop: list[str] = Field( + default_factory=lambda: ["ALL"], + description="Linux capabilities to drop in the execution container.", + ) + tmpfs_size: str = Field( + default="64m", + description="Size of the writable /tmp tmpfs mounted when the rootfs is read-only.", + ) + runtime: str | None = Field( + default=None, + description="Container runtime for the execution step, e.g. 'runsc' for gVisor. " + "None uses Docker's default runtime.", + ) + seccomp_profile: PathLibPath | None = Field( + default=None, + description="Path to a custom seccomp JSON profile applied to the execution " + "container. None keeps Docker's built-in default seccomp profile active.", + ) class _HostConfig(BaseModel): @@ -53,9 +84,21 @@ class _HostConfig(BaseModel): CpuPeriod: int = Field(default=100_000, description="CPU CFS period in microseconds") PidsLimit: int = Field(default=64, description="Max number of PIDs in the container") ReadonlyRootfs: bool = Field(default=False, description="Mount root filesystem as read-only") + CapDrop: list[str] = Field( + default_factory=list, + description="Linux capabilities to drop (e.g. ['ALL'])", + ) + Tmpfs: dict[str, str] = Field( + default_factory=dict, + description="Writable tmpfs mounts (path -> mount options)", + ) SecurityOpt: list[str] = Field( default_factory=lambda: ["no-new-privileges"], - description="Security options (e.g. no-new-privileges)", + description="Security options (e.g. no-new-privileges, seccomp=...)", + ) + Runtime: str | None = Field( + default=None, + description="Container runtime (e.g. 'runsc' for gVisor); None = Docker default", ) @@ -69,19 +112,25 @@ class _ContainerConfig(BaseModel): description="Working directory inside the container", ) NetworkDisabled: bool = Field(default=True, description="Disable network access for isolation") + User: str = Field(default="", description="UID:GID to run as; empty = image default (root)") + Env: list[str] = Field(default_factory=list, description="Environment as KEY=VALUE strings") HostConfig: _HostConfig = Field(description="Host-level resource constraints") class DockerBackend: """Execute scripts in an isolated Docker container via aiodocker. - Provides full OS-level sandboxing: no host filesystem access, no network, - restricted memory/CPU. Requires Docker daemon and the ``aiodocker`` package + Provides full OS-level sandboxing for untrusted codegen output: no host + filesystem access, no network, a non-root user, all capabilities dropped, a + read-only root filesystem (with a small writable ``/tmp`` tmpfs), restricted + memory/CPU/PIDs, no-new-privileges, and Docker's seccomp profile (or a custom + one). Requires the Docker daemon and the ``aiodocker`` package (install with ``pip install xlstruct[docker]``). """ def __init__(self, config: DockerConfig | None = None) -> None: cfg = config or DockerConfig() + self._cfg = cfg self._image = cfg.image self._mem_limit = cfg.mem_limit self._cpu_quota = cfg.cpu_quota @@ -95,6 +144,36 @@ def _prepared_image_name(self) -> str: # ^ e.g. "python:3.11-slim" → "python:3.11-slim-xlstruct-ready" return f"{self._image}-{_PREPARED_IMAGE_TAG}" + def _install_host_config(self) -> _HostConfig: + """Permissive host config for the one-time package install (needs root + writes).""" + mem = _parse_mem_limit(self._mem_limit) + return _HostConfig(Memory=mem, MemorySwap=mem, CpuQuota=self._cpu_quota) + + def _exec_host_config(self) -> _HostConfig: + """Hardened host config for running untrusted scripts.""" + mem = _parse_mem_limit(self._mem_limit) + security_opt = ["no-new-privileges"] + if self._cfg.seccomp_profile is not None: + # ^ aiodocker passes profiles inline as JSON content, not a path. + profile = self._cfg.seccomp_profile.read_text(encoding="utf-8") + security_opt.append(f"seccomp={profile}") + + tmpfs: dict[str, str] = {} + if self._cfg.read_only_rootfs: + # ^ Scripts only print to stdout; a small writable /tmp covers any scratch. + tmpfs["/tmp"] = f"rw,nosuid,nodev,noexec,size={self._cfg.tmpfs_size},mode=1777" + + return _HostConfig( + Memory=mem, + MemorySwap=mem, + CpuQuota=self._cpu_quota, + ReadonlyRootfs=self._cfg.read_only_rootfs, + CapDrop=list(self._cfg.cap_drop), + Tmpfs=tmpfs, + SecurityOpt=security_opt, + Runtime=self._cfg.runtime, + ) + async def _ensure_image(self) -> None: """Prepare a Docker image with dependencies pre-installed. @@ -133,7 +212,7 @@ async def _ensure_image(self) -> None: logger.info("Pulling base image: %s", self._image) await docker.pull(self._image) - # * Stage 1: Install packages with network enabled → commit + # * Stage 1: Install packages with network enabled → commit (permissive: needs root) logger.info("Preparing image: installing %s", ", ".join(DOCKER_PIP_PACKAGES)) install_config = _ContainerConfig( Image=self._image, @@ -144,15 +223,11 @@ async def _ensure_image(self) -> None: *DOCKER_PIP_PACKAGES, ], NetworkDisabled=False, - HostConfig=_HostConfig( - Memory=_parse_mem_limit(self._mem_limit), - MemorySwap=_parse_mem_limit(self._mem_limit), - CpuQuota=self._cpu_quota, - ), + HostConfig=self._install_host_config(), ) container = await docker.containers.create( - config=install_config.model_dump(), + config=install_config.model_dump(exclude_none=True), ) try: @@ -184,7 +259,8 @@ async def execute( ) -> tuple[int, str, str]: """Execute code in a Docker container with full isolation. - Uses the prepared image (packages pre-installed) with network disabled. + Uses the prepared image (packages pre-installed) with network disabled, + a non-root user, dropped capabilities, and a read-only root filesystem. """ try: import aiodocker @@ -201,7 +277,7 @@ async def execute( if not source.exists(): return 1, "", f"Source file not found: {source_path}" - # * Stage 2: Run script with network disabled + # * Stage 2: Run script under the hardened, network-disabled config async with aiodocker.Docker() as docker: container_config = _ContainerConfig( Image=self._ready_image, @@ -211,19 +287,19 @@ async def execute( f"/workspace/{source.name}", ], NetworkDisabled=self._network_disabled, - HostConfig=_HostConfig( - Memory=_parse_mem_limit(self._mem_limit), - MemorySwap=_parse_mem_limit(self._mem_limit), - CpuQuota=self._cpu_quota, - ), + User=self._cfg.user, + # ^ Avoid .pyc writes against the read-only rootfs; flush stdout promptly. + Env=["PYTHONDONTWRITEBYTECODE=1", "PYTHONUNBUFFERED=1"], + HostConfig=self._exec_host_config(), ) container = await docker.containers.create( - config=container_config.model_dump(), + config=container_config.model_dump(exclude_none=True), ) try: - # * Copy files into container via tar archive + # * Copy files into container via tar archive (before start: lands in the + # writable layer, then stays readable once the rootfs is mounted read-only) tar_bytes = _build_tar_archive( ("script.py", code.encode("utf-8")), (source.name, source.read_bytes()), diff --git a/src/xlstruct/codegen/backends/resolver.py b/src/xlstruct/codegen/backends/resolver.py new file mode 100644 index 0000000..85458be --- /dev/null +++ b/src/xlstruct/codegen/backends/resolver.py @@ -0,0 +1,69 @@ +"""Resolve the execution backend for codegen — sandbox-by-default, fail-closed.""" + +import importlib.util +import logging +from typing import Literal + +from xlstruct.codegen.backends.base import ExecutionBackend +from xlstruct.codegen.backends.docker import DockerBackend +from xlstruct.codegen.backends.subprocess import SubprocessBackend +from xlstruct.exceptions import CodegenSecurityError, ErrorCode + +logger = logging.getLogger(__name__) + +CodegenSandbox = Literal["auto", "docker", "subprocess"] + + +def docker_available() -> bool: + """True if ``aiodocker`` is importable. Does not probe the Docker daemon. + + Importability is treated as the signal that the user installed + ``xlstruct[docker]`` intending to sandbox execution. If the daemon is down, + DockerBackend fails clearly at execution time (still fail-closed). + """ + return importlib.util.find_spec("aiodocker") is not None + + +def resolve_execution_backend( + explicit: ExecutionBackend | None, + sandbox: CodegenSandbox, +) -> ExecutionBackend: + """Pick the backend for executing untrusted, LLM-generated codegen output. + + Precedence: an explicitly injected backend always wins. Otherwise ``sandbox`` + decides: + + - ``"auto"``: DockerBackend when ``aiodocker`` is installed, else fail-closed + with a :class:`CodegenSecurityError` (never a silent fallback to subprocess). + - ``"docker"``: DockerBackend (errors at execution time if the daemon/package + is absent). + - ``"subprocess"``: the trusted/dev-only SubprocessBackend, which is NOT a + security boundary. + """ + if explicit is not None: + return explicit + + if sandbox == "subprocess": + logger.warning( + "codegen_sandbox='subprocess' selected — executing untrusted LLM output in a " + "non-isolating subprocess. Use only in trusted/development environments." + ) + return SubprocessBackend(trusted=True) + + if sandbox == "docker": + return DockerBackend() + + # * auto — sandbox when possible, fail-closed otherwise + if docker_available(): + return DockerBackend() + + raise CodegenSecurityError( + "Codegen executes untrusted, LLM-generated Python, and no OS-level sandbox is " + 'available. Install Docker support (`pip install "xlstruct[docker]"` with a running ' + "Docker daemon) to sandbox execution, or — only in a trusted/development environment " + "— opt into the unsandboxed subprocess backend via " + 'ExtractorConfig(codegen_sandbox="subprocess") or ' + "execution_backend=SubprocessBackend(trusted=True). The subprocess backend is NOT a " + "security boundary.", + code=ErrorCode.CODEGEN_NO_SANDBOX, + ) diff --git a/src/xlstruct/codegen/backends/subprocess.py b/src/xlstruct/codegen/backends/subprocess.py index b01fdaa..d8451e5 100644 --- a/src/xlstruct/codegen/backends/subprocess.py +++ b/src/xlstruct/codegen/backends/subprocess.py @@ -1,14 +1,30 @@ -"""SubprocessBackend — hardened subprocess execution with resource limits.""" +"""SubprocessBackend — hardened subprocess execution (trusted/dev-only). + +This backend is NOT a security boundary. The pre-execution AST scan is bypassable +by a code generator, and a subprocess shares the host filesystem, so a hostile +script can still read on-disk secrets. Use ``DockerBackend`` for untrusted output; +keep this backend only for trusted/development environments. The measures below are +defense-in-depth (reduce env/import surface, bound resources, contain timeouts). +""" import asyncio +import functools import logging import os +import signal +import subprocess import sys import tempfile from pathlib import Path as PathLibPath logger = logging.getLogger(__name__) +# ^ Shipped helper that runs the script with a reduced builtins namespace. +_BOOTSTRAP_PATH = PathLibPath(__file__).parent / "_subprocess_bootstrap.py" + +# ^ Emit the "not a security boundary" notice once per process, not per script. +_trusted_warning_emitted = False + # ^ Whitelist approach: only these env vars are passed to the subprocess ALLOWED_ENV_KEYS = frozenset( { @@ -39,42 +55,86 @@ def _build_safe_env() -> dict[str, str]: return {k: v for k, v in os.environ.items() if k in ALLOWED_ENV_KEYS} -def _apply_resource_limits() -> None: - """Apply resource limits to subprocess (Linux/macOS only). +def _apply_resource_limits(timeout: int) -> None: + """Apply resource limits in the forked child before exec (POSIX only). - Called as preexec_fn in subprocess. Limits memory to 512MB, - file descriptors to 64, max file size to 50MB. + Runs as ``preexec_fn``. Critical limits (CPU, file descriptors, file size) + are fail-closed: if they cannot be set the exception propagates, the child + never execs, and the parent's spawn fails. Address-space and process-count + limits are best-effort because their failure is environmental (e.g. an + unstable ``RLIMIT_AS`` on macOS, or a host already near its per-UID process + cap), not a missing boundary — CPU time plus process-group SIGKILL still + bound runaway and fork-bomb damage. """ - try: - import resource + import resource + + cpu = max(1, timeout) + # * Fail-closed limits — failure here aborts the child exec. + resource.setrlimit(resource.RLIMIT_CPU, (cpu, cpu + 5)) + resource.setrlimit(resource.RLIMIT_NOFILE, (256, 256)) + resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 1024**2, 50 * 1024**2)) - # ^ 512MB memory limit - resource.setrlimit(resource.RLIMIT_AS, (512 * 1024**2, 512 * 1024**2)) - # ^ 64 file descriptors - resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64)) - # ^ 50MB max file write size - resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 1024**2, 50 * 1024**2)) - except (ImportError, ValueError, OSError) as e: - logger.warning("Failed to set resource limits: %s", e) + # * Best-effort limits — environmental failure must not break trusted runs. + # Logging is unsafe post-fork, so failures are swallowed silently here. + try: + resource.setrlimit(resource.RLIMIT_NPROC, (512, 512)) # ^ fork-bomb cap + except (ValueError, OSError): + pass + try: + resource.setrlimit(resource.RLIMIT_AS, (512 * 1024**2, 512 * 1024**2)) # ^ 512MB + except (ValueError, OSError): + pass class SubprocessBackend: - """Execute scripts in a hardened subprocess. + """Execute scripts in a hardened subprocess (trusted/dev-only — NOT a boundary). - Security measures: + Defense-in-depth measures: - Credential environment variables stripped (whitelist approach) - - Memory limit: 512MB (via setrlimit) - - File descriptor limit: 64 (via setrlimit) - - Max file write size: 50MB (via setrlimit) + - Isolated interpreter (``python -I``): host env vars, user site-packages, and + the script directory are excluded from the import path + - Reduced builtins via bootstrap (eval/exec/compile/breakpoint/input removed) + - Resource limits: CPU time, file descriptors, file size (fail-closed); memory + and process count (best-effort) + - Timeout kills the whole process group (orphan-safe) with a bounded drain """ + def __init__(self, trusted: bool = False) -> None: + """Create the subprocess backend. + + Args: + trusted: Acknowledge that this backend is NOT a security boundary and + is being used in a trusted/development environment. When False, a + one-time warning is emitted on first use. Set True to suppress it. + """ + self._trusted = trusted + + def _warn_if_untrusted(self) -> None: + global _trusted_warning_emitted + if self._trusted or _trusted_warning_emitted: + return + _trusted_warning_emitted = True + logger.warning( + "SubprocessBackend is NOT a security boundary — the AST scan is bypassable " + "and the script shares the host filesystem. Use DockerBackend for untrusted " + "codegen output, or pass SubprocessBackend(trusted=True) to acknowledge " + "trusted/dev-only use and silence this warning." + ) + async def execute( self, code: str, source_path: str, timeout: int, ) -> tuple[int, str, str]: - """Execute code in subprocess with security hardening.""" + """Execute code in a hardened subprocess. + + Raises: + RuntimeError: Fail-closed — the subprocess could not be spawned with the + required resource limits applied. + """ + self._warn_if_untrusted() + fd, tmp_str = tempfile.mkstemp(suffix=".py", prefix="xlstruct_codegen_") tmp_path = PathLibPath(tmp_str) @@ -83,23 +143,45 @@ async def execute( with os.fdopen(fd, "w", encoding="utf-8") as f: f.write(code) - proc = await asyncio.create_subprocess_exec( - sys.executable, - str(tmp_path), - source_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=_build_safe_env(), - preexec_fn=_apply_resource_limits if sys.platform != "win32" else None, - ) + is_posix = sys.platform != "win32" + if is_posix: + # ^ preexec applies limits (fail-closed); new session enables killpg + preexec = functools.partial(_apply_resource_limits, timeout) + else: + preexec = None + logger.warning( + "Resource limits are not enforced on Windows for SubprocessBackend " + "(preexec_fn unsupported); rely on DockerBackend for real isolation." + ) + + try: + proc = await asyncio.create_subprocess_exec( + sys.executable, + "-I", # ^ isolated: ignore env vars, user site, and script dir + str(_BOOTSTRAP_PATH), + str(tmp_path), + source_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=_build_safe_env(), + preexec_fn=preexec, + start_new_session=is_posix, + ) + except (OSError, ValueError, subprocess.SubprocessError) as e: + # ^ Fail-closed: a preexec failure surfaces as subprocess.SubprocessError + # ("Exception occurred in preexec_fn."); spawn errors as OSError/ValueError. + # Either way the limits were not applied — refuse to run. + raise RuntimeError( + f"Refusing to execute codegen script: subprocess hardening could not " + f"be applied ({e})." + ) from e try: stdout_bytes, stderr_bytes = await asyncio.wait_for( proc.communicate(), timeout=timeout ) except TimeoutError: - proc.kill() - await proc.communicate() + await self._terminate(proc, is_posix) return -1, "", f"Script killed after {timeout}s timeout." stdout = stdout_bytes.decode("utf-8", errors="replace") @@ -108,3 +190,19 @@ async def execute( finally: tmp_path.unlink(missing_ok=True) + + @staticmethod + async def _terminate(proc: "asyncio.subprocess.Process", is_posix: bool) -> None: + """Kill a timed-out process (and its group on POSIX) with a bounded drain.""" + try: + if is_posix: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + else: + proc.kill() + except (ProcessLookupError, PermissionError): + pass + # ^ Bounded drain — a wedged pipe must not hang the caller forever. + try: + await asyncio.wait_for(proc.communicate(), timeout=5) + except (TimeoutError, ProcessLookupError): + pass diff --git a/src/xlstruct/codegen/cache.py b/src/xlstruct/codegen/cache.py index e5b9b46..9771e74 100644 --- a/src/xlstruct/codegen/cache.py +++ b/src/xlstruct/codegen/cache.py @@ -3,11 +3,19 @@ Caches generated scripts by sheet structure signature so that files with the same layout can reuse a previously generated script without additional LLM calls. + +Cached scripts are executed on a hit, so the cache is integrity-protected: the +directory is private (0700), entries are private (0600), and each script carries +an HMAC keyed by a per-user secret. Entries that fail verification are refused +(never executed) and regenerated. """ import hashlib +import hmac import json import logging +import os +import secrets from datetime import UTC, datetime from pathlib import Path as PathLibPath @@ -20,6 +28,9 @@ DEFAULT_CACHE_DIR = PathLibPath.home() / ".xlstruct" / "cache" +# ^ Per-user HMAC key, co-located in the (0700) cache dir so only the owner can read it. +_SECRET_FILENAME = ".hmac_key" + class CacheMetadata(BaseModel): """Metadata stored alongside a cached script.""" @@ -32,6 +43,7 @@ class CacheMetadata(BaseModel): header_sample: list[str] created_at: str explanation: str + mac: str = "" # ^ HMAC over (signature, code); "" marks a legacy/unverified entry def compute_structure_signature( @@ -45,6 +57,9 @@ def compute_structure_signature( - Header cell values (column names) - Column count - Schema field names and types + + Returns the full 256-bit SHA-256 hex digest (not truncated) so the cache + key is not feasibly predictable/collidable from the inputs. """ # * Collect header cell values header_values: list[str] = [] @@ -66,11 +81,27 @@ def compute_structure_signature( "|".join(field_sig), ] - return hashlib.sha256("\n".join(components).encode()).hexdigest()[:16] + return hashlib.sha256("\n".join(components).encode()).hexdigest() + + +def _compute_mac(secret: bytes, signature: str, code: str) -> str: + """HMAC-SHA256 over the signature-bound script bytes.""" + msg = signature.encode("utf-8") + b"\n" + code.encode("utf-8") + return hmac.new(secret, msg, hashlib.sha256).hexdigest() + + +def _is_hex_sha256(value: str) -> bool: + """True if ``value`` is a 64-char lowercase hex digest. + + Guards ``hmac.compare_digest`` against a non-ASCII ``mac`` from attacker-controlled + JSON (it raises TypeError on non-ASCII str), so a malformed value is treated as a + verification failure rather than crashing extraction. + """ + return len(value) == 64 and all(c in "0123456789abcdef" for c in value) class ScriptCache: - """File-based cache for generated codegen scripts.""" + """File-based cache for generated codegen scripts (integrity-protected).""" def __init__(self, cache_dir: PathLibPath | None = None) -> None: self._cache_dir = cache_dir or DEFAULT_CACHE_DIR @@ -79,8 +110,57 @@ def __init__(self, cache_dir: PathLibPath | None = None) -> None: def cache_dir(self) -> PathLibPath: return self._cache_dir + # * Integrity helpers + + def _ensure_dir(self) -> None: + """Create the cache dir and lock it to owner-only (0700).""" + self._cache_dir.mkdir(parents=True, exist_ok=True) + try: + os.chmod(self._cache_dir, 0o700) + except OSError as e: + # ^ Observable boundary: perms may be unsupported (e.g. Windows); HMAC still applies. + logger.warning("Could not restrict cache dir permissions to 0700: %s", e) + + def _get_secret(self) -> bytes: + """Load (or create) the per-user HMAC key stored 0600 in the cache dir.""" + self._ensure_dir() + secret_path = self._cache_dir / _SECRET_FILENAME + if secret_path.exists(): + try: + os.chmod(secret_path, 0o600) + except OSError: + pass + return secret_path.read_bytes() + + key = secrets.token_bytes(32) + try: + fd = os.open(secret_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + except FileExistsError: + # ^ Race: another process created it first — use theirs. + return secret_path.read_bytes() + with os.fdopen(fd, "wb") as f: + f.write(key) + return key + + @staticmethod + def _write_private(path: PathLibPath, data: str) -> None: + """Write text to ``path`` with owner-only (0600) permissions.""" + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(data) + try: + os.chmod(path, 0o600) # ^ enforce 0600 even when the file pre-existed + except OSError: + pass + + # * Public API + def get(self, signature: str) -> GeneratedScript | None: - """Look up a cached script by structure signature.""" + """Look up a cached script by structure signature. + + Returns None on a miss OR when integrity verification fails — a failed + entry is refused (never executed) and the caller regenerates. + """ script_path = self._cache_dir / f"{signature}.py" meta_path = self._cache_dir / f"{signature}.json" @@ -91,12 +171,32 @@ def get(self, signature: str) -> GeneratedScript | None: code = script_path.read_text(encoding="utf-8") meta_raw = json.loads(meta_path.read_text(encoding="utf-8")) meta = CacheMetadata.model_validate(meta_raw) - logger.info("Cache hit: %s (created %s)", signature, meta.created_at) - return GeneratedScript(code=code, explanation=meta.explanation) except Exception as e: logger.warning("Cache read failed for %s: %s", signature, e) return None + # * Integrity gate — refuse to return code we cannot verify. Any failure here + # (malformed/non-ASCII mac, unreadable secret) degrades to a clean miss, never + # a crash, preserving the skip+regenerate contract against hostile cache files. + try: + verified = _is_hex_sha256(meta.mac) and hmac.compare_digest( + meta.mac, _compute_mac(self._get_secret(), signature, code) + ) + except Exception as e: + logger.warning("Cache integrity check errored for %s: %s", signature, e) + verified = False + + if not verified: + logger.warning( + "Cache integrity check FAILED for %s — refusing to load the cached script " + "(it will not be executed). Regenerating.", + signature, + ) + return None + + logger.info("Cache hit: %s (created %s)", signature, meta.created_at) + return GeneratedScript(code=code, explanation=meta.explanation) + def put( self, signature: str, @@ -105,8 +205,9 @@ def put( header_rows: list[int], schema: type[BaseModel], ) -> PathLibPath: - """Store a script in the cache.""" - self._cache_dir.mkdir(parents=True, exist_ok=True) + """Store a script in the cache (private perms + integrity MAC).""" + self._ensure_dir() + secret = self._get_secret() script_path = self._cache_dir / f"{signature}.py" meta_path = self._cache_dir / f"{signature}.json" @@ -126,10 +227,11 @@ def put( header_sample=header_sample, created_at=datetime.now(UTC).isoformat(), explanation=script.explanation, + mac=_compute_mac(secret, signature, script.code), ) - script_path.write_text(script.code, encoding="utf-8") - meta_path.write_text(meta.model_dump_json(indent=2), encoding="utf-8") + self._write_private(script_path, script.code) + self._write_private(meta_path, meta.model_dump_json(indent=2)) logger.info("Cached script: %s → %s", signature, script_path) return script_path diff --git a/src/xlstruct/codegen/engine.py b/src/xlstruct/codegen/engine.py index f584f9b..9744a05 100644 --- a/src/xlstruct/codegen/engine.py +++ b/src/xlstruct/codegen/engine.py @@ -2,10 +2,14 @@ from typing import Any, TypeVar -import instructor from pydantic import BaseModel -from xlstruct.config import ExtractorConfig, apply_cache_control, build_instructor_client +from xlstruct.config import ( + ExtractorConfig, + apply_cache_control, + build_instructor_client, + thinking_call_kwargs, +) from xlstruct.exceptions import ErrorCode, ExtractionError from xlstruct.prompts.codegen import CODEGEN_SYSTEM_PROMPT from xlstruct.schemas.codegen import ( @@ -24,46 +28,8 @@ class CodegenEngine: def __init__(self, config: ExtractorConfig, tracker: UsageTracker | None = None) -> None: self._config = config self._tracker = tracker - self._model: str | None = None - self._client = self._build_client() - - def _build_client(self) -> Any: - """Create async Instructor client with provider-specific kwargs.""" - # ^ Anthropic thinking requires ANTHROPIC_REASONING_TOOLS mode - if self._config.thinking and self._config.provider.startswith("anthropic/"): - from anthropic import AsyncAnthropic # type: ignore - - model = self._config.provider.split("/", 1)[1] - client_kwargs: dict[str, Any] = {} - if self._config.api_key: - client_kwargs["api_key"] = self._config.api_key.get_secret_value() - client = instructor.from_anthropic( # type: ignore - AsyncAnthropic(**client_kwargs), # type: ignore - mode=instructor.Mode.ANTHROPIC_REASONING_TOOLS, - ) - # ^ Store model name for create() calls - self._model = model - return client - - self._model = None - return build_instructor_client(self._config) - - def _thinking_kwargs(self, temperature: float) -> dict[str, Any]: - """Build kwargs for create() with optional extended thinking. - - When thinking is enabled, forces temperature=1 - (required by Anthropic API) and sets a default budget. - """ - if self._config.thinking: - result: dict[str, Any] = { - "temperature": 1, - "thinking": {"type": "enabled", "budget_tokens": 10_000}, - "max_tokens": 16_000, - } - if self._model: - result["model"] = self._model - return result - return {"temperature": temperature} + self._client = build_instructor_client(config) + self._thinking_kwargs = thinking_call_kwargs(config) async def _call_llm( self, @@ -76,12 +42,12 @@ async def _call_llm( ) -> _T: """Execute LLM call with usage tracking and error handling.""" try: - kwargs = self._thinking_kwargs(temperature=temperature) + call_kwargs = {"temperature": temperature, **self._thinking_kwargs} result, completion = await self._client.create_with_completion( response_model=response_model, messages=messages, max_retries=self._config.max_retries, - **kwargs, + **call_kwargs, ) if self._tracker: self._tracker.record(label, completion) diff --git a/src/xlstruct/codegen/orchestrator.py b/src/xlstruct/codegen/orchestrator.py index aad77d8..1182468 100644 --- a/src/xlstruct/codegen/orchestrator.py +++ b/src/xlstruct/codegen/orchestrator.py @@ -14,7 +14,7 @@ from pydantic import BaseModel, ValidationError from xlstruct.codegen.backends.base import ExecutionBackend -from xlstruct.codegen.backends.subprocess import SubprocessBackend +from xlstruct.codegen.backends.resolver import resolve_execution_backend from xlstruct.codegen.engine import CodegenEngine from xlstruct.codegen.schema_utils import get_schema_source from xlstruct.codegen.validation import ScriptValidator @@ -49,7 +49,18 @@ def __init__( ) -> None: self._config = config self._engine = CodegenEngine(config, tracker=tracker) - self._backend: ExecutionBackend = backend or SubprocessBackend() + # ^ Resolve lazily at first script execution — header detection needs no + # backend, so fail-closed only triggers when untrusted code actually runs. + self._explicit_backend = backend + self._resolved_backend: ExecutionBackend | None = None + + def _get_backend(self) -> ExecutionBackend: + """Resolve (and memoize) the execution backend, fail-closed by default.""" + if self._resolved_backend is None: + self._resolved_backend = resolve_execution_backend( + self._explicit_backend, self._config.codegen_sandbox + ) + return self._resolved_backend # * Public API @@ -205,7 +216,9 @@ async def run_extraction( output_schema: type[BaseModel], ) -> list[Any]: """Execute a validated script and parse JSON output into Pydantic models.""" - validator = ScriptValidator(timeout=self._config.codegen_timeout, backend=self._backend) + validator = ScriptValidator( + timeout=self._config.codegen_timeout, backend=self._get_backend() + ) validation = await validator.validate( script.code, source, @@ -235,7 +248,9 @@ async def _validate_and_correct( Uses conversation-based correction: error feedback is appended to the existing messages history instead of re-sending the entire original prompt. """ - validator = ScriptValidator(timeout=self._config.codegen_timeout, backend=self._backend) + validator = ScriptValidator( + timeout=self._config.codegen_timeout, backend=self._get_backend() + ) attempts: list[CodegenAttempt] = [] current_code = result.code current_result = result @@ -373,14 +388,28 @@ def _parse_script_output(stdout: str, schema: type[BaseModel]) -> list[Any]: data = cast(list[dict[str, Any]], raw) results: list[BaseModel] = [] + failures: list[str] = [] for i, item in enumerate(data): # ^ Extract provenance before validation (not part of schema) source_row = item.pop("_source_row", None) try: record = schema.model_validate(item) - if source_row is not None: - object.__setattr__(record, "_source_rows", [source_row]) - results.append(record) except ValidationError as e: - logger.warning("Record %d failed validation, skipping: %s", i, e) + failures.append(f"Record {i}: {e}") + continue + if source_row is not None: + object.__setattr__(record, "_source_rows", [source_row]) + results.append(record) + + # ^ The script already passed full-output validation; invalid records here are a + # ^ real inconsistency — fail loudly instead of silently dropping them. + if failures: + shown = "\n".join(failures[:5]) + if len(failures) > 5: + shown += f"\n(+{len(failures) - 5} more)" + raise ExtractionError( + f"Codegen output failed schema validation: " + f"{len(failures)}/{len(data)} records invalid.\n{shown}", + code=ErrorCode.EXTRACTION_SCHEMA_VALIDATION_FAILED, + ) return results diff --git a/src/xlstruct/codegen/validation.py b/src/xlstruct/codegen/validation.py index aa93d9b..2b06b93 100644 --- a/src/xlstruct/codegen/validation.py +++ b/src/xlstruct/codegen/validation.py @@ -7,7 +7,6 @@ from pydantic import BaseModel, ValidationError from xlstruct.codegen.backends.base import ExecutionBackend -from xlstruct.codegen.backends.subprocess import SubprocessBackend from xlstruct.codegen.executor import scan_blocked_imports logger = logging.getLogger(__name__) @@ -29,11 +28,13 @@ class ScriptValidator: def __init__( self, + backend: ExecutionBackend, timeout: int = 60, - backend: ExecutionBackend | None = None, ) -> None: + # ^ backend is required: no silent SubprocessBackend default — the caller + # (orchestrator) resolves a sandboxed backend fail-closed. self._timeout = timeout - self._backend = backend or SubprocessBackend() + self._backend = backend async def validate( self, @@ -68,8 +69,11 @@ async def validate( "openpyxl, python-calamine, pydantic, json, sys, re, datetime, " "decimal, math, typing, enum, collections, dataclasses, copy, " "itertools, functools, csv, and similar standard data " - "processing libraries. Dangerous builtins (exec, eval, open, " - "getattr, globals, etc.) and dunder escape patterns are also blocked." + "processing libraries. The builtins __import__, exec, eval, compile, " + "open, and breakpoint are blocked, as are the dotted patterns " + "sys.modules/sys.path/sys._getframe/sys.meta_path and dunder escape " + "attributes (__subclasses__, __globals__, __builtins__, __bases__, " + "__mro__, __code__). This scan is a best-effort filter, not a sandbox." ), ) @@ -150,11 +154,11 @@ def _filter_by_required_fields( ] if len(filtered) < original_count: - logger.info( - "Post-filter: %d → %d records (removed %d with null required fields)", - original_count, - len(filtered), + # ^ WARNING (not INFO): dropping records is data loss the caller should see + logger.warning( + "Post-filter dropped %d/%d records with null required fields", original_count - len(filtered), + original_count, ) return json.dumps(filtered, ensure_ascii=False, indent=2, default=str) @@ -163,12 +167,14 @@ def _filter_by_required_fields( def _validate_output( stdout: str, schema: type[BaseModel], - max_sample: int = 5, + max_errors: int = 5, total_data_rows: int | None = None, ) -> str: """Validate stdout JSON against the Pydantic schema. - Parses stdout as JSON array and validates a sample of items. + Validates EVERY record (not a sample) so type errors beyond the first few + rows still fail the script and drive self-correction instead of being + silently dropped downstream. Reports up to ``max_errors`` failures. Returns empty string if valid, error description if invalid. """ stdout_stripped = stdout.strip() @@ -231,20 +237,23 @@ def _validate_output( ) return msg - # * Validate a sample of items against the schema + # * Validate every record against the schema (no sampling) errors: list[str] = [] - sample_indices = list(range(min(max_sample, len(data)))) - for idx in sample_indices: - item = data[idx] + failed = 0 + for idx, item in enumerate(data): try: schema.model_validate(item) except ValidationError as e: - errors.append(f"Record {idx}: {e}") + failed += 1 + if len(errors) < max_errors: + errors.append(f"Record {idx}: {e}") - if errors: + if failed: error_detail = "\n".join(errors) + if failed > len(errors): + error_detail += f"\n(+{failed - len(errors)} more failing records not shown)" return ( - f"OUTPUT VALIDATION ERROR: {len(errors)}/{len(sample_indices)} sampled records " + f"OUTPUT VALIDATION ERROR: {failed}/{len(data)} records " f"failed schema validation ({schema.__name__}).\n{error_detail}" ) diff --git a/src/xlstruct/config.py b/src/xlstruct/config.py index 066dc13..1a6358b 100644 --- a/src/xlstruct/config.py +++ b/src/xlstruct/config.py @@ -2,7 +2,7 @@ from enum import StrEnum from pathlib import Path as PathLibPath -from typing import Any +from typing import Any, Literal import instructor from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator @@ -38,6 +38,14 @@ class ExtractorConfig(BaseModel): max_tokens: int = Field(default=8192, gt=0) max_codegen_retries: int = Field(default=3, ge=0) codegen_timeout: int = Field(default=60, gt=0) + codegen_sandbox: Literal["auto", "docker", "subprocess"] = Field( + default="auto", + description="Sandbox for executing untrusted, LLM-generated codegen scripts. " + "'auto' uses Docker when xlstruct[docker] is installed, else fails closed with a " + "clear error. 'docker' forces the Docker backend. 'subprocess' opts into the " + "non-isolating subprocess backend — NOT a security boundary; trusted/dev only. " + "An explicit execution_backend always overrides this.", + ) thinking: bool = Field( default=False, description="Enable Anthropic extended thinking mode. " @@ -70,6 +78,11 @@ class ExtractorConfig(BaseModel): description="Evaluate formula cells using the formulas library. " "Requires: pip install xlstruct[formulas]", ) + csv_encoding: str = Field( + default="utf-8", + description="Text encoding for CSV files (e.g. 'euc-kr', 'cp1252', 'shift-jis'). " + "'utf-8' transparently strips a BOM if present. Ignored for Excel formats.", + ) min_chunk_rows: int = Field( default=10, gt=0, @@ -80,6 +93,14 @@ class ExtractorConfig(BaseModel): gt=0, description="Force chunking for sheets exceeding this row count.", ) + max_concurrent_chunks: int = Field( + default=5, + gt=0, + description="Max chunk LLM calls run concurrently for a single chunked sheet. " + "Bounds parallelism to respect provider rate limits. Multiplies with " + "extract_workbook(concurrency=...) on multi-sheet runs " + "(worst-case in-flight calls = concurrency × max_concurrent_chunks).", + ) provider_options: dict[str, Any] = Field(default_factory=dict) storage_options: dict[str, Any] = Field(default_factory=dict) @@ -151,9 +172,11 @@ def is_anthropic(provider: str) -> bool: def apply_cache_control(messages: list[dict[str, Any]], provider: str) -> list[dict[str, Any]]: - """Apply Anthropic prompt caching markers to messages. + """Apply the Anthropic prompt-caching marker to the static system prompt only. - Wraps system and first user message content with cache_control ephemeral markers. + The system prompt is stable across calls, so caching it reliably hits. The user + message holds variable sheet data — it changes every call, would essentially never + hit, and tagging it only burns a cache breakpoint. So it is left unmarked. Returns messages unchanged for non-Anthropic providers. """ if not is_anthropic(provider): @@ -164,8 +187,8 @@ def apply_cache_control(messages: list[dict[str, Any]], provider: str) -> list[d role = msg["role"] content = msg["content"] - # ^ Apply cache_control to system prompt and first user message - if role in ("system", "user") and isinstance(content, str): + # ^ Cache only the stable system prompt, not the variable user message + if role == "system" and isinstance(content, str): result.append( { "role": role, @@ -205,11 +228,28 @@ def get_provider_kwargs(config: ExtractorConfig) -> dict[str, Any]: return defaults +def _use_anthropic_thinking(config: "ExtractorConfig") -> bool: + """Whether extended thinking applies — Anthropic-only, requires 'anthropic/'.""" + return config.thinking and config.provider.startswith("anthropic/") + + def build_instructor_client(config: "ExtractorConfig") -> Any: - """Create async Instructor client with provider-specific kwargs. + """Create async Instructor client honoring extended thinking. - Centralizes the common pattern: get_provider_kwargs → inject api_key → from_provider. + For Anthropic + thinking, uses ANTHROPIC_REASONING_TOOLS mode (the model is supplied + per-call via thinking_call_kwargs). Otherwise uses the standard provider client. """ + if _use_anthropic_thinking(config): + from anthropic import AsyncAnthropic # type: ignore + + client_kwargs: dict[str, Any] = {} + if config.api_key: + client_kwargs["api_key"] = config.api_key.get_secret_value() + return instructor.from_anthropic( # type: ignore + AsyncAnthropic(**client_kwargs), # type: ignore + mode=instructor.Mode.ANTHROPIC_REASONING_TOOLS, + ) + kwargs = get_provider_kwargs(config) if config.api_key: kwargs["api_key"] = config.api_key.get_secret_value() @@ -218,3 +258,21 @@ def build_instructor_client(config: "ExtractorConfig") -> Any: async_client=True, **kwargs, ) + + +def thinking_call_kwargs(config: "ExtractorConfig") -> dict[str, Any]: + """Per-call create_with_completion() kwargs for extended thinking. + + Empty when thinking is off. For Anthropic + thinking, forces temperature=1 + (Anthropic requirement), enables the thinking block, and carries max_tokens + model + (from_anthropic does not embed the model the way from_provider does). Callers merge + this last so it overrides any caller-supplied temperature. + """ + if not _use_anthropic_thinking(config): + return {} + return { + "temperature": 1, + "thinking": {"type": "enabled", "budget_tokens": 10_000}, + "max_tokens": 16_000, + "model": config.provider.split("/", 1)[1], + } diff --git a/src/xlstruct/encoder/_formatting.py b/src/xlstruct/encoder/_formatting.py index cab8477..c60d306 100644 --- a/src/xlstruct/encoder/_formatting.py +++ b/src/xlstruct/encoder/_formatting.py @@ -344,14 +344,18 @@ def _generalize_formula(formula: str, row: int) -> str: # * Empty row/column detection -def find_empty_rows(sheet: SheetData) -> set[int]: - """Find row numbers that are completely empty.""" +def find_non_empty_rows(sheet: SheetData) -> set[int]: + """Return row numbers that contain at least one non-empty cell. + + Callers invert the membership test to skip empty rows; returning the + non-empty set avoids materializing a set over the full row range + (1..row_count), which is wasteful for tall sheets. + """ non_empty_rows: set[int] = set() for cell in sheet.cells: if cell.display_value is not None: non_empty_rows.add(cell.row) - all_rows = set(range(1, sheet.row_count + 1)) - return all_rows - non_empty_rows + return non_empty_rows # * Column type summary diff --git a/src/xlstruct/encoder/compressed.py b/src/xlstruct/encoder/compressed.py index 5998333..508576c 100644 --- a/src/xlstruct/encoder/compressed.py +++ b/src/xlstruct/encoder/compressed.py @@ -10,7 +10,7 @@ build_column_headers, build_multi_row_headers, detect_header_row, - find_empty_rows, + find_non_empty_rows, format_cell_value, format_merged_regions, summarize_column_types, @@ -72,9 +72,9 @@ def encode(self, sheet: SheetData, header_rows: list[int] | None = None) -> str: parts.append(", ".join(type_parts)) # * Build markdown table (full or sampled) - empty_rows = find_empty_rows(sheet) + non_empty_rows = find_non_empty_rows(sheet) table_text, total_data_rows, shown_data_rows = self._build_table( - sheet, headers, data_start, empty_rows + sheet, headers, data_start, non_empty_rows ) parts.append("") @@ -92,7 +92,7 @@ def _build_table( sheet: SheetData, headers: dict[int, str], data_start: int, - empty_rows: set[int], + non_empty_rows: set[int], ) -> tuple[str, int, int]: """Build markdown table, optionally sampling rows. @@ -118,7 +118,7 @@ def _build_table( row_num = row_cells[0].row if row_num < data_start: continue - if row_num in empty_rows: + if row_num not in non_empty_rows: continue row_values: dict[int, str] = {} diff --git a/src/xlstruct/exceptions.py b/src/xlstruct/exceptions.py index 19d9936..d794c6f 100644 --- a/src/xlstruct/exceptions.py +++ b/src/xlstruct/exceptions.py @@ -26,6 +26,7 @@ class ErrorCode(StrEnum): CODEGEN_MAX_RETRIES = "CODEGEN_MAX_RETRIES" CODEGEN_SYNTAX_ERROR = "CODEGEN_SYNTAX_ERROR" CODEGEN_EXECUTION_FAILED = "CODEGEN_EXECUTION_FAILED" + CODEGEN_NO_SANDBOX = "CODEGEN_NO_SANDBOX" class XLStructError(Exception): @@ -54,3 +55,12 @@ class CodegenValidationError(XLStructError): def __init__(self, message: str, attempts: "list[Any]", code: ErrorCode | None = None) -> None: super().__init__(message, code=code) self.attempts = attempts + + +class CodegenSecurityError(XLStructError): + """No OS-level sandbox is available to execute untrusted codegen output. + + Raised (fail-closed) instead of silently running LLM-generated code in the + non-isolating subprocess backend. Resolve by installing Docker support or by + explicitly opting into the trusted/dev-only subprocess backend. + """ diff --git a/src/xlstruct/extraction/chunking.py b/src/xlstruct/extraction/chunking.py index b4d5039..c8a0428 100644 --- a/src/xlstruct/extraction/chunking.py +++ b/src/xlstruct/extraction/chunking.py @@ -1,6 +1,6 @@ """ChunkSplitter: Splits large sheets into processable chunks.""" -from xlstruct._tokens import estimate_sheet_tokens +from xlstruct._tokens import estimate_row_token_costs, estimate_sheet_tokens from xlstruct.schemas.core import CellData, SheetData # ^ Minimum rows per chunk to avoid degenerate cases @@ -67,22 +67,30 @@ def split( if not sorted_row_nums: return [sheet] - # ^ Calculate rows per chunk from both token budget and row threshold - data_row_count = len(sorted_row_nums) - total_tokens = estimate_sheet_tokens(sheet) - - if total_tokens > token_budget: - # ^ Token-based: split proportionally - chunk_count = max(1, total_tokens // token_budget) - rows_per_chunk = max(min_chunk_rows, data_row_count // chunk_count) - else: - # ^ Row-based: cap at threshold - rows_per_chunk = max(min_chunk_rows, row_threshold) + # ^ Size chunks by each data row's real token cost (greedy bin-packing), + # ^ reserving the header cost that every chunk re-includes. This replaces + # ^ "sampled total / budget assuming uniform rows": chunks are packed to + # ^ the actual per-row weight, so a chunk fits (token_budget - header) + # ^ regardless of where the heavy rows sit. The row_threshold accuracy cap + # ^ still bounds each chunk's row count. + all_costs = estimate_row_token_costs( + [header_cells, *(data_rows[rn] for rn in sorted_row_nums)] + ) + header_cost = all_costs[0] + row_costs = all_costs[1:] + data_budget = max(1, token_budget - header_cost) + + row_groups = self._pack_rows( + sorted_row_nums, + row_costs, + data_budget=data_budget, + min_chunk_rows=min_chunk_rows, + row_threshold=row_threshold, + ) # * Build chunks chunks: list[SheetData] = [] - for i in range(0, len(sorted_row_nums), rows_per_chunk): - chunk_row_nums = sorted_row_nums[i : i + rows_per_chunk] + for chunk_row_nums in row_groups: chunk_cells = list(header_cells) # ^ Copy header cells into each chunk for rn in chunk_row_nums: chunk_cells.extend(data_rows[rn]) @@ -102,3 +110,40 @@ def split( ) return chunks + + @staticmethod + def _pack_rows( + row_nums: list[int], + row_costs: list[int], + *, + data_budget: int, + min_chunk_rows: int, + row_threshold: int, + ) -> list[list[int]]: + """Greedily group rows into chunks that fit data_budget tokens. + + A chunk closes when it reaches row_threshold rows (the accuracy cap, a + hard upper bound that takes precedence over min_chunk_rows) or when + adding the next row would exceed data_budget — but the budget close only + applies once the chunk already holds min_chunk_rows, so tiny chunks are + avoided. Two degenerate cases can still exceed the budget, both + unavoidable: a single row heavier than data_budget (a row cannot be + split), and the first min_chunk_rows rows when they are collectively + heavier than data_budget (they are appended before the budget gate + activates). + """ + groups: list[list[int]] = [] + current: list[int] = [] + current_cost = 0 + for row_num, cost in zip(row_nums, row_costs): + over_budget = current_cost + cost > data_budget and len(current) >= min_chunk_rows + at_row_cap = len(current) >= row_threshold + if current and (over_budget or at_row_cap): + groups.append(current) + current = [] + current_cost = 0 + current.append(row_num) + current_cost += cost + if current: + groups.append(current) + return groups diff --git a/src/xlstruct/extraction/concurrency.py b/src/xlstruct/extraction/concurrency.py new file mode 100644 index 0000000..8c64ba6 --- /dev/null +++ b/src/xlstruct/extraction/concurrency.py @@ -0,0 +1,78 @@ +"""Bounded-concurrency map shared by Extractor.extract_batch / extract_workbook. + +Centralizes the semaphore + progress-event + ordered-gather boilerplate that was +duplicated as near-identical closures in both methods. The per-item logic differs +(batch calls ``self.extract``; workbook builds a per-sheet engine), so callers pass +a ``worker`` that runs under the semaphore and returns ``(result, status, error)``. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from typing import TypeVar + +from xlstruct.schemas.progress import ProgressEvent, ProgressStatus + +ItemT = TypeVar("ItemT") +ResultT = TypeVar("ResultT") + + +async def run_concurrent( + items: list[ItemT], + worker: Callable[[ItemT], Awaitable[tuple[ResultT, ProgressStatus, str | None]]], + *, + concurrency: int, + label: Callable[[ItemT], str], + on_progress: Callable[[ProgressEvent], None] | None = None, +) -> list[ResultT]: + """Map ``worker`` over ``items`` with bounded concurrency, preserving input order. + + Emits a STARTED ``ProgressEvent`` before acquiring the semaphore and a terminal + event (the status ``worker`` returns) after each item completes. ``worker`` runs + under the semaphore and must not raise — it returns ``(result, status, error)``. + + Args: + items: Work items, processed in order (results returned in the same order). + worker: Async per-item handler. Returns ``(result, terminal status, error)``. + concurrency: Max items processed simultaneously. + label: Maps an item to its ``ProgressEvent.source`` label. + on_progress: Optional callback invoked on STARTED and on completion. + """ + semaphore = asyncio.Semaphore(concurrency) + total = len(items) + completed = 0 + count_lock = asyncio.Lock() + + async def _run(item: ItemT) -> ResultT: + nonlocal completed + + if on_progress: + on_progress( + ProgressEvent( + source=label(item), + status=ProgressStatus.STARTED, + completed=completed, + total=total, + ) + ) + + async with semaphore: + result, status, error = await worker(item) + + async with count_lock: + completed += 1 + current = completed + + if on_progress: + on_progress( + ProgressEvent( + source=label(item), + status=status, + completed=current, + total=total, + error=error, + ) + ) + + return result + + return await asyncio.gather(*[_run(item) for item in items]) diff --git a/src/xlstruct/extraction/engine.py b/src/xlstruct/extraction/engine.py index fa00610..a849976 100644 --- a/src/xlstruct/extraction/engine.py +++ b/src/xlstruct/extraction/engine.py @@ -4,7 +4,12 @@ from pydantic import BaseModel, Field, create_model -from xlstruct.config import ExtractorConfig, apply_cache_control, build_instructor_client +from xlstruct.config import ( + ExtractorConfig, + apply_cache_control, + build_instructor_client, + thinking_call_kwargs, +) from xlstruct.exceptions import ErrorCode, ExtractionError from xlstruct.prompts.extraction import build_extraction_prompt from xlstruct.prompts.system import SYSTEM_PROMPT @@ -93,8 +98,9 @@ def _split_confidence( # ^ Extract and convert confidence scores for name in field_names: conf_key = f"{name}_confidence" - level = data.pop(conf_key, "moderate") - field_confidences[name].append(CONFIDENCE_SCORES.get(level, 0.5)) + # ^ Required by the confidence wrapper schema — a missing key means a structural bug + level = data.pop(conf_key) + field_confidences[name].append(CONFIDENCE_SCORES[level]) record = original_schema.model_validate(data) cleaned.append(record) @@ -109,6 +115,7 @@ def __init__(self, config: ExtractorConfig, tracker: UsageTracker | None = None) self._config = config self._tracker = tracker self._client = build_instructor_client(config) + self._thinking_kwargs = thinking_call_kwargs(config) async def extract( self, @@ -149,50 +156,58 @@ async def extract( response_schema, exclude_fields=provenance_fields ) + # ^ Message prep is our code — keep it outside the LLM try/except below + messages = apply_cache_control( + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + self._config.provider, + ) + call_kwargs: dict[str, Any] = { + "temperature": self._config.temperature, + **self._thinking_kwargs, + } + + # * Only the provider call is wrapped — our post-processing keeps its own exception type try: - messages = apply_cache_control( - [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - self._config.provider, - ) result, completion = await self._client.create_with_completion( response_model=list[response_schema], # type: ignore messages=messages, max_retries=self._config.max_retries, - temperature=self._config.temperature, + **call_kwargs, ) - if self._tracker: - self._tracker.record("extraction", completion) - - items = list(result) - - # ^ Split confidence first (outermost wrapper), then provenance - field_confidences: dict[str, list[float]] | None = None - if include_confidence: - # ^ The schema to split against includes provenance fields if both are enabled - split_schema = _build_provenance_schema(schema) if track_provenance else schema - items, field_confidences = _split_confidence(items, split_schema) - - if track_provenance: - items = self._split_provenance(items, schema) - - # ^ Attach confidence data to each record for collection by Extractor - if field_confidences is not None: - for i, item in enumerate(items): - per_record: dict[str, float] = {} - for field_name, scores in field_confidences.items(): - if i < len(scores): - per_record[field_name] = scores[i] - object.__setattr__(item, "_field_confidences", per_record) - - return items # type: ignore except Exception as e: raise ExtractionError( f"LLM extraction failed: {e}", code=ErrorCode.EXTRACTION_LLM_FAILED ) from e + if self._tracker: + self._tracker.record("extraction", completion) + + items = list(result) + + # ^ Split confidence first (outermost wrapper), then provenance + field_confidences: dict[str, list[float]] | None = None + if include_confidence: + # ^ The schema to split against includes provenance fields if both are enabled + split_schema = _build_provenance_schema(schema) if track_provenance else schema + items, field_confidences = _split_confidence(items, split_schema) + + if track_provenance: + items = self._split_provenance(items, schema) + + # ^ Attach confidence data to each record for collection by Extractor + if field_confidences is not None: + for i, item in enumerate(items): + per_record: dict[str, float] = {} + for field_name, scores in field_confidences.items(): + if i < len(scores): + per_record[field_name] = scores[i] + object.__setattr__(item, "_field_confidences", per_record) + + return items # type: ignore + @staticmethod def _split_provenance(items: list[BaseModel], original_schema: type[T]) -> list[T]: """Strip provenance fields from wrapper records, rebuild as original schema. diff --git a/src/xlstruct/extraction/pipeline.py b/src/xlstruct/extraction/pipeline.py new file mode 100644 index 0000000..31cae11 --- /dev/null +++ b/src/xlstruct/extraction/pipeline.py @@ -0,0 +1,714 @@ +"""ExtractionPipeline: orchestration engine behind the Extractor facade. + +Owns the Storage → Reader → Encoder → Engine pipeline plus codegen routing, +chunked and streaming execution, and multi-sheet/batch concurrency. ``Extractor`` +is a thin public facade that delegates every operation here. +""" + +import asyncio +import logging +import re +from collections.abc import AsyncGenerator, Callable +from pathlib import Path as PathLibPath +from typing import Any, TypeVar + +from pydantic import BaseModel + +from xlstruct._tokens import count_tokens +from xlstruct.codegen.backends.base import ExecutionBackend +from xlstruct.codegen.cache import ScriptCache, compute_structure_signature +from xlstruct.codegen.orchestrator import CodegenOrchestrator +from xlstruct.config import ( + SAMPLE_ROWS, + ExtractionConfig, + ExtractionMode, + ExtractorConfig, + apply_cache_control, + build_instructor_client, + thinking_call_kwargs, +) +from xlstruct.encoder.compressed import CompressedEncoder +from xlstruct.exceptions import ErrorCode, ExtractionError, ReaderError +from xlstruct.extraction.chunking import ChunkSplitter, needs_chunking +from xlstruct.extraction.concurrency import run_concurrent +from xlstruct.extraction.engine import ExtractionEngine +from xlstruct.extraction.report import build_extraction_report +from xlstruct.extraction.result import ExtractionResult +from xlstruct.reader.base import ReaderOptions +from xlstruct.reader.dispatch import read_workbook +from xlstruct.schemas.batch import BatchResult, FileResult +from xlstruct.schemas.codegen import GeneratedScript +from xlstruct.schemas.core import SheetData, WorkbookData +from xlstruct.schemas.progress import ProgressEvent, ProgressStatus +from xlstruct.schemas.report import ExtractionReport +from xlstruct.schemas.usage import UsageTracker +from xlstruct.schemas.workbook import SheetResult, WorkbookResult +from xlstruct.storage import read_file + +logger = logging.getLogger(__name__) + +T = TypeVar("T", bound=BaseModel) + + +class ExtractionPipeline: + """Orchestration engine for XLStruct. Public surface lives on ``Extractor``.""" + + def __init__( + self, + config: ExtractorConfig, + *, + execution_backend: ExecutionBackend | None = None, + ) -> None: + self._config = config + self._execution_backend = execution_backend + self._tracker = UsageTracker() + self._engine = ExtractionEngine(self._config, tracker=self._tracker) + self._codegen: CodegenOrchestrator | None = None + self._chunk_splitter = ChunkSplitter() + self._cache: ScriptCache | None = None + if self._config.cache_enabled: + self._cache = ScriptCache(cache_dir=self._config.cache_dir) + + # * Public read-only state — consumed by the Extractor facade + + @property + def config(self) -> ExtractorConfig: + return self._config + + @property + def engine(self) -> ExtractionEngine: + return self._engine + + @property + def chunk_splitter(self) -> ChunkSplitter: + return self._chunk_splitter + + @property + def cache(self) -> ScriptCache | None: + return self._cache + + # * Script export + + def _export_script(self, source: str, script: GeneratedScript) -> PathLibPath | None: + """Save generated script to export_dir if configured.""" + export_dir = self._config.export_dir + if export_dir is None: + return None + + export_dir.mkdir(parents=True, exist_ok=True) + + # ^ Derive filename from source: "report.xlsx" → "report_codegen.py" + stem = PathLibPath(source.rsplit("/", 1)[-1]).stem + safe_stem = re.sub(r"[^\w\-]", "_", stem) + script_path = export_dir / f"{safe_stem}_codegen.py" + + script_path.write_text(script.code, encoding="utf-8") + logger.info("Exported codegen script: %s", script_path) + return script_path + + # * Lazy codegen orchestrator + + def _get_codegen(self) -> CodegenOrchestrator: + if self._codegen is None: + self._codegen = CodegenOrchestrator( + self._config, backend=self._execution_backend, tracker=self._tracker + ) + return self._codegen + + # * Single-sheet extraction + + async def extract( + self, + source: str, + schema: type[T] | None = None, + *, + extraction_config: ExtractionConfig | None = None, + sheet: str | None = None, + instructions: str | None = None, + **storage_options: Any, + ) -> ExtractionResult[T]: + """Implementation of ``Extractor.extract()``.""" + self._tracker.reset() + + if extraction_config is not None: + items, resolved_mode = await self._run_configured_extraction( + source, extraction_config, **storage_options + ) + elif schema is not None: + workbook = await self.load_workbook(source, sheet_name=sheet, **storage_options) + target_sheet = workbook.sheets[0] + self._require_non_empty_sheet(target_sheet) + items = await self.run_sheet_extraction( + target_sheet, schema, instructions, engine=self._engine + ) + resolved_mode = ExtractionMode.DIRECT + else: + raise ValueError("Either schema or extraction_config must be provided") + + usage = self._tracker.snapshot() + logger.info(usage) + + report = build_extraction_report(items, resolved_mode, usage) + return ExtractionResult(items, report=report) + + async def generate_script( + self, + source: str, + extraction_config: ExtractionConfig, + **storage_options: Any, + ) -> GeneratedScript: + """Implementation of ``Extractor.generate_script()``.""" + workbook = await self.load_workbook( + source, sheet_name=extraction_config.sheet, **storage_options + ) + full_sheet = workbook.sheets[0] + codegen = self._get_codegen() + + # * Auto-detect header rows if not provided + header_rows = extraction_config.header_rows + if header_rows is None: + header_rows = await codegen.detect_header_rows(full_sheet) + + script = await codegen.generate_script(source, full_sheet, header_rows, extraction_config) + self._export_script(source, script) + return script + + # * Streaming extraction + + async def stream( + self, + source: str, + schema: type[T] | None = None, + *, + extraction_config: ExtractionConfig | None = None, + sheet: str | None = None, + instructions: str | None = None, + **storage_options: Any, + ) -> AsyncGenerator[T, None]: + """Implementation of ``Extractor.stream()``.""" + if extraction_config is not None: + async for item in self._stream_configured_extraction( + source, extraction_config, **storage_options + ): + yield item + elif schema is not None: + workbook = await self.load_workbook(source, sheet_name=sheet, **storage_options) + target_sheet = workbook.sheets[0] + async for item in self._stream_sheet_extraction( + target_sheet, schema, instructions, engine=self._engine + ): + yield item + else: + raise ValueError("Either schema or extraction_config must be provided") + + async def suggest_schema( + self, + source: str, + *, + sheet: str | None = None, + instructions: str | None = None, + **storage_options: Any, + ) -> type[BaseModel]: + """Implementation of ``Extractor.suggest_schema()``.""" + from pydantic import Field, create_model + + workbook = await self.load_workbook(source, sheet_name=sheet, **storage_options) + target_sheet = workbook.sheets[0] + + encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) + encoded = encoder.encode(target_sheet) + + hint = "" + if instructions: + hint = f"\nAdditional context: {instructions}\n" + + prompt = ( + "Analyze the following spreadsheet data and suggest a Pydantic model.\n\n" + "Rules:\n" + "- Return a JSON array of field definitions\n" + "- Each field: {name (snake_case), type, nullable, description}\n" + "- type must be one of: str, int, float, bool, date, datetime\n" + "- description should mention the original Excel column name\n" + "- model_name: PascalCase name for the model\n" + f"{hint}\n" + f"Spreadsheet data:\n{encoded}" + ) + + from xlstruct.prompts.system import SYSTEM_PROMPT + from xlstruct.schemas.suggest import SuggestedFields + + client = build_instructor_client(self._config) + call_kwargs: dict[str, Any] = {"temperature": 0.0, **thinking_call_kwargs(self._config)} + + messages = apply_cache_control( + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + self._config.provider, + ) + result, completion = await client.create_with_completion( + response_model=SuggestedFields, + messages=messages, + **call_kwargs, + ) + if self._tracker: + self._tracker.record("suggest_schema", completion) + + # * Build dynamic Pydantic model via create_model() + type_map: dict[str, type] = { + "str": str, + "int": int, + "float": float, + "bool": bool, + "date": __import__("datetime").date, + "datetime": __import__("datetime").datetime, + } + + field_definitions: dict[str, Any] = {} + for f in result.fields: + python_type = type_map.get(f.type, str) + if f.nullable: + python_type = python_type | None # type: ignore + field_definitions[f.name] = ( + python_type, + Field(description=f.description), + ) + + return create_model(result.model_name, **field_definitions) + + # * Multi-sheet extraction + + async def extract_workbook( + self, + source: str, + sheet_schemas: dict[str, type[BaseModel]], + *, + concurrency: int = 5, + instructions: str | None = None, + on_progress: Callable[[ProgressEvent], None] | None = None, + **storage_options: Any, + ) -> WorkbookResult: + """Implementation of ``Extractor.extract_workbook()``.""" + # ^ Load all sheets at once (sheet_name=None) + workbook = await self.load_workbook(source, sheet_name=None, **storage_options) + + async def _worker( + pair: tuple[str, type[BaseModel]], + ) -> tuple[tuple[str, SheetResult[Any]], ProgressStatus, str | None]: + sheet_name, schema = pair + sheet_data = workbook.get_sheet(sheet_name) + if sheet_data is None: + error_msg = f"Sheet '{sheet_name}' not found. Available: {workbook.sheet_names}" + result: SheetResult[Any] = SheetResult( + sheet_name=sheet_name, success=False, error=error_msg + ) + return (sheet_name, result), ProgressStatus.FAILED, error_msg + + try: + tracker = UsageTracker() + engine = ExtractionEngine(self._config, tracker=tracker) + items = await self.run_sheet_extraction( + sheet_data, schema, instructions, engine=engine + ) + result = SheetResult( + sheet_name=sheet_name, + success=True, + records=items, + usage=tracker.snapshot(), + ) + return (sheet_name, result), ProgressStatus.COMPLETED, None + except Exception as e: + logger.warning("Workbook extraction failed for sheet '%s': %s", sheet_name, e) + error_msg = f"{type(e).__name__}: {e}" + result = SheetResult(sheet_name=sheet_name, success=False, error=error_msg) + return (sheet_name, result), ProgressStatus.FAILED, error_msg + + pairs = await run_concurrent( + list(sheet_schemas.items()), + _worker, + concurrency=concurrency, + label=lambda pair: pair[0], + on_progress=on_progress, + ) + return WorkbookResult(results=dict(pairs)) + + # * Cross-sheet extraction + + async def extract_cross_sheet( + self, + source: str, + *, + schema: type[T], + sheets: list[str], + header_rows: dict[str, list[int]] | list[int] | None = None, + instructions: str | None = None, + **storage_options: Any, + ) -> ExtractionResult[T]: + """Implementation of ``Extractor.extract_cross_sheet()``.""" + if len(sheets) < 2: + raise ValueError(f"extract_cross_sheet requires at least 2 sheets, got {len(sheets)}") + + self._tracker.reset() + + # * Load entire workbook (all sheets) + workbook = await self.load_workbook(source, sheet_name=None, **storage_options) + + # * Validate all requested sheets exist + missing = [s for s in sheets if workbook.get_sheet(s) is None] + if missing: + raise ValueError(f"Sheets not found: {missing}. Available: {workbook.sheet_names}") + + # * Encode each sheet separately and concatenate + encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) + encoded_parts: list[str] = [] + for sheet_name in sheets: + sheet_data = workbook.get_sheet(sheet_name) + assert sheet_data is not None # ^ Already validated above + + # * Resolve header_rows for this sheet + sheet_header_rows: list[int] | None + if header_rows is None: + sheet_header_rows = None + elif isinstance(header_rows, list): + sheet_header_rows = header_rows + else: + sheet_header_rows = header_rows.get(sheet_name) + + encoded_parts.append(encoder.encode(sheet_data, header_rows=sheet_header_rows)) + + combined_encoding = "\n\n".join(encoded_parts) + + # * Validate combined encoding fits within token budget + combined_tokens = count_tokens(combined_encoding) + if combined_tokens > self._config.token_budget: + raise ExtractionError( + f"Combined cross-sheet encoding ({combined_tokens:,} tokens) exceeds " + f"token budget ({self._config.token_budget:,}). " + f"Reduce the number of sheets or increase token_budget.", + code=ErrorCode.EXTRACTION_LLM_FAILED, + ) + + # * Send combined encoding to ExtractionEngine + items = await self._engine.extract( + combined_encoding, + schema, + instructions, + ) + + usage = self._tracker.snapshot() + logger.info(usage) + + report = ExtractionReport( + mode=ExtractionMode.DIRECT, + usage=usage, + ) + return ExtractionResult(items, report=report) + + # * Batch extraction + + async def extract_batch( + self, + sources: list[str], + schema: type[T] | None = None, + *, + extraction_config: ExtractionConfig | None = None, + concurrency: int = 5, + sheet: str | None = None, + instructions: str | None = None, + on_progress: Callable[[ProgressEvent], None] | None = None, + **storage_options: Any, + ) -> BatchResult[T]: + """Implementation of ``Extractor.extract_batch()``.""" + + async def _worker(source: str) -> tuple[FileResult[T], ProgressStatus, str | None]: + try: + result = await self.extract( + source, + schema, + extraction_config=extraction_config, + sheet=sheet, + instructions=instructions, + **storage_options, + ) + file_result = FileResult( + source=source, + success=True, + records=list(result), + usage=result.report.usage, + ) + return file_result, ProgressStatus.COMPLETED, None + except Exception as e: + logger.warning("Batch extraction failed for %s: %s", source, e) + error_msg = f"{type(e).__name__}: {e}" + return ( + FileResult[T](source=source, success=False, error=error_msg), + ProgressStatus.FAILED, + error_msg, + ) + + file_results = await run_concurrent( + sources, + _worker, + concurrency=concurrency, + label=lambda source: source, + on_progress=on_progress, + ) + return BatchResult(results=list(file_results)) + + # * Private pipeline methods + + @staticmethod + def _require_non_empty_sheet(sheet: SheetData) -> None: + """Fail fast on a 0-row sheet before spending an LLM header-detection call.""" + if sheet.row_count == 0: + raise ReaderError( + f"Sheet '{sheet.name}' has no rows.", + code=ErrorCode.READER_PARSE_FAILED, + ) + + async def load_workbook( + self, + source: str, + sheet_name: str | None = None, + **storage_options: Any, + ) -> WorkbookData: + """Storage → Reader pipeline.""" + merged_options = {**self._config.storage_options, **storage_options} + file_bytes = await read_file(source, **merged_options) + + options = ReaderOptions( + strict_formulas=self._config.strict_formulas, + evaluate_formulas=self._config.evaluate_formulas, + csv_encoding=self._config.csv_encoding, + ) + workbook = await asyncio.to_thread( + read_workbook, file_bytes, source, sheet_name, options=options + ) + + workbook.file_name = source.rsplit("/", 1)[-1] + workbook.file_size = len(file_bytes) + return workbook + + @staticmethod + def _resolve_auto_mode( + full_sheet: SheetData, + header_rows: list[int], + requested_mode: ExtractionMode, + *, + log_label: str = "Auto-routing", + ) -> ExtractionMode: + """Resolve AUTO to DIRECT/CODEGEN by data-row count; pass other modes through. + + Routing: data rows ≤ SAMPLE_ROWS → DIRECT, > SAMPLE_ROWS → CODEGEN. + """ + if requested_mode != ExtractionMode.AUTO: + return requested_mode + data_rows = full_sheet.row_count - max(header_rows) + mode = ExtractionMode.CODEGEN if data_rows > SAMPLE_ROWS else ExtractionMode.DIRECT + logger.info("%s: %d data rows → mode=%s", log_label, data_rows, mode.value) + return mode + + async def _run_configured_extraction( + self, + source: str, + config: ExtractionConfig, + **storage_options: Any, + ) -> tuple[list[Any], ExtractionMode]: + """Config-based extraction with mode selection. + + - mode=auto: heuristic routing (≤ SAMPLE_ROWS → direct, > SAMPLE_ROWS → codegen). + - mode=direct: always use LLM direct extraction. + - mode=codegen: always use code generation pipeline. + + Returns: + Tuple of (extracted items, resolved extraction mode). + """ + workbook = await self.load_workbook(source, sheet_name=config.sheet, **storage_options) + full_sheet = workbook.sheets[0] + self._require_non_empty_sheet(full_sheet) + codegen = self._get_codegen() + + # * Auto-detect header rows if not provided + header_rows = config.header_rows + if header_rows is None: + header_rows = await codegen.detect_header_rows(full_sheet) + + mode = self._resolve_auto_mode(full_sheet, header_rows, config.mode) + + if mode == ExtractionMode.CODEGEN: + items = await self._run_codegen(source, full_sheet, header_rows, config, codegen) + return items, ExtractionMode.CODEGEN + + items = await self._run_direct(full_sheet, header_rows, config) + return items, ExtractionMode.DIRECT + + async def _run_codegen( + self, + source: str, + full_sheet: SheetData, + header_rows: list[int], + config: ExtractionConfig, + codegen: CodegenOrchestrator, + ) -> list[Any]: + """Code generation pipeline: cache lookup → generate script → execute → parse.""" + script: GeneratedScript | None = None + signature: str | None = None + + # * Cache lookup + if self._cache is not None: + signature = compute_structure_signature(full_sheet, header_rows, config.output_schema) + script = self._cache.get(signature) + + if script is None: + # * Cache miss — generate via LLM + script = await codegen.generate_script(source, full_sheet, header_rows, config) + self._export_script(source, script) + + # * Store in cache + if self._cache is not None and signature is not None: + self._cache.put(signature, script, full_sheet, header_rows, config.output_schema) + + return await codegen.run_extraction(source, script, config.output_schema) + + async def _run_direct( + self, + full_sheet: SheetData, + header_rows: list[int], + config: ExtractionConfig, + ) -> list[Any]: + """Direct LLM extraction: encode → LLM → Pydantic.""" + encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) + encoded = encoder.encode(full_sheet, header_rows=header_rows) + + return await self._engine.extract( + encoded, + config.output_schema, + config.instructions, + is_sampled=True, + total_rows=full_sheet.row_count, + track_provenance=config.track_provenance, + include_confidence=config.include_confidence, + ) + + async def run_sheet_extraction( + self, + sheet: SheetData, + schema: type[T], + instructions: str | None = None, + *, + engine: ExtractionEngine, + ) -> list[T]: + """Encoder → (optional Chunking) → ExtractionEngine pipeline.""" + target_engine = engine + encoder = CompressedEncoder() + + if needs_chunking(sheet, self._config.token_budget, self._config.chunking_row_threshold): + # * Chunked extraction + chunks = self._chunk_splitter.split( + sheet, + self._config.token_budget, + min_chunk_rows=self._config.min_chunk_rows, + row_threshold=self._config.chunking_row_threshold, + ) + # * Parallel chunk extraction with bounded concurrency (rate-limit safe) + semaphore = asyncio.Semaphore(self._config.max_concurrent_chunks) + + async def _extract_chunk(chunk: SheetData) -> list[T]: + async with semaphore: + encoded = encoder.encode(chunk) + return await target_engine.extract(encoded, schema, instructions) + + # ^ gather preserves input order, so records stay in chunk order + chunk_results = await asyncio.gather(*[_extract_chunk(c) for c in chunks]) + all_results: list[T] = [] + for partial in chunk_results: + all_results.extend(partial) + return all_results + else: + # * Single-pass extraction + encoded = encoder.encode(sheet) + return await target_engine.extract(encoded, schema, instructions) + + # * Streaming private helpers + + async def _stream_configured_extraction( + self, + source: str, + config: ExtractionConfig, + **storage_options: Any, + ) -> AsyncGenerator[Any, None]: + """Streaming variant of _run_configured_extraction. + + For codegen mode, yields all records at once (script produces all results + in a single run). For direct mode, yields records as each chunk completes. + """ + workbook = await self.load_workbook(source, sheet_name=config.sheet, **storage_options) + full_sheet = workbook.sheets[0] + self._require_non_empty_sheet(full_sheet) + + # * Auto-detect header rows if not provided (requires codegen orchestrator) + header_rows = config.header_rows + if header_rows is None: + codegen = self._get_codegen() + header_rows = await codegen.detect_header_rows(full_sheet) + + mode = self._resolve_auto_mode( + full_sheet, header_rows, config.mode, log_label="Auto-routing (stream)" + ) + + if mode == ExtractionMode.CODEGEN: + codegen = self._get_codegen() + items = await self._run_codegen(source, full_sheet, header_rows, config, codegen) + for item in items: + yield item + return + + # * Direct mode — single-pass with sampling (no chunking in config mode) + encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) + encoded = encoder.encode(full_sheet, header_rows=header_rows) + items = await self._engine.extract( + encoded, + config.output_schema, + config.instructions, + is_sampled=True, + total_rows=full_sheet.row_count, + track_provenance=config.track_provenance, + ) + for item in items: + yield item + + async def _stream_sheet_extraction( + self, + sheet: SheetData, + schema: type[T], + instructions: str | None = None, + *, + engine: ExtractionEngine, + ) -> AsyncGenerator[T, None]: + """Streaming variant of run_sheet_extraction. + + Yields records incrementally as each chunk's LLM call completes. + For single-chunk sheets, yields all records at once. + """ + encoder = CompressedEncoder() + + if needs_chunking(sheet, self._config.token_budget, self._config.chunking_row_threshold): + # * Chunked extraction — yield from each chunk as it completes + chunks = self._chunk_splitter.split( + sheet, + self._config.token_budget, + min_chunk_rows=self._config.min_chunk_rows, + row_threshold=self._config.chunking_row_threshold, + ) + for chunk in chunks: + encoded = encoder.encode(chunk) + partial = await engine.extract(encoded, schema, instructions) + for item in partial: + yield item + else: + # * Single-pass extraction + encoded = encoder.encode(sheet) + items = await engine.extract(encoded, schema, instructions) + for item in items: + yield item diff --git a/src/xlstruct/extraction/report.py b/src/xlstruct/extraction/report.py new file mode 100644 index 0000000..58033f8 --- /dev/null +++ b/src/xlstruct/extraction/report.py @@ -0,0 +1,56 @@ +"""Assemble an ExtractionReport from extracted records. + +Pulls provenance (``_source_rows`` / ``_source_cells``) and per-field confidence +(``_field_confidences``) that ``ExtractionEngine`` attaches to records, and folds +them into the report. Extracted from ``Extractor.extract`` to keep that method +a thin orchestrator. +""" + +from typing import Any + +from xlstruct.config import ExtractionMode +from xlstruct.schemas.report import ExtractionReport +from xlstruct.schemas.usage import TokenUsage + + +def build_extraction_report( + items: list[Any], + mode: ExtractionMode, + usage: TokenUsage, +) -> ExtractionReport: + """Build an ExtractionReport, folding in provenance/confidence from records. + + Args: + items: Extracted records (each may carry ``_source_rows`` / ``_source_cells`` / + ``_field_confidences`` attributes set by the engine). + mode: The extraction mode that was actually used. + usage: Token usage snapshot for this extraction. + """ + # * Provenance — parallel to items; dropped entirely if nothing was tracked + source_rows: list[list[int]] = [getattr(item, "_source_rows", []) for item in items] + source_cells: list[dict[str, str]] = [getattr(item, "_source_cells", {}) for item in items] + if not any(source_rows): + source_rows = [] + if not any(source_cells): + source_cells = [] + + # * Confidence — only when the engine attached per-field scores + field_confidences: dict[str, list[float]] | None = None + if items and hasattr(items[0], "_field_confidences"): + all_fields: set[str] = set() + for item in items: + per_record = getattr(item, "_field_confidences", {}) + all_fields.update(per_record.keys()) + field_confidences = {name: [] for name in sorted(all_fields)} + for item in items: + per_record = getattr(item, "_field_confidences", {}) + for name in field_confidences: + field_confidences[name].append(per_record.get(name, 0.5)) + + return ExtractionReport( + mode=mode, + usage=usage, + source_rows=source_rows, + source_cells=source_cells, + field_confidences=field_confidences, + ) diff --git a/src/xlstruct/extraction/result.py b/src/xlstruct/extraction/result.py new file mode 100644 index 0000000..a4bd36a --- /dev/null +++ b/src/xlstruct/extraction/result.py @@ -0,0 +1,49 @@ +"""ExtractionResult: list[T] carrying an attached ExtractionReport. + +Lives in its own module so both the Extractor facade and the ExtractionPipeline +can construct it without a circular import. +""" + +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from pandas import DataFrame # type: ignore + +from pydantic import BaseModel + +from xlstruct.schemas.report import ExtractionReport + +T = TypeVar("T", bound=BaseModel) + + +class ExtractionResult(list[T]): # type: ignore + """List of extracted records with an attached extraction report. + + Behaves exactly like list[T] (iteration, indexing, len, etc.) + but also exposes a ``.report`` attribute containing extraction metadata + (mode used, token usage, provenance, etc.). + """ + + report: ExtractionReport + + def __init__(self, items: list[T], report: ExtractionReport) -> None: + super().__init__(items) + self.report = report + + def to_dataframe(self) -> "DataFrame": + """Convert extracted records to a pandas DataFrame. + + Requires pandas to be installed: ``pip install xlstruct[pandas]`` + + Returns: + pandas DataFrame with one row per extracted record. + """ + try: + import pandas as pd + except ImportError: + raise ImportError( + "pandas is required for to_dataframe(). " + "Install it with: pip install xlstruct[pandas]" + ) from None + + return pd.DataFrame([item.model_dump() for item in self]) diff --git a/src/xlstruct/extractor.py b/src/xlstruct/extractor.py index f747315..5a6f950 100644 --- a/src/xlstruct/extractor.py +++ b/src/xlstruct/extractor.py @@ -1,85 +1,34 @@ -"""Extractor: Public API for XLStruct. +"""Extractor: public API facade for XLStruct. -Orchestrates the full pipeline: Storage → Reader → Encoder → Engine. -Delegates code generation to CodegenOrchestrator. +Thin facade over :class:`ExtractionPipeline`. It builds the config, owns a single +pipeline instance, and delegates every operation to it — all orchestration +(Storage → Reader → Encoder → Engine, codegen routing, chunking, concurrency) +lives in the pipeline. """ import asyncio -import logging -import re from collections.abc import AsyncGenerator, Callable, Iterator -from pathlib import Path as PathLibPath -from typing import TYPE_CHECKING, Any, TypeVar - -if TYPE_CHECKING: - from pandas import DataFrame # type: ignore +from typing import Any, TypeVar from pydantic import BaseModel, SecretStr -from xlstruct._tokens import count_tokens from xlstruct.codegen.backends.base import ExecutionBackend -from xlstruct.codegen.cache import ScriptCache, compute_structure_signature -from xlstruct.codegen.orchestrator import CodegenOrchestrator -from xlstruct.config import ( - SAMPLE_ROWS, - ExtractionConfig, - ExtractionMode, - ExtractorConfig, - apply_cache_control, - build_instructor_client, -) -from xlstruct.encoder.compressed import CompressedEncoder -from xlstruct.exceptions import ErrorCode, ExtractionError, ReaderError -from xlstruct.extraction.chunking import ChunkSplitter, needs_chunking +from xlstruct.codegen.cache import ScriptCache +from xlstruct.config import ExtractionConfig, ExtractorConfig +from xlstruct.extraction.chunking import ChunkSplitter from xlstruct.extraction.engine import ExtractionEngine -from xlstruct.reader.hybrid_reader import HybridReader -from xlstruct.schemas.batch import BatchResult, FileResult +from xlstruct.extraction.pipeline import ExtractionPipeline +from xlstruct.extraction.result import ExtractionResult +from xlstruct.reader.dispatch import get_source_ext +from xlstruct.schemas.batch import BatchResult from xlstruct.schemas.codegen import GeneratedScript from xlstruct.schemas.core import SheetData, WorkbookData -from xlstruct.schemas.progress import ProgressEvent, ProgressStatus -from xlstruct.schemas.report import ExtractionReport -from xlstruct.schemas.usage import UsageTracker -from xlstruct.schemas.workbook import SheetResult, WorkbookResult -from xlstruct.storage import read_file - -logger = logging.getLogger(__name__) +from xlstruct.schemas.progress import ProgressEvent +from xlstruct.schemas.workbook import WorkbookResult T = TypeVar("T", bound=BaseModel) -class ExtractionResult(list[T]): # type: ignore - """List of extracted records with an attached extraction report. - - Behaves exactly like list[T] (iteration, indexing, len, etc.) - but also exposes a ``.report`` attribute containing extraction metadata - (mode used, token usage, provenance, etc.). - """ - - report: ExtractionReport - - def __init__(self, items: list[T], report: ExtractionReport) -> None: - super().__init__(items) - self.report = report - - def to_dataframe(self) -> "DataFrame": - """Convert extracted records to a pandas DataFrame. - - Requires pandas to be installed: ``pip install xlstruct[pandas]`` - - Returns: - pandas DataFrame with one row per extracted record. - """ - try: - import pandas as pd - except ImportError: - raise ImportError( - "pandas is required for to_dataframe(). " - "Install it with: pip install xlstruct[pandas]" - ) from None - - return pd.DataFrame([item.model_dump() for item in self]) - - def _run_sync(coro: Any) -> Any: """Run a coroutine synchronously, with Jupyter/notebook compatibility. @@ -124,19 +73,26 @@ def __init__( **kwargs: Any, ) -> None: if config is not None: - self._config = config + resolved_config = config else: secret_key = SecretStr(api_key) if api_key is not None else None - self._config = ExtractorConfig(provider=provider, api_key=secret_key, **kwargs) + resolved_config = ExtractorConfig(provider=provider, api_key=secret_key, **kwargs) - self._execution_backend = execution_backend - self._tracker = UsageTracker() - self._engine = ExtractionEngine(self._config, tracker=self._tracker) - self._codegen: CodegenOrchestrator | None = None - self._chunk_splitter = ChunkSplitter() - self._cache: ScriptCache | None = None - if self._config.cache_enabled: - self._cache = ScriptCache(cache_dir=self._config.cache_dir) + self._pipeline = ExtractionPipeline(resolved_config, execution_backend=execution_backend) + + # * Shared internals — the pipeline owns these; exposed for introspection and tests + + @property + def _config(self) -> ExtractorConfig: + return self._pipeline.config + + @property + def _engine(self) -> ExtractionEngine: + return self._pipeline.engine + + @property + def _chunk_splitter(self) -> ChunkSplitter: + return self._pipeline.chunk_splitter @property def cache(self) -> ScriptCache | None: @@ -145,35 +101,7 @@ def cache(self) -> ScriptCache | None: Returns None if caching is disabled (``cache_enabled=False``). When enabled, provides ``list_entries()``, ``clear()``, ``remove()`` methods. """ - return self._cache - - # * Script export - - def _export_script(self, source: str, script: GeneratedScript) -> PathLibPath | None: - """Save generated script to export_dir if configured.""" - export_dir = self._config.export_dir - if export_dir is None: - return None - - export_dir.mkdir(parents=True, exist_ok=True) - - # ^ Derive filename from source: "report.xlsx" → "report_codegen.py" - stem = PathLibPath(source.rsplit("/", 1)[-1]).stem - safe_stem = re.sub(r"[^\w\-]", "_", stem) - script_path = export_dir / f"{safe_stem}_codegen.py" - - script_path.write_text(script.code, encoding="utf-8") - logger.info("Exported codegen script: %s", script_path) - return script_path - - # * Lazy codegen orchestrator - - def _get_codegen(self) -> CodegenOrchestrator: - if self._codegen is None: - self._codegen = CodegenOrchestrator( - self._config, backend=self._execution_backend, tracker=self._tracker - ) - return self._codegen + return self._pipeline.cache # * Public API @@ -205,55 +133,14 @@ async def extract( Returns: ExtractionResult — list[T] with ``.report`` for extraction metadata. """ - self._tracker.reset() - - if extraction_config is not None: - items, resolved_mode = await self._run_configured_extraction( - source, extraction_config, **storage_options - ) - elif schema is not None: - workbook = await self._load_workbook(source, sheet_name=sheet, **storage_options) - target_sheet = workbook.sheets[0] - items = await self._run_sheet_extraction( - target_sheet, schema, instructions, engine=self._engine - ) - resolved_mode = ExtractionMode.DIRECT - else: - raise ValueError("Either schema or extraction_config must be provided") - - # * Collect provenance from records (set by ExtractionEngine._split_provenance) - source_rows: list[list[int]] = [getattr(item, "_source_rows", []) for item in items] - source_cells: list[dict[str, str]] = [getattr(item, "_source_cells", {}) for item in items] - # ^ Only include if any provenance was actually tracked - if not any(source_rows): - source_rows = [] - if not any(source_cells): - source_cells = [] - - # * Collect confidence from records (set by ExtractionEngine.extract) - field_confidences: dict[str, list[float]] | None = None - if items and hasattr(items[0], "_field_confidences"): - all_fields: set[str] = set() - for item in items: - per_record = getattr(item, "_field_confidences", {}) - all_fields.update(per_record.keys()) - field_confidences = {name: [] for name in sorted(all_fields)} - for item in items: - per_record = getattr(item, "_field_confidences", {}) - for name in field_confidences: - field_confidences[name].append(per_record.get(name, 0.5)) - - usage = self._tracker.snapshot() - logger.info(usage) - - report = ExtractionReport( - mode=resolved_mode, - usage=usage, - source_rows=source_rows, - source_cells=source_cells, - field_confidences=field_confidences, + return await self._pipeline.extract( + source, + schema, + extraction_config=extraction_config, + sheet=sheet, + instructions=instructions, + **storage_options, ) - return ExtractionResult(items, report=report) async def generate_script( self, @@ -271,20 +158,7 @@ async def generate_script( Returns: GeneratedScript with code and explanation. """ - workbook = await self._load_workbook( - source, sheet_name=extraction_config.sheet, **storage_options - ) - full_sheet = workbook.sheets[0] - codegen = self._get_codegen() - - # * Auto-detect header rows if not provided - header_rows = extraction_config.header_rows - if header_rows is None: - header_rows = await codegen.detect_header_rows(full_sheet) - - script = await codegen.generate_script(source, full_sheet, header_rows, extraction_config) - self._export_script(source, script) - return script + return await self._pipeline.generate_script(source, extraction_config, **storage_options) def generate_script_sync( self, @@ -339,20 +213,15 @@ async def stream( Yields: Individual ``T`` instances as they are extracted. """ - if extraction_config is not None: - async for item in self._stream_configured_extraction( - source, extraction_config, **storage_options - ): - yield item - elif schema is not None: - workbook = await self._load_workbook(source, sheet_name=sheet, **storage_options) - target_sheet = workbook.sheets[0] - async for item in self._stream_sheet_extraction( - target_sheet, schema, instructions, engine=self._engine - ): - yield item - else: - raise ValueError("Either schema or extraction_config must be provided") + async for item in self._pipeline.stream( + source, + schema, + extraction_config=extraction_config, + sheet=sheet, + instructions=instructions, + **storage_options, + ): + yield item def stream_sync( self, @@ -400,71 +269,9 @@ async def suggest_schema( Returns: A Pydantic model class built via ``pydantic.create_model()``. """ - from pydantic import Field, create_model - - workbook = await self._load_workbook(source, sheet_name=sheet, **storage_options) - target_sheet = workbook.sheets[0] - - encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) - encoded = encoder.encode(target_sheet) - - hint = "" - if instructions: - hint = f"\nAdditional context: {instructions}\n" - - prompt = ( - "Analyze the following spreadsheet data and suggest a Pydantic model.\n\n" - "Rules:\n" - "- Return a JSON array of field definitions\n" - "- Each field: {name (snake_case), type, nullable, description}\n" - "- type must be one of: str, int, float, bool, date, datetime\n" - "- description should mention the original Excel column name\n" - "- model_name: PascalCase name for the model\n" - f"{hint}\n" - f"Spreadsheet data:\n{encoded}" - ) - - from xlstruct.prompts.system import SYSTEM_PROMPT - from xlstruct.schemas.suggest import SuggestedFields - - client = build_instructor_client(self._config) - - messages = apply_cache_control( - [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - self._config.provider, - ) - result, completion = await client.create_with_completion( - response_model=SuggestedFields, - messages=messages, - temperature=0.0, + return await self._pipeline.suggest_schema( + source, sheet=sheet, instructions=instructions, **storage_options ) - if self._tracker: - self._tracker.record("suggest_schema", completion) - - # * Build dynamic Pydantic model via create_model() - type_map: dict[str, type] = { - "str": str, - "int": int, - "float": float, - "bool": bool, - "date": __import__("datetime").date, - "datetime": __import__("datetime").datetime, - } - - field_definitions: dict[str, Any] = {} - for f in result.fields: - python_type = type_map.get(f.type, str) - if f.nullable: - python_type = python_type | None # type: ignore - field_definitions[f.name] = ( - python_type, - Field(description=f.description), - ) - - return create_model(result.model_name, **field_definitions) def suggest_schema_sync( self, @@ -540,87 +347,14 @@ async def extract_workbook( Returns: WorkbookResult with per-sheet results keyed by sheet name. """ - # ^ Load all sheets at once (sheet_name=None) - workbook = await self._load_workbook(source, sheet_name=None, **storage_options) - - semaphore = asyncio.Semaphore(concurrency) - total = len(sheet_schemas) - completed_count = 0 - count_lock = asyncio.Lock() - - async def _extract_sheet( - sheet_name: str, schema: type[BaseModel] - ) -> tuple[str, SheetResult[Any]]: - nonlocal completed_count - - if on_progress: - on_progress( - ProgressEvent( - source=sheet_name, - status=ProgressStatus.STARTED, - completed=completed_count, - total=total, - ) - ) - - async with semaphore: - sheet_data = workbook.get_sheet(sheet_name) - if sheet_data is None: - error_msg = f"Sheet '{sheet_name}' not found. Available: {workbook.sheet_names}" - sheet_result: SheetResult[Any] = SheetResult( - sheet_name=sheet_name, - success=False, - error=error_msg, - ) - status = ProgressStatus.FAILED - else: - try: - tracker = UsageTracker() - engine = ExtractionEngine(self._config, tracker=tracker) - items = await self._run_sheet_extraction( - sheet_data, schema, instructions, engine=engine - ) - sheet_result = SheetResult( - sheet_name=sheet_name, - success=True, - records=items, - usage=tracker.snapshot(), - ) - status = ProgressStatus.COMPLETED - error_msg = None - except Exception as e: - logger.warning( - "Workbook extraction failed for sheet '%s': %s", sheet_name, e - ) - error_msg = f"{type(e).__name__}: {e}" - sheet_result = SheetResult( - sheet_name=sheet_name, - success=False, - error=error_msg, - ) - status = ProgressStatus.FAILED - - async with count_lock: - completed_count += 1 - current_completed = completed_count - - if on_progress: - on_progress( - ProgressEvent( - source=sheet_name, - status=status, - completed=current_completed, - total=total, - error=error_msg, - ) - ) - - return sheet_name, sheet_result - - pairs = await asyncio.gather( - *[_extract_sheet(name, schema) for name, schema in sheet_schemas.items()] + return await self._pipeline.extract_workbook( + source, + sheet_schemas, + concurrency=concurrency, + instructions=instructions, + on_progress=on_progress, + **storage_options, ) - return WorkbookResult(results=dict(pairs)) def extract_workbook_sync( self, @@ -669,65 +403,15 @@ async def extract_cross_sheet( Raises: ValueError: If fewer than 2 sheets are specified or a sheet is not found. """ - if len(sheets) < 2: - raise ValueError(f"extract_cross_sheet requires at least 2 sheets, got {len(sheets)}") - - self._tracker.reset() - - # * Load entire workbook (all sheets) - workbook = await self._load_workbook(source, sheet_name=None, **storage_options) - - # * Validate all requested sheets exist - missing = [s for s in sheets if workbook.get_sheet(s) is None] - if missing: - raise ValueError(f"Sheets not found: {missing}. Available: {workbook.sheet_names}") - - # * Encode each sheet separately and concatenate - encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) - encoded_parts: list[str] = [] - for sheet_name in sheets: - sheet_data = workbook.get_sheet(sheet_name) - assert sheet_data is not None # ^ Already validated above - - # * Resolve header_rows for this sheet - sheet_header_rows: list[int] | None - if header_rows is None: - sheet_header_rows = None - elif isinstance(header_rows, list): - sheet_header_rows = header_rows - else: - sheet_header_rows = header_rows.get(sheet_name) - - encoded_parts.append(encoder.encode(sheet_data, header_rows=sheet_header_rows)) - - combined_encoding = "\n\n".join(encoded_parts) - - # * Validate combined encoding fits within token budget - combined_tokens = count_tokens(combined_encoding) - if combined_tokens > self._config.token_budget: - raise ExtractionError( - f"Combined cross-sheet encoding ({combined_tokens:,} tokens) exceeds " - f"token budget ({self._config.token_budget:,}). " - f"Reduce the number of sheets or increase token_budget.", - code=ErrorCode.EXTRACTION_LLM_FAILED, - ) - - # * Send combined encoding to ExtractionEngine - items = await self._engine.extract( - combined_encoding, - schema, - instructions, + return await self._pipeline.extract_cross_sheet( + source, + schema=schema, + sheets=sheets, + header_rows=header_rows, + instructions=instructions, + **storage_options, ) - usage = self._tracker.snapshot() - logger.info(usage) - - report = ExtractionReport( - mode=ExtractionMode.DIRECT, - usage=usage, - ) - return ExtractionResult(items, report=report) - def extract_cross_sheet_sync( self, source: str, @@ -773,71 +457,16 @@ async def extract_batch( Returns: BatchResult with per-file results and aggregated usage. """ - semaphore = asyncio.Semaphore(concurrency) - total = len(sources) - completed_count = 0 - count_lock = asyncio.Lock() - - async def _process_one(source: str) -> FileResult[T]: - nonlocal completed_count - - if on_progress: - on_progress( - ProgressEvent( - source=source, - status=ProgressStatus.STARTED, - completed=completed_count, - total=total, - ) - ) - - async with semaphore: - try: - result = await self.extract( - source, - schema, - extraction_config=extraction_config, - sheet=sheet, - instructions=instructions, - **storage_options, - ) - file_result = FileResult( - source=source, - success=True, - records=list(result), - usage=result.report.usage, - ) - status = ProgressStatus.COMPLETED - error_msg = None - except Exception as e: - logger.warning("Batch extraction failed for %s: %s", source, e) - error_msg = f"{type(e).__name__}: {e}" - file_result = FileResult[T]( - source=source, - success=False, - error=error_msg, - ) - status = ProgressStatus.FAILED - - async with count_lock: - completed_count += 1 - current_completed = completed_count - - if on_progress: - on_progress( - ProgressEvent( - source=source, - status=status, - completed=current_completed, - total=total, - error=error_msg, - ) - ) - - return file_result - - file_results = await asyncio.gather(*[_process_one(s) for s in sources]) - return BatchResult(results=list(file_results)) + return await self._pipeline.extract_batch( + sources, + schema, + extraction_config=extraction_config, + concurrency=concurrency, + sheet=sheet, + instructions=instructions, + on_progress=on_progress, + **storage_options, + ) def extract_batch_sync( self, @@ -850,19 +479,12 @@ def extract_batch_sync( self.extract_batch(sources, schema, **kwargs) ) - # * Private pipeline methods + # * Private delegators — kept on the facade for direct introspection/tests @staticmethod def _get_source_ext(source: str) -> str: """Extract and validate file extension from source path/URL.""" - lower = source.lower().rsplit("?", 1)[0] # ^ Strip query params for URLs - for ext in (".xlsm", ".xltx", ".xltm", ".xlsx", ".xls", ".csv"): - if lower.endswith(ext): - return ext - raise ReaderError( - f"Unsupported file format: {source}", - code=ErrorCode.READER_UNSUPPORTED_FORMAT, - ) + return get_source_ext(source) async def _load_workbook( self, @@ -870,125 +492,8 @@ async def _load_workbook( sheet_name: str | None = None, **storage_options: Any, ) -> WorkbookData: - """Storage → Reader pipeline.""" - merged_options = {**self._config.storage_options, **storage_options} - file_bytes = await read_file(source, **merged_options) - - ext = self._get_source_ext(source) - - if ext == ".csv": - from xlstruct.reader.csv_reader import CsvReader - - csv_reader = CsvReader() - workbook = await asyncio.to_thread(csv_reader.read, file_bytes, sheet_name) - else: - reader = HybridReader() - workbook = await asyncio.to_thread( - reader.read, - file_bytes, - sheet_name, - source_ext=ext, - strict_formulas=self._config.strict_formulas, - evaluate_formulas=self._config.evaluate_formulas, - ) - - workbook.file_name = source.rsplit("/", 1)[-1] - workbook.file_size = len(file_bytes) - return workbook - - async def _run_configured_extraction( - self, - source: str, - config: ExtractionConfig, - **storage_options: Any, - ) -> tuple[list[Any], ExtractionMode]: - """Config-based extraction with mode selection. - - - mode=auto: heuristic routing (≤ SAMPLE_ROWS → direct, > SAMPLE_ROWS → codegen). - - mode=direct: always use LLM direct extraction. - - mode=codegen: always use code generation pipeline. - - Returns: - Tuple of (extracted items, resolved extraction mode). - """ - workbook = await self._load_workbook(source, sheet_name=config.sheet, **storage_options) - full_sheet = workbook.sheets[0] - codegen = self._get_codegen() - - # * Auto-detect header rows if not provided - header_rows = config.header_rows - if header_rows is None: - header_rows = await codegen.detect_header_rows(full_sheet) - - # * Resolve mode - mode = config.mode - if mode == ExtractionMode.AUTO: - max_header_row = max(header_rows) - data_rows = full_sheet.row_count - max_header_row - if data_rows > SAMPLE_ROWS: - mode = ExtractionMode.CODEGEN - else: - mode = ExtractionMode.DIRECT - logger.info( - "Auto-routing: %d data rows → mode=%s", - data_rows, - mode.value, - ) - - if mode == ExtractionMode.CODEGEN: - items = await self._run_codegen(source, full_sheet, header_rows, config, codegen) - return items, ExtractionMode.CODEGEN - - items = await self._run_direct(full_sheet, header_rows, config) - return items, ExtractionMode.DIRECT - - async def _run_codegen( - self, - source: str, - full_sheet: SheetData, - header_rows: list[int], - config: ExtractionConfig, - codegen: CodegenOrchestrator, - ) -> list[Any]: - """Code generation pipeline: cache lookup → generate script → execute → parse.""" - script: GeneratedScript | None = None - signature: str | None = None - - # * Cache lookup - if self._cache is not None: - signature = compute_structure_signature(full_sheet, header_rows, config.output_schema) - script = self._cache.get(signature) - - if script is None: - # * Cache miss — generate via LLM - script = await codegen.generate_script(source, full_sheet, header_rows, config) - self._export_script(source, script) - - # * Store in cache - if self._cache is not None and signature is not None: - self._cache.put(signature, script, full_sheet, header_rows, config.output_schema) - - return await codegen.run_extraction(source, script, config.output_schema) - - async def _run_direct( - self, - full_sheet: SheetData, - header_rows: list[int], - config: ExtractionConfig, - ) -> list[Any]: - """Direct LLM extraction: encode → LLM → Pydantic.""" - encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) - encoded = encoder.encode(full_sheet, header_rows=header_rows) - - return await self._engine.extract( - encoded, - config.output_schema, - config.instructions, - is_sampled=True, - total_rows=full_sheet.row_count, - track_provenance=config.track_provenance, - include_confidence=config.include_confidence, - ) + """Storage → Reader pipeline (delegates to ExtractionPipeline).""" + return await self._pipeline.load_workbook(source, sheet_name=sheet_name, **storage_options) async def _run_sheet_extraction( self, @@ -998,118 +503,5 @@ async def _run_sheet_extraction( *, engine: ExtractionEngine, ) -> list[T]: - """Encoder → (optional Chunking) → ExtractionEngine pipeline.""" - target_engine = engine - encoder = CompressedEncoder() - - if needs_chunking(sheet, self._config.token_budget, self._config.chunking_row_threshold): - # * Chunked extraction - chunks = self._chunk_splitter.split( - sheet, - self._config.token_budget, - min_chunk_rows=self._config.min_chunk_rows, - row_threshold=self._config.chunking_row_threshold, - ) - all_results: list[T] = [] - for chunk in chunks: - encoded = encoder.encode(chunk) - partial = await target_engine.extract(encoded, schema, instructions) - all_results.extend(partial) - return all_results - else: - # * Single-pass extraction - encoded = encoder.encode(sheet) - return await target_engine.extract(encoded, schema, instructions) - - # * Streaming private helpers - - async def _stream_configured_extraction( - self, - source: str, - config: ExtractionConfig, - **storage_options: Any, - ) -> AsyncGenerator[Any, None]: - """Streaming variant of _run_configured_extraction. - - For codegen mode, yields all records at once (script produces all results - in a single run). For direct mode, yields records as each chunk completes. - """ - workbook = await self._load_workbook(source, sheet_name=config.sheet, **storage_options) - full_sheet = workbook.sheets[0] - - # * Auto-detect header rows if not provided (requires codegen orchestrator) - header_rows = config.header_rows - if header_rows is None: - codegen = self._get_codegen() - header_rows = await codegen.detect_header_rows(full_sheet) - - # * Resolve mode - mode = config.mode - if mode == ExtractionMode.AUTO: - max_header_row = max(header_rows) - data_rows = full_sheet.row_count - max_header_row - if data_rows > SAMPLE_ROWS: - mode = ExtractionMode.CODEGEN - else: - mode = ExtractionMode.DIRECT - logger.info( - "Auto-routing (stream): %d data rows → mode=%s", - data_rows, - mode.value, - ) - - if mode == ExtractionMode.CODEGEN: - codegen = self._get_codegen() - items = await self._run_codegen(source, full_sheet, header_rows, config, codegen) - for item in items: - yield item - return - - # * Direct mode — single-pass with sampling (no chunking in config mode) - encoder = CompressedEncoder(sample_size=SAMPLE_ROWS) - encoded = encoder.encode(full_sheet, header_rows=header_rows) - items = await self._engine.extract( - encoded, - config.output_schema, - config.instructions, - is_sampled=True, - total_rows=full_sheet.row_count, - track_provenance=config.track_provenance, - ) - for item in items: - yield item - - async def _stream_sheet_extraction( - self, - sheet: SheetData, - schema: type[T], - instructions: str | None = None, - *, - engine: ExtractionEngine, - ) -> AsyncGenerator[T, None]: - """Streaming variant of _run_sheet_extraction. - - Yields records incrementally as each chunk's LLM call completes. - For single-chunk sheets, yields all records at once. - """ - encoder = CompressedEncoder() - - if needs_chunking(sheet, self._config.token_budget, self._config.chunking_row_threshold): - # * Chunked extraction — yield from each chunk as it completes - chunks = self._chunk_splitter.split( - sheet, - self._config.token_budget, - min_chunk_rows=self._config.min_chunk_rows, - row_threshold=self._config.chunking_row_threshold, - ) - for chunk in chunks: - encoded = encoder.encode(chunk) - partial = await engine.extract(encoded, schema, instructions) - for item in partial: - yield item - else: - # * Single-pass extraction - encoded = encoder.encode(sheet) - items = await engine.extract(encoded, schema, instructions) - for item in items: - yield item + """Encoder → (optional Chunking) → ExtractionEngine pipeline (delegates).""" + return await self._pipeline.run_sheet_extraction(sheet, schema, instructions, engine=engine) diff --git a/src/xlstruct/mcp_server.py b/src/xlstruct/mcp_server.py index 1298324..a816658 100644 --- a/src/xlstruct/mcp_server.py +++ b/src/xlstruct/mcp_server.py @@ -24,7 +24,7 @@ from xlstruct.codegen.cache import ScriptCache from xlstruct.config import ExtractionConfig, ExtractionMode from xlstruct.extractor import ExtractionResult, Extractor -from xlstruct.reader.hybrid_reader import HybridReader +from xlstruct.reader.dispatch import read_workbook from xlstruct.schemas.batch import BatchResult from xlstruct.schemas.core import SheetData from xlstruct.storage import read_file @@ -156,17 +156,7 @@ def _validate_source(source: str) -> None: async def _load_sheet(source: str, sheet: str | None = None) -> SheetData: """Load a single sheet from a file for inspection.""" file_bytes = await read_file(source) - ext = Extractor._get_source_ext(source) # pyright: ignore[reportPrivateUsage] - - if ext == ".csv": - from xlstruct.reader.csv_reader import CsvReader - - csv_reader = CsvReader() - workbook = await asyncio.to_thread(csv_reader.read, file_bytes, sheet) - else: - reader = HybridReader() - workbook = await asyncio.to_thread(reader.read, file_bytes, sheet, source_ext=ext) - + workbook = await asyncio.to_thread(read_workbook, file_bytes, source, sheet) return workbook.sheets[0] diff --git a/src/xlstruct/reader/base.py b/src/xlstruct/reader/base.py index d241267..2100228 100644 --- a/src/xlstruct/reader/base.py +++ b/src/xlstruct/reader/base.py @@ -1,32 +1,45 @@ -"""ExcelReader protocol definition.""" +"""Reader contracts: ReaderOptions and the WorkbookReader dispatch protocol.""" from typing import Protocol +from pydantic import BaseModel + from xlstruct.schemas.core import WorkbookData -class ExcelReader(Protocol): - """Protocol for reading Excel files into WorkbookData. +class ReaderOptions(BaseModel): + """Options shared across reader implementations. - Implementations are sync (openpyxl is sync). - The Extractor layer wraps calls with asyncio.to_thread(). + Each reader uses only the subset it understands: HybridReader reads + ``strict_formulas`` / ``evaluate_formulas``; CsvReader reads ``csv_encoding``. + The unused fields are simply ignored by the other reader. """ - def read( + strict_formulas: bool = True + evaluate_formulas: bool = False + csv_encoding: str = "utf-8" + + +class WorkbookReader(Protocol): + """Contract for READER_REGISTRY entries: bytes + options → WorkbookData. + + Implementations are sync (calamine/openpyxl are sync). The Extractor and + MCP server wrap calls with ``asyncio.to_thread()``. + """ + + def __call__( self, file_bytes: bytes, - sheet_name: str | None = None, - *, - source_ext: str = ".xlsx", + sheet_name: str | None, + source_ext: str, + options: ReaderOptions, ) -> WorkbookData: - """Read an Excel file from bytes. + """Read raw file bytes into WorkbookData. Args: - file_bytes: Raw bytes of the .xlsx file. + file_bytes: Raw bytes of the file. sheet_name: If provided, read only this sheet. None = all sheets. - source_ext: File extension hint for format detection. - - Returns: - WorkbookData with one or more SheetData entries. + source_ext: Resolved file extension (e.g. ".xlsx", ".csv"). + options: Reader options (formula handling, CSV encoding). """ ... diff --git a/src/xlstruct/reader/csv_reader.py b/src/xlstruct/reader/csv_reader.py index bab9dab..fc65b62 100644 --- a/src/xlstruct/reader/csv_reader.py +++ b/src/xlstruct/reader/csv_reader.py @@ -12,6 +12,7 @@ from openpyxl.utils import get_column_letter +from xlstruct.exceptions import ErrorCode, ReaderError from xlstruct.schemas.core import CellData, SheetData, WorkbookData log = logging.getLogger(__name__) @@ -60,7 +61,15 @@ def read( """ # ^ utf-8-sig strips BOM if present; identical to utf-8 otherwise effective_encoding = "utf-8-sig" if encoding == "utf-8" else encoding - text = file_bytes.decode(effective_encoding) + try: + text = file_bytes.decode(effective_encoding) + except (UnicodeDecodeError, LookupError) as e: + # ^ UnicodeDecodeError: wrong encoding; LookupError: unknown encoding name + raise ReaderError( + f"Failed to decode CSV as '{encoding}': {e}. " + "Set the correct encoding via ExtractorConfig(csv_encoding=...).", + code=ErrorCode.READER_PARSE_FAILED, + ) from e # * Dialect auto-detection dialect = self._detect_dialect(text) diff --git a/src/xlstruct/reader/dispatch.py b/src/xlstruct/reader/dispatch.py new file mode 100644 index 0000000..07ad909 --- /dev/null +++ b/src/xlstruct/reader/dispatch.py @@ -0,0 +1,82 @@ +"""Reader dispatch shared by Extractor and the MCP server. + +A single ``READER_REGISTRY`` maps a source extension to a reader adapter, so a +new format is added in one place instead of being branched on at every call +site. Replaces the duplicated ``.csv`` / else if-blocks that previously lived in +both extractor.py and mcp_server.py (the latter via private-member access). +""" + +from xlstruct.exceptions import ErrorCode, ReaderError +from xlstruct.reader.base import ReaderOptions, WorkbookReader +from xlstruct.reader.csv_reader import CsvReader +from xlstruct.reader.hybrid_reader import HybridReader +from xlstruct.schemas.core import WorkbookData + + +def _read_csv( + file_bytes: bytes, + sheet_name: str | None, + source_ext: str, + options: ReaderOptions, +) -> WorkbookData: + return CsvReader().read(file_bytes, sheet_name, encoding=options.csv_encoding) + + +def _read_hybrid( + file_bytes: bytes, + sheet_name: str | None, + source_ext: str, + options: ReaderOptions, +) -> WorkbookData: + return HybridReader().read( + file_bytes, + sheet_name, + source_ext=source_ext, + strict_formulas=options.strict_formulas, + evaluate_formulas=options.evaluate_formulas, + ) + + +READER_REGISTRY: dict[str, WorkbookReader] = { + ".csv": _read_csv, + ".xlsx": _read_hybrid, + ".xlsm": _read_hybrid, + ".xltx": _read_hybrid, + ".xltm": _read_hybrid, + ".xls": _read_hybrid, +} + + +def get_source_ext(source: str) -> str: + """Extract and validate the file extension from a source path/URL.""" + lower = source.lower().rsplit("?", 1)[0] # ^ Strip query params for URLs + # ^ Longest extension first so a future suffix-overlapping format can't shadow + for ext in sorted(READER_REGISTRY, key=len, reverse=True): + if lower.endswith(ext): + return ext + raise ReaderError( + f"Unsupported file format: {source}", + code=ErrorCode.READER_UNSUPPORTED_FORMAT, + ) + + +def read_workbook( + file_bytes: bytes, + source: str, + sheet_name: str | None = None, + *, + options: ReaderOptions | None = None, +) -> WorkbookData: + """Dispatch raw file bytes to the right reader by source extension. + + Synchronous — callers wrap this in ``asyncio.to_thread()``. + + Args: + file_bytes: Raw bytes of the file. + source: File path or URL (used only to resolve the extension). + sheet_name: If provided, read only this sheet. None = all sheets. + options: Reader options. Defaults to ``ReaderOptions()`` when omitted. + """ + ext = get_source_ext(source) + reader = READER_REGISTRY[ext] + return reader(file_bytes, sheet_name, ext, options or ReaderOptions()) diff --git a/tests/test_backend_resolver.py b/tests/test_backend_resolver.py new file mode 100644 index 0000000..3523b3b --- /dev/null +++ b/tests/test_backend_resolver.py @@ -0,0 +1,59 @@ +"""Tests for execution backend resolution (S1): sandbox-by-default, fail-closed.""" + +import pytest + +from xlstruct.codegen.backends import resolver as resolver_mod +from xlstruct.codegen.backends.docker import DockerBackend +from xlstruct.codegen.backends.resolver import resolve_execution_backend +from xlstruct.codegen.backends.subprocess import SubprocessBackend +from xlstruct.config import ExtractorConfig +from xlstruct.exceptions import CodegenSecurityError, ErrorCode + + +class TestResolveExecutionBackend: + def test_explicit_backend_always_wins(self): + """An injected backend is returned verbatim regardless of sandbox mode.""" + explicit = SubprocessBackend(trusted=True) + for mode in ("auto", "docker", "subprocess"): + assert resolve_execution_backend(explicit, mode) is explicit + + def test_subprocess_mode_returns_trusted_subprocess(self): + backend = resolve_execution_backend(None, "subprocess") + assert isinstance(backend, SubprocessBackend) + assert backend._trusted is True # ^ explicit opt-in suppresses the warning + + def test_docker_mode_returns_docker_backend(self): + # ^ DockerBackend construction is lazy; no daemon needed here + backend = resolve_execution_backend(None, "docker") + assert isinstance(backend, DockerBackend) + + def test_auto_uses_docker_when_available(self, monkeypatch): + monkeypatch.setattr(resolver_mod, "docker_available", lambda: True) + backend = resolve_execution_backend(None, "auto") + assert isinstance(backend, DockerBackend) + + def test_auto_fails_closed_without_docker(self, monkeypatch): + """The core S1 guarantee: no silent fallback to the non-isolating subprocess.""" + monkeypatch.setattr(resolver_mod, "docker_available", lambda: False) + with pytest.raises(CodegenSecurityError) as exc_info: + resolve_execution_backend(None, "auto") + assert exc_info.value.code == ErrorCode.CODEGEN_NO_SANDBOX + # ^ Error must point the user at both remedies + msg = str(exc_info.value) + assert "xlstruct[docker]" in msg + assert "subprocess" in msg + + +class TestConfigDefault: + def test_codegen_security_error_is_public(self): + # ^ new public exception must be catchable from the top-level package + import xlstruct + + assert xlstruct.CodegenSecurityError is CodegenSecurityError + + def test_codegen_sandbox_defaults_to_auto(self): + assert ExtractorConfig().codegen_sandbox == "auto" + + def test_codegen_sandbox_rejects_unknown_value(self): + with pytest.raises(ValueError): + ExtractorConfig(codegen_sandbox="yolo") # type: ignore[arg-type] diff --git a/tests/test_cache.py b/tests/test_cache.py index 95bbd6a..ead57e8 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -1,6 +1,9 @@ """Tests for codegen script caching.""" import json +import os +import stat +import sys import pytest from pydantic import BaseModel @@ -8,6 +11,7 @@ from xlstruct.codegen.cache import ( CacheMetadata, ScriptCache, + _compute_mac, compute_structure_signature, ) from xlstruct.schemas.codegen import GeneratedScript @@ -72,9 +76,9 @@ def test_deterministic(self, sample_sheet): assert sig1 == sig2 def test_length(self, sample_sheet): - """Signature is 16-char hex string.""" + """Signature is the full 64-char (256-bit) SHA-256 hex digest.""" sig = compute_structure_signature(sample_sheet, [1], Invoice) - assert len(sig) == 16 + assert len(sig) == 64 assert all(c in "0123456789abcdef" for c in sig) def test_different_schema_different_signature(self, sample_sheet): @@ -229,3 +233,107 @@ def test_missing_script_file_returns_none(self, cache, sample_sheet, sample_scri result = cache.get(sig) assert result is None + + +# * Integrity protection (S3) + + +class TestCacheIntegrity: + def test_valid_entry_round_trips(self, cache, sample_sheet, sample_script): + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + got = cache.get(sig) + assert got is not None + assert got.code == sample_script.code + + def test_tampered_script_is_refused(self, cache, sample_sheet, sample_script): + """A modified cached script must NOT be returned (never executed).""" + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + # ^ Inject code at the predicted path (harmless probe; refused before any run) + script_path = cache.cache_dir / f"{sig}.py" + script_path.write_text("import os\nos.system('exit 0') # injected", encoding="utf-8") + + assert cache.get(sig) is None + + def test_legacy_entry_without_mac_is_refused(self, cache, sample_sheet, sample_script): + """An entry with no MAC (pre-integrity cache) is refused and regenerated.""" + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + meta_path = cache.cache_dir / f"{sig}.json" + meta = json.loads(meta_path.read_text()) + meta["mac"] = "" + meta_path.write_text(json.dumps(meta)) + + assert cache.get(sig) is None + + def test_mac_from_different_secret_is_refused(self, cache, sample_sheet, sample_script): + """A MAC forged under a different key must not verify.""" + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + meta_path = cache.cache_dir / f"{sig}.json" + meta = json.loads(meta_path.read_text()) + meta["mac"] = _compute_mac(b"attacker-key", sig, sample_script.code) + meta_path.write_text(json.dumps(meta)) + + assert cache.get(sig) is None + + def test_mac_binds_signature_blocking_cross_entry_reuse( + self, cache, sample_sheet, sample_script + ): + """A validly-signed entry cannot be relocated to a different signature path.""" + sig_a = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig_a, sample_script, sample_sheet, [1], Invoice) + + sig_b = "b" * 64 + # ^ Copy A's code + A's (valid) mac into B's paths + (cache.cache_dir / f"{sig_b}.py").write_text(sample_script.code, encoding="utf-8") + meta = json.loads((cache.cache_dir / f"{sig_a}.json").read_text()) + (cache.cache_dir / f"{sig_b}.json").write_text(json.dumps(meta)) + + # ^ MAC was computed over sig_a, so it fails to verify under sig_b + assert cache.get(sig_b) is None + + def test_non_ascii_mac_is_refused_not_crash(self, cache, sample_sheet, sample_script): + """A hostile non-ASCII mac must degrade to skip+regenerate, not crash extraction.""" + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + meta_path = cache.cache_dir / f"{sig}.json" + meta = json.loads(meta_path.read_text()) + meta["mac"] = "café" * 16 # ^ non-ASCII would raise in hmac.compare_digest + meta_path.write_text(json.dumps(meta)) + + assert cache.get(sig) is None # ^ returns cleanly (no TypeError) + + def test_malformed_mac_is_refused(self, cache, sample_sheet, sample_script): + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + meta_path = cache.cache_dir / f"{sig}.json" + meta = json.loads(meta_path.read_text()) + meta["mac"] = "deadbeef" # ^ valid hex but wrong length + meta_path.write_text(json.dumps(meta)) + + assert cache.get(sig) is None + + def test_clear_preserves_secret(self, cache, sample_sheet, sample_script): + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + secret_path = cache.cache_dir / ".hmac_key" + assert secret_path.exists() + cache.clear() + assert secret_path.exists() + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX permission bits") + def test_private_permissions(self, cache, sample_sheet, sample_script): + sig = compute_structure_signature(sample_sheet, [1], Invoice) + cache.put(sig, sample_script, sample_sheet, [1], Invoice) + + assert stat.S_IMODE(os.stat(cache.cache_dir).st_mode) == 0o700 + assert stat.S_IMODE(os.stat(cache.cache_dir / f"{sig}.py").st_mode) == 0o600 + assert stat.S_IMODE(os.stat(cache.cache_dir / f"{sig}.json").st_mode) == 0o600 + assert stat.S_IMODE(os.stat(cache.cache_dir / ".hmac_key").st_mode) == 0o600 diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 13ea7e7..61070ed 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -1,6 +1,12 @@ """Tests for extraction/chunking.py: needs_chunking() and ChunkSplitter.""" -from xlstruct.extraction.chunking import ChunkSplitter, needs_chunking +from xlstruct._tokens import estimate_row_token_costs, estimate_sheet_tokens +from xlstruct.extraction.chunking import ( + _CHUNKING_ROW_THRESHOLD, + _MIN_CHUNK_ROWS, + ChunkSplitter, + needs_chunking, +) from xlstruct.schemas.core import CellData, SheetData # * Fixtures @@ -212,3 +218,167 @@ def test_split_custom_thresholds_headers_in_every_chunk(self): for chunk in chunks: header_cells = [c for c in chunk.cells if c.row == 1] assert len(header_cells) > 0 + + +# * Token-based chunking (header-aware budgeting) + + +def _make_uniform_sheet(cols: int, data_rows: int) -> SheetData: + """Header + data rows where every data cell is a constant-width value. + + Flat per-row density makes the chunk count predictable, so this fixture is + used to check that the row_threshold accuracy cap still bounds chunk size + even when the token budget is loose. + """ + cells = [CellData(row=1, col=c, value=f"Col{c}", data_type="s") for c in range(1, cols + 1)] + for r in range(2, data_rows + 2): + for c in range(1, cols + 1): + cells.append(CellData(row=r, col=c, value="val", data_type="s")) + return SheetData( + name="Uniform", + dimensions=f"A1:Z{data_rows + 1}", + cells=cells, + merged_ranges=[], + row_count=data_rows + 1, + col_count=cols, + ) + + +def _make_header_heavy_sheet(cols: int, data_rows: int) -> SheetData: + """Long-label header row over compact numeric data — the header is a large + fraction of a small token budget, so reserving it per chunk matters.""" + label = "Very Long Descriptive Column Header Label Number {c} With Extra Words" + cells = [ + CellData(row=1, col=c, value=label.format(c=c), data_type="s") for c in range(1, cols + 1) + ] + for r in range(2, data_rows + 2): + for c in range(1, cols + 1): + cells.append(CellData(row=r, col=c, value=(r * c) % 9, data_type="n")) + return SheetData( + name="HeaderHeavy", + dimensions=f"A1:Z{data_rows + 1}", + cells=cells, + merged_ranges=[], + row_count=data_rows + 1, + col_count=cols, + ) + + +def _make_heavy_tail_sheet(cols: int, light_rows: int, heavy_rows: int) -> SheetData: + """Light rows followed by much heavier rows. + + A proportional (equal-row-count) split groups the heavy tail into one + over-budget chunk; greedy packing by real per-row cost keeps every chunk + within budget. This is the case the old sampled-total formula got wrong. + """ + heavy = "lorem ipsum dolor sit amet consectetur adipiscing elit sed" + cells = [CellData(row=1, col=c, value=f"Col{c}", data_type="s") for c in range(1, cols + 1)] + row = 2 + for _ in range(light_rows): + for c in range(1, cols + 1): + cells.append(CellData(row=row, col=c, value="x", data_type="s")) + row += 1 + for _ in range(heavy_rows): + for c in range(1, cols + 1): + cells.append(CellData(row=row, col=c, value=heavy, data_type="s")) + row += 1 + return SheetData( + name="HeavyTail", + dimensions=f"A1:Z{row - 1}", + cells=cells, + merged_ranges=[], + row_count=row - 1, + col_count=cols, + ) + + +def _exact_chunk_cost(chunk: SheetData, header_row_count: int = 1) -> int: + """The splitter's own cost model for a chunk: header cost + sum of row costs.""" + header_cells = [c for c in chunk.cells if c.row <= header_row_count] + rows: dict[int, list[CellData]] = {} + for c in chunk.cells: + if c.row > header_row_count: + rows.setdefault(c.row, []).append(c) + groups = [header_cells, *(rows[r] for r in sorted(rows))] + return sum(estimate_row_token_costs(groups)) + + +def _naive_proportional_max_cost(sheet: SheetData, budget: int, header_row_count: int = 1) -> int: + """Worst chunk cost under the pre-fix proportional split (equal row counts, + sampled total, no header reservation). Shows the old approach overflowed.""" + header_cells = [c for c in sheet.cells if c.row <= header_row_count] + data_rows: dict[int, list[CellData]] = {} + for c in sheet.cells: + if c.row > header_row_count: + data_rows.setdefault(c.row, []).append(c) + srn = sorted(data_rows) + total = estimate_sheet_tokens(sheet) + if total > budget: + chunk_count = max(1, total // budget) + rows_per_chunk = max(_MIN_CHUNK_ROWS, len(srn) // chunk_count) + else: + rows_per_chunk = max(_MIN_CHUNK_ROWS, _CHUNKING_ROW_THRESHOLD) + header_cost = estimate_row_token_costs([header_cells])[0] + worst = 0 + for i in range(0, len(srn), rows_per_chunk): + group = srn[i : i + rows_per_chunk] + cost = header_cost + sum(estimate_row_token_costs([data_rows[r] for r in group])) + worst = max(worst, cost) + return worst + + +class TestTokenBasedChunking: + """The token-based path sizes chunks by each row's real token cost (greedy + bin-packing, header reserved) so every chunk fits token_budget regardless of + where the heavy rows sit — the case the old proportional split got wrong.""" + + def test_header_aware_chunks_fit_budget(self): + sheet = _make_header_heavy_sheet(cols=20, data_rows=200) + budget = 2_000 + assert estimate_sheet_tokens(sheet) > budget # ^ token path is exercised + + chunks = ChunkSplitter().split(sheet, token_budget=budget) + assert len(chunks) > 1 + for chunk in chunks: + assert _exact_chunk_cost(chunk) <= budget, f"{chunk.name} exceeds budget" + # ^ The pre-fix proportional split (no header reservation) overflowed here + assert _naive_proportional_max_cost(sheet, budget) > budget + + def test_variable_row_sizes_each_chunk_fits(self): + # ^ Heavy rows concentrated in the tail: a proportional split overflows + # ^ the tail chunk; greedy packing by real cost keeps every chunk in budget. + sheet = _make_heavy_tail_sheet(cols=8, light_rows=300, heavy_rows=100) + budget = 2_000 + + chunks = ChunkSplitter().split(sheet, token_budget=budget) + assert len(chunks) > 1 + for chunk in chunks: + assert _exact_chunk_cost(chunk) <= budget, f"{chunk.name} exceeds budget" + # ^ Prove the fix matters: the old proportional split overflows here + assert _naive_proportional_max_cost(sheet, budget) > budget + + def test_row_threshold_caps_chunk_row_count(self): + # ^ Even with a loose budget, the accuracy cap bounds each chunk's rows + sheet = _make_uniform_sheet(cols=10, data_rows=500) + chunks = ChunkSplitter().split(sheet, token_budget=3_000) + assert len(chunks) > 1 + for chunk in chunks: + data_row_nums = {c.row for c in chunk.cells if c.row != 1} + assert len(data_row_nums) <= _CHUNKING_ROW_THRESHOLD + + def test_token_branch_covers_all_data_rows(self): + # ^ No data loss: every data row appears in exactly one chunk + sheet = _make_heavy_tail_sheet(cols=8, light_rows=300, heavy_rows=100) + chunks = ChunkSplitter().split(sheet, token_budget=2_000) + all_rows: list[int] = [] + for chunk in chunks: + rows = {c.row for c in chunk.cells if c.row != 1} + all_rows.extend(rows) + # ^ rows 2..401 each appear exactly once (sorted equality rules out dupes) + assert sorted(all_rows) == list(range(2, 402)) + + def test_header_in_every_chunk_token_branch(self): + sheet = _make_heavy_tail_sheet(cols=8, light_rows=300, heavy_rows=100) + chunks = ChunkSplitter().split(sheet, token_budget=2_000) + for chunk in chunks: + assert any(c.row == 1 for c in chunk.cells), f"{chunk.name} missing header" diff --git a/tests/test_codegen.py b/tests/test_codegen.py index 4b8be79..c454cfb 100644 --- a/tests/test_codegen.py +++ b/tests/test_codegen.py @@ -4,6 +4,7 @@ import logging from unittest.mock import AsyncMock +import pytest from pydantic import BaseModel from xlstruct.codegen.executor import ( @@ -12,8 +13,10 @@ _build_safe_env, scan_blocked_imports, ) +from xlstruct.codegen.orchestrator import CodegenOrchestrator from xlstruct.codegen.schema_utils import get_schema_source from xlstruct.codegen.validation import ScriptValidator +from xlstruct.exceptions import ErrorCode, ExtractionError from xlstruct.schemas.codegen import ColumnMapping, MappingPlan # * Test schemas @@ -221,6 +224,31 @@ def test_getattr_non_dunder_string_clean(self): assert scan_blocked_imports(code) == [] +# * Characterization: the AST scan is a best-effort filter, NOT a boundary (S1) + + +class TestScanIsNotABoundary: + """Pin the known scanner bypasses that justify sandbox-by-default. + + scan_blocked_imports only parses the AST — it never executes — so these + probes are harmless. They demonstrate why untrusted output needs Docker + isolation, not just the static scan. If the scanner is ever hardened to + catch one of these, update the corresponding assertion. + """ + + def test_alias_import_evades_sys_modules_pattern(self): + # ^ chain is 's.modules', not 'sys.modules', so the pattern misses it + assert scan_blocked_imports("import sys as s\nx = s.modules") == [] + + def test_globals_subscript_is_not_flagged(self): + # ^ globals() is not a blocked builtin; subscript access is not analyzed + assert scan_blocked_imports('g = globals()\nx = g["__bui" + "ltins__"]') == [] + + def test_getattr_with_concatenated_string_evades_detection(self): + # ^ 2nd arg is a BinOp, not an ast.Constant, so it slips past the check + assert scan_blocked_imports('cls = getattr((), "__cla" + "ss__")') == [] + + # * _build_safe_env @@ -447,7 +475,8 @@ def test_nested_model_includes_both_in_dependency_order(self): class TestScriptValidatorValidate: async def test_disallowed_imports_returns_failure_without_execution(self): - validator = ScriptValidator(timeout=10) + # ^ backend is required now; the security scan rejects before it is ever called + validator = ScriptValidator(timeout=10, backend=AsyncMock()) code = "import socket\nprint(socket.gethostname())" result = await validator.validate(code, source_path="/fake/path.xlsx") assert result.success is False @@ -629,3 +658,56 @@ def test_multiple_duplicates_all_deduplicated(self): "value", "note", ] + + +# * C2 — codegen records that fail validation are surfaced, not silently dropped + + +class TestValidateOutputAllRecords: + def test_validates_beyond_first_five_records(self): + # ^ 5 valid rows then a 6th with a type error — must NOT pass as "success" + data = [{"name": f"i{n}", "value": n} for n in range(5)] + data.append({"name": "bad", "value": "NOT_AN_INT"}) + result = ScriptValidator._validate_output(json.dumps(data), SampleRecord) + assert "VALIDATION ERROR" in result + assert "Record 5" in result + + def test_reports_total_failure_count_with_capped_detail(self): + data = [{"name": "x", "value": f"bad{n}"} for n in range(8)] + result = ScriptValidator._validate_output(json.dumps(data), SampleRecord) + assert "8/8 records" in result + assert "more failing records not shown" in result + + def test_all_valid_returns_empty(self): + data = [{"name": f"i{n}", "value": n} for n in range(20)] + result = ScriptValidator._validate_output(json.dumps(data), SampleRecord) + assert result == "" + + +class TestParseScriptOutput: + def test_valid_records_parsed(self): + stdout = json.dumps([{"name": "a", "value": 1}, {"name": "b", "value": 2}]) + result = CodegenOrchestrator._parse_script_output(stdout, SampleRecord) + assert len(result) == 2 + + def test_invalid_record_raises_instead_of_dropping(self): + # ^ A record that fails validation must fail loudly, not vanish silently + stdout = json.dumps([{"name": "a", "value": 1}, {"name": "b", "value": "x"}]) + with pytest.raises(ExtractionError) as exc: + CodegenOrchestrator._parse_script_output(stdout, SampleRecord) + assert exc.value.code == ErrorCode.EXTRACTION_SCHEMA_VALIDATION_FAILED + + def test_provenance_source_row_stripped(self): + stdout = json.dumps([{"name": "a", "value": 1, "_source_row": 5}]) + result = CodegenOrchestrator._parse_script_output(stdout, SampleRecord) + assert len(result) == 1 + assert getattr(result[0], "_source_rows") == [5] + + +class TestFilterLogsWarning: + def test_drop_emits_warning(self, caplog): + data = [{"name": "ok", "value": 1}, {"name": None, "value": 2}] + with caplog.at_level(logging.WARNING, logger="xlstruct.codegen.validation"): + ScriptValidator._filter_by_required_fields(json.dumps(data), SampleRecord) + assert any(rec.levelno == logging.WARNING for rec in caplog.records) + assert "dropped" in caplog.text.lower() diff --git a/tests/test_config.py b/tests/test_config.py index 2d81ba9..ba6d963 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,7 +3,7 @@ import pytest from pydantic import ValidationError -from xlstruct.config import ExtractorConfig, get_provider_kwargs +from xlstruct.config import ExtractorConfig, apply_cache_control, get_provider_kwargs class TestExtractorConfig: @@ -73,3 +73,28 @@ def test_provider_options_override(self): ) kwargs = get_provider_kwargs(config) assert kwargs["max_tokens"] == 4096 + + +class TestApplyCacheControl: + def test_marks_system_only_for_anthropic(self): + messages = [ + {"role": "system", "content": "STATIC SYSTEM PROMPT"}, + {"role": "user", "content": "VARIABLE SHEET DATA"}, + ] + result = apply_cache_control(messages, "anthropic/claude-sonnet-4-6") + + # ^ system prompt is wrapped with a cache_control marker + sys_content = result[0]["content"] + assert isinstance(sys_content, list) + assert sys_content[0]["cache_control"] == {"type": "ephemeral"} + + # ^ variable user message is left as a plain string (no cache breakpoint) + assert result[1]["content"] == "VARIABLE SHEET DATA" + + def test_non_anthropic_returns_messages_unchanged(self): + messages = [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "DATA"}, + ] + result = apply_cache_control(messages, "openai/gpt-4o") + assert result == messages diff --git a/tests/test_csv_reader.py b/tests/test_csv_reader.py index ed29013..b67fa51 100644 --- a/tests/test_csv_reader.py +++ b/tests/test_csv_reader.py @@ -3,6 +3,9 @@ import csv from unittest.mock import patch +import pytest + +from xlstruct.exceptions import ErrorCode, ReaderError from xlstruct.reader.csv_reader import CsvReader reader = CsvReader() @@ -178,3 +181,28 @@ def test_non_date_string_remains_string() -> None: assert values[(2, 2)] == "2024-hello" assert types[(2, 2)] == "s" + + +# * C3 — non-UTF8 decode failure is wrapped as ReaderError; encoding is honored + + +class TestCsvEncoding: + EUCKR = "이름,값\n사과,10\n".encode("euc-kr") + + def test_non_utf8_raises_reader_error(self): + # ^ euc-kr bytes are invalid utf-8 — must surface as ReaderError, not raw UnicodeDecodeError + with pytest.raises(ReaderError) as exc: + reader.read(self.EUCKR) + assert exc.value.code == ErrorCode.READER_PARSE_FAILED + + def test_explicit_encoding_parses(self): + wb = reader.read(self.EUCKR, encoding="euc-kr") + values = [c.value for c in wb.sheets[0].cells] + assert "이름" in values + assert "사과" in values + + def test_unknown_encoding_name_raises_reader_error(self): + # ^ bogus encoding name raises LookupError — must also surface as ReaderError + with pytest.raises(ReaderError) as exc: + reader.read(b"a,b\n1,2\n", encoding="nonsense") + assert exc.value.code == ErrorCode.READER_PARSE_FAILED diff --git a/tests/test_docker_backend.py b/tests/test_docker_backend.py new file mode 100644 index 0000000..7a00dcf --- /dev/null +++ b/tests/test_docker_backend.py @@ -0,0 +1,148 @@ +"""Tests for the hardened DockerBackend (S2). + +Config-shape tests always run (no daemon needed). Container probes run only when +Docker is available and use harmless probes: read the uid, attempt a write to the +read-only rootfs, and attempt a DNS-port connection that isolation should block. +""" + +import asyncio +import json + +import pytest + +from xlstruct.codegen.backends.docker import DockerBackend, DockerConfig + + +def _docker_available() -> bool: + try: + import aiodocker + except ImportError: + return False + + async def _ping() -> bool: + docker = aiodocker.Docker() + try: + await docker.version() + return True + finally: + await docker.close() + + try: + return asyncio.run(_ping()) + except Exception: + return False + + +# * Always-on: the hardening must be present in the execution container config + + +class TestDockerExecConfigHardening: + def test_exec_host_config_is_hardened(self): + hc = DockerBackend()._exec_host_config().model_dump(exclude_none=True) + assert hc["ReadonlyRootfs"] is True + assert hc["CapDrop"] == ["ALL"] + assert "/tmp" in hc["Tmpfs"] + assert "no-new-privileges" in hc["SecurityOpt"] + # ^ existing controls must be preserved + assert hc["PidsLimit"] == 64 + assert hc["Memory"] > 0 + assert hc["MemorySwap"] == hc["Memory"] # swap disabled + # ^ no runtime override by default + assert "Runtime" not in hc + + def test_exec_container_runs_as_nonroot_with_no_bytecode_writes(self): + from xlstruct.codegen.backends.docker import _ContainerConfig, _HostConfig + + cfg = _ContainerConfig( + Image="x", + Cmd=["python"], + User=DockerBackend()._cfg.user, + Env=["PYTHONDONTWRITEBYTECODE=1"], + HostConfig=_HostConfig(Memory=1, MemorySwap=1, CpuQuota=1), + ).model_dump(exclude_none=True) + assert cfg["User"] == "65534:65534" + assert "PYTHONDONTWRITEBYTECODE=1" in cfg["Env"] + + def test_install_host_config_is_permissive(self): + """The one-time package install must NOT be hardened (needs root + writes).""" + hc = DockerBackend()._install_host_config().model_dump(exclude_none=True) + assert hc["ReadonlyRootfs"] is False + assert hc["CapDrop"] == [] + assert hc["Tmpfs"] == {} + + def test_gvisor_runtime_surfaces(self): + hc = ( + DockerBackend(DockerConfig(runtime="runsc")) + ._exec_host_config() + .model_dump(exclude_none=True) + ) + assert hc["Runtime"] == "runsc" + + def test_custom_seccomp_profile_surfaces(self, tmp_path): + profile = tmp_path / "seccomp.json" + profile.write_text(json.dumps({"defaultAction": "SCMP_ACT_ALLOW"})) + hc = ( + DockerBackend(DockerConfig(seccomp_profile=profile)) + ._exec_host_config() + .model_dump(exclude_none=True) + ) + assert any(opt.startswith("seccomp=") for opt in hc["SecurityOpt"]) + + def test_default_seccomp_is_not_disabled(self): + """With no custom profile, Docker's built-in default profile stays active.""" + hc = DockerBackend()._exec_host_config().model_dump(exclude_none=True) + assert not any("seccomp=unconfined" in opt for opt in hc["SecurityOpt"]) + + +# * Gated: real container isolation probes + + +@pytest.fixture +def dummy_source(tmp_path): + p = tmp_path / "source.xlsx" + p.write_bytes(b"PK\x03\x04") # ^ minimal bytes; probes do not read it + return str(p) + + +@pytest.mark.skipif(not _docker_available(), reason="Docker daemon not available") +class TestDockerContainerProbes: + async def test_runs_as_nonroot(self, dummy_source): + backend = DockerBackend() + code = "import os\nprint(os.getuid())" + exit_code, stdout, stderr = await backend.execute(code, dummy_source, timeout=60) + assert exit_code == 0, stderr + assert stdout.strip().endswith("65534") + + async def test_root_filesystem_is_read_only(self, dummy_source): + backend = DockerBackend() + code = ( + "try:\n" + " open('/probe_root.txt', 'w').write('x')\n" + " print('WROTE')\n" + "except OSError:\n" + " print('READONLY')\n" + ) + exit_code, stdout, stderr = await backend.execute(code, dummy_source, timeout=60) + assert exit_code == 0, stderr + assert "READONLY" in stdout + + async def test_tmp_is_writable(self, dummy_source): + backend = DockerBackend() + code = "open('/tmp/probe.txt', 'w').write('ok')\nprint('TMP_OK')" + exit_code, stdout, stderr = await backend.execute(code, dummy_source, timeout=60) + assert exit_code == 0, stderr + assert "TMP_OK" in stdout + + async def test_network_is_blocked(self, dummy_source): + backend = DockerBackend() + code = ( + "import socket\n" + "try:\n" + " socket.create_connection(('8.8.8.8', 53), timeout=3)\n" + " print('NETWORK_OK')\n" + "except OSError:\n" + " print('NETWORK_BLOCKED')\n" + ) + exit_code, stdout, stderr = await backend.execute(code, dummy_source, timeout=60) + assert exit_code == 0, stderr + assert "NETWORK_BLOCKED" in stdout diff --git a/tests/test_extractor.py b/tests/test_extractor.py index 3ea266e..c9f7997 100644 --- a/tests/test_extractor.py +++ b/tests/test_extractor.py @@ -1,5 +1,6 @@ """Integration tests for Extractor class with mocked LLM extraction.""" +import asyncio import io from unittest.mock import AsyncMock, MagicMock, patch @@ -7,8 +8,10 @@ import pytest from pydantic import BaseModel -from xlstruct.config import ExtractionMode, ExtractorConfig +from xlstruct.config import ExtractionConfig, ExtractionMode, ExtractorConfig +from xlstruct.exceptions import ErrorCode, ReaderError from xlstruct.extractor import ExtractionResult, Extractor +from xlstruct.schemas.core import SheetData, WorkbookData from xlstruct.schemas.report import ExtractionReport from xlstruct.schemas.usage import TokenUsage @@ -295,3 +298,138 @@ async def test_storage_options_merged(self, product_xlsx_file): workbook = await extractor._load_workbook(product_xlsx_file) assert len(workbook.sheets) == 1 + + +# * P1 — chunked single-sheet extraction runs concurrently (bounded), order preserved + + +@pytest.fixture +def large_xlsx_file(tmp_path) -> str: + """Xlsx with enough rows (> 100) to trigger chunking.""" + wb = openpyxl.Workbook() + ws = wb.active + ws.title = "Products" + ws["A1"], ws["B1"], ws["C1"] = "Name", "Price", "Stock" + for i in range(2, 153): # ^ 151 data rows + ws[f"A{i}"], ws[f"B{i}"], ws[f"C{i}"] = f"Product{i}", float(i), i * 10 + buf = io.BytesIO() + wb.save(buf) + path = tmp_path / "large.xlsx" + path.write_bytes(buf.getvalue()) + return str(path) + + +class TestChunkedConcurrency: + async def test_chunks_run_concurrently_and_bounded(self, large_xlsx_file): + # ^ concurrency=2 + multiple chunks → peak in-flight is exactly 2 + extractor = Extractor(max_concurrent_chunks=2) + active = 0 + peak = 0 + + async def mock_extract(encoded, schema, instructions=None): + nonlocal active, peak + active += 1 + peak = max(peak, active) + await asyncio.sleep(0.02) # ^ hold the slot so peers can enter + active -= 1 + return [Product(name="x", price=1.0, stock=1)] + + with patch.object(extractor._engine, "extract", side_effect=mock_extract): + await extractor.extract(large_xlsx_file, Product) + + assert peak == 2 # > 1 proves parallel; == 2 proves the semaphore bound holds + + async def test_chunked_order_preserved_despite_reverse_completion(self): + extractor = Extractor(max_concurrent_chunks=5) + chunks = [SheetData(name=f"c{i}", row_count=1) for i in range(4)] + + async def mock_extract(encoded, schema, instructions=None): + idx = int(encoded) + # ^ later chunks finish first; gather must still restore input order + await asyncio.sleep(0.01 * (4 - idx)) + return [Product(name=f"Chunk{idx}", price=float(idx), stock=idx)] + + with ( + patch("xlstruct.extraction.pipeline.needs_chunking", return_value=True), + patch.object(extractor._chunk_splitter, "split", return_value=chunks), + patch( + "xlstruct.extraction.pipeline.CompressedEncoder.encode", + side_effect=lambda chunk: chunk.name[1:], + ), + patch.object(extractor._engine, "extract", side_effect=mock_extract), + ): + results = await extractor._run_sheet_extraction( + chunks[0], Product, engine=extractor._engine + ) + + assert [r.name for r in results] == ["Chunk0", "Chunk1", "Chunk2", "Chunk3"] + + +# * B3 — empty sheet fast-fails before the header-detection LLM call + + +class TestEmptySheetFastFail: + @staticmethod + def _empty_workbook() -> WorkbookData: + return WorkbookData(sheets=[SheetData(name="Empty", row_count=0, col_count=0)]) + + async def test_empty_sheet_raises_reader_error(self): + extractor = Extractor() + config = ExtractionConfig(output_schema=Product, mode=ExtractionMode.DIRECT) + with patch.object( + extractor._pipeline, + "load_workbook", + new_callable=AsyncMock, + return_value=self._empty_workbook(), + ): + with pytest.raises(ReaderError) as exc: + await extractor.extract("x.xlsx", extraction_config=config) + assert exc.value.code == ErrorCode.READER_PARSE_FAILED + + async def test_empty_sheet_raises_in_streaming(self): + extractor = Extractor() + config = ExtractionConfig(output_schema=Product, mode=ExtractionMode.DIRECT) + with patch.object( + extractor._pipeline, + "load_workbook", + new_callable=AsyncMock, + return_value=self._empty_workbook(), + ): + with pytest.raises(ReaderError): + async for _ in extractor.stream("x.xlsx", extraction_config=config): + pass + + async def test_empty_sheet_raises_in_legacy_schema_path(self): + extractor = Extractor() + with patch.object( + extractor._pipeline, + "load_workbook", + new_callable=AsyncMock, + return_value=self._empty_workbook(), + ): + with pytest.raises(ReaderError) as exc: + await extractor.extract("x.xlsx", Product) + assert exc.value.code == ErrorCode.READER_PARSE_FAILED + + +# * C3 — csv_encoding is threaded from config through _load_workbook to the CSV reader + + +class TestCsvEncodingThreading: + @staticmethod + def _euckr_csv(tmp_path) -> str: + path = tmp_path / "k.csv" + path.write_bytes("이름,값\n사과,10\n".encode("euc-kr")) + return str(path) + + async def test_default_utf8_raises_reader_error_on_euckr(self, tmp_path): + extractor = Extractor() + with pytest.raises(ReaderError) as exc: + await extractor._load_workbook(self._euckr_csv(tmp_path)) + assert exc.value.code == ErrorCode.READER_PARSE_FAILED + + async def test_config_encoding_parses_euckr(self, tmp_path): + extractor = Extractor(csv_encoding="euc-kr") + workbook = await extractor._load_workbook(self._euckr_csv(tmp_path)) + values = [c.value for c in workbook.sheets[0].cells] + assert "이름" in values diff --git a/tests/test_formatting.py b/tests/test_formatting.py index b3619c2..b45df19 100644 --- a/tests/test_formatting.py +++ b/tests/test_formatting.py @@ -3,7 +3,7 @@ from xlstruct.encoder._formatting import ( build_column_headers, detect_header_row, - find_empty_rows, + find_non_empty_rows, format_cell_value, format_merged_regions, summarize_column_types, @@ -75,11 +75,22 @@ def test_with_value(self, merged_sheet: SheetData): assert "Invoice #2024-001" in regions[0] -class TestFindEmptyRows: - def test_no_empty_rows(self, simple_sheet: SheetData): - empty = find_empty_rows(simple_sheet) +class TestFindNonEmptyRows: + def test_all_rows_non_empty(self, simple_sheet: SheetData): + non_empty = find_non_empty_rows(simple_sheet) # ^ All rows 1-6 have data - assert len(empty) == 0 + assert non_empty == {1, 2, 3, 4, 5, 6} + + def test_skips_empty_rows(self): + # ^ Rows 1 and 3 hold data; row 2 is a gap that must be excluded + cells = [ + CellData(row=1, col=1, value="a"), + CellData(row=3, col=1, value="b"), + ] + sheet = SheetData(name="gap", row_count=3, col_count=1, cells=cells) + non_empty = find_non_empty_rows(sheet) + assert non_empty == {1, 3} + assert 2 not in non_empty class TestSummarizeColumnTypes: diff --git a/tests/test_stream.py b/tests/test_stream.py index cf5370b..2714de1 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -211,8 +211,10 @@ async def test_stream_with_codegen_mode(self, product_xlsx_file): extractor = Extractor() with ( - patch.object(extractor, "_get_codegen", return_value=MagicMock()), - patch.object(extractor, "_run_codegen", new_callable=AsyncMock) as mock_codegen, + patch.object(extractor._pipeline, "_get_codegen", return_value=MagicMock()), + patch.object( + extractor._pipeline, "_run_codegen", new_callable=AsyncMock + ) as mock_codegen, ): mock_codegen.return_value = expected results = [ diff --git a/tests/test_subprocess_backend.py b/tests/test_subprocess_backend.py new file mode 100644 index 0000000..5d35357 --- /dev/null +++ b/tests/test_subprocess_backend.py @@ -0,0 +1,110 @@ +"""Tests for the hardened (trusted/dev-only) SubprocessBackend (S1). + +All probes are harmless: they read an env var the test sets, check which builtins +are present, print argv, or sleep. None perform destructive, credential-reading, or +exfiltrating actions. +""" + +import logging +import sys + +import pytest + +from xlstruct.codegen.backends import subprocess as subprocess_mod +from xlstruct.codegen.backends.subprocess import SubprocessBackend + +# ^ preexec_fn / killpg / setrlimit are POSIX-only; the hardening behaves differently +# on Windows (warns, no limits), so these execution tests target POSIX. +pytestmark = pytest.mark.skipif( + sys.platform == "win32", reason="subprocess hardening probes are POSIX-only" +) + + +@pytest.fixture +def source(tmp_path): + """A throwaway source path passed to scripts as argv[1].""" + p = tmp_path / "source.xlsx" + p.write_bytes(b"") + return str(p) + + +class TestSubprocessHardening: + async def test_legitimate_script_runs(self, source): + backend = SubprocessBackend(trusted=True) + code = "import openpyxl\nprint(len([1, 2, 3]))" + exit_code, stdout, stderr = await backend.execute(code, source, timeout=10) + assert exit_code == 0, stderr + assert "3" in stdout + + async def test_credential_env_is_stripped(self, source, monkeypatch): + """A secret in the parent env must not reach the child (whitelist).""" + monkeypatch.setenv("XLSTRUCT_PROBE_SECRET", "topsecret") + backend = SubprocessBackend(trusted=True) + code = "import os\nprint(os.environ.get('XLSTRUCT_PROBE_SECRET', 'ABSENT'))" + exit_code, stdout, stderr = await backend.execute(code, source, timeout=10) + assert exit_code == 0, stderr + assert "ABSENT" in stdout + assert "topsecret" not in stdout + + async def test_dynamic_exec_builtins_are_stripped(self, source): + """The bootstrap removes eval/exec/compile/breakpoint/input from builtins.""" + backend = SubprocessBackend(trusted=True) + code = ( + "for name in ('eval', 'exec', 'compile', 'breakpoint', 'input'):\n" + " try:\n" + " __builtins__[name] if isinstance(__builtins__, dict) " + "else getattr(__builtins__, name)\n" + " print(name, 'PRESENT')\n" + " except (KeyError, AttributeError):\n" + " print(name, 'STRIPPED')\n" + ) + exit_code, stdout, stderr = await backend.execute(code, source, timeout=10) + assert exit_code == 0, stderr + for name in ("eval", "exec", "compile", "breakpoint", "input"): + assert f"{name} STRIPPED" in stdout + + async def test_source_path_passed_as_argv(self, source): + backend = SubprocessBackend(trusted=True) + code = "import sys\nprint(sys.argv[1])" + exit_code, stdout, stderr = await backend.execute(code, source, timeout=10) + assert exit_code == 0, stderr + assert source in stdout + + async def test_timeout_is_contained(self, source): + """A sleeping script is killed (process group) and returns promptly.""" + backend = SubprocessBackend(trusted=True) + code = "import time\ntime.sleep(30)" + exit_code, stdout, stderr = await backend.execute(code, source, timeout=1) + assert exit_code == -1 + assert "timeout" in stderr.lower() + + async def test_fail_closed_when_limits_cannot_be_applied(self, source, monkeypatch): + """If resource limits cannot be applied, refuse to run (no unsandboxed exec).""" + + def _boom(_timeout): + raise ValueError("simulated setrlimit failure") + + monkeypatch.setattr(subprocess_mod, "_apply_resource_limits", _boom) + backend = SubprocessBackend(trusted=True) + code = "print('should not run')" + with pytest.raises(RuntimeError, match="hardening could not be applied"): + await backend.execute(code, source, timeout=10) + + +class TestTrustedWarning: + def test_untrusted_backend_warns_once(self, monkeypatch, caplog): + monkeypatch.setattr(subprocess_mod, "_trusted_warning_emitted", False) + backend = SubprocessBackend() + with caplog.at_level(logging.WARNING): + backend._warn_if_untrusted() + backend._warn_if_untrusted() + hits = [r for r in caplog.records if "NOT a security boundary" in r.getMessage()] + assert len(hits) == 1 + + def test_trusted_backend_does_not_warn(self, monkeypatch, caplog): + monkeypatch.setattr(subprocess_mod, "_trusted_warning_emitted", False) + backend = SubprocessBackend(trusted=True) + with caplog.at_level(logging.WARNING): + backend._warn_if_untrusted() + hits = [r for r in caplog.records if "NOT a security boundary" in r.getMessage()] + assert hits == [] diff --git a/tests/test_thinking.py b/tests/test_thinking.py new file mode 100644 index 0000000..d932d15 --- /dev/null +++ b/tests/test_thinking.py @@ -0,0 +1,155 @@ +"""ExtractionEngine: thinking unification (C1), confidence fail-fast (B2), error scope (C4).""" + +import io +from unittest.mock import AsyncMock, MagicMock, patch + +import openpyxl +import pytest +from pydantic import BaseModel + +from xlstruct.config import ExtractorConfig, thinking_call_kwargs +from xlstruct.exceptions import ErrorCode, ExtractionError +from xlstruct.extraction.engine import ExtractionEngine, _split_confidence +from xlstruct.extractor import Extractor +from xlstruct.schemas.suggest import FieldDef, SuggestedFields + +THINKING_BLOCK = {"type": "enabled", "budget_tokens": 10_000} + + +class Sample(BaseModel): + name: str + value: int + + +def _mock_client(return_value: object) -> MagicMock: + # ^ usage=None → UsageTracker.record() treats it as zero usage and no-ops + completion = MagicMock(usage=None) + client = MagicMock() + client.create_with_completion = AsyncMock(return_value=(return_value, completion)) + return client + + +# * thinking_call_kwargs + + +class TestThinkingCallKwargs: + def test_empty_when_thinking_off(self): + cfg = ExtractorConfig(provider="anthropic/claude-sonnet-4-6", thinking=False) + assert thinking_call_kwargs(cfg) == {} + + def test_empty_for_non_anthropic_even_when_on(self): + cfg = ExtractorConfig(provider="openai/gpt-4o", thinking=True) + assert thinking_call_kwargs(cfg) == {} + + def test_block_for_anthropic_thinking(self): + cfg = ExtractorConfig(provider="anthropic/claude-sonnet-4-6", thinking=True) + kw = thinking_call_kwargs(cfg) + assert kw["temperature"] == 1 + assert kw["thinking"] == THINKING_BLOCK + assert kw["max_tokens"] == 16_000 + assert kw["model"] == "claude-sonnet-4-6" + + +# * Direct extraction honors thinking (the C1 bug) + + +class TestDirectExtractionThinking: + async def test_thinking_kwargs_reach_llm(self): + cfg = ExtractorConfig(provider="anthropic/claude-sonnet-4-6", thinking=True) + client = _mock_client([Sample(name="a", value=1)]) + with patch("xlstruct.extraction.engine.build_instructor_client", return_value=client): + engine = ExtractionEngine(cfg) + await engine.extract("sheet data", Sample) + + kwargs = client.create_with_completion.await_args.kwargs + assert kwargs["thinking"] == THINKING_BLOCK + assert kwargs["temperature"] == 1 + assert kwargs["model"] == "claude-sonnet-4-6" + + async def test_no_thinking_kwargs_when_disabled(self): + cfg = ExtractorConfig( + provider="anthropic/claude-sonnet-4-6", thinking=False, temperature=0.0 + ) + client = _mock_client([Sample(name="a", value=1)]) + with patch("xlstruct.extraction.engine.build_instructor_client", return_value=client): + engine = ExtractionEngine(cfg) + await engine.extract("sheet data", Sample) + + kwargs = client.create_with_completion.await_args.kwargs + assert "thinking" not in kwargs + assert kwargs["temperature"] == 0.0 + + +# * suggest_schema honors thinking (the C1 bug) + + +class TestSuggestSchemaThinking: + async def test_thinking_kwargs_reach_llm(self, tmp_path): + wb = openpyxl.Workbook() + ws = wb.active + ws["A1"], ws["B1"] = "Name", "Price" + ws["A2"], ws["B2"] = "Apple", 1.5 + buf = io.BytesIO() + wb.save(buf) + path = tmp_path / "s.xlsx" + path.write_bytes(buf.getvalue()) + + suggested = SuggestedFields( + model_name="Row", + fields=[FieldDef(name="name", type="str", nullable=False, description="A")], + ) + suggest_client = _mock_client(suggested) + with ( + # ^ Engine built in Extractor.__init__ must not hit a real client + patch("xlstruct.extraction.engine.build_instructor_client", return_value=MagicMock()), + patch( + "xlstruct.extraction.pipeline.build_instructor_client", return_value=suggest_client + ), + ): + extractor = Extractor(provider="anthropic/claude-sonnet-4-6", thinking=True) + await extractor.suggest_schema(str(path)) + + kwargs = suggest_client.create_with_completion.await_args.kwargs + assert kwargs["thinking"] == THINKING_BLOCK + assert kwargs["temperature"] == 1 + + +# * B2 — confidence default removed (fail-fast on structural invariant break) + + +class TestConfidenceFailFast: + def test_missing_confidence_key_raises(self): + # ^ A bare record lacks the "{field}_confidence" keys the wrapper guarantees + with pytest.raises(KeyError): + _split_confidence([Sample(name="a", value=1)], Sample) + + +# * C4 — only the provider call is wrapped as ExtractionError; our post-processing is not + + +class TestExtractionErrorScope: + async def test_llm_call_error_wrapped_as_extraction_error(self): + cfg = ExtractorConfig(provider="openai/gpt-4o") + client = MagicMock() + client.create_with_completion = AsyncMock(side_effect=RuntimeError("api down")) + with patch("xlstruct.extraction.engine.build_instructor_client", return_value=client): + engine = ExtractionEngine(cfg) + with pytest.raises(ExtractionError) as exc: + await engine.extract("data", Sample) + assert exc.value.code == ErrorCode.EXTRACTION_LLM_FAILED + + async def test_postprocessing_error_propagates_unwrapped(self): + # ^ A bug in our own post-processing must NOT be mislabeled as an LLM failure + cfg = ExtractorConfig(provider="openai/gpt-4o") + client = _mock_client([Sample(name="a", value=1)]) + with patch("xlstruct.extraction.engine.build_instructor_client", return_value=client): + engine = ExtractionEngine(cfg) + with ( + patch.object( + ExtractionEngine, + "_split_provenance", + side_effect=RuntimeError("post-proc boom"), + ), + pytest.raises(RuntimeError, match="post-proc boom"), + ): + await engine.extract("data", Sample, track_provenance=True) diff --git a/tests/test_tokens.py b/tests/test_tokens.py index 7956937..0b4d7e0 100644 --- a/tests/test_tokens.py +++ b/tests/test_tokens.py @@ -1,7 +1,12 @@ """Tests for token counting utilities.""" -from xlstruct._tokens import estimate_sheet_tokens -from xlstruct.schemas.core import SheetData +from xlstruct._tokens import ( + count_tokens, + estimate_cells_tokens, + estimate_row_token_costs, + estimate_sheet_tokens, +) +from xlstruct.schemas.core import CellData, SheetData class TestEstimateSheetTokens: @@ -14,3 +19,31 @@ def test_simple_sheet(self, simple_sheet: SheetData): assert tokens > 0 # ^ 24 cells with short values: should be reasonable assert tokens < 1000 + + +class TestSpecialTokenContent: + """Cell values may contain literal special-token text (e.g. '<|endoftext|>'). + The strict encode() raises on it; estimation must use encode_ordinary and + never crash on valid spreadsheet content.""" + + _SPECIAL = "before <|endoftext|> after" + + def test_count_tokens_does_not_raise(self): + assert count_tokens(self._SPECIAL) > 0 + + def test_estimate_cells_tokens_does_not_raise(self): + cells = [CellData(row=1, col=1, value="<|endoftext|>")] + assert estimate_cells_tokens(cells) > 0 + + def test_estimate_sheet_tokens_does_not_raise(self): + sheet = SheetData( + name="s", + row_count=1, + col_count=1, + cells=[CellData(row=1, col=1, value=self._SPECIAL)], + ) + assert estimate_sheet_tokens(sheet) > 0 + + def test_estimate_row_token_costs_does_not_raise(self): + rows = [[CellData(row=2, col=1, value=self._SPECIAL)]] + assert estimate_row_token_costs(rows)[0] > 0 diff --git a/uv.lock b/uv.lock index f534cb5..da0099b 100644 --- a/uv.lock +++ b/uv.lock @@ -3259,7 +3259,7 @@ wheels = [ [[package]] name = "xlstruct" -version = "0.6.0" +version = "0.7.0" source = { editable = "." } dependencies = [ { name = "fsspec" },