diff --git a/openrag/core/config/infrastructure.py b/openrag/core/config/infrastructure.py index ee4e9e897..fea9a078f 100644 --- a/openrag/core/config/infrastructure.py +++ b/openrag/core/config/infrastructure.py @@ -139,3 +139,4 @@ class PromptsConfig(ConfigMixin): spoken_style_answer: str = "spoken_style_answer_tmpl.txt" hyde: str = "hyde.txt" multi_query: str = "multi_query_pmpt_tmpl.txt" + topic_tagger: str = "topic_tagger_tmpl.txt" diff --git a/openrag/core/indexing/topic_tags.py b/openrag/core/indexing/topic_tags.py new file mode 100644 index 000000000..cf25cea96 --- /dev/null +++ b/openrag/core/indexing/topic_tags.py @@ -0,0 +1,160 @@ +"""Document-level topic tagging for indexed chunks.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import re +from collections.abc import Sequence + +from core.llm import LLM +from core.models.chunk import Chunk + +logger = logging.getLogger(__name__) + +_MAX_PROMPT_CHARS = 8000 + + +class TopicTagger: + """Extract a compact set of document topic labels with an LLM.""" + + def __init__( + self, + llm: LLM, + system_prompt: str, + *, + timeout_seconds: float | None = None, + ) -> None: + self._llm = llm + self._system_prompt = system_prompt + self._timeout = timeout_seconds + + async def tag( + self, + chunks: Sequence[Chunk], + *, + filename: str = "", + max_tags: int = 7, + lang: str = "en", + ) -> list[str]: + """Return normalized, unique topic tags for a document.""" + chunks = list(chunks) + if not chunks: + return [] + if max_tags <= 0: + return [] + + try: + messages = _build_messages( + system_prompt=self._system_prompt, + chunks=chunks, + filename=filename, + max_tags=max_tags, + lang=lang, + ) + operation = self._llm.chat(messages) + response = ( + await asyncio.wait_for(operation, timeout=self._timeout) + if self._timeout is not None + else await operation + ) + return _parse_topic_tags(_chat_response_text(response), max_tags=max_tags) + except (TimeoutError, OSError, RuntimeError, ValueError, TypeError) as exc: + logger.warning("Error extracting topic tags for %s: %s", filename, exc) + return [] + + +def _build_messages( + *, + system_prompt: str, + chunks: Sequence[Chunk], + filename: str, + max_tags: int, + lang: str, +) -> list[dict[str, str]]: + chunk_text = "\n\n".join( + f"Chunk {index + 1}:\n{chunk.text}" for index, chunk in enumerate(chunks[:12]) if chunk.text.strip() + ) + chunk_text = chunk_text[:_MAX_PROMPT_CHARS] + user_prompt = ( + f"Filename: {filename or 'unknown'}\n" + f"Language: {lang}\n" + f"Maximum topics: {max_tags}\n\n" + "Document chunks:\n" + f"{chunk_text}\n\n" + "Return only a JSON array of short topic strings." + ) + return [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + + +def _parse_topic_tags(text: str, *, max_tags: int) -> list[str]: + parsed = _load_json_array(text) + if parsed is None: + return [] + + tags: list[str] = [] + seen: set[str] = set() + for item in parsed: + if not isinstance(item, str): + continue + tag = _normalize_display_tag(item) + key = tag.casefold() + if not tag or key in seen: + continue + tags.append(tag) + seen.add(key) + if len(tags) >= max_tags: + break + return tags + + +def _load_json_array(text: str) -> list[object] | None: + try: + value = json.loads(text) + except json.JSONDecodeError: + match = re.search(r"\[[\s\S]*?\]", text) + if not match: + return None + try: + value = json.loads(match.group(0)) + except json.JSONDecodeError: + return None + + if isinstance(value, list): + return value + if isinstance(value, dict) and isinstance(value.get("topics"), list): + return value["topics"] + if isinstance(value, dict) and isinstance(value.get("tags"), list): + return value["tags"] + return None + + +def _normalize_display_tag(value: str) -> str: + return re.sub(r"\s+", " ", value).strip()[:80] + + +def _chat_response_text(response: object) -> str: + if isinstance(response, str): + return response + if not isinstance(response, dict): + return "" + + choices = response.get("choices") + if isinstance(choices, list) and choices: + first = choices[0] + if isinstance(first, dict): + message = first.get("message") + if isinstance(message, dict) and isinstance(message.get("content"), str): + return message["content"] + if isinstance(first.get("text"), str): + return first["text"] + + content = response.get("content") + return content if isinstance(content, str) else "" + + +__all__ = ["TopicTagger"] diff --git a/openrag/core/models/prompt.py b/openrag/core/models/prompt.py index 23f2b38a0..970ff83cc 100644 --- a/openrag/core/models/prompt.py +++ b/openrag/core/models/prompt.py @@ -17,6 +17,7 @@ class PromptType(str, Enum): HYDE = "hyde" MULTI_QUERY = "multi_query" SPOKEN_STYLE_ANSWER = "spoken_style_answer" + TOPIC_TAGGER = "topic_tagger" class Prompt(BaseModel): diff --git a/openrag/core/ports/topic_tag_repo.py b/openrag/core/ports/topic_tag_repo.py index 99287f46c..faa27b7c6 100644 --- a/openrag/core/ports/topic_tag_repo.py +++ b/openrag/core/ports/topic_tag_repo.py @@ -12,10 +12,10 @@ class TopicTagRepository(ABC): async def bulk_insert(self, tags: list[dict]) -> int: ... @abstractmethod - async def get_by_document(self, document_id: str) -> list[dict]: ... + async def get_by_document(self, document_id: str, partition: str) -> list[dict]: ... @abstractmethod - async def delete_by_document(self, document_id: str) -> int: ... + async def delete_by_document(self, document_id: str, partition: str) -> int: ... @abstractmethod async def search(self, partition: str, tag: str, top_k: int = 10) -> list[dict]: ... diff --git a/openrag/prompts/templates/topic_tagger_tmpl.txt b/openrag/prompts/templates/topic_tagger_tmpl.txt new file mode 100644 index 000000000..fc96829c0 --- /dev/null +++ b/openrag/prompts/templates/topic_tagger_tmpl.txt @@ -0,0 +1,12 @@ +You extract concise topic labels from document chunks for search and filtering. + +Guidelines: +- Return only a JSON array of strings. +- Use the same language as the document when possible. +- Prefer durable topics over generic words. +- Keep each topic short, normally 1-4 words. +- Do not include explanations, markdown, numbering, or confidence scores. + +Examples: +["climate finance", "portfolio risk", "carbon markets"] +["budget departemental", "subventions", "decision publique"] diff --git a/openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py b/openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py new file mode 100644 index 000000000..1f6376a04 --- /dev/null +++ b/openrag/services/persistence/migrations/alembic/versions/b7c8d9e0f1a2_add_topic_tags.py @@ -0,0 +1,67 @@ +"""add topic_tags table + +Revision ID: b7c8d9e0f1a2 +Revises: 06dd2101ea3a +Create Date: 2026-06-19 00:00:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from schema_helpers import index_exists, table_exists + +# revision identifiers, used by Alembic. +revision: str = "b7c8d9e0f1a2" +down_revision: str | Sequence[str] | None = "06dd2101ea3a" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + if not table_exists("topic_tags"): + op.create_table( + "topic_tags", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("document_id", sa.String(), nullable=False), + sa.Column("partition", sa.String(), nullable=False), + sa.Column("tag", sa.String(), nullable=False), + sa.Column("normalized_tag", sa.String(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["document_id", "partition"], + ["files.file_id", "files.partition_name"], + ondelete="CASCADE", + name="fk_topic_tags_file", + ), + sa.UniqueConstraint( + "document_id", + "partition", + "normalized_tag", + name="uix_topic_tags_document_partition_tag", + ), + ) + + if not index_exists("topic_tags", "ix_topic_tags_partition"): + op.create_index("ix_topic_tags_partition", "topic_tags", ["partition"]) + if not index_exists("topic_tags", "ix_topic_tags_document_id"): + op.create_index("ix_topic_tags_document_id", "topic_tags", ["document_id"]) + if not index_exists("topic_tags", "ix_topic_tags_partition_tag"): + op.create_index("ix_topic_tags_partition_tag", "topic_tags", ["partition", "normalized_tag"]) + + +def downgrade() -> None: + if index_exists("topic_tags", "ix_topic_tags_partition_tag"): + op.drop_index("ix_topic_tags_partition_tag", table_name="topic_tags") + if index_exists("topic_tags", "ix_topic_tags_document_id"): + op.drop_index("ix_topic_tags_document_id", table_name="topic_tags") + if index_exists("topic_tags", "ix_topic_tags_partition"): + op.drop_index("ix_topic_tags_partition", table_name="topic_tags") + if table_exists("topic_tags"): + op.drop_table("topic_tags") diff --git a/openrag/services/persistence/schema.py b/openrag/services/persistence/schema.py index 618946427..a5d1467e5 100644 --- a/openrag/services/persistence/schema.py +++ b/openrag/services/persistence/schema.py @@ -23,6 +23,7 @@ DateTime, Float, ForeignKey, + ForeignKeyConstraint, Index, Integer, LargeBinary, @@ -145,6 +146,32 @@ ) +topic_tags = Table( + "topic_tags", + metadata, + Column("id", Integer, primary_key=True), + Column("document_id", String, nullable=False), + Column("partition", String, nullable=False, index=True), + Column("tag", String, nullable=False), + Column("normalized_tag", String, nullable=False), + Column( + "created_at", + DateTime(timezone=True), + server_default=text("now()"), + nullable=False, + ), + ForeignKeyConstraint( + ["document_id", "partition"], + ["files.file_id", "files.partition_name"], + ondelete="CASCADE", + name="fk_topic_tags_file", + ), + UniqueConstraint("document_id", "partition", "normalized_tag", name="uix_topic_tags_document_partition_tag"), + Index("ix_topic_tags_document_id", "document_id"), + Index("ix_topic_tags_partition_tag", "partition", "normalized_tag"), +) + + users = Table( "users", metadata, @@ -270,6 +297,7 @@ "metadata", "model_endpoints", "pipeline_presets", + "topic_tags", "partitions", "files", "users", diff --git a/openrag/services/persistence/topic_tag_repo.py b/openrag/services/persistence/topic_tag_repo.py index c6398004a..3c0e457d2 100644 --- a/openrag/services/persistence/topic_tag_repo.py +++ b/openrag/services/persistence/topic_tag_repo.py @@ -1,30 +1,145 @@ -"""Stub :class:`TopicTagRepository`. - -Topic/tag attachment per document is a future feature — useful for -faceted search and for "show me docs about X" UIs. No table exists -today. -""" +"""asyncpg-backed :class:`TopicTagRepository`.""" from __future__ import annotations +import re +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + from core.ports.topic_tag_repo import TopicTagRepository -from services.persistence._stubs import _StubRepositoryBase, stub_not_implemented +if TYPE_CHECKING: + import asyncpg + + +class PgTopicTagRepository(TopicTagRepository): + """Store document-level topic tags in Postgres.""" -class PgTopicTagRepository(_StubRepositoryBase, TopicTagRepository): - """TODO: real impl once the ``topic_tags`` table is added.""" + def __init__(self, pool_getter: Callable[[], asyncpg.Pool]) -> None: + self._pool_getter = pool_getter + + @property + def pool(self) -> asyncpg.Pool: + return self._pool_getter() + + @staticmethod + def _row_to_dict(row: asyncpg.Record) -> dict: + result = { + "document_id": row["document_id"], + "partition": row["partition"], + "tag": row["tag"], + } + try: + result["created_at"] = row["created_at"] + except KeyError: + pass + return result async def bulk_insert(self, tags: list[dict]) -> int: - raise stub_not_implemented("Topic / tag storage") + rows = _normalize_rows(tags) + if not rows: + return 0 - async def get_by_document(self, document_id: str) -> list[dict]: - raise stub_not_implemented("Topic / tag storage") + return await self.pool.fetchval( + """ + WITH rows AS ( + SELECT * + FROM unnest($1::text[], $2::text[], $3::text[], $4::text[]) + AS t(document_id, partition, tag, normalized_tag) + ), + inserted AS ( + INSERT INTO topic_tags (document_id, partition, tag, normalized_tag) + SELECT document_id, partition, tag, normalized_tag + FROM rows + ON CONFLICT (document_id, partition, normalized_tag) + DO UPDATE SET tag = EXCLUDED.tag + RETURNING 1 + ) + SELECT COUNT(*)::int FROM inserted + """, + [row["document_id"] for row in rows], + [row["partition"] for row in rows], + [row["tag"] for row in rows], + [row["normalized_tag"] for row in rows], + ) - async def delete_by_document(self, document_id: str) -> int: - raise stub_not_implemented("Topic / tag storage") + async def get_by_document(self, document_id: str, partition: str) -> list[dict]: + rows = await self.pool.fetch( + """ + SELECT document_id, partition, tag, created_at + FROM topic_tags + WHERE document_id = $1 AND partition = $2 + ORDER BY normalized_tag + """, + document_id, + partition, + ) + return [self._row_to_dict(row) for row in rows] + + async def delete_by_document(self, document_id: str, partition: str) -> int: + result = await self.pool.execute( + "DELETE FROM topic_tags WHERE document_id = $1 AND partition = $2", + document_id, + partition, + ) + return _delete_count(result) async def search(self, partition: str, tag: str, top_k: int = 10) -> list[dict]: - raise stub_not_implemented("Topic / tag storage") + rows = await self.pool.fetch( + """ + SELECT document_id, partition, tag, created_at + FROM topic_tags + WHERE partition = $1 AND normalized_tag = $2 + ORDER BY document_id + LIMIT $3 + """, + partition, + _normalize_key(tag), + max(1, top_k), + ) + return [self._row_to_dict(row) for row in rows] + + +def _normalize_rows(tags: list[dict]) -> list[dict[str, str]]: + rows: list[dict[str, str]] = [] + seen: set[tuple[str, str, str]] = set() + for tag in tags: + document_id = str(tag.get("document_id") or "").strip() + partition = str(tag.get("partition") or "").strip() + display_tag = _normalize_display_tag(tag.get("tag")) + normalized_tag = _normalize_key(display_tag) + if not document_id or not partition or not display_tag: + continue + key = (document_id, partition, normalized_tag) + if key in seen: + continue + seen.add(key) + rows.append( + { + "document_id": document_id, + "partition": partition, + "tag": display_tag, + "normalized_tag": normalized_tag, + } + ) + return rows + + +def _normalize_display_tag(value: Any) -> str: + if not isinstance(value, str): + return "" + return re.sub(r"\s+", " ", value).strip()[:80] + + +def _normalize_key(value: str) -> str: + return _normalize_display_tag(value).casefold() + + +def _delete_count(result: str) -> int: + try: + return int(result.rsplit(" ", 1)[1]) + except (IndexError, ValueError): + return 0 __all__ = ["PgTopicTagRepository"] diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index d162e35b3..d582a788a 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -32,10 +32,12 @@ def __init__( pipeline: IndexingPipeline, task_state_manager: Any, document_repo: Any = None, + topic_tag_repo: Any = None, ) -> None: self._pipeline = pipeline self._tsm = task_state_manager self._document_repo = document_repo + self._topic_tag_repo = topic_tag_repo async def process_file( self, @@ -80,6 +82,14 @@ async def process_file( replace=replace, indexation_config=indexation_config, ) + if self._topic_tag_repo is not None: + await _replace_topic_tags_if_needed( + topic_tag_repo=self._topic_tag_repo, + row=row, + metadata=metadata, + partition=partition, + indexation_config=indexation_config, + ) await self._tsm.set_state.remote(task_id, "COMPLETED") return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} except Exception: @@ -122,6 +132,43 @@ async def _write_catalog_record( ) +async def _replace_topic_tags_if_needed( + *, + topic_tag_repo: Any, + row: dict[str, Any], + metadata: dict[str, Any], + partition: str, + indexation_config: dict[str, Any] | None, +) -> None: + file_id = metadata.get("file_id", "") + if not file_id: + return + + has_topic_tags = "topic_tags" in row + topic_tagging_disabled = indexation_config is not None and indexation_config.get("enable_topic_tagging") is False + if not has_topic_tags and not topic_tagging_disabled: + return + + raw_tags = row.get("topic_tags", []) + if not isinstance(raw_tags, list): + raise TypeError("topic_tags must be a list of strings") + + tags = [tag for tag in raw_tags if isinstance(tag, str) and tag.strip()] + await topic_tag_repo.delete_by_document(file_id, partition=partition) + if not tags: + return + await topic_tag_repo.bulk_insert( + [ + { + "document_id": file_id, + "partition": partition, + "tag": tag, + } + for tag in tags + ] + ) + + def _load_document( path: str, metadata: dict[str, Any], diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index 863dfb6ee..badd52dad 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -31,6 +31,7 @@ def __init__(self) -> None: chunker = _build_chunker(cfg) embedder_factory = _build_embedder_factory(cfg) contextualizer_factory = _build_contextualizer_factory(cfg) + topic_tagger_factory = _build_topic_tagger_factory(cfg) embed_cfg = cfg.embedder embedder = embedder_registry.create( @@ -53,6 +54,7 @@ def __init__(self) -> None: chunker_factory=_build_chunker_from_config, embedder_factory=embedder_factory, contextualizer_factory=contextualizer_factory, + topic_tagger_factory=topic_tagger_factory, ) rdb_cfg = cfg.rdb.model_copy(update={"database": f"partitions_for_collection_{cfg.vectordb.collection_name}"}) self._catalog_store = PostgresStore(rdb_cfg, run_migrations=False) @@ -62,6 +64,7 @@ def __init__(self) -> None: pipeline=pipeline, task_state_manager=task_state_manager, document_repo=self._catalog_store.document_repo, + topic_tag_repo=self._catalog_store.topic_tag_repo, ) async def _ensure_catalog(self) -> None: @@ -237,6 +240,51 @@ def factory(name: str = "default") -> ChunkContextualizer: return factory +def _build_topic_tagger_factory(cfg: Settings) -> Any: + """Build a cached factory yielding ``TopicTagger`` instances.""" + import services.inference.ollama_client # noqa: F401 + import services.inference.vllm_client # noqa: F401 + from core.indexing.topic_tags import TopicTagger + from core.llm import llm_registry + from core.prompts import load_template_by_key + + named_llms = getattr(getattr(cfg, "models", None), "llm", {}) or {} + fallback_cfg = _global_llm_endpoint_config(cfg) + if not named_llms and fallback_cfg is None: + return None + + system_prompt = load_template_by_key(cfg.paths.prompts_dir, cfg.prompts, "topic_tagger") + cache: dict[str, TopicTagger] = {} + lock = threading.Lock() + + def factory(name: str = "default") -> TopicTagger: + if name in cache: + return cache[name] + with lock: + if name in cache: + return cache[name] + model_cfg = named_llms.get(name) + if model_cfg is None: + if name == "default" and fallback_cfg is not None: + model_cfg = fallback_cfg + else: + raise KeyError(f"Unknown llm '{name}'. Available: {list(named_llms)}") + impl_kwargs = {key: value for key, value in model_cfg.extra.items() if key != "implementation"} + impl = model_cfg.extra.get("implementation", "vllm") + llm = llm_registry.create( + impl, + endpoint=model_cfg.endpoint, + model_name=model_cfg.model_name, + timeout=model_cfg.timeout, + **impl_kwargs, + ) + tagger = TopicTagger(llm, system_prompt, timeout_seconds=model_cfg.timeout) + cache[name] = tagger + return tagger + + return factory + + def _global_llm_endpoint_config(cfg: Any) -> Any | None: """Adapt the legacy/global ``cfg.llm`` block into a ``ModelEndpointConfig``. diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index bcb507fdb..f4d8b8b13 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -9,6 +9,7 @@ from core.embeddings.embedder import Embedder from core.indexing.contextualize import ChunkContextualizer from core.indexing.parsers.document_parser import DocumentParser +from core.indexing.topic_tags import TopicTagger from core.vector_stores.vector_store import VectorStore from core.vlm.vlm import VLM from services.workers.stages.caption import caption_stage @@ -17,6 +18,7 @@ from services.workers.stages.embed import embed_stage from services.workers.stages.parse import parse_stage from services.workers.stages.store import store_stage +from services.workers.stages.topic_tag import topic_tag_stage @dataclass(slots=True, frozen=True) @@ -33,6 +35,7 @@ class PipelineTimeouts: embed_per_chunk: float = 0.0 store: float | None = None store_per_chunk: float = 0.0 + topic_tag: float | None = None @dataclass(slots=True, frozen=True) @@ -45,6 +48,7 @@ class IndexingPipeline: vector_store: VectorStore vlm: VLM | None = None contextualizer: ChunkContextualizer | None = None + topic_tagger: TopicTagger | None = None timeouts: PipelineTimeouts = PipelineTimeouts() indexation_config: IndexationPipelineConfig | None = None parser_factory: Callable[[str], DocumentParser] | None = None @@ -52,6 +56,7 @@ class IndexingPipeline: embedder_factory: Callable[[str], Embedder] | None = None vlm_factory: Callable[[str], VLM] | None = None contextualizer_factory: Callable[[str], ChunkContextualizer] | None = None + topic_tagger_factory: Callable[[str], TopicTagger] | None = None async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: """Run a single row through parse, optional enrichments, embed, and store.""" @@ -62,6 +67,7 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: embedder = self._select_embedder(row) vlm = self._select_vlm(config) contextualizer = self._select_contextualizer(config) + topic_tagger = self._select_topic_tagger(config) await parse_stage(row, parser, timeout=self.timeouts.parse) if vlm is not None: @@ -79,6 +85,14 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: timeout=self.timeouts.contextualize, per_chunk_timeout=self.timeouts.contextualize_per_chunk, ) + if topic_tagger is not None: + max_tags = config.max_topic_tags if config is not None else 7 + await topic_tag_stage( + row, + topic_tagger, + max_tags=max_tags, + timeout=self.timeouts.topic_tag, + ) await embed_stage( row, embedder, @@ -135,6 +149,14 @@ def _select_contextualizer(self, config: IndexationPipelineConfig | None) -> Chu return self.contextualizer_factory(config.contextualization_llm or "default") return self.contextualizer + def _select_topic_tagger(self, config: IndexationPipelineConfig | None) -> TopicTagger | None: + if config is not None: + if not config.enable_topic_tagging: + return None + if self.topic_tagger_factory is not None: + return self.topic_tagger_factory(config.topic_tagging_llm or "default") + return self.topic_tagger + def build_indexing_pipeline( *, @@ -144,6 +166,7 @@ def build_indexing_pipeline( vector_store: VectorStore, vlm: VLM | None = None, contextualizer: ChunkContextualizer | None = None, + topic_tagger: TopicTagger | None = None, timeouts: PipelineTimeouts | None = None, indexation_config: IndexationPipelineConfig | None = None, parser_factory: Callable[[str], DocumentParser] | None = None, @@ -151,6 +174,7 @@ def build_indexing_pipeline( embedder_factory: Callable[[str], Embedder] | None = None, vlm_factory: Callable[[str], VLM] | None = None, contextualizer_factory: Callable[[str], ChunkContextualizer] | None = None, + topic_tagger_factory: Callable[[str], TopicTagger] | None = None, ) -> IndexingPipeline: """Build the default sequential indexing pipeline.""" @@ -161,6 +185,7 @@ def build_indexing_pipeline( vector_store=vector_store, vlm=vlm, contextualizer=contextualizer, + topic_tagger=topic_tagger, timeouts=timeouts or PipelineTimeouts(), indexation_config=indexation_config, parser_factory=parser_factory, @@ -168,6 +193,7 @@ def build_indexing_pipeline( embedder_factory=embedder_factory, vlm_factory=vlm_factory, contextualizer_factory=contextualizer_factory, + topic_tagger_factory=topic_tagger_factory, ) diff --git a/openrag/services/workers/stages/topic_tag.py b/openrag/services/workers/stages/topic_tag.py new file mode 100644 index 000000000..265ea6799 --- /dev/null +++ b/openrag/services/workers/stages/topic_tag.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.indexing.topic_tags import TopicTagger +from core.models.chunk import Chunk +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials + + +async def topic_tag_stage( + row: MutableMapping[str, Any], + topic_tagger: TopicTagger, + *, + max_tags: int = 7, + timeout: float | None = None, +) -> MutableMapping[str, Any]: + """Extract document-level topic tags into ``row["topic_tags"]``.""" + try: + chunks = row.get("chunks") + if not _is_chunk_list(chunks): + raise ValueError("topic_tag_stage row must contain a list[Chunk] under 'chunks'") + + filename = str(row.get("filename") or "") + language = str(row.get("language") or row.get("lang") or "en") + row["topic_tags"] = await run_with_optional_timeout( + lambda: topic_tagger.tag(chunks, filename=filename, max_tags=max_tags, lang=language), + timeout, + ) + row["stage"] = "topic_tagged" + row.pop("error", None) + return row + except Exception as exc: + row["stage"] = "topic_tag_failed" + row["error"] = str(exc) + raise + finally: + scrub_credentials(row) + + +def _is_chunk_list(value: Any) -> bool: + return isinstance(value, list) and all(isinstance(chunk, Chunk) for chunk in value) diff --git a/tests/unit/core/indexing/test_topic_tags.py b/tests/unit/core/indexing/test_topic_tags.py new file mode 100644 index 000000000..264b7bdf2 --- /dev/null +++ b/tests/unit/core/indexing/test_topic_tags.py @@ -0,0 +1,69 @@ +import asyncio + +import pytest +from core.indexing.topic_tags import TopicTagger +from core.models.chunk import Chunk + + +class FakeLLM: + def __init__(self, content: str) -> None: + self.content = content + self.messages: list[list[dict[str, str]]] = [] + + async def chat(self, messages: list[dict[str, str]], **kwargs): + self.messages.append(messages) + return {"choices": [{"message": {"content": self.content}}]} + + +@pytest.mark.asyncio +async def test_topic_tagger_extracts_normalized_unique_tags(): + llm = FakeLLM('["Finance", "climate risk", "Finance", ""]') + tagger = TopicTagger(llm, "extract topics") + + tags = await tagger.tag( + [Chunk(id="c1", text="green finance"), Chunk(id="c2", text="portfolio risk")], + filename="report.pdf", + max_tags=3, + lang="en", + ) + + assert tags == ["Finance", "climate risk"] + assert "report.pdf" in llm.messages[0][1]["content"] + assert "green finance" in llm.messages[0][1]["content"] + + +@pytest.mark.asyncio +async def test_topic_tagger_falls_back_to_empty_list_on_bad_response(): + llm = FakeLLM("this is not structured") + tagger = TopicTagger(llm, "extract topics") + + assert await tagger.tag([Chunk(id="c1", text="hello")], max_tags=5) == [] + + +@pytest.mark.asyncio +async def test_topic_tagger_returns_empty_when_max_tags_is_not_positive(): + llm = FakeLLM('["finance"]') + tagger = TopicTagger(llm, "extract topics") + + assert await tagger.tag([Chunk(id="c1", text="hello")], max_tags=0) == [] + assert llm.messages == [] + + +@pytest.mark.asyncio +async def test_topic_tagger_timeout_zero_uses_timeout_path(): + class SlowLLM: + async def chat(self, messages: list[dict[str, str]], **kwargs): + await asyncio.sleep(0) + return {"choices": [{"message": {"content": '["finance"]'}}]} + + tagger = TopicTagger(SlowLLM(), "extract topics", timeout_seconds=0) + + assert await tagger.tag([Chunk(id="c1", text="hello")], max_tags=5) == [] + + +@pytest.mark.asyncio +async def test_topic_tagger_parses_json_array_before_bracket_suffix(): + llm = FakeLLM('["finance"] [Sources: 1]') + tagger = TopicTagger(llm, "extract topics") + + assert await tagger.tag([Chunk(id="c1", text="hello")], max_tags=5) == ["finance"] diff --git a/tests/unit/services/persistence/test_topic_tag_repo.py b/tests/unit/services/persistence/test_topic_tag_repo.py new file mode 100644 index 000000000..e6ef315ab --- /dev/null +++ b/tests/unit/services/persistence/test_topic_tag_repo.py @@ -0,0 +1,104 @@ +import pytest + + +class _FakePool: + def __init__(self): + self.executed: list[tuple[str, tuple]] = [] + self._fetch_result: list = [] + self._fetchval_result: int = 0 + + async def fetch(self, query: str, *params): + self.executed.append((query, params)) + return self._fetch_result + + async def fetchval(self, query: str, *params): + self.executed.append((query, params)) + return self._fetchval_result + + async def execute(self, query: str, *params): + self.executed.append((query, params)) + return "DELETE 2" + + +def _row(**kwargs): + base = { + "document_id": "file-1", + "partition": "tenant-a", + "tag": "finance", + } + base.update(kwargs) + return base + + +@pytest.mark.asyncio +async def test_bulk_insert_normalizes_and_deduplicates_tags(): + from services.persistence.topic_tag_repo import PgTopicTagRepository + + pool = _FakePool() + pool._fetchval_result = 2 + repo = PgTopicTagRepository(pool_getter=lambda: pool) + + inserted = await repo.bulk_insert( + [ + {"document_id": "file-1", "partition": "tenant-a", "tag": "Finance"}, + {"document_id": "file-1", "partition": "tenant-a", "tag": " finance "}, + {"document_id": "file-1", "partition": "tenant-a", "tag": "risk"}, + ] + ) + + assert inserted == 2 + query, params = pool.executed[0] + assert "INSERT INTO topic_tags" in query + assert params[0] == ["file-1", "file-1"] + assert params[1] == ["tenant-a", "tenant-a"] + assert params[2] == ["Finance", "risk"] + assert params[3] == ["finance", "risk"] + + +@pytest.mark.asyncio +async def test_get_by_document_returns_rows_ordered_by_tag(): + from services.persistence.topic_tag_repo import PgTopicTagRepository + + pool = _FakePool() + pool._fetch_result = [_row(tag="finance"), _row(tag="risk")] + repo = PgTopicTagRepository(pool_getter=lambda: pool) + + result = await repo.get_by_document("file-1", partition="tenant-a") + + assert [row["tag"] for row in result] == ["finance", "risk"] + query, params = pool.executed[0] + assert "WHERE document_id = $1 AND partition = $2" in query + assert "ORDER BY normalized_tag" in query + assert params == ("file-1", "tenant-a") + + +@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 + 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") + + +@pytest.mark.asyncio +async def test_search_is_partition_scoped_and_case_insensitive(): + from services.persistence.topic_tag_repo import PgTopicTagRepository + + pool = _FakePool() + pool._fetch_result = [_row()] + repo = PgTopicTagRepository(pool_getter=lambda: pool) + + await repo.search("tenant-a", "Finance", top_k=3) + + query, params = pool.executed[0] + assert "WHERE partition = $1 AND normalized_tag = $2" in query + assert "LIMIT $3" in query + assert params == ("tenant-a", "finance", 3) diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index 0803886d6..e4e11d475 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -4,6 +4,7 @@ from types import SimpleNamespace import pytest +from core.config.model_endpoints import ModelEndpointConfig class _NativeChunker: @@ -100,6 +101,40 @@ def fake_options(**kwargs): assert calls["max_concurrency"] == 4 +def test_build_topic_tagger_factory_resolves_named_llm(monkeypatch: pytest.MonkeyPatch) -> None: + from core.llm import llm_registry + from services.workers.indexer_pool import _build_topic_tagger_factory + + class ProbeLLM: + def __init__(self, **kwargs): + self.kwargs = kwargs + + llm_registry.register("topic-probe")(ProbeLLM) + cfg = SimpleNamespace( + models=SimpleNamespace( + llm={ + "topic-a": ModelEndpointConfig( + endpoint="http://llm:8000/v1", + model_name="topic-model", + timeout=9.0, + extra={"implementation": "topic-probe", "temperature": 0.1}, + ) + } + ), + llm=SimpleNamespace(base_url="", model=""), + paths=SimpleNamespace(prompts_dir="/tmp/prompts"), + prompts=SimpleNamespace(topic_tagger="topic.txt"), + ) + monkeypatch.setattr("core.prompts.load_template_by_key", lambda *_args: "extract topics") + + factory = _build_topic_tagger_factory(cfg) + tagger = factory("topic-a") + + assert tagger._llm.kwargs["endpoint"] == "http://llm:8000/v1" + assert tagger._llm.kwargs["model_name"] == "topic-model" + assert tagger._llm.kwargs["temperature"] == 0.1 + + def test_build_contextualizer_factory_returns_none_without_llm_config(tmp_path) -> None: from services.workers.indexer_pool import _build_contextualizer_factory @@ -206,6 +241,7 @@ def test_indexer_pool_wires_contextualizer_factory(monkeypatch: pytest.MonkeyPat captured = {} contextualizer_factory = object() + topic_tagger_factory = object() class RDBConfig: def model_copy(self, *, update): @@ -227,6 +263,7 @@ def model_copy(self, *, update): class Store: document_repo = object() + topic_tag_repo = object() class Worker: def __init__(self, **kwargs): @@ -240,15 +277,26 @@ def fake_build_pipeline(**kwargs): monkeypatch.setattr(module, "_build_chunker", lambda _cfg: object()) monkeypatch.setattr(module, "_build_embedder_factory", lambda _cfg: object()) monkeypatch.setattr(module, "_build_contextualizer_factory", lambda _cfg: contextualizer_factory) + monkeypatch.setattr(module, "_build_topic_tagger_factory", lambda _cfg: topic_tagger_factory) monkeypatch.setattr(core.embeddings.embedder_registry, "create", lambda *args, **kwargs: object()) monkeypatch.setattr(milvus_store, "MilvusVectorStore", lambda _cfg: object()) monkeypatch.setattr(postgres_store, "PostgresStore", lambda *args, **kwargs: Store()) monkeypatch.setattr(parser_bridge, "DocSerializerBridgeParser", lambda **kwargs: object()) monkeypatch.setattr(pipeline_builder, "build_indexing_pipeline", fake_build_pipeline) - monkeypatch.setattr(module.ray, "get_actor", lambda *args, **kwargs: object()) + actor_calls = [] + + def fake_get_actor(*args, **kwargs): + 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_factory diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 7d2ae1a3a..8a871be6d 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -92,6 +92,20 @@ async def update_file_in_partition(self, **kwargs: Any) -> bool: return True +class FakeTopicTagRepo: + def __init__(self) -> None: + self.deleted: list[tuple[str, str]] = [] + self.inserted: list[list[dict[str, str]]] = [] + + async def delete_by_document(self, document_id: str, partition: str) -> int: + self.deleted.append((document_id, partition)) + return 0 + + async def bulk_insert(self, tags: list[dict]) -> int: + self.inserted.append(tags) + return len(tags) + + # --------------------------------------------------------------------------- # Tests — _load_document helper # --------------------------------------------------------------------------- @@ -388,3 +402,101 @@ async def add_file_to_partition(self, **kwargs: Any) -> bool: tsm.set_failed_if_not_cancelled.remote.assert_called_once() completed_calls = [call for call in tsm.set_state.remote.call_args_list if call.args == ("t-fail", "COMPLETED")] assert completed_calls == [] + + +@pytest.mark.asyncio +async def test_process_file_replaces_topic_tags_after_successful_pipeline(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + + class TaggingPipeline: + async def run(self, row: dict[str, Any]) -> dict[str, Any]: + row["topic_tags"] = ["finance", "risk"] + row["stored_count"] = 1 + row["stage"] = "stored" + return row + + repo = FakeTopicTagRepo() + worker = IndexerWorker( + pipeline=TaggingPipeline(), + task_state_manager=_fake_tsm(), + topic_tag_repo=repo, + ) + + await worker.process_file( + task_id="t-tags", + path=str(path), + metadata={"file_id": "f1"}, + partition="tenant-a", + ) + + assert repo.deleted == [("f1", "tenant-a")] + assert repo.inserted == [ + [ + {"document_id": "f1", "partition": "tenant-a", "tag": "finance"}, + {"document_id": "f1", "partition": "tenant-a", "tag": "risk"}, + ] + ] + + +@pytest.mark.asyncio +async def test_process_file_deletes_topic_tags_when_tagging_is_disabled(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + + class UntaggedPipeline: + async def run(self, row: dict[str, Any]) -> dict[str, Any]: + row["stored_count"] = 1 + row["stage"] = "stored" + return row + + repo = FakeTopicTagRepo() + worker = IndexerWorker( + pipeline=UntaggedPipeline(), + task_state_manager=_fake_tsm(), + topic_tag_repo=repo, + ) + + await worker.process_file( + task_id="t-disabled-tags", + path=str(path), + metadata={"file_id": "f1"}, + partition="tenant-a", + indexation_config={"enable_topic_tagging": False}, + ) + + assert repo.deleted == [("f1", "tenant-a")] + assert repo.inserted == [] + + +@pytest.mark.asyncio +async def test_process_file_rejects_malformed_topic_tags_before_delete(tmp_path: Path) -> None: + path = tmp_path / "doc.txt" + path.write_bytes(b"content") + tsm = _fake_tsm() + + class BrokenTaggingPipeline: + async def run(self, row: dict[str, Any]) -> dict[str, Any]: + row["topic_tags"] = "finance" + row["stored_count"] = 1 + row["stage"] = "stored" + return row + + repo = FakeTopicTagRepo() + worker = IndexerWorker( + pipeline=BrokenTaggingPipeline(), + task_state_manager=tsm, + topic_tag_repo=repo, + ) + + with pytest.raises(TypeError, match="topic_tags"): + await worker.process_file( + task_id="t-bad-tags", + path=str(path), + metadata={"file_id": "f1"}, + partition="tenant-a", + ) + + assert repo.deleted == [] + assert repo.inserted == [] + tsm.set_failed_if_not_cancelled.remote.assert_called_once() diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 7c30cd59b..3cedd161b 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -72,6 +72,23 @@ async def contextualize(self, chunks, *, filename: str = "", lang: str = "en") - return [chunk.model_copy(update={"text": f"ctx {chunk.text}", "context": "ctx"}) for chunk in chunks] +class FakeTopicTagger: + def __init__(self, tags: list[str] | None = None) -> None: + self.tags = tags or ["finance", "risk"] + self.calls: list[tuple[list[Chunk], str, int, str]] = [] + + async def tag( + self, + chunks, + *, + filename: str = "", + max_tags: int = 7, + lang: str = "en", + ) -> list[str]: + self.calls.append((list(chunks), filename, max_tags, lang)) + return self.tags + + @pytest.mark.asyncio async def test_pipeline_runs_required_stages_in_order_and_keeps_row_object(): document = Document(filename="note.txt", text="hello", partition="tenant-a") @@ -173,11 +190,13 @@ async def test_pipeline_row_indexation_config_selects_components(): selected_embedder = FakeEmbedder([[0.5]]) selected_vlm = FakeVLM() selected_contextualizer = FakeContextualizer() + selected_topic_tagger = FakeTopicTagger(["portfolio"]) parser_calls: list[str] = [] chunker_calls: list[object] = [] embedder_calls: list[str] = [] vlm_calls: list[str] = [] contextualizer_calls: list[str] = [] + topic_tagger_calls: list[str] = [] pipeline = build_indexing_pipeline( parser=FakeParser(default_processed), @@ -189,6 +208,7 @@ async def test_pipeline_row_indexation_config_selects_components(): embedder_factory=lambda name: embedder_calls.append(name) or selected_embedder, vlm_factory=lambda name: vlm_calls.append(name) or selected_vlm, contextualizer_factory=lambda name: contextualizer_calls.append(name) or selected_contextualizer, + topic_tagger_factory=lambda name: topic_tagger_calls.append(name) or selected_topic_tagger, ) row = { "document": document, @@ -202,6 +222,9 @@ async def test_pipeline_row_indexation_config_selects_components(): "vlm": "vlm-fast", "enable_contextualization": True, "contextualization_llm": "llm-context", + "enable_topic_tagging": True, + "topic_tagging_llm": "llm-topic", + "max_topic_tags": 3, }, } @@ -212,8 +235,35 @@ async def test_pipeline_row_indexation_config_selects_components(): assert embedder_calls == ["embed-fast"] assert vlm_calls == ["vlm-fast"] assert contextualizer_calls == ["llm-context"] + assert topic_tagger_calls == ["llm-topic"] assert selected_parser.calls == [document] assert selected_chunker.calls == [(row["processed_document"], "tenant-a")] assert selected_vlm.calls == [b"png"] assert selected_contextualizer.calls[0][1:] == ("note.txt", "en") + assert selected_topic_tagger.calls == [ + ([row["chunks"][0].model_copy(update={"embedding": None})], "note.txt", 3, "en") + ] + assert row["topic_tags"] == ["portfolio"] assert row["chunks"][0].embedding == [0.5] + + +@pytest.mark.asyncio +async def test_pipeline_indexation_config_disables_topic_tagging(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="hello")]) + chunks = [Chunk(id="c1", text="hello", partition="tenant-a")] + topic_tagger = FakeTopicTagger() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker(chunks), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + topic_tagger=topic_tagger, + indexation_config=IndexationPipelineConfig(enable_topic_tagging=False), + ) + + row = {"document": document, "partition": "tenant-a", "filename": "note.txt"} + await pipeline.run(row) + + assert topic_tagger.calls == [] + assert row.get("topic_tags") is None