feat(indexing): implement topic tagging pipeline - #528
Merged
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughImplements end-to-end document-level topic tagging for the indexing pipeline. Adds a ChangesTopic Tagging Feature
Sequence Diagram(s)sequenceDiagram
participant IndexerPool
participant IndexingPipeline
participant topic_tag_stage
participant TopicTagger
participant LLM
participant IndexerWorker
participant PgTopicTagRepository
rect rgba(70, 130, 180, 0.5)
note over IndexerPool: Build phase
IndexerPool->>IndexerPool: _build_topic_tagger_factory(cfg)
IndexerPool->>IndexingPipeline: build_indexing_pipeline(topic_tagger_factory=...)
IndexerPool->>IndexerWorker: __init__(topic_tag_repo=pg_repo)
end
rect rgba(60, 179, 113, 0.5)
note over IndexerWorker: process_file per document
IndexerWorker->>IndexingPipeline: run(row, config)
IndexingPipeline->>topic_tag_stage: row["chunks"]
topic_tag_stage->>TopicTagger: tag(chunks, filename, max_tags)
TopicTagger->>LLM: chat(system+user messages)
LLM-->>TopicTagger: JSON array response
TopicTagger-->>topic_tag_stage: list[str] normalized tags
topic_tag_stage-->>IndexingPipeline: row["topic_tags"] populated
IndexingPipeline-->>IndexerWorker: completed row
IndexerWorker->>PgTopicTagRepository: delete_by_document(document_id, partition)
IndexerWorker->>PgTopicTagRepository: bulk_insert(tag_records)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
openrag/services/persistence/schema.py (1)
296-307: 💤 Low valueConsider adding
topic_tagsto__all__for consistency.Other tables like
files,users, andpartition_membershipsare exported via__all__. Omittingtopic_tagscreates an inconsistency, though it doesn't break functionality since the table can still be imported directly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/schema.py` around lines 296 - 307, The `topic_tags` table is missing from the `__all__` list in the schema.py file, creating an inconsistency with other exported tables like `files`, `users`, and `partition_memberships`. Add `topic_tags` to the `__all__` list to ensure it is properly exported alongside the other table definitions, maintaining consistency in the module's public API.openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py (1)
22-63: 💤 Low valueRedundant unique constraint creation block.
The
UniqueConstraintis already defined withincreate_table(lines 43-48), so the separate check and creation at lines 57-62 will always find the constraint already exists. This block is a no-op in normal operation and can be removed for clarity.♻️ Proposed simplification
if not index_exists("topic_tags", "ix_topic_tags_partition_tag"): op.create_index("ix_topic_tags_partition_tag", "topic_tags", ["partition", "normalized_tag"]) - if not unique_constraint_exists("topic_tags", "uix_topic_tags_document_partition_tag"): - op.create_unique_constraint( - "uix_topic_tags_document_partition_tag", - "topic_tags", - ["document_id", "partition", "normalized_tag"], - )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py` around lines 22 - 63, The upgrade function contains a redundant unique constraint creation block. The UniqueConstraint with name uix_topic_tags_document_partition_tag is already defined within the op.create_table call for the topic_tags table, so the separate check using unique_constraint_exists and the subsequent op.create_unique_constraint call at the end of the upgrade function is unnecessary and will always be a no-op. Remove the entire block that checks if the unique constraint exists and attempts to create it, as the constraint will be automatically created as part of the table definition.tests/unit/services/persistence/test_topic_tag_repo.py (1)
75-88: ⚡ Quick winConsider adding test coverage for
delete_by_document.The test file covers
bulk_insert,get_by_document, andsearch, butdelete_by_documentis untested. A simple test verifying the query structure and_delete_countparsing would complete the coverage.💚 Proposed test
`@pytest.mark.asyncio` async def test_delete_by_document_returns_affected_count(): from services.persistence.topic_tag_repo import PgTopicTagRepository pool = _FakePool() repo = PgTopicTagRepository(pool_getter=lambda: pool) count = await repo.delete_by_document("file-1", partition="tenant-a") assert count == 2 # _FakePool.execute returns "DELETE 2" query, params = pool.executed[0] assert "DELETE FROM topic_tags" in query assert "WHERE document_id = $1 AND partition = $2" in query assert params == ("file-1", "tenant-a")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/services/persistence/test_topic_tag_repo.py` around lines 75 - 88, The test file is missing coverage for the delete_by_document method of PgTopicTagRepository. Add a new test function named test_delete_by_document_returns_affected_count that creates a _FakePool instance, instantiates PgTopicTagRepository with the pool, calls the delete_by_document method with parameters like "file-1" and partition="tenant-a", and then verifies the executed query contains the DELETE FROM statement, has the correct WHERE clause with document_id and partition conditions, and that the parameters are passed in the correct order as the method expects.openrag/services/persistence/topic_tag_repo.py (1)
145-152: 💤 Low valueDuplicate normalization logic across layers.
_normalize_display_taghere duplicates similar logic inopenrag/core/indexing/topic_tags.py. Both collapse whitespace and truncate to 80 characters. If the normalization rules evolve, both locations must be updated in sync.Consider extracting a shared helper to
openrag/core/utils/if this becomes a maintenance concern, though the current duplication is acceptable given the layer separation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/topic_tag_repo.py` around lines 145 - 152, The normalization logic in _normalize_display_tag (collapsing whitespace and truncating to 80 characters) is duplicated in openrag/core/indexing/topic_tags.py. If maintaining synchronized normalization rules across both locations becomes problematic during future changes, extract a shared helper function to openrag/core/utils/ that both _normalize_display_tag in topic_tag_repo.py and the similar logic in topic_tags.py can import and use, ensuring consistent behavior across layers.tests/unit/core/indexing/test_topic_tags.py (1)
16-39: ⚡ Quick winAdd regression tests for timeout-zero and bracket-suffix parsing.
Please add focused cases for:
timeout_seconds=0still using timeout path, and- parsing output like
["a"] [Sources: 1]to keep["a"].🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/core/indexing/test_topic_tags.py` around lines 16 - 39, Add two new test functions to the file following the same pattern as the existing tests. First, create a test for the timeout_seconds=0 edge case that verifies the TopicTagger still applies timeout handling even when timeout_seconds is set to zero. Second, create a test for parsing responses with bracket suffixes where the TopicTagger should correctly extract only the JSON array portion and ignore trailing content like "[Sources: 1]", using FakeLLM to simulate an LLM response in the format ["a"] [Sources: 1] and verify the tagger returns only ["a"].openrag/core/ports/topic_tag_repo.py (1)
15-18: ⚡ Quick winMake
partitionrequired on document-scoped read/delete methods.Line 15 and Line 18 currently allow
partition=None, which makes accidental cross-partition reads/deletes easier in future callers. Given the rest of this feature is partition-scoped, requiringpartitionin the port is safer.Suggested contract tightening
- async def get_by_document(self, document_id: str, partition: str | None = None) -> list[dict]: ... + async def get_by_document(self, document_id: str, partition: str) -> list[dict]: ... - async def delete_by_document(self, document_id: str, partition: str | None = None) -> int: ... + async def delete_by_document(self, document_id: str, partition: str) -> int: ...🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/core/ports/topic_tag_repo.py` around lines 15 - 18, The `get_by_document` and `delete_by_document` methods in TopicTagRepo currently have optional partition parameters with `partition: str | None = None`, which increases the risk of accidental cross-partition operations. Remove the `| None` union type and the `= None` default value from the partition parameter in both methods to make partition a required argument, ensuring all callers must explicitly specify which partition to operate on.tests/unit/services/workers/test_indexer_worker.py (1)
408-439: ⚡ Quick winAdd a branch test for “disabled topic tagging triggers delete-only cleanup.”
Please add a focused case where
indexation_config={"enable_topic_tagging": False}and norow["topic_tags"]is produced, then assertdelete_by_documentis called andbulk_insertis not. This protects the stale-tag cleanup contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/services/workers/test_indexer_worker.py` around lines 408 - 439, Add a new test function after test_process_file_replaces_topic_tags_after_successful_pipeline that validates the delete-only cleanup behavior when topic tagging is disabled. Create a test case where IndexerWorker is initialized with indexation_config set to disable topic tagging (enable_topic_tagging False), use a TaggingPipeline that does not produce topic_tags in the row output, call worker.process_file with the same parameters, and then assert that the FakeTopicTagRepo's deleted list contains the document cleanup entry while the inserted list remains empty, verifying that stale tags are removed when topic tagging is disabled.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/core/indexing/topic_tags.py`:
- Around line 33-45: The tag method needs an early guard clause to handle
non-positive max_tags values. Currently, when max_tags is 0 or negative, the
function can still return one tag because validation happens after appending at
line 104. Add a check right after the chunks list conversion and empty chunks
check in the tag method that returns an empty list if max_tags is less than or
equal to 0, ensuring the function returns deterministically for invalid max_tags
values before any tagging logic executes.
- Around line 113-118: The regex pattern in the tag extraction logic uses a
greedy quantifier that matches from the first opening bracket to the last
closing bracket in the text, causing multi-bracket responses like ["finance"]
[Sources: 1] to be treated as a single invalid JSON string. Fix this by changing
the greedy quantifier in the regex pattern from `*` to `*?` in the
`re.search(r"\[[\s\S]*\]", text)` call to make it non-greedy, so it will match
only the shortest bracket-enclosed segment and correctly extract valid JSON tags
instead of discarding them.
- Line 55: The timeout check on line 55 uses a truthy evaluation (`if
self._timeout`) which treats 0 as falsy and skips the asyncio.wait_for timeout
enforcement entirely. Replace the truthy check with an explicit None check by
changing `if self._timeout` to `if self._timeout is not None` in the conditional
expression. This ensures that a timeout_seconds value of 0 is properly
recognized as a valid timeout configuration and the timeout enforcement is
applied correctly.
In `@openrag/services/workers/indexer_actor.py`:
- Around line 152-154: The code at line 153 iterates over row.get("topic_tags",
[]) without validating its type first, and the delete_by_document call happens
at line 152 before any validation occurs. This means if row["topic_tags"] is a
string instead of a list, it will iterate character-by-character creating
invalid tags, and if an error occurs during iteration, the existing tags are
already deleted with no rollback. Add type validation for row["topic_tags"]
before the delete_by_document call on line 152 to ensure it is a list, either by
checking isinstance(row.get("topic_tags"), list) or coercing the value to a list
type, so that validation happens before any destructive operations.
---
Nitpick comments:
In `@openrag/core/ports/topic_tag_repo.py`:
- Around line 15-18: The `get_by_document` and `delete_by_document` methods in
TopicTagRepo currently have optional partition parameters with `partition: str |
None = None`, which increases the risk of accidental cross-partition operations.
Remove the `| None` union type and the `= None` default value from the partition
parameter in both methods to make partition a required argument, ensuring all
callers must explicitly specify which partition to operate on.
In
`@openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py`:
- Around line 22-63: The upgrade function contains a redundant unique constraint
creation block. The UniqueConstraint with name
uix_topic_tags_document_partition_tag is already defined within the
op.create_table call for the topic_tags table, so the separate check using
unique_constraint_exists and the subsequent op.create_unique_constraint call at
the end of the upgrade function is unnecessary and will always be a no-op.
Remove the entire block that checks if the unique constraint exists and attempts
to create it, as the constraint will be automatically created as part of the
table definition.
In `@openrag/services/persistence/schema.py`:
- Around line 296-307: The `topic_tags` table is missing from the `__all__` list
in the schema.py file, creating an inconsistency with other exported tables like
`files`, `users`, and `partition_memberships`. Add `topic_tags` to the `__all__`
list to ensure it is properly exported alongside the other table definitions,
maintaining consistency in the module's public API.
In `@openrag/services/persistence/topic_tag_repo.py`:
- Around line 145-152: The normalization logic in _normalize_display_tag
(collapsing whitespace and truncating to 80 characters) is duplicated in
openrag/core/indexing/topic_tags.py. If maintaining synchronized normalization
rules across both locations becomes problematic during future changes, extract a
shared helper function to openrag/core/utils/ that both _normalize_display_tag
in topic_tag_repo.py and the similar logic in topic_tags.py can import and use,
ensuring consistent behavior across layers.
In `@tests/unit/core/indexing/test_topic_tags.py`:
- Around line 16-39: Add two new test functions to the file following the same
pattern as the existing tests. First, create a test for the timeout_seconds=0
edge case that verifies the TopicTagger still applies timeout handling even when
timeout_seconds is set to zero. Second, create a test for parsing responses with
bracket suffixes where the TopicTagger should correctly extract only the JSON
array portion and ignore trailing content like "[Sources: 1]", using FakeLLM to
simulate an LLM response in the format ["a"] [Sources: 1] and verify the tagger
returns only ["a"].
In `@tests/unit/services/persistence/test_topic_tag_repo.py`:
- Around line 75-88: The test file is missing coverage for the
delete_by_document method of PgTopicTagRepository. Add a new test function named
test_delete_by_document_returns_affected_count that creates a _FakePool
instance, instantiates PgTopicTagRepository with the pool, calls the
delete_by_document method with parameters like "file-1" and
partition="tenant-a", and then verifies the executed query contains the DELETE
FROM statement, has the correct WHERE clause with document_id and partition
conditions, and that the parameters are passed in the correct order as the
method expects.
In `@tests/unit/services/workers/test_indexer_worker.py`:
- Around line 408-439: Add a new test function after
test_process_file_replaces_topic_tags_after_successful_pipeline that validates
the delete-only cleanup behavior when topic tagging is disabled. Create a test
case where IndexerWorker is initialized with indexation_config set to disable
topic tagging (enable_topic_tagging False), use a TaggingPipeline that does not
produce topic_tags in the row output, call worker.process_file with the same
parameters, and then assert that the FakeTopicTagRepo's deleted list contains
the document cleanup entry while the inserted list remains empty, verifying that
stale tags are removed when topic tagging is disabled.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 0b31e42f-fafd-43ef-b9ff-b6ecb6313495
📒 Files selected for processing (17)
openrag/core/config/infrastructure.pyopenrag/core/indexing/topic_tags.pyopenrag/core/models/prompt.pyopenrag/core/ports/topic_tag_repo.pyopenrag/prompts/templates/topic_tagger_tmpl.txtopenrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.pyopenrag/services/persistence/schema.pyopenrag/services/persistence/topic_tag_repo.pyopenrag/services/workers/indexer_actor.pyopenrag/services/workers/indexer_pool.pyopenrag/services/workers/pipeline_builder.pyopenrag/services/workers/stages/topic_tag.pytests/unit/core/indexing/test_topic_tags.pytests/unit/services/persistence/test_topic_tag_repo.pytests/unit/services/workers/test_indexer_pool.pytests/unit/services/workers/test_indexer_worker.pytests/unit/services/workers/test_pipeline_builder.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/services/workers/test_indexer_pool.py (1)
286-290: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAssert
ray.get_actornamespace in this wiring test.The stub currently ignores call arguments, so a regression dropping
namespace="openrag"would still pass.Suggested test hardening
- monkeypatch.setattr(module.ray, "get_actor", lambda *args, **kwargs: object()) + actor_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + def _fake_get_actor(*args: Any, **kwargs: Any) -> object: + actor_calls.append((args, kwargs)) + return object() + monkeypatch.setattr(module.ray, "get_actor", _fake_get_actor) monkeypatch.setattr(module, "IndexerWorker", Worker) actor_class = module.IndexerPool.__ray_metadata__.modified_class actor_class() + assert actor_calls + assert actor_calls[0][0][0] == "TaskStateManager" + assert actor_calls[0][1].get("namespace") == "openrag" assert captured["contextualizer_factory"] is contextualizer_factory assert captured["topic_tagger_factory"] is topic_tagger_factoryAs per coding guidelines, Ray actors should be accessed via
ray.get_actor()using namespace"openrag".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/services/workers/test_indexer_pool.py` around lines 286 - 290, The monkeypatch for module.ray.get_actor currently uses a lambda that ignores all arguments, making it impossible to verify that the correct namespace parameter is being passed. Replace the lambda mock with a proper mock object (using unittest.mock.MagicMock or similar) that can capture call arguments, then add an assertion after the actor_class() call to verify that ray.get_actor was called with the namespace parameter set to "openrag". This ensures that any future regression where the namespace argument is accidentally removed from the ray.get_actor call will be caught by the test.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unit/services/workers/test_indexer_pool.py`:
- Around line 286-290: The monkeypatch for module.ray.get_actor currently uses a
lambda that ignores all arguments, making it impossible to verify that the
correct namespace parameter is being passed. Replace the lambda mock with a
proper mock object (using unittest.mock.MagicMock or similar) that can capture
call arguments, then add an assertion after the actor_class() call to verify
that ray.get_actor was called with the namespace parameter set to "openrag".
This ensures that any future regression where the namespace argument is
accidentally removed from the ray.get_actor call will be caught by the test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 04bea148-b114-43a3-8839-81f835b2b6d7
📒 Files selected for processing (11)
openrag/core/indexing/topic_tags.pyopenrag/core/ports/topic_tag_repo.pyopenrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.pyopenrag/services/persistence/schema.pyopenrag/services/persistence/topic_tag_repo.pyopenrag/services/workers/indexer_actor.pyopenrag/services/workers/indexer_pool.pytests/unit/core/indexing/test_topic_tags.pytests/unit/services/persistence/test_topic_tag_repo.pytests/unit/services/workers/test_indexer_pool.pytests/unit/services/workers/test_indexer_worker.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tests/unit/services/persistence/test_topic_tag_repo.py
- openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py
- openrag/core/indexing/topic_tags.py
- openrag/services/persistence/topic_tag_repo.py
- openrag/services/workers/indexer_pool.py
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
Topic tagging was exposed through indexation config, but it did not actually generate or store any tags. That made the preset flag misleading and left retrieval/facet work without usable data.
Change
This PR turns the scaffold into a working backend slice: topic tags are generated during indexing when enabled, normalized and deduplicated, then stored in Postgres with partition-scoped cleanup on re-index or file deletion.
Retrieval ranking and UI filtering are intentionally left for follow-up work; this PR makes the indexed data exist first.
Closes #523.
Validation
uv run --no-env-file pytest tests/unituv run --no-env-file ruff check openrag tests/unit/core/indexing/test_topic_tags.py tests/unit/services/persistence/test_topic_tag_repo.py tests/unit/services/workers/test_pipeline_builder.py tests/unit/services/workers/test_indexer_worker.py tests/unit/services/workers/test_indexer_pool.pyuv run --no-env-file ruff format --check openrag tests/unit/core/indexing/test_topic_tags.py tests/unit/services/persistence/test_topic_tag_repo.py tests/unit/services/workers/test_pipeline_builder.py tests/unit/services/workers/test_indexer_worker.py tests/unit/services/workers/test_indexer_pool.pycd openrag/services/persistence/migrations/alembic && uv run --no-env-file alembic headsSummary by CodeRabbit
topic_tagsper document/partition.