mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-31 21:13:49 +08:00
Feature/table parser column roles (#13710)
### What problem does this PR solve? The table file parser (CSV/Excel) currently treats all columns identically — every column is both vectorized (embedded in chunk text) and stored as filterable metadata. There's no way for users to control which columns should be searchable by semantic meaning versus which should only be filterable attributes. For example, when ingesting a news articles CSV with columns like title, content, country, category, source, etc., the embedding includes metadata fields like country: Brazil and source: Reuters in the chunk text, which dilutes the semantic quality of the embedding without adding retrieval value. The RDBMS connector (MySQL/PostgreSQL) already supports content_columns / metadata_columns, but this capability was missing for file-based table ingestion. This PR adds column-level control (vectorize / metadata / both) for the table file parser, following RAGFlow's existing patterns. Backward compatible: Datasets without table_column_roles or with table_column_mode: auto behave exactly as before (all columns = both). ### Type of change - [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
@@ -18,14 +18,15 @@
|
||||
|
||||
from unittest.mock import Mock
|
||||
from api.utils.validation_utils import (
|
||||
validate_immutable_fields,
|
||||
ParserConfig,
|
||||
UpdateDocumentReq,
|
||||
validate_chunk_method,
|
||||
validate_document_name,
|
||||
validate_chunk_method
|
||||
validate_immutable_fields,
|
||||
)
|
||||
from api.constants import FILE_NAME_LEN_LIMIT
|
||||
from api.db import FileType
|
||||
from common.constants import RetCode
|
||||
from api.utils.validation_utils import UpdateDocumentReq
|
||||
|
||||
|
||||
def test_validate_immutable_fields_no_changes():
|
||||
@@ -299,4 +300,15 @@ def test_validate_chunk_method_other_extensions_still_valid():
|
||||
|
||||
error_msg, error_code = validate_chunk_method(doc)
|
||||
assert error_msg is None
|
||||
assert error_code is None
|
||||
assert error_code is None
|
||||
|
||||
|
||||
def test_parser_config_normalizes_legacy_vectorize_table_column_role():
|
||||
p = ParserConfig(
|
||||
table_column_roles={"title": "vectorize", "country": "metadata", "x": "both"},
|
||||
)
|
||||
assert p.table_column_roles == {
|
||||
"title": "indexing",
|
||||
"country": "metadata",
|
||||
"x": "both",
|
||||
}
|
||||
0
test/unit_test/rag/app/__init__.py
Normal file
0
test/unit_test/rag/app/__init__.py
Normal file
235
test/unit_test/rag/app/test_table_chunk_column_roles.py
Normal file
235
test/unit_test/rag/app/test_table_chunk_column_roles.py
Normal file
@@ -0,0 +1,235 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
#
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License on an "AS IS" BASIS, WITHOUT WARRANTIES
|
||||
# OR CONDITIONS OF ANY KIND, either express or implied. See the License
|
||||
# for the specific language governing permissions and limitations under
|
||||
# the License.
|
||||
#
|
||||
|
||||
"""Integration-style tests for rag.app.table.chunk() column roles (mocked KB + tokenizer)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Mock heavy modules that trigger ONNX model loading at import time
|
||||
# table.py -> deepdoc.parser.figure_parser -> rag.app.picture -> OCR()
|
||||
for mod in [
|
||||
"deepdoc.vision.ocr",
|
||||
"deepdoc.parser.figure_parser",
|
||||
"rag.app.picture",
|
||||
]:
|
||||
if mod not in sys.modules:
|
||||
sys.modules[mod] = MagicMock()
|
||||
|
||||
import warnings
|
||||
|
||||
# Importing rag.app.table pulls api -> rag.llm -> deepdoc -> xgboost; xgboost may warn on
|
||||
# pkg_resources in a way that breaks its compat shim unless pkg_resources loads first.
|
||||
warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*", category=UserWarning)
|
||||
import pkg_resources # noqa: F401 — stabilize xgboost import during collection
|
||||
|
||||
import pytest
|
||||
|
||||
import common.settings as settings
|
||||
from rag.app.table import chunk
|
||||
|
||||
# chunk() removes columns named id, _id, index, idx — use row_id instead of id.
|
||||
TEST_CSV = b"""row_id,title,content,country,category
|
||||
1,Earthquake hits Turkey,A 5.8 magnitude earthquake struck Konya,Turkey,Disaster
|
||||
2,Oil prices surge,Brent crude jumped 4.2 percent,Global,Economy
|
||||
3,AI regulation proposed,EU unveiled a draft regulation,EU,Technology
|
||||
"""
|
||||
|
||||
FILENAME = "test.csv"
|
||||
KB_ID = "test_kb_id"
|
||||
|
||||
|
||||
def _noop_callback(*_a, **_k):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _es_doc_engine(monkeypatch):
|
||||
monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False)
|
||||
monkeypatch.setattr(settings, "DOC_ENGINE_OCEANBASE", False)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stub_rag_tokenizer(monkeypatch):
|
||||
"""Avoid NLTK / infinity tokenizer deps; keep string content inspectable."""
|
||||
|
||||
def fake_tokenize(line):
|
||||
return str(line)
|
||||
|
||||
monkeypatch.setattr("rag.nlp.rag_tokenizer.tokenize", fake_tokenize)
|
||||
monkeypatch.setattr("rag.nlp.rag_tokenizer.fine_grained_tokenize", fake_tokenize)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_update_kb():
|
||||
with patch("rag.app.table.KnowledgebaseService.update_parser_config") as m:
|
||||
yield m
|
||||
|
||||
|
||||
def _run_chunk(parser_config: dict, mock_update_kb: MagicMock):
|
||||
return chunk(
|
||||
FILENAME,
|
||||
binary=TEST_CSV,
|
||||
callback=_noop_callback,
|
||||
kb_id=KB_ID,
|
||||
parser_config=parser_config,
|
||||
lang="Chinese",
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_auto_mode_all_columns_in_text_and_stored(mock_update_kb: MagicMock):
|
||||
parser_config: dict = {}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
assert len(chunks) == 3
|
||||
first = chunks[0]
|
||||
cww = first["content_with_weight"]
|
||||
assert "Earthquake hits Turkey" in cww
|
||||
assert "Konya" in cww
|
||||
assert "Turkey" in cww
|
||||
assert "Disaster" in cww
|
||||
assert "1" in cww or "row_id" in cww
|
||||
# ES path: stored typed fields for text columns include *_tks and *_raw; row_id is int -> *_long
|
||||
assert "row_id_long" in first
|
||||
assert "title_raw" in first and "country_raw" in first
|
||||
|
||||
|
||||
def test_chunk_manual_mode_indexing_only(mock_update_kb: MagicMock):
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {
|
||||
"title": "indexing",
|
||||
"content": "indexing",
|
||||
"row_id": "metadata",
|
||||
"country": "metadata",
|
||||
"category": "metadata",
|
||||
},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
cww = first["content_with_weight"]
|
||||
assert "- title:" in cww and "Earthquake" in cww
|
||||
assert "- content:" in cww and "Konya" in cww
|
||||
assert "- country:" not in cww
|
||||
assert "- category:" not in cww
|
||||
assert "- row_id:" not in cww
|
||||
# Column title/content not stored as table fields
|
||||
assert "title_raw" not in first
|
||||
assert "content_raw" not in first
|
||||
assert "country_raw" in first and "category_raw" in first
|
||||
assert "row_id_long" in first
|
||||
|
||||
|
||||
def test_chunk_manual_mode_legacy_vectorize_role(mock_update_kb: MagicMock):
|
||||
"""Stored configs may still use role *vectorize*; chunking treats it like *indexing*."""
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {
|
||||
"title": "vectorize",
|
||||
"content": "indexing",
|
||||
"row_id": "metadata",
|
||||
"country": "metadata",
|
||||
"category": "metadata",
|
||||
},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
cww = first["content_with_weight"]
|
||||
assert "- title:" in cww and "Earthquake" in cww
|
||||
assert "- content:" in cww and "Konya" in cww
|
||||
assert "- country:" not in cww
|
||||
|
||||
|
||||
def test_chunk_manual_mode_metadata_only(mock_update_kb: MagicMock):
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {
|
||||
"title": "metadata",
|
||||
"content": "metadata",
|
||||
"row_id": "metadata",
|
||||
"country": "metadata",
|
||||
"category": "metadata",
|
||||
},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
assert (first.get("content_with_weight") or "").strip() == ""
|
||||
assert "country_raw" in first and "title_raw" in first
|
||||
|
||||
|
||||
def test_chunk_manual_mode_both(mock_update_kb: MagicMock):
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {c: "both" for c in ["title", "content", "country", "category", "row_id"]},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
cww = first["content_with_weight"]
|
||||
assert "Earthquake hits Turkey" in cww
|
||||
assert "Turkey" in cww
|
||||
assert "Disaster" in cww
|
||||
assert "row_id_long" in first
|
||||
assert "title_raw" in first and "country_raw" in first
|
||||
|
||||
|
||||
def test_chunk_manual_mode_partial_roles_default_to_both(mock_update_kb: MagicMock):
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {
|
||||
"title": "indexing",
|
||||
"country": "metadata",
|
||||
},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
cww = first["content_with_weight"]
|
||||
assert "- title:" in cww and "Earthquake" in cww
|
||||
assert "- country:" not in cww
|
||||
assert "- row_id:" in cww
|
||||
assert "- content:" in cww
|
||||
assert "- category:" in cww
|
||||
assert "title_raw" not in first
|
||||
assert "country_raw" in first and "country_tks" in first
|
||||
assert "content_raw" in first and "category_raw" in first
|
||||
|
||||
|
||||
def test_chunk_manual_mode_raw_fields_for_es(mock_update_kb: MagicMock):
|
||||
parser_config = {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {c: "both" for c in ["title", "content", "country", "category", "row_id"]},
|
||||
}
|
||||
chunks = _run_chunk(parser_config, mock_update_kb)
|
||||
first = chunks[0]
|
||||
for col in ("title", "content", "country", "category"):
|
||||
assert f"{col}_raw" in first
|
||||
assert f"{col}_tks" in first
|
||||
|
||||
|
||||
def test_chunk_updates_table_column_names(mock_update_kb: MagicMock):
|
||||
_run_chunk({}, mock_update_kb)
|
||||
mock_update_kb.assert_called_once()
|
||||
args, kwargs = mock_update_kb.call_args
|
||||
assert args[0] == KB_ID
|
||||
payload = args[1]
|
||||
names = payload["table_column_names"]
|
||||
assert names == ["row_id", "title", "content", "country", "category"]
|
||||
|
||||
|
||||
def test_chunk_count_matches_row_count(mock_update_kb: MagicMock):
|
||||
chunks = _run_chunk({}, mock_update_kb)
|
||||
assert len(chunks) == 3
|
||||
1
test/unit_test/rag/svr/__init__.py
Normal file
1
test/unit_test/rag/svr/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# Unit tests for rag/svr
|
||||
132
test/unit_test/rag/svr/test_table_column_roles_helpers.py
Normal file
132
test/unit_test/rag/svr/test_table_column_roles_helpers.py
Normal file
@@ -0,0 +1,132 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
"""Unit tests for ES table metadata helpers (rag.utils.table_es_metadata)."""
|
||||
|
||||
from rag.utils.table_es_metadata import (
|
||||
_es_field_value_to_doc_metadata,
|
||||
_es_raw_field_key_from_typed,
|
||||
_probe_es_typed_key_for_column,
|
||||
_resolve_es_chunk_field_key,
|
||||
merge_table_parser_config_from_kb,
|
||||
table_parser_strip_doc_metadata_keys,
|
||||
)
|
||||
|
||||
|
||||
class TestProbeEsTypedKeyForColumn:
|
||||
def test_probe_es_typed_key_tks(self):
|
||||
chunk = {"country_tks": "tok", "other": 1}
|
||||
assert _probe_es_typed_key_for_column("country", chunk) == "country_tks"
|
||||
|
||||
def test_probe_es_typed_key_dt(self):
|
||||
chunk = {"published_date_dt": "2024-01-01"}
|
||||
assert _probe_es_typed_key_for_column("published_date", chunk) == "published_date_dt"
|
||||
|
||||
def test_probe_es_typed_key_raw(self):
|
||||
# Only raw field present (no _tks) — probe returns the raw key
|
||||
chunk = {"country_raw": "Brazil"}
|
||||
assert _probe_es_typed_key_for_column("country", chunk) == "country_raw"
|
||||
|
||||
def test_probe_es_typed_key_no_match(self):
|
||||
chunk = {"other_kwd": "x"}
|
||||
assert _probe_es_typed_key_for_column("country", chunk) is None
|
||||
|
||||
def test_probe_es_typed_key_empty_col(self):
|
||||
assert _probe_es_typed_key_for_column("", {"a_tks": "x"}) is None
|
||||
assert _probe_es_typed_key_for_column(None, {"a_tks": "x"}) is None
|
||||
|
||||
|
||||
class TestResolveEsChunkFieldKey:
|
||||
def test_resolve_es_field_empty_fieldmap_uses_probe(self):
|
||||
sample = {"country_tks": ["tok"]}
|
||||
tk, src = _resolve_es_chunk_field_key("country", {}, sample)
|
||||
assert tk == "country_tks"
|
||||
assert src == "probe"
|
||||
|
||||
def test_resolve_es_field_fieldmap_priority(self):
|
||||
fm = {"guojia_tks": "country"}
|
||||
sample = {"guojia_tks": ["x"], "country_tks": ["y"]}
|
||||
tk, src = _resolve_es_chunk_field_key("country", fm, sample)
|
||||
assert tk == "guojia_tks"
|
||||
assert src == "field_map"
|
||||
|
||||
|
||||
class TestEsRawFieldKeyFromTyped:
|
||||
def test_es_raw_field_key_from_tks(self):
|
||||
assert _es_raw_field_key_from_typed("country_tks") == "country_raw"
|
||||
|
||||
def test_es_raw_field_key_from_non_tks(self):
|
||||
assert _es_raw_field_key_from_typed("country_dt") is None
|
||||
|
||||
def test_es_raw_field_key_from_none(self):
|
||||
assert _es_raw_field_key_from_typed(None) is None
|
||||
|
||||
|
||||
class TestEsFieldValueToDocMetadata:
|
||||
def test_es_field_value_string(self):
|
||||
assert _es_field_value_to_doc_metadata("Brazil", from_tks_fallback=False) == "Brazil"
|
||||
|
||||
def test_es_field_value_list_joined(self):
|
||||
assert (
|
||||
_es_field_value_to_doc_metadata(["hello", "world"], from_tks_fallback=True)
|
||||
== "hello world"
|
||||
)
|
||||
|
||||
def test_es_field_value_empty(self):
|
||||
assert _es_field_value_to_doc_metadata(None, from_tks_fallback=True) is None
|
||||
assert _es_field_value_to_doc_metadata("", from_tks_fallback=True) is None
|
||||
assert _es_field_value_to_doc_metadata([], from_tks_fallback=True) is None
|
||||
|
||||
|
||||
class TestMergeTableParserConfigFromKb:
|
||||
def test_merge_table_parser_config_from_kb(self):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {"llm_id": "x"},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"a": "metadata"},
|
||||
"table_column_names": ["a", "b"],
|
||||
},
|
||||
}
|
||||
merged = merge_table_parser_config_from_kb(task)
|
||||
assert merged["table_column_mode"] == "manual"
|
||||
assert merged["table_column_roles"] == {"a": "metadata"}
|
||||
assert merged["table_column_names"] == ["a", "b"]
|
||||
assert merged["llm_id"] == "x"
|
||||
|
||||
def test_merge_table_parser_config_auto_default(self):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {"foo": 1},
|
||||
"kb_parser_config": {"llm_id": "abc"},
|
||||
}
|
||||
merged = merge_table_parser_config_from_kb(task)
|
||||
assert merged == {"foo": 1} # no table_* keys copied from kb without kb_parser_config keys
|
||||
|
||||
|
||||
class TestTableParserStripDocMetadataKeys:
|
||||
def test_uses_table_column_names_when_present(self):
|
||||
eff = {"table_column_names": ["Region", " SKU "]}
|
||||
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"Region", "SKU"})
|
||||
|
||||
def test_falls_back_to_role_keys_when_no_names(self):
|
||||
eff = {"table_column_roles": {"x": "metadata", "y": "indexing"}}
|
||||
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"x", "y"})
|
||||
|
||||
def test_empty_names_falls_back_to_roles(self):
|
||||
eff = {"table_column_names": [], "table_column_roles": {"only": "both"}}
|
||||
assert table_parser_strip_doc_metadata_keys(eff) == frozenset({"only"})
|
||||
230
test/unit_test/rag/svr/test_table_metadata_aggregation.py
Normal file
230
test/unit_test/rag/svr/test_table_metadata_aggregation.py
Normal file
@@ -0,0 +1,230 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
"""Unit tests for aggregate_table_manual_doc_metadata."""
|
||||
|
||||
import pytest
|
||||
|
||||
from rag.utils.table_es_metadata import aggregate_table_manual_doc_metadata, merge_table_parser_config_from_kb
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def es_engine(monkeypatch):
|
||||
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_INFINITY", False)
|
||||
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_OCEANBASE", False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def infinity_engine(monkeypatch):
|
||||
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_INFINITY", True)
|
||||
monkeypatch.setattr("rag.utils.table_es_metadata.settings.DOC_ENGINE_OCEANBASE", False)
|
||||
|
||||
|
||||
def _table_task(**kb_extra):
|
||||
return {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"country": "metadata", "category": "metadata"},
|
||||
"table_column_names": ["country", "category"],
|
||||
"field_map": {
|
||||
"country_tks": "country",
|
||||
"category_tks": "category",
|
||||
},
|
||||
**kb_extra,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestAggregateTableManualDocMetadata:
|
||||
def test_aggregate_manual_mode_happy_path(self, es_engine):
|
||||
task = _table_task()
|
||||
chunks = [
|
||||
{
|
||||
"country_raw": "Brazil",
|
||||
"category_raw": "Economy",
|
||||
"country_tks": "x",
|
||||
"category_tks": "y",
|
||||
},
|
||||
{
|
||||
"country_raw": "Turkey",
|
||||
"category_raw": "Disaster",
|
||||
"country_tks": "x",
|
||||
"category_tks": "y",
|
||||
},
|
||||
{
|
||||
"country_raw": "Brazil",
|
||||
"category_raw": "Economy",
|
||||
"country_tks": "x",
|
||||
"category_tks": "y",
|
||||
},
|
||||
]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out["country"] == ["Brazil", "Turkey"]
|
||||
assert out["category"] == ["Economy", "Disaster"]
|
||||
|
||||
def test_aggregate_auto_mode_returns_empty(self, es_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "auto",
|
||||
"table_column_roles": {"country": "metadata"},
|
||||
},
|
||||
}
|
||||
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
|
||||
|
||||
def test_aggregate_no_mode_returns_empty(self, es_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_roles": {"country": "metadata"},
|
||||
},
|
||||
}
|
||||
assert aggregate_table_manual_doc_metadata([{}], task) == {}
|
||||
|
||||
def test_aggregate_no_metadata_columns(self, es_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"country": "indexing"},
|
||||
"table_column_names": ["country"],
|
||||
},
|
||||
}
|
||||
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
|
||||
|
||||
def test_aggregate_prefers_raw_over_tks(self, es_engine):
|
||||
task = _table_task()
|
||||
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
|
||||
task["kb_parser_config"]["table_column_names"] = ["country"]
|
||||
chunks = [{"country_raw": "Brazil", "country_tks": ["brazil"]}]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out == {"country": ["Brazil"]}
|
||||
|
||||
def test_aggregate_tks_fallback(self, es_engine):
|
||||
task = _table_task()
|
||||
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
|
||||
task["kb_parser_config"]["table_column_names"] = ["country"]
|
||||
chunks = [{"country_tks": ["brazil"]}]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out == {"country": ["brazil"]}
|
||||
|
||||
def test_aggregate_partial_roles_defaults_to_both(self, es_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"country": "indexing"},
|
||||
"table_column_names": ["country", "city"],
|
||||
"field_map": {"city_tks": "city"},
|
||||
},
|
||||
}
|
||||
chunks = [{"city_raw": "SP", "city_tks": "t", "country_tks": "x"}]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out == {"city": ["SP"]}
|
||||
assert "country" not in out
|
||||
|
||||
def test_aggregate_empty_roles_all_columns_both(self, es_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {},
|
||||
"table_column_names": ["country", "city"],
|
||||
"field_map": {"country_tks": "country", "city_tks": "city"},
|
||||
},
|
||||
}
|
||||
chunks = [
|
||||
{"country_raw": "BR", "city_raw": "SP", "country_tks": "x", "city_tks": "y"},
|
||||
]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert "country" in out and "city" in out
|
||||
|
||||
def test_aggregate_deduplicates_values(self, es_engine):
|
||||
task = _table_task()
|
||||
task["kb_parser_config"]["table_column_roles"] = {"country": "metadata"}
|
||||
task["kb_parser_config"]["table_column_names"] = ["country"]
|
||||
chunks = [
|
||||
{"country_raw": "US", "country_tks": "x"},
|
||||
{"country_raw": "UK", "country_tks": "y"},
|
||||
{"country_raw": "US", "country_tks": "x"},
|
||||
]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out["country"] == ["US", "UK"]
|
||||
|
||||
def test_aggregate_kb_reload_field_map(self, es_engine, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
class MockKBS:
|
||||
@staticmethod
|
||||
def get_by_id(kid):
|
||||
kb = MagicMock()
|
||||
kb.parser_config = {"field_map": {"country_tks": "country"}}
|
||||
return True, kb
|
||||
|
||||
monkeypatch.setattr(
|
||||
"rag.utils.table_es_metadata._knowledgebase_service_cls",
|
||||
lambda: MockKBS,
|
||||
)
|
||||
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"country": "metadata"},
|
||||
"table_column_names": ["country"],
|
||||
},
|
||||
"kb_id": "kb-1",
|
||||
}
|
||||
chunks = [{"country_raw": "X", "country_tks": "t"}]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out == {"country": ["X"]}
|
||||
|
||||
def test_merge_infinity_chunk_data(self, infinity_engine):
|
||||
task = {
|
||||
"parser_id": "table",
|
||||
"parser_config": {},
|
||||
"kb_parser_config": {
|
||||
"table_column_mode": "manual",
|
||||
"table_column_roles": {"country": "both"},
|
||||
"table_column_names": ["country"],
|
||||
},
|
||||
}
|
||||
chunks = [
|
||||
{"chunk_data": {"country": "US"}},
|
||||
{"chunk_data": {"country": "UK"}},
|
||||
]
|
||||
out = aggregate_table_manual_doc_metadata(chunks, task)
|
||||
assert out == {"country": ["US", "UK"]}
|
||||
|
||||
|
||||
class TestMergeTableParserConfigFromKbExtra:
|
||||
"""Merge tests also covered in helpers file; keep one explicit case for aggregation module."""
|
||||
|
||||
def test_merge_preserves_parser_config_when_parser_not_table(self):
|
||||
task = {
|
||||
"parser_id": "naive",
|
||||
"parser_config": {"a": 1},
|
||||
"kb_parser_config": {"table_column_mode": "manual"},
|
||||
}
|
||||
assert merge_table_parser_config_from_kb(task) == {"a": 1}
|
||||
Reference in New Issue
Block a user