feat(indexing): implement topic tagging pipeline

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"

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"]

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):

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]: ...

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"]

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")

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",