fix: support auto mode in table parser document metadata aggregation (#15780)

### What problem does this PR solve?

Table parser metadata aggregation previously only ran when
`table_column_mode` was set to `manual`. In auto mode (default), all
columns default to `"both"` role, meaning they should also be aggregated
into document-level metadata for UI/chat filters. Additionally, the task
snapshot could be stale — `table_column_names` are written to KB
`parser_config` during `chunk()` but the task may have been created
before that.

Changes:
- Renames `aggregate_table_manual_doc_metadata` →
`aggregate_table_doc_metadata`
- Supports both `"manual"` and `"auto"` `table_column_mode` (defaults to
`"auto"`)
- Reloads `table_column_names` from KB DB when missing from task
snapshot
- Removes the manual-only guard in `task_executor` and refactored
`post_processor`
- Updates all tests with new function name and adds auto mode test cases

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
euvre
2026-06-08 04:08:23 -07:00
committed by GitHub
parent 2c64febc93
commit d9a04ef702
5 changed files with 277 additions and 289 deletions

View File

@@ -69,12 +69,13 @@ class TestPostProcessorProcessTableParserMetadata:
ctx.write_interceptor = None
chunks = [{"col_key": "val"}]
with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \
patch("rag.svr.task_executor_refactor.post_processor.update_metadata_to") as mock_update, \
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta:
with (
patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge,
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_doc_metadata") as mock_agg,
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip,
patch("rag.svr.task_executor_refactor.post_processor.update_metadata_to") as mock_update,
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta,
):
mock_merge.return_value = {"table_column_mode": "manual"}
mock_agg.return_value = {"col_key": ["val1", "val2"]}
mock_strip.return_value = set()
@@ -95,11 +96,12 @@ class TestPostProcessorProcessTableParserMetadata:
ctx.raw_task = {}
ctx.write_interceptor = MagicMock()
with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta:
with (
patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge,
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_doc_metadata") as mock_agg,
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip,
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta,
):
mock_merge.return_value = {"table_column_mode": "manual"}
mock_agg.return_value = {"key": ["v"]}
mock_strip.return_value = set()
@@ -108,9 +110,7 @@ class TestPostProcessorProcessTableParserMetadata:
service = PostProcessor(ctx=ctx)
await service.process_table_parser_metadata("doc_1", [])
ctx.write_interceptor.intercept.assert_called_once_with(
"DocMetadataService.update_document_metadata"
)
ctx.write_interceptor.intercept.assert_called_once_with("DocMetadataService.update_document_metadata")
class TestPostProcessorInsertTocChunk:
@@ -160,9 +160,7 @@ class TestPostProcessorInsertTocChunk:
result = await service.insert_toc_chunk(toc_chunk, chunk_service)
assert result is True
chunk_service.insert_chunks.assert_called_once_with(
"task_1", "tenant_1", "kb_1", [toc_chunk]
)
chunk_service.insert_chunks.assert_called_once_with("task_1", "tenant_1", "kb_1", [toc_chunk])
@pytest.mark.asyncio
async def test_handles_insert_failure(self):
@@ -179,4 +177,4 @@ class TestPostProcessorInsertTocChunk:
result = await service.insert_toc_chunk(toc_chunk, chunk_service)
assert result is False
assert result is False

View File

@@ -14,11 +14,11 @@
# limitations under the License.
#
"""Unit tests for aggregate_table_manual_doc_metadata."""
"""Unit tests for aggregate_table_doc_metadata."""
import pytest
from rag.utils.table_es_metadata import aggregate_table_manual_doc_metadata, merge_table_parser_config_from_kb
from rag.utils.table_es_metadata import aggregate_table_doc_metadata, merge_table_parser_config_from_kb
@pytest.fixture
@@ -73,30 +73,64 @@ class TestAggregateTableManualDocMetadata:
"category_tks": "y",
},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert out["country"] == ["Brazil", "Turkey"]
assert out["category"] == ["Economy", "Disaster"]
def test_aggregate_auto_mode_returns_empty(self, es_engine):
def test_aggregate_auto_mode_returns_data(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_mode": "auto",
"table_column_roles": {"country": "metadata"},
"table_column_names": ["country"],
"field_map": {"country_tks": "country"},
},
}
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
chunks = [{"country_raw": "Brazil", "country_tks": "x"}]
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["Brazil"]}
def test_aggregate_no_mode_returns_empty(self, es_engine):
def test_aggregate_auto_mode_all_columns_both(self, es_engine):
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_roles": {"country": "metadata"},
"table_column_mode": "auto",
"table_column_names": ["country", "category"],
"field_map": {"country_tks": "country", "category_tks": "category"},
},
}
assert aggregate_table_manual_doc_metadata([{}], task) == {}
chunks = [
{"country_raw": "Brazil", "country_tks": "x", "category_raw": "Economy", "category_tks": "y"},
{"country_raw": "Turkey", "country_tks": "x", "category_raw": "Disaster", "category_tks": "y"},
]
out = aggregate_table_doc_metadata(chunks, task)
assert out["country"] == ["Brazil", "Turkey"]
assert out["category"] == ["Economy", "Disaster"]
def test_aggregate_no_mode_defaults_to_auto(self, es_engine):
"""When table_column_mode is missing, it defaults to 'auto'."""
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {
"table_column_names": ["country"],
"field_map": {"country_tks": "country"},
},
}
chunks = [{"country_raw": "Brazil", "country_tks": "x"}]
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["Brazil"]}
def test_aggregate_no_mode_no_columns_returns_empty(self, es_engine):
"""No mode and no column names/roles -> empty (nothing to aggregate)."""
task = {
"parser_id": "table",
"parser_config": {},
"kb_parser_config": {},
}
assert aggregate_table_doc_metadata([{}], task) == {}
def test_aggregate_no_metadata_columns(self, es_engine):
task = {
@@ -108,14 +142,14 @@ class TestAggregateTableManualDocMetadata:
"table_column_names": ["country"],
},
}
assert aggregate_table_manual_doc_metadata([{"country_tks": "x"}], task) == {}
assert aggregate_table_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)
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["Brazil"]}
def test_aggregate_tks_fallback(self, es_engine):
@@ -123,7 +157,7 @@ class TestAggregateTableManualDocMetadata:
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)
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["brazil"]}
def test_aggregate_partial_roles_defaults_to_both(self, es_engine):
@@ -138,7 +172,7 @@ class TestAggregateTableManualDocMetadata:
},
}
chunks = [{"city_raw": "SP", "city_tks": "t", "country_tks": "x"}]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"city": ["SP"]}
assert "country" not in out
@@ -156,7 +190,7 @@ class TestAggregateTableManualDocMetadata:
chunks = [
{"country_raw": "BR", "city_raw": "SP", "country_tks": "x", "city_tks": "y"},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert "country" in out and "city" in out
def test_aggregate_deduplicates_values(self, es_engine):
@@ -168,7 +202,7 @@ class TestAggregateTableManualDocMetadata:
{"country_raw": "UK", "country_tks": "y"},
{"country_raw": "US", "country_tks": "x"},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert out["country"] == ["US", "UK"]
def test_aggregate_kb_reload_field_map(self, es_engine, monkeypatch):
@@ -197,7 +231,7 @@ class TestAggregateTableManualDocMetadata:
"kb_id": "kb-1",
}
chunks = [{"country_raw": "X", "country_tks": "t"}]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["X"]}
def test_merge_infinity_chunk_data(self, infinity_engine):
@@ -214,7 +248,7 @@ class TestAggregateTableManualDocMetadata:
{"chunk_data": {"country": "US"}},
{"chunk_data": {"country": "UK"}},
]
out = aggregate_table_manual_doc_metadata(chunks, task)
out = aggregate_table_doc_metadata(chunks, task)
assert out == {"country": ["US", "UK"]}