mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-19 01:50:51 -05:00
feat(indexing): add SKIP_SPREADSHEET_CONTENT to index spreadsheets by title only (#16075)
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
7cddaf2278
commit
81b2e2bed8
21 files changed
+324
-33
No files matched your search
@@ -48,7 +48,11 @@ from onyx.background.indexing.index_attempt_utils import (
|
||||
cleanup_index_attempts,
|
||||
get_old_index_attempt_ids,
|
||||
)
|
||||
from onyx.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP, PERSISTENT_INDEXING
|
||||
from onyx.configs.app_configs import (
|
||||
DISABLE_INDEX_UPDATE_ON_SWAP,
|
||||
PERSISTENT_INDEXING,
|
||||
SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
from onyx.configs.constants import (
|
||||
CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
|
||||
CELERY_INDEXING_LOCK_TIMEOUT,
|
||||
@@ -2080,6 +2084,7 @@ def _docprocessing_task(
|
||||
request_id=index_attempt_metadata.request_id,
|
||||
adapter=adapter,
|
||||
from_beginning=from_beginning,
|
||||
skip_spreadsheet_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
|
||||
# Track chunk indexing usage for cloud usage limits
|
||||
|
||||
@@ -351,6 +351,8 @@ def _load_user_file_documents(
|
||||
file_names=[file_name] if file_name else None,
|
||||
)
|
||||
connector.load_credentials({})
|
||||
# User files keep their spreadsheet content regardless of the setting.
|
||||
connector.skip_spreadsheet_content = False
|
||||
|
||||
# User files aren't attempt-scoped, so the docfetching staging reapers don't cover them.
|
||||
staging_callback, staged_csv_ids = build_tracking_raw_file_callback(
|
||||
|
||||
@@ -22,7 +22,7 @@ from collections.abc import Iterable, Sequence
|
||||
from more_itertools import chunked
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from onyx.configs.app_configs import INDEX_BATCH_SIZE
|
||||
from onyx.configs.app_configs import INDEX_BATCH_SIZE, SKIP_SPREADSHEET_CONTENT
|
||||
from onyx.connectors.factory import instantiate_connector
|
||||
from onyx.connectors.interfaces import Resolver
|
||||
from onyx.connectors.models import (
|
||||
@@ -158,6 +158,7 @@ def _flush_batch(
|
||||
document_batch=documents,
|
||||
request_id=metadata.request_id,
|
||||
adapter=adapter,
|
||||
skip_spreadsheet_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
|
||||
failed_ids: set[str] = set()
|
||||
|
||||
@@ -1554,6 +1554,17 @@ INDEXING_EMBEDDING_MODEL_NUM_THREADS = int(
|
||||
# Default is 512MB worth of characters (536,870,912). Configurable via MAX_DOCUMENT_CHARS env var.
|
||||
MAX_DOCUMENT_CHARS = int(os.environ.get("MAX_DOCUMENT_CHARS") or 536_870_912)
|
||||
|
||||
# When true, connectors replace the content of spreadsheets and other tabular
|
||||
# files (xlsx, Google Sheets, CSV, TSV) with a "<content-skipped>" placeholder
|
||||
# at fetch time, so each file indexes as its title only. The sheets are not
|
||||
# parsed, but the raw file is still downloaded and stored. Docprocessing applies
|
||||
# the same replacement to tabular sections that still reach it, e.g. in batches
|
||||
# staged before this was set. Bounds the embedding cost of large spreadsheets.
|
||||
# User uploads keep their content. Set it on docfetching and docprocessing.
|
||||
SKIP_SPREADSHEET_CONTENT = (
|
||||
os.environ.get("SKIP_SPREADSHEET_CONTENT", "").lower() == "true"
|
||||
)
|
||||
|
||||
# Max RSS (in MB) for a spawned indexing worker process. When exceeded, the
|
||||
# docfetching watchdog terminates the worker and fails the attempt with a clear
|
||||
# error, pre-empting a kernel OOM kill of the whole pod — which would also kill
|
||||
|
||||
@@ -14,7 +14,11 @@ from botocore.credentials import RefreshableCredentials
|
||||
from botocore.exceptions import ClientError, NoCredentialsError, PartialCredentialsError
|
||||
from botocore.session import get_session
|
||||
|
||||
from onyx.configs.app_configs import BLOB_STORAGE_SIZE_THRESHOLD, INDEX_BATCH_SIZE
|
||||
from onyx.configs.app_configs import (
|
||||
BLOB_STORAGE_SIZE_THRESHOLD,
|
||||
INDEX_BATCH_SIZE,
|
||||
SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
from onyx.configs.constants import BlobType, DocumentSource, FileOrigin
|
||||
from onyx.connectors.cross_connector_utils.miscellaneous_utils import (
|
||||
process_onyx_metadata,
|
||||
@@ -486,7 +490,7 @@ class BlobStorageConnector(LoadConnector, PollConnector):
|
||||
downloaded_file = self._download_object(key)
|
||||
if downloaded_file is None:
|
||||
continue
|
||||
tabular_sections: list[TabularSection] = []
|
||||
tabular_sections: list[TabularSection | TextSection] = []
|
||||
staged_file_id: str | None = None
|
||||
if self.raw_file_callback is not None:
|
||||
content_type, _ = mimetypes.guess_type(file_name)
|
||||
@@ -496,6 +500,7 @@ class BlobStorageConnector(LoadConnector, PollConnector):
|
||||
content_type=content_type or "application/octet-stream",
|
||||
raw_file_callback=self.raw_file_callback,
|
||||
link=link,
|
||||
skip_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
tabular_sections = result.sections
|
||||
staged_file_id = result.staged_file_id
|
||||
|
||||
@@ -4,7 +4,7 @@ from typing import IO
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from onyx.connectors.models import TabularSection
|
||||
from onyx.connectors.models import TabularSection, TextSection
|
||||
from onyx.file_processing.extract_file_text import file_io_to_text, stage_xlsx_sheets
|
||||
from onyx.file_processing.file_types import OnyxFileExtensions
|
||||
from onyx.file_store.staging import RawFileCallback
|
||||
@@ -12,9 +12,11 @@ from onyx.utils.logger import setup_logger
|
||||
|
||||
logger = setup_logger()
|
||||
|
||||
CONTENT_SKIPPED_PLACEHOLDER = "<content-skipped>"
|
||||
|
||||
|
||||
class TabularExtractionResult(BaseModel):
|
||||
sections: list[TabularSection]
|
||||
sections: list[TabularSection | TextSection]
|
||||
staged_file_id: str
|
||||
|
||||
|
||||
@@ -38,18 +40,24 @@ def tabular_file_to_sections(
|
||||
file_name: str,
|
||||
stage: RawFileCallback,
|
||||
link: str = "",
|
||||
) -> list[TabularSection]:
|
||||
skip_content: bool = False,
|
||||
) -> list[TabularSection | TextSection]:
|
||||
"""Convert a tabular file into one or more TabularSections.
|
||||
|
||||
- .xlsx → one staged TabularSection per non-empty sheet.
|
||||
- .csv / .tsv → one staged TabularSection for the whole file.
|
||||
- empty input → `[]`.
|
||||
- skip_content → one placeholder TextSection, so the file indexes as its
|
||||
title only. Nothing is parsed or staged.
|
||||
"""
|
||||
lowered = file_name.lower()
|
||||
|
||||
if not lowered.endswith(tuple(OnyxFileExtensions.TABULAR_EXTENSIONS)):
|
||||
raise ValueError(f"{file_name!r} is not a tabular file")
|
||||
|
||||
if skip_content:
|
||||
return [TextSection(text=CONTENT_SKIPPED_PLACEHOLDER, link=link or file_name)]
|
||||
|
||||
if lowered.endswith(tuple(OnyxFileExtensions.SPREADSHEET_EXTENSIONS)):
|
||||
return [
|
||||
TabularSection(
|
||||
@@ -84,13 +92,16 @@ def extract_and_stage_tabular_file(
|
||||
content_type: str,
|
||||
raw_file_callback: RawFileCallback,
|
||||
link: str = "",
|
||||
skip_content: bool = False,
|
||||
) -> TabularExtractionResult:
|
||||
"""Extract tabular sections AND stage the raw bytes via the callback."""
|
||||
"""Extract tabular sections AND stage the raw bytes via the callback. The
|
||||
raw bytes are staged even with skip_content, so the file stays readable."""
|
||||
sections = tabular_file_to_sections(
|
||||
file=file,
|
||||
file_name=file_name,
|
||||
stage=raw_file_callback,
|
||||
link=link,
|
||||
skip_content=skip_content,
|
||||
)
|
||||
# rewind so the callback can re-read what extraction consumed
|
||||
file.seek(0)
|
||||
|
||||
@@ -9,6 +9,7 @@ from onyx.configs.app_configs import (
|
||||
CONTINUE_ON_CONNECTOR_FAILURE,
|
||||
DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD,
|
||||
INDEX_BATCH_SIZE,
|
||||
SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
from onyx.configs.constants import DocumentSource, FileOrigin
|
||||
from onyx.connectors.cross_connector_utils.miscellaneous_utils import (
|
||||
@@ -307,6 +308,7 @@ class DrupalWikiConnector(
|
||||
file_name=file_name,
|
||||
stage=self.raw_file_callback,
|
||||
link=download_url,
|
||||
skip_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
|
||||
@@ -4,7 +4,7 @@ from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import IO, Any
|
||||
|
||||
from onyx.configs.app_configs import INDEX_BATCH_SIZE
|
||||
from onyx.configs.app_configs import INDEX_BATCH_SIZE, SKIP_SPREADSHEET_CONTENT
|
||||
from onyx.configs.constants import DocumentSource, FileOrigin
|
||||
from onyx.connectors.cross_connector_utils.miscellaneous_utils import (
|
||||
process_onyx_metadata,
|
||||
@@ -26,7 +26,11 @@ from onyx.connectors.models import (
|
||||
TabularSection,
|
||||
TextSection,
|
||||
)
|
||||
from onyx.file_processing.extract_file_text import extract_text_and_images, get_file_ext
|
||||
from onyx.file_processing.extract_file_text import (
|
||||
ExtractionResult,
|
||||
extract_text_and_images,
|
||||
get_file_ext,
|
||||
)
|
||||
from onyx.file_processing.file_types import OnyxFileExtensions
|
||||
from onyx.file_processing.image_utils import store_image_and_create_section
|
||||
from onyx.file_store.file_store import get_default_file_store
|
||||
@@ -87,6 +91,7 @@ def _process_file(
|
||||
pdf_pass: str | None,
|
||||
file_type: str | None,
|
||||
stage: RawFileCallback | None,
|
||||
skip_spreadsheet_content: bool = False,
|
||||
) -> list[Document]:
|
||||
"""
|
||||
Process a file and return a list of Documents.
|
||||
@@ -156,13 +161,24 @@ def _process_file(
|
||||
# 2) Otherwise: text-based approach. Possibly with embedded images.
|
||||
file.seek(0)
|
||||
|
||||
# Extract text and images from the file
|
||||
extraction_result = extract_text_and_images(
|
||||
file=file,
|
||||
file_name=file_name,
|
||||
pdf_pass=pdf_pass,
|
||||
content_type=file_type,
|
||||
)
|
||||
# Extract text and images from the file. A skipped spreadsheet is not
|
||||
# parsed: its content is replaced, and xlsx extraction yields no metadata.
|
||||
# CSV/TSV still go through extraction, which reads their ONYX_METADATA.
|
||||
extraction_result: ExtractionResult
|
||||
if (
|
||||
skip_spreadsheet_content
|
||||
and extension in OnyxFileExtensions.SPREADSHEET_EXTENSIONS
|
||||
):
|
||||
extraction_result = ExtractionResult(
|
||||
text_content="", embedded_images=[], metadata={}
|
||||
)
|
||||
else:
|
||||
extraction_result = extract_text_and_images(
|
||||
file=file,
|
||||
file_name=file_name,
|
||||
pdf_pass=pdf_pass,
|
||||
content_type=file_type,
|
||||
)
|
||||
|
||||
# Each file may have file-specific ONYX_METADATA https://docs.onyx.app/admins/connectors/official/file
|
||||
# If so, we should add it to any metadata processed so far
|
||||
@@ -227,6 +243,7 @@ def _process_file(
|
||||
file_name=file_name,
|
||||
stage=stage,
|
||||
link=link or "",
|
||||
skip_content=skip_spreadsheet_content,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -304,6 +321,9 @@ class LocalFileConnector(LoadConnector):
|
||||
batch_size: int = INDEX_BATCH_SIZE,
|
||||
) -> None:
|
||||
self.file_locations = [str(loc) for loc in file_locations]
|
||||
# Not an __init__ parameter, so it stays out of the connector config.
|
||||
# The user file loader sets it to False to keep spreadsheet content.
|
||||
self.skip_spreadsheet_content: bool = SKIP_SPREADSHEET_CONTENT
|
||||
self.batch_size = batch_size
|
||||
self.pdf_pass: str | None = None
|
||||
self._zip_metadata_file_id = zip_metadata_file_id
|
||||
@@ -345,6 +365,7 @@ class LocalFileConnector(LoadConnector):
|
||||
pdf_pass=self.pdf_pass,
|
||||
file_type=file_record.file_type,
|
||||
stage=self.raw_file_callback,
|
||||
skip_spreadsheet_content=self.skip_spreadsheet_content,
|
||||
)
|
||||
documents.extend(new_docs)
|
||||
|
||||
|
||||
@@ -7,7 +7,10 @@ from urllib.parse import urlparse, urlunparse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from onyx.access.models import ExternalAccess
|
||||
from onyx.configs.app_configs import GOOGLE_DRIVE_ADVANCED_PARSE_MAX_BYTES
|
||||
from onyx.configs.app_configs import (
|
||||
GOOGLE_DRIVE_ADVANCED_PARSE_MAX_BYTES,
|
||||
SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
from onyx.configs.constants import DocumentSource, FileOrigin
|
||||
from onyx.connectors.cross_connector_utils.section_utils import cap_sections_text
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
@@ -299,6 +302,7 @@ def _download_and_extract_sections_basic(
|
||||
content_type=content_type,
|
||||
raw_file_callback=raw_file_callback,
|
||||
link=link,
|
||||
skip_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
tabular_sections = result.sections
|
||||
staged_file_id = result.staged_file_id
|
||||
|
||||
@@ -23,7 +23,7 @@ from office365.onedrive.driveitems.driveItem import DriveItem
|
||||
from office365.runtime.paths.resource_path import ResourcePath
|
||||
from pydantic import AliasChoices, BaseModel, Field
|
||||
|
||||
from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS
|
||||
from onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS, SKIP_SPREADSHEET_CONTENT
|
||||
from onyx.configs.constants import FileOrigin
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
extract_and_stage_tabular_file,
|
||||
@@ -655,6 +655,7 @@ def extract_drive_item_content(
|
||||
content_type=mime_type or "application/octet-stream",
|
||||
raw_file_callback=raw_file_callback,
|
||||
link=driveitem.web_url or "",
|
||||
skip_content=SKIP_SPREADSHEET_CONTENT,
|
||||
)
|
||||
sections.extend(result.sections)
|
||||
staged_file_id = result.staged_file_id
|
||||
|
||||
@@ -21,6 +21,9 @@ from onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled
|
||||
from onyx.connectors.cross_connector_utils.miscellaneous_utils import (
|
||||
get_experts_stores_representations,
|
||||
)
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
CONTENT_SKIPPED_PLACEHOLDER,
|
||||
)
|
||||
from onyx.connectors.models import (
|
||||
ConnectorFailure,
|
||||
ConnectorStopSignal,
|
||||
@@ -31,6 +34,7 @@ from onyx.connectors.models import (
|
||||
IndexingDocument,
|
||||
Section,
|
||||
SectionType,
|
||||
TabularSection,
|
||||
TextSection,
|
||||
)
|
||||
from onyx.db.connector_credential_pair import get_connector_credential_pair
|
||||
@@ -505,6 +509,7 @@ def index_doc_batch_with_handler(
|
||||
llm_enrichment_allowed: bool = True,
|
||||
image_summarization_llm: LLM | None = None,
|
||||
llm: LLM | None = None,
|
||||
skip_spreadsheet_content: bool = False,
|
||||
) -> IndexingPipelineResult:
|
||||
# Documents whose prepare() committed before the batch raised. Their change
|
||||
# gates would now see nothing to do, so a retry of them forces the update.
|
||||
@@ -533,6 +538,7 @@ def index_doc_batch_with_handler(
|
||||
llm=llm,
|
||||
force_update=force_update,
|
||||
on_prepared=prepared.update,
|
||||
skip_spreadsheet_content=skip_spreadsheet_content,
|
||||
)
|
||||
except ConnectorStopSignal as e:
|
||||
logger.warning(
|
||||
@@ -751,6 +757,34 @@ def index_doc_batch_prepare(
|
||||
)
|
||||
|
||||
|
||||
def skip_tabular_content(documents: list[Document]) -> list[Document]:
|
||||
"""Replace each document's tabular sections with one placeholder text
|
||||
section, so a spreadsheet indexes as its title plus the placeholder. Other
|
||||
sections stay. Connectors already skip tabular content at extraction; this
|
||||
catches sections that still reach docprocessing (batches staged before
|
||||
SKIP_SPREADSHEET_CONTENT was set, or a connector that builds them itself).
|
||||
The result matches extraction, so both give the same content hash."""
|
||||
result: list[Document] = []
|
||||
for document in documents:
|
||||
if not any(isinstance(s, TabularSection) for s in document.sections):
|
||||
result.append(document)
|
||||
continue
|
||||
|
||||
sections: list[TextSection | ImageSection | TabularSection] = []
|
||||
placeholder_added: bool = False
|
||||
for section in document.sections:
|
||||
if not isinstance(section, TabularSection):
|
||||
sections.append(section)
|
||||
elif not placeholder_added:
|
||||
sections.append(
|
||||
TextSection(text=CONTENT_SKIPPED_PLACEHOLDER, link=section.link)
|
||||
)
|
||||
placeholder_added = True
|
||||
logger.info("Skipped tabular content of document %s", document.id)
|
||||
result.append(document.model_copy(update={"sections": sections}))
|
||||
return result
|
||||
|
||||
|
||||
def filter_documents(
|
||||
document_batch: list[Document],
|
||||
) -> tuple[list[Document], list[ConnectorFailure]]:
|
||||
@@ -1468,6 +1502,7 @@ def index_doc_batch(
|
||||
filter_fnc: Callable[
|
||||
[list[Document]], tuple[list[Document], list[ConnectorFailure]]
|
||||
] = filter_documents,
|
||||
skip_spreadsheet_content: bool = False,
|
||||
) -> IndexingPipelineResult:
|
||||
"""End-to-end indexing for a pre-batched set of documents."""
|
||||
"""Takes different pieces of the indexing pipeline and applies it to a batch of documents
|
||||
@@ -1500,6 +1535,8 @@ def index_doc_batch(
|
||||
_attempt_metadata.attempt_id if _attempt_metadata is not None else None
|
||||
)
|
||||
|
||||
if skip_spreadsheet_content:
|
||||
document_batch = skip_tabular_content(document_batch)
|
||||
filtered_documents, filter_failures = filter_fnc(document_batch)
|
||||
filtered_documents = _apply_document_ingestion_hook(filtered_documents)
|
||||
with time_stage_if_set(IndexAttemptStage.DOC_DB_PREPARE, attempt_id):
|
||||
@@ -1729,11 +1766,16 @@ def run_indexing_pipeline(
|
||||
ignore_time_skip: bool = False,
|
||||
index_to_secondary: bool = False,
|
||||
from_beginning: bool = False,
|
||||
skip_spreadsheet_content: bool = False,
|
||||
) -> IndexingPipelineResult:
|
||||
"""Builds a pipeline which takes in a list (batch) of docs and indexes them.
|
||||
|
||||
index_to_secondary disables the content-hash gate + stamp for FUTURE writes so
|
||||
they don't cross-suppress the PRESENT index's shared hash.
|
||||
|
||||
skip_spreadsheet_content indexes each tabular file as its title plus a
|
||||
placeholder (see skip_tabular_content). Only connector indexing sets it:
|
||||
user files and the ingestion API need the content they were given.
|
||||
"""
|
||||
if db_session is not None:
|
||||
all_search_settings = get_active_search_settings(db_session)
|
||||
@@ -1799,4 +1841,5 @@ def run_indexing_pipeline(
|
||||
ignore_time_skip=ignore_time_skip,
|
||||
index_to_secondary=index_to_secondary,
|
||||
from_beginning=from_beginning,
|
||||
skip_spreadsheet_content=skip_spreadsheet_content,
|
||||
)
|
||||
+46
-5
@@ -7,9 +7,12 @@ import pytest
|
||||
from openpyxl.worksheet.worksheet import Worksheet
|
||||
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
CONTENT_SKIPPED_PLACEHOLDER,
|
||||
extract_and_stage_tabular_file,
|
||||
is_tabular_file,
|
||||
tabular_file_to_sections,
|
||||
)
|
||||
from onyx.connectors.models import TabularSection, TextSection
|
||||
|
||||
|
||||
def _make_stage_callback() -> tuple[
|
||||
@@ -76,8 +79,10 @@ class TestTabularFileToSections:
|
||||
stage=fake_stage,
|
||||
)
|
||||
assert len(sections) == 1
|
||||
assert "Alice" in staged[sections[0].csv_file_id][0].decode("utf-8")
|
||||
assert sections[0].heading == "budget.xlsm :: Sheet1"
|
||||
section = sections[0]
|
||||
assert isinstance(section, TabularSection)
|
||||
assert "Alice" in staged[section.csv_file_id][0].decode("utf-8")
|
||||
assert section.heading == "budget.xlsm :: Sheet1"
|
||||
|
||||
def test_unknown_extension_raises(self) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
@@ -98,6 +103,41 @@ class TestTabularFileToSections:
|
||||
|
||||
assert sections == []
|
||||
|
||||
def test_skip_content_returns_placeholder_without_parsing(self) -> None:
|
||||
"""The bytes are not a valid workbook, so any parse would raise."""
|
||||
staged, fake_stage = _make_stage_callback()
|
||||
|
||||
sections = tabular_file_to_sections(
|
||||
io.BytesIO(b"not a workbook"),
|
||||
file_name="budget.xlsx",
|
||||
stage=fake_stage,
|
||||
skip_content=True,
|
||||
)
|
||||
|
||||
assert sections == [
|
||||
TextSection(text=CONTENT_SKIPPED_PLACEHOLDER, link="budget.xlsx")
|
||||
]
|
||||
assert staged == {}
|
||||
|
||||
def test_skip_content_still_stages_the_raw_file(self) -> None:
|
||||
staged, fake_stage = _make_stage_callback()
|
||||
|
||||
result = extract_and_stage_tabular_file(
|
||||
io.BytesIO(b"a,b\n1,2"),
|
||||
file_name="data.csv",
|
||||
content_type="text/csv",
|
||||
raw_file_callback=fake_stage,
|
||||
link="https://example.com/data.csv",
|
||||
skip_content=True,
|
||||
)
|
||||
|
||||
assert result.sections == [
|
||||
TextSection(
|
||||
text=CONTENT_SKIPPED_PLACEHOLDER, link="https://example.com/data.csv"
|
||||
)
|
||||
]
|
||||
assert staged == {result.staged_file_id: (b"a,b\n1,2", "text/csv")}
|
||||
|
||||
|
||||
class TestFileBackedXlsx:
|
||||
"""Every non-empty workbook sheet is staged as CSV and referenced by
|
||||
@@ -114,7 +154,7 @@ class TestFileBackedXlsx:
|
||||
|
||||
assert len(sections) == 1
|
||||
section = sections[0]
|
||||
assert section.csv_file_id is not None
|
||||
assert isinstance(section, TabularSection)
|
||||
assert section.heading == "big.xlsx :: Sheet1"
|
||||
|
||||
csv_bytes, content_type = staged[section.csv_file_id]
|
||||
@@ -134,5 +174,6 @@ class TestFileBackedXlsx:
|
||||
xlsx, file_name="small.xlsx", stage=fake_stage
|
||||
)
|
||||
assert len(sections) == 1
|
||||
assert sections[0].csv_file_id is not None
|
||||
assert staged[sections[0].csv_file_id][0].decode("utf-8") == "a,b\n1,2\n"
|
||||
section = sections[0]
|
||||
assert isinstance(section, TabularSection)
|
||||
assert staged[section.csv_file_id][0].decode("utf-8") == "a,b\n1,2\n"
|
||||
@@ -0,0 +1,41 @@
|
||||
import io
|
||||
from typing import IO
|
||||
from unittest.mock import patch
|
||||
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
CONTENT_SKIPPED_PLACEHOLDER,
|
||||
)
|
||||
from onyx.connectors.file.connector import _process_file
|
||||
from onyx.connectors.models import TextSection
|
||||
|
||||
|
||||
def test_skipped_spreadsheet_is_not_parsed() -> None:
|
||||
"""The bytes are not a valid workbook, so any parse would fail."""
|
||||
staged: list[str] = []
|
||||
|
||||
def stage(_content: IO[bytes], content_type: str) -> str:
|
||||
staged.append(content_type)
|
||||
return "staged-file"
|
||||
|
||||
with patch(
|
||||
"onyx.connectors.file.connector.extract_text_and_images"
|
||||
) as extract_text_and_images:
|
||||
documents = _process_file(
|
||||
file_id="file-1",
|
||||
file_name="Budget.xlsx",
|
||||
file=io.BytesIO(b"not a workbook"),
|
||||
metadata=None,
|
||||
pdf_pass=None,
|
||||
file_type=None,
|
||||
stage=stage,
|
||||
skip_spreadsheet_content=True,
|
||||
)
|
||||
|
||||
extract_text_and_images.assert_not_called()
|
||||
assert staged == []
|
||||
assert len(documents) == 1
|
||||
assert documents[0].sections == [
|
||||
TextSection(text=CONTENT_SKIPPED_PLACEHOLDER, link="Budget.xlsx")
|
||||
]
|
||||
# The original file stays readable by the Python tool.
|
||||
assert documents[0].file_id == "file-1"
|
||||
@@ -4,6 +4,7 @@ import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import nullcontext
|
||||
from datetime import datetime, timezone
|
||||
from io import BytesIO
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, List, cast
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
@@ -11,6 +12,10 @@ from unittest.mock import MagicMock, Mock, patch
|
||||
import pytest
|
||||
|
||||
from onyx.access.models import ExternalAccess
|
||||
from onyx.connectors.cross_connector_utils.tabular_section_utils import (
|
||||
CONTENT_SKIPPED_PLACEHOLDER,
|
||||
tabular_file_to_sections,
|
||||
)
|
||||
from onyx.connectors.models import (
|
||||
Document,
|
||||
DocumentSource,
|
||||
@@ -18,6 +23,9 @@ from onyx.connectors.models import (
|
||||
TabularSection,
|
||||
TextSection,
|
||||
)
|
||||
from onyx.document_index.chunk_content_enrichment import (
|
||||
generate_enriched_content_for_chunk_embedding,
|
||||
)
|
||||
from onyx.error_handling.error_codes import OnyxErrorCode
|
||||
from onyx.error_handling.exceptions import OnyxError
|
||||
from onyx.hooks.executor import HookSkipped, HookSoftFailed
|
||||
@@ -1165,6 +1173,78 @@ def _make_tabular_doc(doc_id: str, section: TabularSection) -> Document:
|
||||
)
|
||||
|
||||
|
||||
def test_index_batch_skips_tabular_sections_that_reach_it() -> None:
|
||||
"""A tabular section that reaches docprocessing (for example, in a batch
|
||||
staged before the option was set) becomes one placeholder before the change
|
||||
gate. Other sections and documents stay."""
|
||||
workbook = Document(
|
||||
id="workbook",
|
||||
title="Budget.xlsx",
|
||||
semantic_identifier="Budget.xlsx",
|
||||
sections=[
|
||||
TextSection(text="Notes", link="notes-link"),
|
||||
TabularSection(csv_file_id="sheet-1", link="sheet-1-link"),
|
||||
TabularSection(csv_file_id="sheet-2", link="sheet-2-link"),
|
||||
],
|
||||
source=DocumentSource.GOOGLE_DRIVE,
|
||||
metadata={},
|
||||
)
|
||||
text_doc = _make_doc("text-doc")
|
||||
adapter = MagicMock()
|
||||
adapter.index_attempt_metadata = None
|
||||
adapter.prepare.return_value = None
|
||||
|
||||
with patch(
|
||||
f"{_PATCH_PREFIX}._apply_document_ingestion_hook",
|
||||
side_effect=lambda documents: documents,
|
||||
):
|
||||
index_doc_batch(
|
||||
document_batch=[workbook, text_doc],
|
||||
chunker=MagicMock(),
|
||||
embedder=MagicMock(),
|
||||
document_index=MagicMock(),
|
||||
request_id=None,
|
||||
tenant_id="tenant",
|
||||
adapter=adapter,
|
||||
skip_spreadsheet_content=True,
|
||||
)
|
||||
|
||||
prepared_workbook, prepared_text_doc = adapter.prepare.call_args.args[0]
|
||||
assert prepared_workbook.sections == [
|
||||
TextSection(text="Notes", link="notes-link"),
|
||||
TextSection(text=CONTENT_SKIPPED_PLACEHOLDER, link="sheet-1-link"),
|
||||
]
|
||||
assert prepared_text_doc == text_doc
|
||||
|
||||
|
||||
def test_skipped_spreadsheet_embeds_as_its_title(
|
||||
embedder: DefaultIndexingEmbedder,
|
||||
) -> None:
|
||||
workbook = Document(
|
||||
id="workbook",
|
||||
title="Budget.xlsx",
|
||||
semantic_identifier="Budget.xlsx",
|
||||
sections=tabular_file_to_sections(
|
||||
BytesIO(b"unused"),
|
||||
file_name="Budget.xlsx",
|
||||
stage=lambda _content, _content_type: "unused",
|
||||
skip_content=True,
|
||||
),
|
||||
source=DocumentSource.GOOGLE_DRIVE,
|
||||
metadata={},
|
||||
)
|
||||
chunker = Chunker(tokenizer=embedder.embedding_model.tokenizer)
|
||||
|
||||
chunks = chunker.chunk(process_image_sections([workbook]))
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].content == CONTENT_SKIPPED_PLACEHOLDER
|
||||
assert chunks[0].title_prefix.startswith("Budget.xlsx")
|
||||
assert generate_enriched_content_for_chunk_embedding(chunks[0]).startswith(
|
||||
chunks[0].title_prefix + CONTENT_SKIPPED_PLACEHOLDER
|
||||
)
|
||||
|
||||
|
||||
class TestProcessImageSections:
|
||||
"""Validate that parallel image summarization places results in the
|
||||
correct section positions — especially under concurrent execution."""
|
||||
|
||||
Reference in new issue
Block a user