diff --git a/FORWARD_PORT_LOG.md b/FORWARD_PORT_LOG.md
new file mode 100644
index 000000000..e30c56c74
--- /dev/null
+++ b/FORWARD_PORT_LOG.md
@@ -0,0 +1,27 @@
+# Forward Port Log
+
+Tracks `dev` changes during MODE 2 isolation (Phases 5-9).
+Each entry: what changed on `dev`, whether it was forward-ported or deferred
+to the cutover re-implementation queue.
+
+> Created retroactively at the start of Phase 5 (2026-04-29). The Phase 0-4
+> Mode 1 work merged from `dev` cleanly so this log starts empty.
+
+---
+
+## Forward-ported (critical)
+
+_Security fixes, data-loss bugs, production outages re-implemented against
+the new architecture. Each entry pairs the dev commit with the refactor
+commit so a reviewer can audit equivalence._
+
+(none yet)
+
+## Deferred to cutover (features)
+
+_Non-critical changes that landed on `dev` during MODE 2. These will be
+re-implemented directly in the new architecture during MODE 3 (Phases
+10-12) or post-cutover. List the dev PR number / commit and the target
+location in the new layout._
+
+(none yet)
diff --git a/REFACTORING_DECISION_LOG.md b/REFACTORING_DECISION_LOG.md
index 8cf2a8003..ddf668a9c 100644
--- a/REFACTORING_DECISION_LOG.md
+++ b/REFACTORING_DECISION_LOG.md
@@ -55,7 +55,141 @@ match reality, then record the reasoning here.
---
-## Template for future entries
+## Phase 5 — Core domain logic for retrieval, chunking, prompts (2026-04-29)
+
+**1. New `RetrievalSearcher` port in `core/retrieval/searcher.py`, separate from
+the narrow `VectorStore` ABC.**
+The retriever needs four operations (search by query string, multi-query
+search, related-chunk lookup, ancestor lookup) that the Phase-4
+`VectorStore` ABC does not cover — that ABC is intentionally narrow
+(`search(embedding, top_k)`). We added a transitional ABC the retriever
+depends on, implemented by `services/storage/milvus_ray_shim.py` over the
+legacy Ray actor.
+- Why: STRATEGY §5A says the retriever should "call `VectorStore.search()`
+ (port method), not `vectordb.async_search.remote()`". But the legacy Ray
+ actor's `async_search` takes a query *string* and embeds internally; the
+ narrow `VectorStore.search(embedding, ...)` ABC doesn't fit. Pre-embedding
+ in the shim before calling Ray is impractical because the actor also owns
+ BM25 and surrounding-chunks semantics. A retrieval-facing port keeps the
+ retriever clean of Ray today and survives Phase 7 — when the Vectordb
+ god object is decomposed, these methods either move onto a richer
+ `VectorStore` or split between `VectorStore` and `ChunkRepository`.
+- Alternative considered: extend `VectorStore` with the four legacy methods.
+ Rejected — bloats the ABC with operations that should not exist past
+ Phase 7. Also considered: skip the new core port and have the retriever
+ call the Ray actor through the shim with the legacy method names —
+ rejected because that leaks legacy method names into core/ and makes the
+ retriever harder to test.
+
+**2. Skipped: bringing up integration tests for the new code.**
+Phase 5 ships pure-domain unit tests only (50 new tests in `core/`, no Ray /
+Milvus / real LLM). The new pipeline is dormant until Phase 8 wires it.
+- Why: Mode 2 forbids touching the legacy wiring; the new pipeline has
+ nowhere to be plugged in yet. Integration coverage will land with Phase 8
+ orchestrators (or Phase 7 storage if it goes first).
+- Alternative considered: stand up a fake searcher in an integration
+ fixture and run a full retriever-pipeline-RRF round trip. Defers the
+ same coverage to Phase 8 with less code; not worth the extra fixtures.
+
+**3. `Query`, `SearchQueries`, `TemporalPredicate` lifted into
+`core/models/query.py`.**
+The legacy `components/pipeline.py` defined these inline. The new
+`RetrieverPipeline` consumes them — they're domain types, not pipeline
+internals.
+- Why: STRATEGY §2 calls these out as `pipeline.py SearchQueries → core/models/query.py`.
+- Alternative considered: keep them in `core/retrieval/`. Rejected — they
+ describe a query in the abstract; the orchestrator (Phase 8) and the API
+ layer will both use them, not just retrieval.
+
+**4. Phase 5.15 (re-export shims) deferred to follow-up — then completed.**
+The first Phase 5 commits (5A/5B/5C, 2026-04-29) created the new core/
+modules but left the legacy `components/` files intact. STRATEGY §4.1
+mandates a three-step move — create new file, update old file to
+re-export from new, update consumers — and Phase 5 step 5.15 says
+"Update old files to re-export from core/". We skipped that.
+The follow-up sweep (2026-05-05) replaced six legacy files with shims:
+
+| Legacy file | Shim strategy |
+|---|---|
+| `components/indexer/chunker/utils.py` | Plain re-export from `core.chunking.markdown_utils` |
+| `components/prompts/prompts.py` | `load_prompt(key)` adapter calling `core.prompts.template_loader.load_template_by_key` |
+| `components/utils.py:format_context` + `format_web_context` | Adapters into `core.prompts.chat_prompt_builder` (rest of utils stays — Phase 6+ scope) |
+| `components/indexer/chunker/chunker.py` | `BaseChunker` / `RecursiveSplitter` delegate to `core.chunking.RecursiveSplitter` via Document↔ProcessedDocument↔Chunk conversion. `ChunkContextualizer` + `ChunkerFactory` retained (5D + Phase 8). |
+| `components/retriever.py` | `Single`/`MultiQuery`/`HyDe` retrievers wrap `core.retrieval.retriever` strategies. Ray actor → `MilvusRayShim`; `ChatOpenAI` → `_LangChainLLMAdapter`. `RetrieverFactory` retained. |
+| `components/pipeline.py` | `Query`/`SearchQueries`/`TemporalPredicate` re-exported from `core.models.query`. `RetrieverPipeline` delegates to `core.retrieval.pipeline.RetrieverPipeline` via a `_LegacyRerankerAdapter` bridging the legacy reranker (Document-in / Document-out) to the core ABC (str-in / `(idx, score)`-out). `RagPipeline` + `RAGMODE` retained (Phase 8). |
+
+- Why: STRATEGY §4.1 is explicit ("Update old file to re-export from new
+ location"); leaving the duplication in place would let the codepaths
+ drift. Three CodeRabbit fixes from PR #352 (image_caption ChunkType,
+ page-marker semantics, chunk_table header-only flush) had to be applied
+ twice or only fixed in core — exactly the failure mode 5.15 prevents.
+ The shim pattern matches the prior-art shims for config (1329cc18) and
+ exceptions (a0f3d9f2).
+- Alternative considered: leave the copies until Phase 8 cutover. Rejected
+ — the doc explicitly puts 5.15 *inside* Phase 5, and one round of
+ drift already happened.
+- Side effect: a noqa side-effect import in `chunker.py` keeps the legacy
+ `components.utils` ↔ `components.indexer.utils.files` circular-import
+ resolving in the right order. Removed once `components.utils` is split
+ in Phase 6+.
+- Behavioral note: image elements in the legacy chunker now stamp
+ `chunk_type=image_caption` (matching `core.models.chunk.ChunkType`)
+ instead of the previous raw `image`. No legacy reader filters on this
+ value, so the change is invisible to consumers.
+
+**5. `core/chunking/recursive.py` keeps `langchain.text_splitter.RecursiveCharacterTextSplitter` as a dependency.**
+STRATEGY §3 lists core as "stdlib + pydantic + pure libs" and §4.7 limits
+LangChain in core to boundary converters (`from_langchain` / `to_langchain`
+on domain models). The chunker imports `RecursiveCharacterTextSplitter`
+directly, which is neither stdlib nor a boundary converter.
+- Why: `RecursiveCharacterTextSplitter` is a self-contained recursive
+ separator-based string splitter — no IO, no LLM client, no Document
+ semantics. Reimplementing it in core/ would be a meaningful chunk of
+ pure code with no behavior change, and the legacy chunker has been
+ using it for two years with stable output. Keeping it in for Phase 5
+ preserves byte-for-byte chunk equivalence, which the strangler-fig
+ shim relies on for behavior parity. The import is also deferred
+ (inside the constructor / `split_text`) so importing the module
+ without LangChain installed doesn't fail.
+- Alternative considered: write a stdlib-only recursive splitter as part
+ of Phase 5. Rejected as scope creep — would couple a behaviorally
+ risky rewrite (chunk boundaries shift, downstream embedding output
+ changes) to the additive Phase 5 cut, breaking the parity guarantee
+ the legacy shims rely on. Tracked as a Phase 12 / post-cutover
+ follow-up: replace the splitter with a stdlib implementation behind
+ the same `Callable[[str], int]` length-function injection point.
+- Scope: limited to `RecursiveCharacterTextSplitter`. No other LangChain
+ symbol leaks into core; `langchain_core.documents.Document` only
+ appears inside `Chunk.from_langchain` / `Chunk.to_langchain` /
+ `Document.from_langchain` / `Document.to_langchain`, all with deferred
+ imports, exactly as §4.7 prescribes.
+
+**6. Two file-layout deviations from STRATEGY §3 / §5A / §5B.**
+- `core/chunking/markdown_section.py` (§5B, line 1038; §3, line 387)
+ → renamed to `core/chunking/markdown_utils.py`. The contents are pure
+ parsing helpers — `MDElement`, `split_md_elements`, `chunk_table`,
+ `parse_markdown_table`, `get_chunk_page_number` — not a
+ section-aware chunker strategy. The name `markdown_utils.py` matches
+ the module's role (utilities consumed by `RecursiveSplitter`); the
+ separate `markdown_section.py` / `markdown_layout.py` *strategies*
+ listed in §3's tree aren't built in Phase 5 and remain available
+ filenames if/when those strategies land.
+- `core/retrieval/hydration.py` (§5A, line 1027) → kept as private
+ `_expand_with_related_chunks` in `retriever.py`. The function is
+ ~60 LOC, only invoked by `BaseRetriever.expand_search_results`, and
+ splitting it would add an import + test fixture surface without any
+ reuse benefit. If a second consumer ever appears (Phase 8 likely),
+ promoting it to a module-level public function in `hydration.py` is
+ a one-commit move.
+- Why: both deviations make the module names track the actual contents
+ rather than the strategy doc's pre-write naming guess. Recording so
+ future readers don't grep for files that aren't there.
+- Alternative considered: rename to match the strategy doc verbatim.
+ Rejected — the strategy filenames anticipated different content
+ (a section/layout chunker, a standalone hydration entry point) than
+ what Phase 5 actually produced.
+
+---
## Phase 1 — Registry + Exceptions (2026-04-21)
diff --git a/openrag/components/indexer/chunker/chunker.py b/openrag/components/indexer/chunker/chunker.py
index de9f8934c..a99c688f0 100644
--- a/openrag/components/indexer/chunker/chunker.py
+++ b/openrag/components/indexer/chunker/chunker.py
@@ -1,7 +1,28 @@
+"""Backward-compatibility shim — chunking primitives delegate to `openrag.core.chunking`.
+
+Phase 5B/5.15 status:
+
+* `BaseChunker._get_chunks` / `_prepare_md_elements` / `split_text` →
+ delegated to a held `core.chunking.recursive.RecursiveSplitter` instance
+ via a `Document` ↔ `ProcessedDocument` ↔ `Chunk` adapter.
+* `ChunkContextualizer` and the `split_document()` orchestration sit
+ outside Phase 5B (they belong to 5D / Phase 8). They stay here until
+ Phase 5D ships `core/indexing/contextualize.py`.
+* `ChunkerFactory` is config-driven; the new code uses
+ `chunking_registry`. Both coexist until Phase 8 cutover.
+
+Scheduled for removal in Phase 12.
+"""
+
from typing import Literal
import openai
-from components.indexer.utils.text_sanitizer import sanitize_text
+
+# Side-effect import: pre-loads the indexer-utils submodule so the legacy
+# circular import between `components.utils` and `components.indexer.utils.files`
+# resolves in the correct order. Removing this line breaks chunker collection.
+# Slated to disappear when `components.utils` is split (Phase 6+).
+from components.indexer.utils import text_sanitizer as _text_sanitizer # noqa: F401
from components.prompts import CHUNK_CONTEXTUALIZER_PROMPT
from components.utils import detect_language, get_vlm_semaphore, load_config
from langchain_core.documents.base import Document
@@ -10,17 +31,15 @@
from tqdm.asyncio import tqdm
from utils.logger import get_logger
-from openrag.consts import IMAGE_PLACEHOLDER
+from openrag.core.chunking.recursive import RecursiveSplitter as _CoreRecursiveSplitter
+from openrag.core.models.document import ProcessedDocument, TextBlock
from ..embeddings import BaseEmbedding
-from .utils import MDElement, chunk_table, get_chunk_page_number, split_md_elements
logger = get_logger()
config = load_config()
-# Timeout for individual chunk contextualization LLM calls (in seconds)
CONTEXTUALIZATION_TIMEOUT = config.chunker.contextualization_timeout
-# Maximum concurrent contextualization tasks to prevent system overload
MAX_CONCURRENT_CONTEXTUALIZATION = config.chunker.max_concurrent_contextualization
BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]"
@@ -28,7 +47,12 @@
class ChunkContextualizer:
- """Handles contextualization of document chunks."""
+ """Handles contextualization of document chunks.
+
+ Stays in `components/` until Phase 5D moves the orchestration into
+ `core/indexing/contextualize.py`. The pure prompt-builders for the
+ LLM call already live at `core.prompts.contextualization_builder`.
+ """
def __init__(self, llm_config: dict):
llm_config: dict = dict(llm_config)
@@ -97,7 +121,6 @@ async def contextualize_chunks(
contexts = []
batch_size = MAX_CONCURRENT_CONTEXTUALIZATION
- # Process chunks in batches to limit concurrent LLM calls
for batch_start in range(0, len(chunks), batch_size):
batch_end = min(batch_start + batch_size, len(chunks))
batch_tasks = [
@@ -134,8 +157,36 @@ async def contextualize_chunks(
return chunks
+def _chunks_to_documents(chunks: list, base_metadata: dict) -> list[Document]:
+ """Convert a list of core domain Chunks into legacy LangChain Documents.
+
+ Reproduces the legacy metadata shape: `page`, `chunk_type`, plus the
+ document/partition keys the legacy code stamps onto every chunk.
+ """
+ out: list[Document] = []
+ for c in chunks:
+ meta = dict(c.metadata)
+ meta.update(
+ {
+ "file_id": c.document_id,
+ "partition": c.partition,
+ "page": c.page_number,
+ "chunk_type": c.chunk_type.value,
+ }
+ )
+ # Preserve legacy keys that weren't lifted into core fields.
+ for k, v in base_metadata.items():
+ meta.setdefault(k, v)
+ out.append(Document(page_content=c.text, metadata=meta))
+ return out
+
+
class BaseChunker:
- """Base class for document chunkers with built-in contextualization capability."""
+ """Legacy chunker shell — markdown-aware splitting delegated to core.
+
+ Subclasses configure ``self._core_splitter`` (a
+ `core.chunking.recursive.RecursiveSplitter`) in their `__init__`.
+ """
def __init__(
self,
@@ -152,11 +203,9 @@ def __init__(
self.llm = ChatOpenAI(**llm_config)
self._length_function = self.llm.get_num_tokens
- self.text_splitter = None
+ self._core_splitter: _CoreRecursiveSplitter | None = None
self.contextual_retrieval = contextual_retrieval
-
- # Initialize contextualizer only if needed
self.contextualizer = ChunkContextualizer(llm_config) if contextual_retrieval else None
async def _apply_contextualization(
@@ -177,112 +226,21 @@ async def _apply_contextualization(
return await self.contextualizer.contextualize_chunks(chunks, lang=lang, filename=filename)
- def _prepare_md_elements(self, content: str) -> tuple[list[MDElement], list[MDElement]]:
- """Prepare and combine markdown elements from raw content."""
- md_elements: list[MDElement] = split_md_elements(content)
-
- tables_and_images, texts = [], []
-
- for e in md_elements:
- if e.type in ("table", "image"):
- if e.type == "image" and IMAGE_PLACEHOLDER.lower() in e.content.lower(): # skip placeholder images
- continue
-
- if self._length_function(e.content) <= 100: # do not isolate small tables/images
- texts.append(e)
- else:
- tables_and_images.append(e)
- else:
- texts.append(e)
-
- return texts, tables_and_images
-
- def split_text(self, text: str) -> list[str]:
- """Split text into chunks using the text splitter."""
- if not self.text_splitter:
- logger.warning("Text splitter not initialized. Initializing with default RecursiveCharacterTextSplitter.")
- from langchain.text_splitter import RecursiveCharacterTextSplitter
-
- self.text_splitter = RecursiveCharacterTextSplitter(
- chunk_size=self.chunk_size,
- chunk_overlap=self.chunk_overlap,
- length_function=self._length_function,
- )
-
- return self.text_splitter.split_text(text)
-
def _get_chunks(self, content: str, metadata: dict | None = None, log=None) -> list[Document]:
log = log or logger
- texts, tables_and_images = self._prepare_md_elements(content=content)
- combined_texts = "\n".join([e.content for e in texts])
-
- # Sanitize the combined text before chunking to remove excessive whitespace
- # and useless characters, which saves tokens and improves quality
- sanitized_texts = sanitize_text(
- combined_texts,
- normalize_whitespace=True,
- remove_control_chars=True,
- remove_zero_width_chars=True,
- max_consecutive_newlines=2,
- normalize_unicode=True,
- )
-
- text_chunks = self.split_text(sanitized_texts)
-
- # Manage tables and images as separate chunks
- chunks = []
- for e in tables_and_images:
- if e.type == "table" and self._length_function(e.content) > self.chunk_size:
- # Chunk large tables separately
- subtables = chunk_table(
- table_element=e,
- chunk_size=self.chunk_size,
- length_function=self._length_function,
- )
-
- s = [
- Document(
- page_content=subtable.content.strip(),
- metadata={
- **metadata,
- "page": subtable.page_number,
- "chunk_type": "table",
- },
- )
- for subtable in subtables
- ]
+ metadata = metadata or {}
+ partition = metadata.get("partition", "default")
- else:
- s = [
- Document(
- page_content=e.content.strip(),
- metadata={
- **metadata,
- "page": e.page_number,
- "chunk_type": e.type,
- },
- )
- ]
- chunks.extend(s)
-
- prev_page_num = 1
- for c in text_chunks:
- page_info = get_chunk_page_number(chunk_str=c, previous_chunk_ending_page=prev_page_num)
- start_page = page_info["start_page"]
- prev_page_num = page_info["end_page"]
- chunks.append(
- Document(
- page_content=c.strip(),
- metadata={**metadata, "page": start_page, "chunk_type": "text"},
- )
- )
-
- if chunks:
- chunks.sort(key=lambda d: d.metadata.get("page"))
- return chunks
- else:
+ doc = ProcessedDocument(
+ document_id=metadata.get("file_id", ""),
+ text_blocks=[TextBlock(text=content)],
+ metadata=metadata,
+ )
+ chunks = self._core_splitter.chunk(doc, partition=partition)
+ if not chunks:
log.warning("No chunks created. Content is empty or image is not informative.")
return []
+ return _chunks_to_documents(chunks, base_metadata=metadata)
async def split_document(self, doc: Document, task_id: str | None = None) -> list[Document]:
"""Split document into chunks with optional contextualization."""
@@ -297,11 +255,9 @@ async def split_document(self, doc: Document, task_id: str | None = None) -> lis
detected_lang = detect_language(text=doc.page_content)
- # Process document through pipeline
chunks = self._get_chunks(doc.page_content.strip(), metadata, log=log)
if chunks:
- # Apply contextualization if enabled
log.info(
"Contextualizing chunks",
apply_contextualization=self.contextual_retrieval,
@@ -323,15 +279,10 @@ def __init__(
**kwargs,
):
super().__init__(chunk_size, chunk_overlap_rate, llm_config, contextual_retrieval, **kwargs)
-
- from langchain.text_splitter import RecursiveCharacterTextSplitter
-
- self.text_splitter = RecursiveCharacterTextSplitter(
+ self._core_splitter = _CoreRecursiveSplitter(
chunk_size=self.chunk_size,
- chunk_overlap=self.chunk_overlap,
+ chunk_overlap_rate=self.chunk_overlap_rate,
length_function=self._length_function,
- is_separator_regex=True,
- separators=["\n", r"(?<=[\.\?\!])"],
)
@@ -345,11 +296,9 @@ def create_chunker(
config,
embedder: BaseEmbedding | None = None,
) -> BaseChunker:
- # Extract parameters
chunker_params = config.chunker.model_dump()
name = chunker_params.pop("name")
- # Initialize and return the chunker
chunker_cls: BaseChunker = ChunkerFactory.CHUNKERS.get(name)
if not chunker_cls:
diff --git a/openrag/components/indexer/chunker/utils.py b/openrag/components/indexer/chunker/utils.py
index 2529b616d..02e3973d8 100644
--- a/openrag/components/indexer/chunker/utils.py
+++ b/openrag/components/indexer/chunker/utils.py
@@ -1,252 +1,33 @@
-import re
-from collections.abc import Callable
-from typing import Literal
-
-from components.indexer.utils.text_sanitizer import clean_markdown_table_spacing
-
-# Regex to match a Markdown table (header + delimiter + at least one row)
-TABLE_RE = re.compile(
- r"((?:^|\n)\|.*?\|\r?\n\|\s*[:-]+(?:\s*\|[:-]+)*\|\r?\n(?:\|.*?\|\r?\n)+)",
- re.DOTALL | re.MULTILINE,
+"""Backward-compatibility shim — re-exports from `openrag.core.chunking.markdown_utils`.
+
+The implementation moved to `openrag/core/chunking/markdown_utils.py` in
+Phase 5B. New code should import from there directly. This file is kept
+so existing legacy imports keep working until the consumers migrate;
+scheduled for removal in Phase 12.
+"""
+
+from openrag.core.chunking.markdown_utils import (
+ IMAGE_RE,
+ PAGE_RE,
+ TABLE_RE,
+ MDElement,
+ chunk_table,
+ get_chunk_page_number,
+ get_page_number,
+ parse_markdown_table,
+ span_inside,
+ split_md_elements,
)
-# Regex to match image descriptions
-IMAGE_RE = re.compile(r"((.*?))", re.DOTALL)
-
-# Regex to match page markers
-PAGE_RE = re.compile(r"\[PAGE_(\d+)\]")
-
-
-class MDElement:
- """Class representing a segment of markdown content."""
-
- def __init__(
- self,
- type: Literal["text", "table", "image"],
- content: str,
- page_number: int | None = None,
- ):
- self.type = type # 'text', 'table', 'image'
- self.content = content
- self.page_number = page_number
-
- def __repr__(self):
- return f"Element(type={self.type}, page_number={self.page_number}, content={self.content[:100]}...)"
-
-
-def span_inside(span: tuple[int, int], container: tuple[int, int]) -> bool:
- return container[0] <= span[0] and span[1] <= container[1]
-
-
-def get_page_number(position, page_markers):
- """
- Given a position in the text and list of (position, page_number) tuples,
- return the page number for that position.
- Content after [PAGE_N] marker belongs to page N+1.
- """
- current_page = 1 # Default to page 1 if before any markers
- for marker_pos, page_num in page_markers:
- if position >= marker_pos:
- current_page = page_num + 1 # Content after [PAGE_N] is on page N+1
- else:
- break
- return current_page
-
-
-def split_md_elements(md_text: str) -> list[MDElement]:
- """
- Split markdown text into segments of text, tables, and images.
- Returns a list of tuples:
- - ('text', content) for text segments
- - ('table', content, page_number) for tables
- - ('image', content, page_number) for images
- """
- # Find all page markers
- page_markers = []
- for match in PAGE_RE.finditer(md_text):
- page_markers.append((match.start(), int(match.group(1))))
- page_markers.sort() # Ensure they're in order
-
- all_matches = []
-
- # Find image matches first and record their spans
- image_spans = []
- for match in IMAGE_RE.finditer(md_text):
- span = match.span()
- page_num = get_page_number(span[0], page_markers)
- all_matches.append((span, "image", match.group(1).strip(), page_num))
- image_spans.append(span)
-
- # Find table matches, but skip those that are fully inside an image description
- for match in TABLE_RE.finditer(md_text):
- span = match.span()
- if not any(span_inside(span, image_span) for image_span in image_spans):
- page_num = get_page_number(span[0], page_markers)
- all_matches.append((span, "table", match.group(1).strip(), page_num))
-
- # Sort matches by start position
- all_matches.sort(key=lambda x: x[0][0])
-
- parts = []
- last = 0
-
- for (start, end), match_type, content, page_num in all_matches:
- # Add text segment before this match if there is any
- if start > last:
- text_segment = md_text[last:start]
- if text_segment.strip(): # Only add non-empty text segments
- parts.append(("text", text_segment.strip()))
-
- # Add the matched segment with page number
- parts.append((match_type, content, page_num))
- last = end
-
- # Add remaining text after the last match
- if last < len(md_text):
- remaining_text = md_text[last:]
- if remaining_text.strip(): # Only add non-empty text segments
- parts.append(("text", remaining_text.strip()))
-
- return [MDElement(*p) for p in parts]
-
-
-def get_chunk_page_number(chunk_str: str, previous_chunk_ending_page=1):
- """
- Determine the start and end pages for a text chunk containing [PAGE_N] separators.
- PAGE_N marks the end of page N - text before separator is on page N.
- """
- # Find all page separator matches in the chunk
- matches = list(PAGE_RE.finditer(chunk_str))
-
- if not matches:
- # No separators found - entire chunk is on previous page
- return {
- "start_page": previous_chunk_ending_page,
- "end_page": previous_chunk_ending_page,
- }
-
- first_match = matches[0]
- last_match = matches[-1]
- last_char_idx = len(chunk_str) - 1
-
- # Determine start page
- if first_match.start() == 0:
- # Chunk starts with a separator - begins on next page
- start_page = int(first_match.group(1)) + 1
- else:
- # Text precedes first separator - starts on previous page
- start_page = previous_chunk_ending_page
-
- # Determine end page
- if last_match.end() - 1 == last_char_idx:
- # Chunk ends exactly at a separator - ends on that page
- end_page = int(last_match.group(1))
- else:
- # Chunk ends after separator - ends on next page
- end_page = int(last_match.group(1)) + 1
-
- return {"start_page": start_page, "end_page": end_page}
-
-
-def parse_markdown_table(markdown_table):
- """
- Parse a markdown table and extract header and groups based on Domain column.
-
- Returns:
- tuple: (header_lines, groups)
- - header_lines: list of [header_row, separator_row]
- - groups: list of lists, each containing rows belonging to one domain
- """
- lines = markdown_table.strip().split("\n")
-
- # Extract header (first 2 lines)
- header_lines = lines[:2]
- data_rows = lines[2:]
-
- # Group rows by Domain (first column)
- groups = []
- current_group = []
-
- for row in data_rows:
- # Parse first column (Domain)
- cells = [cell.strip() for cell in row.split("|")[1:-1]]
- if not cells:
- continue # skip malformed rows
-
- domain = cells[0]
-
- # If Domain is not empty, start a new group
- if domain:
- if current_group: # Save previous group
- groups.append(current_group)
- current_group = [row] # Start new group
- else:
- # Domain is empty, continue current group
- current_group.append(row)
-
- # Don't forget the last group
- if current_group:
- groups.append(current_group)
-
- return header_lines, groups
-
-
-def chunk_table(
- table_element: MDElement,
- chunk_size: int = 512,
- length_function: Callable[[str], int] | None = None,
-) -> list[MDElement]:
- txt = clean_markdown_table_spacing(table_element.content)
- header_lines, groups = parse_markdown_table(txt)
-
- # Convert header lines → text block
- header_text = "\n".join(header_lines)
-
- # Convert group lists → text blocks
- group_texts = ["\n".join(g) for g in groups]
-
- # Precompute token length
- header_ntoks = length_function(header_text)
- groups_ntoks = [length_function(g) for g in group_texts]
-
- subtables = []
- current_rows = [header_text]
- current_size = header_ntoks
-
- prev_last_row = None # for overlap
-
- for group_txt, g_ntoks in zip(group_texts, groups_ntoks, strict=True):
- # If adding this group exceeds the chunk limit
- if current_size + g_ntoks > chunk_size:
- # ---- finalize current subtable ----
- subtables.append("\n".join(current_rows))
-
- # ---- start new subtable with OVERLAP ----
- current_rows = [header_text] # always restart headers
- if prev_last_row:
- current_rows.append(prev_last_row) # add overlapping row
-
- current_rows.append(group_txt)
- current_size = header_ntoks + (length_function(prev_last_row) if prev_last_row else 0) + g_ntoks
-
- else:
- # fits → just append normally
- current_rows.append(group_txt)
- current_size += g_ntoks
-
- # track last row for overlap
- prev_last_row = group_txt
-
- # finalize last subtable
- if current_rows:
- subtables.append("\n".join(current_rows))
-
- # wrap into MDElement list
- return [
- MDElement(
- type="table",
- content=subtable,
- page_number=table_element.page_number,
- )
- for subtable in subtables
- ]
+__all__ = [
+ "IMAGE_RE",
+ "MDElement",
+ "PAGE_RE",
+ "TABLE_RE",
+ "chunk_table",
+ "get_chunk_page_number",
+ "get_page_number",
+ "parse_markdown_table",
+ "span_inside",
+ "split_md_elements",
+]
diff --git a/openrag/components/pipeline.py b/openrag/components/pipeline.py
index 8187392d0..b43888e1b 100644
--- a/openrag/components/pipeline.py
+++ b/openrag/components/pipeline.py
@@ -2,7 +2,6 @@
import copy
from datetime import datetime
from enum import Enum
-from typing import Literal
import openai
import ray
@@ -18,9 +17,16 @@
from langchain_core.documents.base import Document
from langchain_core.exceptions import OutputParserException
from langchain_openai import ChatOpenAI
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import ValidationError
from utils.logger import get_logger
+# Phase 5/5.15: domain query model + RetrieverPipeline live in core/. This file
+# re-exports them and shims the legacy RetrieverPipeline as an adapter.
+from openrag.core.models.chunk import Chunk
+from openrag.core.models.query import Query, SearchQueries, TemporalPredicate
+from openrag.core.rerankers.reranker import Reranker as _CoreReranker
+from openrag.core.retrieval.pipeline import RetrieverPipeline as _CoreRetrieverPipeline
+
from .llm import LLM
from .map_reduce import RAGMapReduce
from .reranker import BaseReranker, RerankerFactory
@@ -31,93 +37,77 @@
config = load_config()
VECTORDB_TIMEOUT = config.ray.indexer.vectordb_timeout
+__all__ = [
+ "Query",
+ "RAGMODE",
+ "RagPipeline",
+ "RetrieverPipeline",
+ "SearchQueries",
+ "TemporalPredicate",
+]
+
class RAGMODE(Enum):
SIMPLERAG = "SimpleRag"
CHATBOTRAG = "ChatBotRag"
-class TemporalPredicate(BaseModel):
- """A single constraint on a document's creation date.
+class _LegacyRerankerAdapter(_CoreReranker):
+ """Wraps a legacy ``BaseReranker`` (Document-in / Document-out) so it
+ satisfies the core ``Reranker`` ABC (str-in / (idx, score)-out).
- Multiple predicates on the same `Query` are combined with logical AND.
- Use two predicates to express a closed range (e.g. last month):
- [{op: ">=", value: "2026-03-01..."}, {op: "<=", value: "2026-03-31..."}]
+ The legacy reranker only returns reordered documents; we tag each
+ incoming text with its original index via the wrapping Document's
+ metadata, then read it back to produce ``(idx, rank-score)`` tuples.
+ Score values are synthetic (``1 / (rank+1)``) — only the order is
+ consumed by the core pipeline.
"""
- field: Literal["created_at"] = Field(
- default="created_at",
- description="Document metadata field to filter on. Always `created_at` for now.",
- )
- operator: Literal[">", "<", ">=", "<="] = Field(
- description="Comparison operator applied to the date field.",
- )
- value: str = Field(
- description='ISO 8601 datetime with timezone, e.g. "2026-03-15T00:00:00+00:00".',
- )
+ def __init__(self, legacy: BaseReranker) -> None:
+ self._legacy = legacy
+ async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]:
+ tagged = [Document(page_content=t, metadata={"_legacy_rerank_idx": i}) for i, t in enumerate(documents)]
+ reordered = await self._legacy.rerank(query=query, documents=tagged, top_k=top_k)
+ return [(d.metadata["_legacy_rerank_idx"], 1.0 / (rank + 1)) for rank, d in enumerate(reordered)]
-class Query(BaseModel):
- """A single vector database search query with optional temporal filters on document creation date.
-
- Predicates in `temporal_filters` are AND-combined. To express an exclusion
- (e.g. "last year except March"), emit TWO `Query` objects, each with its own
- AND-combined predicates covering one side of the gap.
- """
-
- query: str = Field(description="A semantically enriched, descriptive query for vector similarity search.")
- temporal_filters: list[TemporalPredicate] | None = Field(
- default=None,
- description="Date predicates on `created_at`, AND-combined. Null when no temporal reference in the query.",
- )
-
- def to_milvus_filter(self) -> str | None:
- """The temporal_filters attributes are already checked through the Pydantic types, except for date value that is kept as string,
- as LLM sometimes give correct but not entirely complete date
- """
-
- if not self.temporal_filters:
- return None
- parts = []
- for p in self.temporal_filters:
- try:
- datetime.fromisoformat(p.value)
- except (TypeError, ValueError):
- logger.warning(
- "Dropping temporal predicate with non-ISO value",
- field=p.field,
- operator=p.operator,
- value=p.value,
- )
- continue
- parts.append(f'{p.field} {p.operator} ISO "{p.value}"')
- if not parts:
- return None
- return " and ".join(parts)
- def __str__(self) -> str:
- return f"Query: {self.query}, Filter: {self.to_milvus_filter()}"
+def _to_documents(chunks: list[Chunk]) -> list[Document]:
+ return [c.to_langchain() for c in chunks]
-class SearchQueries(BaseModel):
- query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.")
-
- def __str__(self) -> str:
- return " --- ".join(str(q) for q in self.query_list)
+class RetrieverPipeline:
+ """Backward-compat adapter — delegates to ``core.retrieval.pipeline.RetrieverPipeline``.
+ The legacy retriever shim already provides a core ``Retriever``; we
+ wrap the legacy reranker in ``_LegacyRerankerAdapter`` and hand both
+ to the core pipeline. Outputs are converted ``Chunk → Document`` so
+ legacy callers (``RagPipeline``) keep working unchanged.
+ """
-class RetrieverPipeline:
def __init__(self) -> None:
- # retriever
self.retriever: BaseRetriever = RetrieverFactory.create_retriever(config=config)
self.allow_filterless_fallback = config.retriever.allow_filterless_fallback
- # reranker
self.reranker_enabled = config.reranker.enabled
self.reranker: BaseReranker = RerankerFactory.get_reranker(config)
logger.debug("Reranker", enabled=self.reranker_enabled, provider=config.reranker.provider)
self.reranker_top_k = config.reranker.top_k
+ self._core_pipeline: _CoreRetrieverPipeline | None = None
+
+ def _ensure_core_pipeline(self) -> _CoreRetrieverPipeline:
+ # Built lazily so the underlying Ray actor only needs to exist at
+ # first request time, not at module import / pipeline construction.
+ if self._core_pipeline is None:
+ self._core_pipeline = _CoreRetrieverPipeline(
+ retriever=self.retriever._build_core_retriever(),
+ reranker=_LegacyRerankerAdapter(self.reranker) if self.reranker_enabled else None,
+ reranker_top_k=self.reranker_top_k,
+ allow_filterless_fallback=self.allow_filterless_fallback,
+ )
+ return self._core_pipeline
+
async def retrieve_docs(
self,
partition: list[str],
@@ -125,54 +115,11 @@ async def retrieve_docs(
top_k: int | None = None,
filter_params: dict | None = None,
) -> list[Document]:
- milvus_filter = query.to_milvus_filter()
- docs = await self.retriever.retrieve(
- partition=partition, query=query.query, filter=milvus_filter, filter_params=filter_params
+ chunks = await self._ensure_core_pipeline().retrieve_docs(
+ partition=partition, query=query, top_k=top_k, filter_params=filter_params
)
-
- # Fallback: drop temporal filter if it wiped out all candidates.
- # Gated by `retriever.allow_filterless_fallback` so deployments that
- # prefer strict temporal retrieval can opt out (returns no docs
- # rather than temporally-incorrect ones).
- if not docs and milvus_filter and self.allow_filterless_fallback:
- logger.warning(
- "Temporal filter dropped: no documents matched, retrying without filter",
- query=str(query.query),
- filter=milvus_filter,
- partition=partition,
- )
- docs = await self.retriever.retrieve(
- partition=partition, query=query.query, filter=None, filter_params=filter_params
- )
-
- logger.debug("Documents retreived", document_count=len(docs))
-
- if docs:
- # 1. rerank all the docs
- if self.reranker_enabled:
- docs = await self.reranker.rerank(query=query.query, documents=docs, top_k=None)
- logger.debug("Documents reranked", document_count=len(docs))
-
- # 2. expand the docs with related documents
- if self.retriever.expansion_enabled:
- # Limit the number of docs to expand
- top_k = max(self.reranker_top_k, top_k) if top_k else self.reranker_top_k
- docs2expand = copy.deepcopy(docs[:top_k])
-
- logger.debug("Documents to expand", document_count=len(docs2expand))
- expanded_docs = await self.retriever.expand_search_results(results=docs2expand)
- if len(docs2expand) == len(expanded_docs): # no expansion found, keep the original docs
- return docs
-
- logger.debug("Documents expanded", document_count=len(expanded_docs))
- docs = expanded_docs
-
- # rerank again after expansion if reranker is enabled
- if self.reranker_enabled:
- docs = await self.reranker.rerank(query=query.query, documents=docs, top_k=None)
- logger.debug("Documents after expansion and reranking", document_count=len(docs))
-
- return docs
+ logger.debug("Documents retrieved", document_count=len(chunks))
+ return _to_documents(chunks)
async def get_relevant_docs(
self,
@@ -181,16 +128,11 @@ async def get_relevant_docs(
top_k: int | None = None,
filter_params: dict | None = None,
) -> list[Document]:
- tasks = [
- self.retrieve_docs(partition=partition, query=q, top_k=top_k, filter_params=filter_params)
- for q in search_queries.query_list
- ]
- results = await asyncio.gather(*tasks)
- results = self.reranker.rrf_reranking(doc_lists=results)
- if top_k is not None:
- results = results[:top_k]
- logger.debug("Final relevant documents after RRF reranking", document_count=len(results))
- return results
+ chunks = await self._ensure_core_pipeline().get_relevant_docs(
+ partition=partition, search_queries=search_queries, top_k=top_k, filter_params=filter_params
+ )
+ logger.debug("Final relevant documents after RRF reranking", document_count=len(chunks))
+ return _to_documents(chunks)
class RagPipeline:
diff --git a/openrag/components/prompts/prompts.py b/openrag/components/prompts/prompts.py
index afb2c3727..76d67a5b6 100644
--- a/openrag/components/prompts/prompts.py
+++ b/openrag/components/prompts/prompts.py
@@ -1,7 +1,22 @@
+"""Backward-compatibility shim — delegates to `openrag.core.prompts.template_loader`.
+
+The disk-based template loader moved to
+`openrag/core/prompts/template_loader.py` in Phase 5C. This module is
+retained for legacy imports of `load_prompt(...)` and the eagerly-loaded
+SYS_PROMPT_TMPLT / *_PROMPT constants until consumers migrate;
+scheduled for removal in Phase 12.
+
+The new function takes (prompts_dir, mapping, key) explicitly; this
+shim's `load_prompt(key)` resolves the first two from the cached
+config, matching the legacy call site shape.
+"""
+
from pathlib import Path
from config import load_config
+from openrag.core.prompts.template_loader import load_template_by_key
+
config = load_config()
prompts_dir: Path = config.paths.prompts_dir
@@ -13,29 +28,18 @@ def load_prompt(
prompts_dir: Path = prompts_dir,
prompt_mapping=prompt_mapping,
) -> str:
- file_name = getattr(prompt_mapping, prompt_name, None)
- if not file_name:
- raise ValueError(f"No associated file name found for prompt: `{prompt_name}`")
-
- file_path = prompts_dir / file_name
-
- if not file_path.exists():
- raise FileNotFoundError(f"Prompt file not found: `{file_path}`")
-
- with open(file_path) as f:
- sys_msg = f.read()
- return sys_msg
+ return load_template_by_key(prompts_dir, prompt_mapping, prompt_name)
-# Load prompts
+# Eagerly-loaded prompt strings — preserved for legacy callers that
+# import these names directly. New code should call `load_template_by_key`
+# (or `load_template`) on demand instead.
SYS_PROMPT_TMPLT = load_prompt("sys_prompt")
QUERY_CONTEXTUALIZER_PROMPT = load_prompt("query_contextualizer")
CHUNK_CONTEXTUALIZER_PROMPT = load_prompt("chunk_contextualizer")
IMAGE_DESCRIBER = load_prompt("image_describer")
-# Retrievers prompts
HYDE_PROMPT = load_prompt("hyde")
MULTI_QUERY_PROMPT = load_prompt("multi_query")
-# Short answer prompt
SPOKEN_STYLE_ANSWER_PROMPT = load_prompt("spoken_style_answer")
diff --git a/openrag/components/retriever.py b/openrag/components/retriever.py
index 3f074adaf..667f3c21f 100644
--- a/openrag/components/retriever.py
+++ b/openrag/components/retriever.py
@@ -1,20 +1,91 @@
-# Import necessary modules and classes
-import asyncio
+"""Backward-compatibility shim — retriever strategies delegate to `openrag.core.retrieval`.
+
+Phase 5A/5.15 status:
+
+* `BaseRetriever` / `SingleRetriever` / `MultiQueryRetriever` / `HyDeRetriever`
+ → adapters wrapping the corresponding `core.retrieval.retriever` strategies.
+ The Ray actor is wrapped in a `MilvusRayShim` so the core retriever talks
+ to a `RetrievalSearcher` port. The LLM (legacy `ChatOpenAI`) is wrapped in
+ a `_LangChainLLMAdapter` so it fits the core `LLM` ABC.
+* Output is converted from domain `Chunk` back to LangChain `Document` so
+ legacy callers (RetrieverPipeline, RagPipeline) keep working unchanged.
+* `_expand_with_related_chunks` → delegates to
+ `core.retrieval.retriever._expand_with_related_chunks` with the same
+ conversion at the boundary.
+* `RetrieverFactory` is config-driven; the new code uses
+ `retriever_registry`. Both coexist until Phase 8 cutover.
+
+Scheduled for removal in Phase 12.
+"""
+
from abc import ABC, abstractmethod
-from itertools import chain
-from typing import ClassVar
+from typing import Any, ClassVar
from components.prompts import HYDE_PROMPT, MULTI_QUERY_PROMPT
from langchain_core.documents.base import Document
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from utils.dependencies import get_vectordb
from utils.logger import get_logger
+from openrag.core.llm.llm import LLM as _CoreLLM
+from openrag.core.models.chunk import Chunk
+from openrag.core.retrieval.retriever import (
+ HyDeRetriever as _CoreHyDeRetriever,
+)
+from openrag.core.retrieval.retriever import (
+ MultiQueryRetriever as _CoreMultiQueryRetriever,
+)
+from openrag.core.retrieval.retriever import (
+ SingleRetriever as _CoreSingleRetriever,
+)
+from openrag.core.retrieval.retriever import (
+ _expand_with_related_chunks as _core_expand,
+)
+from openrag.services.storage.milvus_ray_shim import MilvusRayShim
+
logger = get_logger()
+# ---------------------------------------------------------------------------
+# Adapters bridging legacy types (ChatOpenAI, Ray actor, Document) to core.
+# ---------------------------------------------------------------------------
+class _LangChainLLMAdapter(_CoreLLM):
+ """Wraps a LangChain ``ChatOpenAI`` so it satisfies the core ``LLM`` ABC."""
+
+ _ROLE_MAP: ClassVar[dict] = {"user": HumanMessage, "system": SystemMessage, "assistant": AIMessage}
+
+ def __init__(self, lc_llm: ChatOpenAI) -> None:
+ self._llm = lc_llm
+
+ async def generate(self, prompt: str, **kwargs) -> str:
+ out = await self._llm.ainvoke(prompt)
+ return out.content if hasattr(out, "content") else str(out)
+
+ async def chat(self, messages: list[dict[str, str]], **kwargs) -> str:
+ lc_msgs = [self._ROLE_MAP[m["role"]](content=m["content"]) for m in messages]
+ out = await self._llm.ainvoke(lc_msgs)
+ return out.content
+
+
+def _searcher() -> MilvusRayShim:
+ """Wrap the legacy Vectordb Ray actor as a core ``RetrievalSearcher``."""
+ return MilvusRayShim(get_vectordb())
+
+
+def _to_documents(chunks: list[Chunk]) -> list[Document]:
+ """Convert core ``Chunk`` objects back to LangChain ``Document``s for legacy callers."""
+ return [c.to_langchain() for c in chunks]
+
+
+def _from_documents(docs: list[Document]) -> list[Chunk]:
+ """Convert legacy ``Document``s into core ``Chunk``s for the expansion helper."""
+ return [Chunk.from_langchain(d) for d in docs]
+
+
+# ---------------------------------------------------------------------------
+# Legacy ABCs — preserved so existing isinstance / type hints keep working.
+# ---------------------------------------------------------------------------
class ABCRetriever(ABC):
"""Abstract class for the base retriever."""
@@ -41,28 +112,23 @@ async def expand_search_results(self, results: list[Document]) -> list[Document]
pass
-# Define the Simple Retriever class
class BaseRetriever(ABCRetriever):
+ """Common adapter — instantiates a core retriever and converts I/O at the boundary."""
+
+ _CORE_CLS: type = _CoreSingleRetriever
+
def __init__(
self,
- top_k=6,
- similarity_threshold=0.95,
- with_surrounding_chunks=True,
- include_related=False,
- include_ancestors=False,
- related_limit=10,
+ top_k: int = 6,
+ similarity_threshold: float = 0.95,
+ with_surrounding_chunks: bool = True,
+ include_related: bool = False,
+ include_ancestors: bool = False,
+ related_limit: int = 10,
max_ancestor_depth: int | None = None,
**kwargs,
- ):
- super().__init__(
- top_k,
- similarity_threshold,
- include_related=include_related,
- include_ancestors=include_ancestors,
- related_limit=related_limit,
- max_ancestor_depth=max_ancestor_depth,
- **kwargs,
- )
+ ) -> None:
+ # Mirror legacy attributes so external callers reading them still work.
self.top_k = top_k
self.similarity_threshold = similarity_threshold
self.with_surrounding_chunks = with_surrounding_chunks
@@ -71,6 +137,23 @@ def __init__(
self.related_limit = related_limit
self.max_ancestor_depth = max_ancestor_depth
self.expansion_enabled = include_related or include_ancestors
+ self._core_kwargs = self._build_core_kwargs(kwargs)
+
+ def _build_core_kwargs(self, extra: dict[str, Any]) -> dict[str, Any]:
+ """Kwargs handed to the core retriever's constructor."""
+ return {
+ "top_k": self.top_k,
+ "similarity_threshold": self.similarity_threshold,
+ "with_surrounding_chunks": self.with_surrounding_chunks,
+ "include_related": self.include_related,
+ "include_ancestors": self.include_ancestors,
+ "related_limit": self.related_limit,
+ "max_ancestor_depth": self.max_ancestor_depth,
+ }
+
+ def _build_core_retriever(self):
+ """Late-bind the core retriever so the Ray actor is only resolved on use."""
+ return self._CORE_CLS(searcher=_searcher(), **self._core_kwargs)
async def retrieve(
self,
@@ -79,152 +162,116 @@ async def retrieve(
filter: str | None = None,
filter_params: dict | None = None,
) -> list[Document]:
- db = get_vectordb()
- chunks = await db.async_search.remote(
- query=query,
- partition=partition,
- top_k=self.top_k,
- filter=filter,
- filter_params=filter_params,
- similarity_threshold=self.similarity_threshold,
- with_surrounding_chunks=self.with_surrounding_chunks,
+ chunks = await self._build_core_retriever().retrieve(
+ partition=partition, query=query, filter=filter, filter_params=filter_params
)
- return chunks
+ return _to_documents(chunks)
async def expand_search_results(self, results: list[Document]) -> list[Document]:
- """Expand search results with related and ancestor chunks."""
- db = get_vectordb()
- return await _expand_with_related_chunks(
- db=db,
- results=results,
+ expanded = await _core_expand(
+ searcher=_searcher(),
+ results=_from_documents(results),
include_related=self.include_related,
include_ancestors=self.include_ancestors,
related_limit=self.related_limit,
max_ancestor_depth=self.max_ancestor_depth,
)
+ return _to_documents(expanded)
class SingleRetriever(BaseRetriever):
- pass
+ _CORE_CLS = _CoreSingleRetriever
class MultiQueryRetriever(BaseRetriever):
+ _CORE_CLS = _CoreMultiQueryRetriever
+
def __init__(
self,
- top_k=6,
- similarity_threshold=0.95,
- with_surrounding_chunks=True,
- include_related=False,
- include_ancestors=False,
- related_limit=10,
- max_ancestor_depth=None,
+ top_k: int = 6,
+ similarity_threshold: float = 0.95,
+ with_surrounding_chunks: bool = True,
+ include_related: bool = False,
+ include_ancestors: bool = False,
+ related_limit: int = 10,
+ max_ancestor_depth: int | None = None,
k_queries: int = 3,
- llm: ChatOpenAI = None,
+ llm: ChatOpenAI | None = None,
**kwargs,
- ):
+ ) -> None:
+ if llm is None:
+ raise ValueError("llm must be provided for MultiQueryRetriever")
+ self.k_queries = k_queries
+ self.llm = llm
super().__init__(
- top_k,
- similarity_threshold,
- with_surrounding_chunks,
- include_related,
- include_ancestors,
- related_limit,
- max_ancestor_depth,
+ top_k=top_k,
+ similarity_threshold=similarity_threshold,
+ with_surrounding_chunks=with_surrounding_chunks,
+ include_related=include_related,
+ include_ancestors=include_ancestors,
+ related_limit=related_limit,
+ max_ancestor_depth=max_ancestor_depth,
**kwargs,
)
- self.k_queries = k_queries
- self.llm = llm
-
- if llm is None:
- raise ValueError("llm must be provided for MultiQueryRetriever")
-
- prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(MULTI_QUERY_PROMPT)
- self.generate_queries = prompt | llm | StrOutputParser() | (lambda x: x.split("[SEP]"))
-
- async def retrieve(
- self, partition: list[str], query: str, filter: str | None = None, filter_params: dict | None = None
- ):
- db = get_vectordb()
- logger.debug("Generating multiple queries", k_queries=self.k_queries)
- generated_queries = await self.generate_queries.ainvoke(
- {
- "query": query,
- "k_queries": self.k_queries,
- }
- )
- chunks = await db.async_multi_query_search.remote(
- queries=generated_queries,
- partition=partition,
- top_k_per_query=self.top_k,
- filter=filter,
- filter_params=filter_params,
- similarity_threshold=self.similarity_threshold,
- with_surrounding_chunks=self.with_surrounding_chunks,
+ def _build_core_kwargs(self, extra: dict[str, Any]) -> dict[str, Any]:
+ kw = super()._build_core_kwargs(extra)
+ kw.update(
+ llm=_LangChainLLMAdapter(self.llm),
+ multi_query_template=MULTI_QUERY_PROMPT,
+ k_queries=self.k_queries,
)
- return chunks
+ return kw
class HyDeRetriever(BaseRetriever):
+ _CORE_CLS = _CoreHyDeRetriever
+
def __init__(
self,
- top_k=6,
- similarity_threshold=0.95,
- with_surrounding_chunks=True,
- include_related=False,
- include_ancestors=False,
- related_limit=10,
- max_ancestor_depth=None,
- llm: ChatOpenAI = None,
+ top_k: int = 6,
+ similarity_threshold: float = 0.95,
+ with_surrounding_chunks: bool = True,
+ include_related: bool = False,
+ include_ancestors: bool = False,
+ related_limit: int = 10,
+ max_ancestor_depth: int | None = None,
+ llm: ChatOpenAI | None = None,
combine: bool = False,
**kwargs,
- ):
- super().__init__(
- top_k,
- similarity_threshold,
- with_surrounding_chunks,
- include_related,
- include_ancestors,
- related_limit,
- max_ancestor_depth,
- **kwargs,
- )
-
- super().__init__(top_k, similarity_threshold, **kwargs)
+ ) -> None:
if llm is None:
raise ValueError("llm must be provided for HyDeRetriever")
-
self.combine = combine
self.llm = llm
+ super().__init__(
+ top_k=top_k,
+ similarity_threshold=similarity_threshold,
+ with_surrounding_chunks=with_surrounding_chunks,
+ include_related=include_related,
+ include_ancestors=include_ancestors,
+ related_limit=related_limit,
+ max_ancestor_depth=max_ancestor_depth,
+ **kwargs,
+ )
- prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(HYDE_PROMPT)
- self.hyde_generator = prompt | llm | StrOutputParser()
-
- async def get_hyde(self, query: str):
- logger.debug("Generating HyDe Document")
- hyde_document = await self.hyde_generator.ainvoke({"query": query})
- return hyde_document
-
- async def retrieve(
- self, partition: list[str], query: str, filter: str | None = None, filter_params: dict | None = None
- ) -> list[Document]:
- db = get_vectordb()
- hyde = await self.get_hyde(query)
- queries = [hyde]
- if self.combine:
- queries.append(query)
-
- return await db.async_multi_query_search.remote(
- queries=queries,
- partition=partition,
- top_k_per_query=self.top_k,
- filter=filter,
- filter_params=filter_params,
- similarity_threshold=self.similarity_threshold,
- with_surrounding_chunks=self.with_surrounding_chunks,
+ def _build_core_kwargs(self, extra: dict[str, Any]) -> dict[str, Any]:
+ kw = super()._build_core_kwargs(extra)
+ kw.update(
+ llm=_LangChainLLMAdapter(self.llm),
+ hyde_template=HYDE_PROMPT,
+ combine=self.combine,
)
+ return kw
+ async def get_hyde(self, query: str) -> str:
+ # Preserved for legacy callers / tests that introspect this method.
+ return await self._build_core_retriever().get_hyde(query)
+
+# ---------------------------------------------------------------------------
+# Legacy free function — preserved for external callers; routes through core.
+# ---------------------------------------------------------------------------
async def _expand_with_related_chunks(
db,
results: list[Document],
@@ -233,83 +280,20 @@ async def _expand_with_related_chunks(
related_limit: int = 10,
max_ancestor_depth: int | None = None,
) -> list[Document]:
- """Expand results with related and/or ancestor chunks."""
- if not results or (not include_related and not include_ancestors):
- return results
-
- # Track what we already have to avoid duplicates
- seen_ids = {doc.metadata.get("_id") for doc in results}
- expanded_results = list(results)
-
- # Collect unique relationship_ids and file_ids from results
- relationship_ids = set()
- file_infos = [] # List of (partition, file_id) tuples
-
- for doc in results:
- metadata = doc.metadata
- if include_related and metadata.get("relationship_id"):
- relationship_ids.add((metadata.get("partition"), metadata.get("relationship_id")))
- if include_ancestors:
- file_infos.append((metadata.get("partition"), metadata.get("file_id")))
-
- # Create tasks for parallel fetching
- async def fetch_related(partition: str, rel_id: str) -> list[Document]:
- """Fetch related chunks with error handling."""
- try:
- return await db.get_related_chunks.remote(
- partition=partition,
- relationship_id=rel_id,
- limit=related_limit,
- )
- except Exception as e:
- logger.warning(
- "Failed to fetch related chunks",
- relationship_id=rel_id,
- error=str(e),
- )
- return []
-
- async def fetch_ancestors(partition: str, file_id: str) -> list[Document]:
- """Fetch ancestor chunks with error handling."""
- try:
- return await db.get_ancestor_chunks.remote(
- partition=partition,
- file_id=file_id,
- limit=related_limit,
- max_ancestor_depth=max_ancestor_depth,
- )
- except Exception as e:
- logger.warning(
- "Failed to fetch ancestor chunks",
- file_id=file_id,
- error=str(e),
- )
- return []
-
- # Build list of tasks for parallel execution
- tasks = []
-
- if include_related:
- tasks.extend(fetch_related(partition, rel_id) for partition, rel_id in relationship_ids if partition and rel_id)
-
- if include_ancestors:
- tasks.extend(fetch_ancestors(partition, file_id) for partition, file_id in file_infos if partition and file_id)
-
- # Execute all tasks in parallel
- if tasks:
- all_results = await asyncio.gather(*tasks)
- for chunk in chain.from_iterable(all_results):
- chunk_id = chunk.metadata.get("_id")
- if chunk_id and chunk_id not in seen_ids:
- seen_ids.add(chunk_id)
- expanded_results.append(chunk)
-
- logger.debug(
- "Expanded results with related/ancestor chunks",
- original_count=len(results),
- expanded_count=len(expanded_results),
+ """Backward-compat free-function — delegates to the core expansion helper.
+
+ The legacy callers pass a Ray actor via ``db``; we wrap it in
+ ``MilvusRayShim`` to satisfy the core ``RetrievalSearcher`` port.
+ """
+ expanded = await _core_expand(
+ searcher=MilvusRayShim(db),
+ results=_from_documents(results),
+ include_related=include_related,
+ include_ancestors=include_ancestors,
+ related_limit=related_limit,
+ max_ancestor_depth=max_ancestor_depth,
)
- return expanded_results
+ return _to_documents(expanded)
class RetrieverFactory:
@@ -322,7 +306,6 @@ class RetrieverFactory:
@classmethod
def create_retriever(cls, config) -> ABCRetriever:
retrieverConfig = config.retriever.model_dump()
-
retriever_type = retrieverConfig.pop("type")
retriever_cls = RetrieverFactory.RETRIEVERS.get(retriever_type, None)
diff --git a/openrag/components/utils.py b/openrag/components/utils.py
index 3e3055508..adea40b53 100644
--- a/openrag/components/utils.py
+++ b/openrag/components/utils.py
@@ -7,7 +7,6 @@
from typing import ClassVar
import ray
-from components.indexer.utils.text_sanitizer import sanitize_text
from config import load_config
from fast_langdetect import LangDetectConfig, LangDetector
from langchain_core.documents.base import Document
@@ -98,28 +97,24 @@ def get_num_tokens():
def format_context(
docs: list[Document], max_context_tokens: int = 4096, number_sources: bool = True
) -> tuple[str, list[int]]:
- if not docs:
- return "No document found from the database", []
+ """Backward-compat shim — delegates to `core.prompts.chat_prompt_builder.format_context`.
- _length_function = get_num_tokens()
-
- reduced_docs = []
- included_indices = []
- total_tokens = 0
-
- for i, doc in enumerate(docs):
- prefix = f"[Source {len(reduced_docs) + 1}]\n" if number_sources else ""
- n_tokens = _length_function(doc.page_content)
- if prefix:
- n_tokens += _length_function(prefix)
- if total_tokens + n_tokens > max_context_tokens:
- break
- reduced_docs.append(f"{prefix}{doc.page_content}")
- included_indices.append(i)
- total_tokens += n_tokens
-
- logger.debug("Context formatted", total_tokens=total_tokens, doc_count=len(reduced_docs))
- return SOURCE_SEPARATOR.join(reduced_docs), included_indices
+ The legacy signature took LangChain Documents and resolved a tokenizer
+ internally; the core version takes raw strings + an injected
+ length_function. We adapt by extracting page_content and threading
+ the cached tokenizer through.
+ """
+ from openrag.core.prompts.chat_prompt_builder import format_context as _core_format_context
+
+ texts = [doc.page_content for doc in docs]
+ text, included = _core_format_context(
+ texts,
+ max_context_tokens=max_context_tokens,
+ length_function=get_num_tokens(),
+ number_sources=number_sources,
+ )
+ logger.debug("Context formatted", doc_count=len(included))
+ return text, included
def format_web_context(
@@ -127,41 +122,21 @@ def format_web_context(
start_index: int = 1,
max_tokens: int = 2000,
) -> tuple[str, list[int], int]:
- """Format web results as numbered [Source N] blocks within a token budget.
+ """Backward-compat shim — delegates to `core.prompts.chat_prompt_builder.format_web_context`.
- Uses fetched page content when available, falling back to the search snippet.
-
- Args:
- web_results: Results from web search provider (list of WebResult)
- start_index: First source number (continues numbering after RAG sources)
- max_tokens: Maximum token budget for all web sources combined
-
- Returns:
- (formatted_string, list_of_source_numbers_used, total_tokens_used)
+ Same adaptation pattern as `format_context`: legacy resolved the
+ tokenizer internally, core takes it as a parameter.
"""
- if not web_results:
- return "", [], 0
-
- _length_function = get_num_tokens()
-
- parts = []
- source_numbers = []
- total_tokens = 0
-
- for i, result in enumerate(web_results):
- n = start_index + i
- title = sanitize_text(result.title)
- body = sanitize_text(result.content) if result.content else sanitize_text(result.snippet)
- block = f"[Source {n}]\n{title}\n{body}"
- block_tokens = _length_function(block)
- if total_tokens + block_tokens > max_tokens and parts:
- break
- parts.append(block)
- source_numbers.append(n)
- total_tokens += block_tokens
-
- logger.debug("Web context formatted", total_tokens=total_tokens, source_count=len(parts))
- return SOURCE_SEPARATOR.join(parts), source_numbers, total_tokens
+ from openrag.core.prompts.chat_prompt_builder import format_web_context as _core_format_web_context
+
+ text, source_numbers, total_tokens = _core_format_web_context(
+ web_results,
+ length_function=get_num_tokens(),
+ start_index=start_index,
+ max_tokens=max_tokens,
+ )
+ logger.debug("Web context formatted", total_tokens=total_tokens, source_count=len(source_numbers))
+ return text, source_numbers, total_tokens
_SOURCES_NONE_RE = re.compile(r"\n?\[?Sources?\]?\s*:\s*\[?\s*none\s*\]?\s*$", re.IGNORECASE)
diff --git a/openrag/core/chunking/__init__.py b/openrag/core/chunking/__init__.py
index 5a9e8413f..a699bc2eb 100644
--- a/openrag/core/chunking/__init__.py
+++ b/openrag/core/chunking/__init__.py
@@ -1,6 +1,24 @@
-"""ChunkingStrategy ABC + registry."""
+"""ChunkingStrategy ABC + registry + concrete strategies."""
from .chunking_strategy import ChunkingStrategy
+from .markdown_utils import (
+ MDElement,
+ chunk_table,
+ get_chunk_page_number,
+ parse_markdown_table,
+ split_md_elements,
+)
+from .recursive import BaseChunker, RecursiveSplitter
from .registry import chunking_registry
-__all__ = ["ChunkingStrategy", "chunking_registry"]
+__all__ = [
+ "ChunkingStrategy",
+ "chunking_registry",
+ "BaseChunker",
+ "RecursiveSplitter",
+ "MDElement",
+ "chunk_table",
+ "get_chunk_page_number",
+ "parse_markdown_table",
+ "split_md_elements",
+]
diff --git a/openrag/core/chunking/markdown_utils.py b/openrag/core/chunking/markdown_utils.py
new file mode 100644
index 000000000..6dd030875
--- /dev/null
+++ b/openrag/core/chunking/markdown_utils.py
@@ -0,0 +1,224 @@
+"""Markdown parsing primitives used by chunking strategies.
+
+Pure functions extracted from ``components/indexer/chunker/utils.py``. They
+recognize page markers, image-description blocks, and tables; split a
+markdown document into typed elements; and split oversize tables along
+their semantic groups.
+
+This module has no IO and no config dependency.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+from openrag.core.utils.text import clean_markdown_table_spacing
+
+# Header + delimiter + at least one row.
+TABLE_RE = re.compile(
+ r"((?:^|\n)\|.*?\|\r?\n\|\s*[:-]+(?:\s*\|[:-]+)*\|\r?\n(?:\|.*?\|\r?\n)+)",
+ re.DOTALL | re.MULTILINE,
+)
+
+# `...` block injected by the VLM step.
+IMAGE_RE = re.compile(r"((.*?))", re.DOTALL)
+
+# `[PAGE_N]` page-boundary markers — content BEFORE [PAGE_N] is on page N.
+PAGE_RE = re.compile(r"\[PAGE_(\d+)\]")
+
+
+ElementType = Literal["text", "table", "image"]
+
+
+@dataclass
+class MDElement:
+ """A typed segment of markdown content with optional source page number."""
+
+ type: ElementType
+ content: str
+ page_number: int | None = None
+
+ def __repr__(self) -> str:
+ return f"MDElement(type={self.type}, page_number={self.page_number}, content={self.content[:100]}...)"
+
+
+def span_inside(span: tuple[int, int], container: tuple[int, int]) -> bool:
+ """Return True if ``span`` is fully contained within ``container``."""
+ return container[0] <= span[0] and span[1] <= container[1]
+
+
+def get_page_number(position: int, page_markers: list[tuple[int, int]]) -> int:
+ """Look up the page number for a position in the source markdown.
+
+ ``page_markers`` is a sorted list of ``(offset, page_n)`` tuples taken
+ from ``[PAGE_N]`` matches. Content AFTER ``[PAGE_N]`` belongs to page
+ ``N + 1``; content before any marker is page 1.
+ """
+ current_page = 1
+ for marker_pos, page_num in page_markers:
+ if position >= marker_pos:
+ current_page = page_num + 1
+ else:
+ break
+ return current_page
+
+
+def split_md_elements(md_text: str) -> list[MDElement]:
+ """Split markdown into ``MDElement`` segments of text, table, and image.
+
+ Tables nested inside an ```` block are NOT extracted
+ as separate elements — they belong to the image.
+ """
+ page_markers: list[tuple[int, int]] = []
+ for match in PAGE_RE.finditer(md_text):
+ page_markers.append((match.start(), int(match.group(1))))
+ page_markers.sort()
+
+ all_matches: list[tuple[tuple[int, int], ElementType, str, int | None]] = []
+ image_spans: list[tuple[int, int]] = []
+
+ for match in IMAGE_RE.finditer(md_text):
+ span = match.span()
+ page_num = get_page_number(span[0], page_markers)
+ all_matches.append((span, "image", match.group(1).strip(), page_num))
+ image_spans.append(span)
+
+ for match in TABLE_RE.finditer(md_text):
+ span = match.span()
+ if not any(span_inside(span, image_span) for image_span in image_spans):
+ page_num = get_page_number(span[0], page_markers)
+ all_matches.append((span, "table", match.group(1).strip(), page_num))
+
+ all_matches.sort(key=lambda x: x[0][0])
+
+ parts: list[MDElement] = []
+ last = 0
+
+ for (start, end), match_type, content, page_num in all_matches:
+ if start > last:
+ text_segment = md_text[last:start]
+ if text_segment.strip():
+ parts.append(MDElement(type="text", content=text_segment.strip()))
+ parts.append(MDElement(type=match_type, content=content, page_number=page_num))
+ last = end
+
+ if last < len(md_text):
+ remaining = md_text[last:]
+ if remaining.strip():
+ parts.append(MDElement(type="text", content=remaining.strip()))
+
+ return parts
+
+
+def get_chunk_page_number(chunk_str: str, previous_chunk_ending_page: int = 1) -> dict[str, int]:
+ """Resolve start and end pages for a text chunk containing ``[PAGE_N]`` markers.
+
+ Returns ``{"start_page": int, "end_page": int}``.
+ """
+ matches = list(PAGE_RE.finditer(chunk_str))
+
+ if not matches:
+ return {
+ "start_page": previous_chunk_ending_page,
+ "end_page": previous_chunk_ending_page,
+ }
+
+ first_match = matches[0]
+ last_match = matches[-1]
+ last_char_idx = len(chunk_str) - 1
+
+ if first_match.start() == 0:
+ start_page = int(first_match.group(1)) + 1
+ else:
+ start_page = previous_chunk_ending_page
+
+ if last_match.end() - 1 == last_char_idx:
+ end_page = int(last_match.group(1))
+ else:
+ end_page = int(last_match.group(1)) + 1
+
+ return {"start_page": start_page, "end_page": end_page}
+
+
+def parse_markdown_table(markdown_table: str) -> tuple[list[str], list[list[str]]]:
+ """Parse a markdown table into header lines + groups of rows.
+
+ Rows are grouped by the first column ("Domain"): a non-empty Domain
+ starts a new group, an empty Domain continues the current group. This
+ preserves the document's logical structure when chunking large tables.
+ """
+ lines = markdown_table.strip().split("\n")
+ header_lines = lines[:2]
+ data_rows = lines[2:]
+
+ groups: list[list[str]] = []
+ current_group: list[str] = []
+
+ for row in data_rows:
+ cells = [cell.strip() for cell in row.split("|")[1:-1]]
+ if not cells:
+ continue
+ domain = cells[0]
+ if domain:
+ if current_group:
+ groups.append(current_group)
+ current_group = [row]
+ else:
+ current_group.append(row)
+
+ if current_group:
+ groups.append(current_group)
+
+ return header_lines, groups
+
+
+def chunk_table(
+ table_element: MDElement,
+ chunk_size: int,
+ length_function: Callable[[str], int],
+) -> list[MDElement]:
+ """Split an oversize markdown table into multiple ``MDElement`` chunks.
+
+ Each chunk repeats the table header. When a new chunk starts, the LAST
+ row of the previous chunk is replayed as overlap so context is preserved
+ across the boundary.
+ """
+ txt = clean_markdown_table_spacing(table_element.content)
+ header_lines, groups = parse_markdown_table(txt)
+
+ header_text = "\n".join(header_lines)
+ group_texts = ["\n".join(g) for g in groups]
+
+ header_ntoks = length_function(header_text)
+ groups_ntoks = [length_function(g) for g in group_texts]
+
+ subtables: list[str] = []
+ body_rows: list[str] = [] # rows under the current chunk, header excluded
+ body_size = 0
+ prev_last_row: str | None = None
+
+ for group_txt, g_ntoks in zip(group_texts, groups_ntoks, strict=True):
+ # Only flush when we actually have body content to flush — otherwise an
+ # oversized first group would emit a header-only chunk.
+ if body_rows and header_ntoks + body_size + g_ntoks > chunk_size:
+ subtables.append("\n".join([header_text, *body_rows]))
+ body_rows = []
+ body_size = 0
+ # Replay only the last row of the previous chunk as overlap
+ # (matches the docstring contract; prev_last_row is the trailing
+ # line of the last admitted group).
+ if prev_last_row:
+ body_rows.append(prev_last_row)
+ body_size += length_function(prev_last_row)
+ body_rows.append(group_txt)
+ body_size += g_ntoks
+ # The "last row" is the trailing line of this group, not the whole group.
+ prev_last_row = group_txt.rsplit("\n", 1)[-1]
+
+ if body_rows:
+ subtables.append("\n".join([header_text, *body_rows]))
+
+ return [MDElement(type="table", content=subtable, page_number=table_element.page_number) for subtable in subtables]
diff --git a/openrag/core/chunking/recursive.py b/openrag/core/chunking/recursive.py
new file mode 100644
index 000000000..0b395d0f5
--- /dev/null
+++ b/openrag/core/chunking/recursive.py
@@ -0,0 +1,277 @@
+"""Recursive markdown-aware chunking strategy.
+
+Pure domain logic — no LLM client, no Ray, no LangChain ``Document``.
+The token-counting function is injected (``length_function``); the actual
+text splitter is ``langchain.text_splitter.RecursiveCharacterTextSplitter``,
+a pure utility kept until a stdlib-only replacement is in place.
+
+Contextualization (the LLM-driven [CONTEXT] block prepended to each chunk)
+lives in ``core/indexing/contextualize.py`` (Phase 5D) and is applied as a
+separate stage by the orchestrator — not from inside the chunker.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from openrag.core.chunking.chunking_strategy import ChunkingStrategy
+from openrag.core.chunking.markdown_utils import (
+ MDElement,
+ chunk_table,
+ get_chunk_page_number,
+ split_md_elements,
+)
+from openrag.core.chunking.registry import chunking_registry
+from openrag.core.models.chunk import Chunk, ChunkType
+from openrag.core.models.document import ProcessedDocument
+from openrag.core.utils.text import sanitize_text
+
+# Substring (case-insensitive) marking a "no useful content" image caption.
+# Detection logic mirrors the legacy chunker, which skips these elements so
+# they don't pollute the index.
+_IMAGE_PLACEHOLDER_MARKER = "[image placeholder]"
+
+# Tables/images smaller than this token count are inlined with surrounding
+# text rather than emitted as standalone chunks.
+_INLINE_ELEMENT_TOKEN_THRESHOLD = 100
+
+
+class BaseChunker(ChunkingStrategy):
+ """Base markdown-aware chunker.
+
+ Subclasses must set ``self.text_splitter`` to an object with a
+ ``.split_text(str) -> list[str]`` method (e.g. LangChain's
+ ``RecursiveCharacterTextSplitter``).
+ """
+
+ def __init__(
+ self,
+ chunk_size: int = 200,
+ chunk_overlap_rate: float = 0.2,
+ length_function: Callable[[str], int] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ if length_function is None:
+ raise ValueError("length_function is required (e.g. tokenizer.count_tokens)")
+ self.chunk_size = chunk_size
+ self.chunk_overlap_rate = chunk_overlap_rate
+ self.chunk_overlap = int(self.chunk_size * self.chunk_overlap_rate)
+ self.length_function = length_function
+ self.text_splitter: Any = None
+
+ # ------------------------------------------------------------------
+ # ChunkingStrategy contract
+ # ------------------------------------------------------------------
+ def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]:
+ """Split a processed document into ``Chunk`` objects."""
+ content = self._content_from(document)
+ if not content.strip():
+ return []
+
+ metadata = self._chunk_metadata_base(document, partition)
+ md_chunks = self._get_chunks(content=content.strip(), metadata=metadata)
+
+ return [
+ Chunk(
+ document_id=metadata.get("file_id", ""),
+ text=md_chunks_meta["page_content"],
+ chunk_index=i,
+ chunk_type=ChunkType(md_chunks_meta["chunk_type"]),
+ metadata={k: v for k, v in md_chunks_meta.items() if k not in ("page_content", "chunk_type", "page")},
+ partition=partition,
+ page_number=md_chunks_meta.get("page"),
+ )
+ for i, md_chunks_meta in enumerate(md_chunks)
+ ]
+
+ # ------------------------------------------------------------------
+ # Helpers
+ # ------------------------------------------------------------------
+ @staticmethod
+ def _content_from(document: ProcessedDocument) -> str:
+ """Reconstruct chunkable markdown from a ProcessedDocument.
+
+ Single-block documents on page 1 (or with no page metadata) flow
+ through unchanged. Anything else gets synthetic ``[PAGE_N]`` markers
+ injected so downstream chunk-page resolution works correctly.
+
+ Marker semantics: a ``[PAGE_N]`` marker means "everything BEFORE this
+ marker was on page N" (see ``markdown_utils.get_page_number``). So we
+ emit the marker for the *outgoing* page just before content from a
+ new page begins, and we also prepend a marker for the first block if
+ it doesn't start on page 1.
+ """
+ if not document.text_blocks:
+ return ""
+ if len(document.text_blocks) == 1 and document.text_blocks[0].page_number in (None, 1):
+ return document.text_blocks[0].text
+
+ parts: list[str] = []
+ last_page: int | None = None
+ for index, block in enumerate(document.text_blocks):
+ if block.page_number is not None:
+ # Emit `[PAGE_{block.page_number - 1}]` immediately *before*
+ # this block's text so downstream resolution lands on
+ # block.page_number. Using `block.page_number - 1` (rather
+ # than `last_page`) handles non-sequential pages (1 -> 5)
+ # and a first block already on page > 1.
+ needs_marker = (index == 0 and block.page_number > 1) or (
+ last_page is not None and block.page_number != last_page
+ )
+ if needs_marker:
+ parts.append(f"[PAGE_{block.page_number - 1}]")
+ parts.append(block.text)
+ last_page = block.page_number
+ return "\n\n".join(parts)
+
+ @staticmethod
+ def _chunk_metadata_base(document: ProcessedDocument, partition: str) -> dict[str, Any]:
+ # Reserved identity fields must win — `chunk()` later reads
+ # metadata["file_id"] to set Chunk.document_id, so a stray key in
+ # `document.metadata` would silently reassign chunks to the wrong doc.
+ return {
+ **document.metadata,
+ "file_id": document.document_id,
+ "partition": partition,
+ }
+
+ def split_text(self, text: str) -> list[str]:
+ """Split a text string with the configured text splitter.
+
+ Lazy-initializes a ``RecursiveCharacterTextSplitter`` if a subclass
+ forgot to set one — preserves legacy behavior.
+ """
+ if self.text_splitter is None:
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+ self.text_splitter = RecursiveCharacterTextSplitter(
+ chunk_size=self.chunk_size,
+ chunk_overlap=self.chunk_overlap,
+ length_function=self.length_function,
+ )
+ return self.text_splitter.split_text(text)
+
+ def _prepare_md_elements(self, content: str) -> tuple[list[MDElement], list[MDElement]]:
+ """Separate markdown into (inline-able texts) and (standalone tables/images)."""
+ md_elements = split_md_elements(content)
+ tables_and_images: list[MDElement] = []
+ texts: list[MDElement] = []
+
+ for element in md_elements:
+ if element.type in ("table", "image"):
+ if element.type == "image" and _IMAGE_PLACEHOLDER_MARKER in element.content.lower():
+ continue
+ if self.length_function(element.content) <= _INLINE_ELEMENT_TOKEN_THRESHOLD:
+ texts.append(element)
+ else:
+ tables_and_images.append(element)
+ else:
+ texts.append(element)
+
+ return texts, tables_and_images
+
+ def _get_chunks(self, content: str, metadata: dict[str, Any]) -> list[dict[str, Any]]:
+ """Produce per-chunk dicts with ``page_content`` + metadata fields.
+
+ The dict shape is intentional — it lets ``chunk()`` build ``Chunk``
+ objects without leaking domain types into the lower-level helpers.
+ """
+ texts, tables_and_images = self._prepare_md_elements(content=content)
+ combined_texts = "\n".join(e.content for e in texts)
+
+ sanitized = sanitize_text(
+ combined_texts,
+ normalize_whitespace=True,
+ remove_control_chars=True,
+ remove_zero_width_chars=True,
+ max_consecutive_newlines=2,
+ normalize_unicode=True,
+ )
+ text_chunks = self.split_text(sanitized)
+
+ chunks: list[dict[str, Any]] = []
+
+ # Reserved per-chunk keys must win over arbitrary `metadata` values —
+ # a stray "chunk_type" / "page" / "page_content" in the document's
+ # metadata would otherwise clobber the resolved value (and crash
+ # `chunk()` when ChunkType(...) is fed an out-of-enum string). Same
+ # defensive pattern as `_chunk_metadata_base`.
+ for element in tables_and_images:
+ if element.type == "table" and self.length_function(element.content) > self.chunk_size:
+ subtables = chunk_table(
+ table_element=element,
+ chunk_size=self.chunk_size,
+ length_function=self.length_function,
+ )
+ chunks.extend(
+ {
+ **metadata,
+ "page_content": subtable.content.strip(),
+ "page": subtable.page_number,
+ "chunk_type": "table",
+ }
+ for subtable in subtables
+ )
+ else:
+ # MDElement.type is the source-markdown literal ("image"/"table");
+ # ChunkType uses "image_caption" for image blocks.
+ ct = "image_caption" if element.type == "image" else element.type
+ chunks.append(
+ {
+ **metadata,
+ "page_content": element.content.strip(),
+ "page": element.page_number,
+ "chunk_type": ct,
+ }
+ )
+
+ prev_page = 1
+ for c in text_chunks:
+ page_info = get_chunk_page_number(chunk_str=c, previous_chunk_ending_page=prev_page)
+ prev_page = page_info["end_page"]
+ chunks.append(
+ {
+ **metadata,
+ "page_content": c.strip(),
+ "page": page_info["start_page"],
+ "chunk_type": "text",
+ }
+ )
+
+ if not chunks:
+ return []
+ chunks.sort(key=lambda d: d.get("page") or 0)
+ return chunks
+
+
+@chunking_registry.register("recursive_splitter")
+class RecursiveSplitter(BaseChunker):
+ """Markdown-aware chunker backed by ``RecursiveCharacterTextSplitter``.
+
+ Splits on paragraph boundaries first, then sentence terminators, then
+ smaller separators.
+ """
+
+ def __init__(
+ self,
+ chunk_size: int = 200,
+ chunk_overlap_rate: float = 0.2,
+ length_function: Callable[[str], int] | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(
+ chunk_size=chunk_size,
+ chunk_overlap_rate=chunk_overlap_rate,
+ length_function=length_function,
+ **kwargs,
+ )
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+ self.text_splitter = RecursiveCharacterTextSplitter(
+ chunk_size=self.chunk_size,
+ chunk_overlap=self.chunk_overlap,
+ length_function=self.length_function,
+ is_separator_regex=True,
+ separators=["\n", r"(?<=[\.\?\!])"],
+ )
diff --git a/openrag/core/chunking/test_markdown_utils.py b/openrag/core/chunking/test_markdown_utils.py
new file mode 100644
index 000000000..549d9ce52
--- /dev/null
+++ b/openrag/core/chunking/test_markdown_utils.py
@@ -0,0 +1,182 @@
+"""Tests for core.chunking.markdown_utils.
+
+Mirrors components/indexer/chunker/test_chunking.py to verify behavior is
+preserved through the move into core/.
+"""
+
+from __future__ import annotations
+
+from openrag.core.chunking.markdown_utils import (
+ MDElement,
+ chunk_table,
+ get_chunk_page_number,
+ parse_markdown_table,
+ span_inside,
+ split_md_elements,
+)
+
+
+def _mock_length(text: str) -> int:
+ """Estimate token count at ~4 chars per token (matches legacy tests)."""
+ return len(text) // 4
+
+
+class TestSplitMdElements:
+ def test_simple_text_only(self):
+ md = "This is a simple paragraph.\n\nAnother paragraph here."
+ elems = split_md_elements(md)
+ assert len(elems) == 1
+ assert elems[0].type == "text"
+ assert elems[0].content == md
+
+ def test_single_table(self):
+ md = (
+ "Some text before.\n\n| Header 1 | Header 2 |\n|----------|----------|\n"
+ "| Cell 1 | Cell 2 |\n| Cell 3 | Cell 4 |\n\nSome text after."
+ )
+ elems = split_md_elements(md)
+ assert [e.type for e in elems] == ["text", "table", "text"]
+ assert "Header 1" in elems[1].content
+
+ def test_single_image(self):
+ md = (
+ "\nText before image.\n\n\nA beautiful sunset over the ocean.\n"
+ "\n\nText after image."
+ )
+ elems = split_md_elements(md)
+ assert [e.type for e in elems] == ["text", "image", "text"]
+ assert "sunset" in elems[1].content
+
+ def test_table_inside_image_description_is_ignored(self):
+ md = (
+ "\n\nThis image contains a table:\n| Col 1 | Col 2 |\n"
+ "|-------|-------|\n| A | B |\n\n\n"
+ "Outside table:\n| Real 1 | Real 2 |\n|--------|--------|\n| X | Y |\n"
+ )
+ elems = split_md_elements(md)
+ tables = [e for e in elems if e.type == "table"]
+ assert len(tables) == 1
+ assert "Real 1" in tables[0].content
+
+ def test_page_markers_with_table(self):
+ md = (
+ "text on page 1.\n[PAGE_1]\nText on page 2.\n\n"
+ "| Header 1 | Header 2 |\n|----------|----------|\n| Data 1 | Data 2 |\n\n"
+ "[PAGE_2]\nMore content.\n"
+ )
+ elems = split_md_elements(md)
+ tables = [e for e in elems if e.type == "table"]
+ assert len(tables) == 1
+ assert tables[0].page_number == 2
+
+ def test_page_markers_with_images(self):
+ md = "\n[PAGE_1]\n[PAGE_2]\n\nImage on page 3.\n\n"
+ elems = split_md_elements(md)
+ images = [e for e in elems if e.type == "image"]
+ assert len(images) == 1
+ assert images[0].page_number == 3
+
+
+class TestGetChunkPageNumber:
+ def test_no_markers_returns_previous_page(self):
+ result = get_chunk_page_number("Just some plain text content.", previous_chunk_ending_page=1)
+ assert result == {"start_page": 1, "end_page": 1}
+
+ def test_chunk_starts_with_marker(self):
+ result = get_chunk_page_number("[PAGE_2]Content on page 3.", previous_chunk_ending_page=1)
+ assert result == {"start_page": 3, "end_page": 3}
+
+ def test_chunk_ends_with_marker(self):
+ result = get_chunk_page_number("Content on page 1.[PAGE_1]", previous_chunk_ending_page=1)
+ assert result == {"start_page": 1, "end_page": 1}
+
+ def test_marker_in_middle(self):
+ result = get_chunk_page_number("Start on page 1.[PAGE_1]End on page 2.", previous_chunk_ending_page=1)
+ assert result == {"start_page": 1, "end_page": 2}
+
+
+class TestChunkTable:
+ def test_small_table_no_chunking(self):
+ content = "| Name | Age |\n|------|-----|\n| John | 30 |\n| Jane | 25 |"
+ elem = MDElement(type="table", content=content, page_number=1)
+ chunks = chunk_table(elem, chunk_size=1000, length_function=_mock_length)
+ assert len(chunks) == 1
+ assert chunks[0].type == "table"
+ assert chunks[0].page_number == 1
+ assert "John" in chunks[0].content
+ assert "Jane" in chunks[0].content
+
+ def test_chunking_preserves_groups(self):
+ header = "| Country | Strategy | Goals |"
+ g1 = "| USA | Cyber | Goal 1 |\n| | | Goal 2 |\n| | | Goal 3 |"
+ g2 = "| Mexico | Defense | Goal X |\n| | | Goal Y |\n| | | Goal Z |"
+ table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n"
+ elem = MDElement(type="table", content=table, page_number=2)
+ chunk_size = _mock_length(table) // 2
+ chunks = chunk_table(elem, chunk_size=chunk_size, length_function=_mock_length)
+ assert len(chunks) == 2
+ assert all(c.type == "table" for c in chunks)
+ assert all(header in c.content for c in chunks)
+ assert "USA" in chunks[0].content
+
+ def test_oversized_first_group_does_not_emit_header_only_chunk(self):
+ """When the very first group is already larger than chunk_size, the
+ old code flushed a header-only chunk (CodeRabbit #1)."""
+ header = "| Country | Strategy | Goals |"
+ g1 = "| USA | Cyber | Goal 1 |\n| | | Goal 2 |\n| | | Goal 3 |"
+ g2 = "| Mexico | Defense | Goal X |"
+ table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n"
+ # Tight budget: g1 alone already exceeds.
+ chunks = chunk_table(
+ MDElement(type="table", content=table, page_number=1), chunk_size=2, length_function=_mock_length
+ )
+ # No chunk may contain only the header.
+ for c in chunks:
+ body = c.content.replace(header, "").strip()
+ assert body, f"header-only chunk emitted: {c.content!r}"
+
+ def test_overlap_replays_only_last_row_not_full_group(self):
+ """The docstring promises last-row overlap; the old code stored the
+ whole previous group (CodeRabbit #1)."""
+ header = "| Country | Strategy | Goal |"
+ g1 = "| USA | Cyber | first |\n| | | second |\n| | | LAST_ROW_OF_G1 |"
+ g2 = "| Mexico | Defense | only |"
+ table = f"{header}\n|----|----|----|\n{g1}\n{g2}\n"
+ # Force a split between g1 and g2.
+ chunk_size = _mock_length(g1) + _mock_length(header)
+ chunks = chunk_table(
+ MDElement(type="table", content=table, page_number=1), chunk_size=chunk_size, length_function=_mock_length
+ )
+ assert len(chunks) >= 2
+ second_chunk = chunks[1].content
+ assert "LAST_ROW_OF_G1" in second_chunk, "last row should be replayed as overlap"
+ # The earlier rows of g1 must NOT appear in the second chunk.
+ assert "first" not in second_chunk
+ assert "second" not in second_chunk
+
+
+def test_md_element_repr_truncates_content():
+ elem = MDElement(type="text", content="x" * 500, page_number=3)
+ rendered = repr(elem)
+ assert "type=text" in rendered
+ assert "page_number=3" in rendered
+ # Long content is truncated to <=100 chars + ellipsis.
+ assert "x" * 200 not in rendered
+
+
+def test_span_inside_helper():
+ assert span_inside((10, 20), (5, 30)) is True
+ assert span_inside((5, 30), (10, 20)) is False
+ assert span_inside((10, 20), (10, 20)) is True
+
+
+def test_parse_markdown_table_skips_blank_data_rows():
+ """A pipe-only row (e.g. an extra blank `|`) yields no cells; it must be
+ skipped without erroring or starting a phantom group."""
+ header = "| Country | Goal |"
+ delim = "|---------|------|"
+ table = f"{header}\n{delim}\n| USA | A |\n|\n| Mexico | B |"
+ headers, groups = parse_markdown_table(table)
+ assert headers == [header, delim]
+ # Two non-empty rows -> two groups (each row has a non-empty Domain).
+ assert len(groups) == 2
diff --git a/openrag/core/chunking/test_recursive.py b/openrag/core/chunking/test_recursive.py
new file mode 100644
index 000000000..69d496a41
--- /dev/null
+++ b/openrag/core/chunking/test_recursive.py
@@ -0,0 +1,251 @@
+"""End-to-end tests for the RecursiveSplitter chunker."""
+
+from __future__ import annotations
+
+from openrag.core.chunking.recursive import RecursiveSplitter
+from openrag.core.chunking.registry import chunking_registry
+from openrag.core.models.chunk import ChunkType
+from openrag.core.models.document import ProcessedDocument, TextBlock
+
+
+def _word_tokens(text: str) -> int:
+ return len(text.split())
+
+
+def test_recursive_splitter_is_registered():
+ assert "recursive_splitter" in chunking_registry
+
+
+def test_recursive_splitter_chunks_simple_document():
+ splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text="alpha beta gamma\ndelta epsilon zeta\neta theta iota.", page_number=1)],
+ metadata={"source": "test.md"},
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks
+ assert all(c.partition == "p1" for c in chunks)
+ assert all(c.document_id == "d1" for c in chunks)
+ assert all(c.chunk_type == ChunkType.TEXT for c in chunks)
+
+
+def test_recursive_splitter_emits_table_chunks():
+ table = "| Col | Val |\n|-----|-----|\n" + "\n".join(f"| Group{i} | {' '.join(['x'] * 50)} |" for i in range(6))
+ splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text=f"Some prose here.\n\n{table}\n\nMore prose.", page_number=1)],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ table_chunks = [c for c in chunks if c.chunk_type == ChunkType.TABLE]
+ assert table_chunks, "expected at least one table-type chunk"
+
+
+def test_recursive_splitter_skips_image_placeholder():
+ placeholder_md = (
+ "Real text first.\n\n\n\n[Image Placeholder]\n\n\n\nReal text after."
+ )
+ splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text=placeholder_md, page_number=1)],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert all(c.chunk_type != ChunkType.IMAGE_CAPTION for c in chunks)
+ for c in chunks:
+ assert "[image placeholder]" not in c.text.lower()
+
+
+def test_recursive_splitter_metadata_passthrough():
+ splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text="alpha beta gamma delta epsilon", page_number=1)],
+ metadata={"source": "test.md", "filename": "test.md", "tag": "v1"},
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks
+ assert chunks[0].metadata.get("source") == "test.md"
+ assert chunks[0].metadata.get("tag") == "v1"
+
+
+def test_recursive_splitter_empty_document_returns_empty():
+ splitter = RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(document_id="d1", text_blocks=[])
+ assert splitter.chunk(doc, partition="p1") == []
+
+
+def test_recursive_splitter_requires_length_function():
+ import pytest
+
+ with pytest.raises(ValueError, match="length_function"):
+ RecursiveSplitter(chunk_size=10, chunk_overlap_rate=0.0)
+
+
+def test_recursive_splitter_joins_multi_block_document_with_synthetic_page_markers():
+ """Multi-block docs need synthetic [PAGE_N] markers so chunks downstream
+ of a page boundary report the right page. Cover that injection path in
+ BaseChunker._content_from with a chunk_size small enough to force a split
+ across pages."""
+ block_text = " ".join([f"word{i}" for i in range(20)])
+ splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[
+ TextBlock(text=block_text, page_number=1),
+ TextBlock(text=block_text, page_number=2),
+ TextBlock(text=block_text, page_number=3),
+ ],
+ metadata={"source": "multi.md"},
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks
+ pages = {c.page_number for c in chunks}
+ # First chunk(s) stay on page 1; once a [PAGE_N] marker lands inside a
+ # chunk's content the next chunk resolves to >=2.
+ assert 1 in pages
+ assert any((p or 0) >= 2 for p in pages)
+
+
+def test_recursive_splitter_inlines_small_table():
+ """Tables under the inline threshold (<=100 length-function tokens) flow
+ through the text path rather than emitting a standalone TABLE chunk."""
+ table = "| A | B |\n|---|---|\n| 1 | 2 |"
+ splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text=f"Lead-in.\n\n{table}\n\nTrailing.", page_number=1)],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks
+ assert all(c.chunk_type != ChunkType.TABLE for c in chunks)
+
+
+def test_recursive_splitter_image_caption_chunk_emitted_when_above_threshold():
+ """Image_description blocks above the inline threshold land as their own
+ chunks (chunk_type=image_caption) — covers the standalone-element path in
+ _get_chunks's else branch."""
+ long_caption = "lorem ipsum dolor sit amet " * 60 # well above inline threshold
+ md = f"Some text.\n\n\n{long_caption}\n\n\nAfter."
+ splitter = RecursiveSplitter(chunk_size=400, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text=md, page_number=1)],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ image_chunks = [c for c in chunks if c.chunk_type == ChunkType.IMAGE_CAPTION]
+ assert image_chunks, "expected at least one image_caption chunk"
+
+
+def test_recursive_splitter_returns_empty_when_only_image_placeholder():
+ """When the only element is a skipped image placeholder, _get_chunks
+ produces nothing — exercise the `if not chunks: return []` guard."""
+ splitter = RecursiveSplitter(chunk_size=200, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[
+ TextBlock(text="\n[Image Placeholder]\n", page_number=1),
+ ],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks == []
+
+
+def test_base_chunker_lazy_initializes_text_splitter():
+ """A BaseChunker subclass that forgets to set self.text_splitter still
+ works — split_text lazy-builds a default RecursiveCharacterTextSplitter."""
+ from openrag.core.chunking.recursive import BaseChunker
+
+ class BareChunker(BaseChunker):
+ pass
+
+ bare = BareChunker(chunk_size=12, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ assert bare.text_splitter is None
+ pieces = bare.split_text("alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu")
+ assert pieces
+ assert bare.text_splitter is not None # cached after first call
+
+
+def test_document_metadata_cannot_override_file_id_or_partition():
+ """Reserved identity fields must win over arbitrary metadata keys
+ (CodeRabbit #3)."""
+ splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="real-doc-id",
+ text_blocks=[TextBlock(text="alpha beta gamma delta", page_number=1)],
+ metadata={
+ "file_id": "MALICIOUS_OVERRIDE",
+ "partition": "MALICIOUS_PARTITION",
+ "source": "ok.md",
+ },
+ )
+ chunks = splitter.chunk(doc, partition="real-partition")
+ assert chunks
+ for c in chunks:
+ assert c.document_id == "real-doc-id"
+ assert c.partition == "real-partition"
+ # Other metadata keys still flow through.
+ assert c.metadata.get("source") == "ok.md"
+
+
+def test_document_metadata_cannot_override_chunk_type_or_page():
+ """Per-chunk reserved keys (chunk_type, page, page_content) must win
+ over `document.metadata`. A poison `chunk_type` value would otherwise
+ crash `chunk()` when ChunkType(...) is constructed (ultrareview)."""
+ splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[TextBlock(text="alpha beta gamma delta epsilon", page_number=2)],
+ metadata={
+ "chunk_type": "POISON",
+ "page": 999,
+ "page_content": "REPLACED",
+ "tag": "v1",
+ },
+ )
+ # Must not raise ValueError("'POISON' is not a valid ChunkType").
+ chunks = splitter.chunk(doc, partition="p1")
+ assert chunks
+ for c in chunks:
+ assert c.chunk_type == ChunkType.TEXT
+ assert c.page_number != 999
+ assert c.text != "REPLACED"
+ # Other metadata keys still flow through.
+ assert c.metadata.get("tag") == "v1"
+
+
+def test_recursive_splitter_first_block_on_page_three_resolves_correctly():
+ """When the first block already starts on page>1, every chunk used to be
+ tagged page 1. Now it should land on the actual block page (CodeRabbit #2)."""
+ block_text = " ".join([f"word{i}" for i in range(20)])
+ splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[
+ TextBlock(text=block_text, page_number=3),
+ TextBlock(text=block_text, page_number=4),
+ ],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ pages = {c.page_number for c in chunks}
+ assert 1 not in pages, f"chunks tagged page 1 despite first block on page 3: {pages}"
+ assert any((p or 0) >= 3 for p in pages)
+
+
+def test_recursive_splitter_skips_pages_get_correct_marker():
+ """Block pages 1 -> 5 (skipping 2/3/4) — the second block's chunks must
+ resolve to page 5, not page 2 (= last_page+1)."""
+ block_text = " ".join([f"word{i}" for i in range(40)])
+ splitter = RecursiveSplitter(chunk_size=8, chunk_overlap_rate=0.0, length_function=_word_tokens)
+ doc = ProcessedDocument(
+ document_id="d1",
+ text_blocks=[
+ TextBlock(text=block_text, page_number=1),
+ TextBlock(text=block_text, page_number=5),
+ ],
+ )
+ chunks = splitter.chunk(doc, partition="p1")
+ pages = sorted({c.page_number for c in chunks if c.page_number is not None})
+ assert 1 in pages
+ assert 5 in pages, f"page 5 missing despite second block on page 5: {pages}"
diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py
index fbebc34cc..9f2d13e87 100644
--- a/openrag/core/config/indexation.py
+++ b/openrag/core/config/indexation.py
@@ -2,7 +2,9 @@
from __future__ import annotations
-from pydantic import Field
+from typing import Any
+
+from pydantic import Field, field_validator
from .base import ConfigMixin
@@ -10,6 +12,20 @@
# Transcriber (nested under loader)
# ---------------------------------------------------------------------------
+# Audio formats the transcription endpoint accepts as-is — uploads of these
+# extensions skip the WAV-conversion pre-step. Configurable via env var
+# `TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES` (pipe-delimited string).
+_DEFAULT_DIRECT_UPLOAD_SUFFIXES = frozenset(
+ {".wav", ".flac", ".ogg", ".mp3", ".mp4", ".m4a", ".webm", ".mpeg", ".mpga"}
+)
+
+
+def _normalize_suffix(s: str) -> str:
+ s = s.strip().lower()
+ if not s:
+ return ""
+ return s if s.startswith(".") else f".{s}"
+
class TranscriberConfig(ConfigMixin):
base_url: str = "http://transcriber:8000/v1"
@@ -18,6 +34,14 @@ class TranscriberConfig(ConfigMixin):
timeout: int = 3600
max_concurrent_chunks: int = 20
use_whisper_lang_detector: bool = True
+ direct_upload_suffixes: set[str] = Field(default_factory=lambda: set(_DEFAULT_DIRECT_UPLOAD_SUFFIXES))
+
+ @field_validator("direct_upload_suffixes", mode="before")
+ @classmethod
+ def _split_suffixes(cls, v: Any) -> Any:
+ if isinstance(v, str):
+ return {n for raw in v.split("|") if (n := _normalize_suffix(raw))}
+ return v
# ---------------------------------------------------------------------------
diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py
index 03510b0c6..b35e99864 100644
--- a/openrag/core/config/loader.py
+++ b/openrag/core/config/loader.py
@@ -104,6 +104,7 @@
("TRANSCRIBER_MODEL", "loader.transcriber.model_name", str),
("TRANSCRIBER_TIMEOUT", "loader.transcriber.timeout", int),
("TRANSCRIBER_MAX_CONCURRENT_CHUNKS", "loader.transcriber.max_concurrent_chunks", int),
+ ("TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES", "loader.transcriber.direct_upload_suffixes", str),
("USE_WHISPER_LANG_DETECTOR", "loader.transcriber.use_whisper_lang_detector", bool),
("OPENAI_LOADER_BASE_URL", "loader.openai.base_url", str),
("OPENAI_LOADER_API_KEY", "loader.openai.api_key", str),
diff --git a/openrag/core/config/test_indexation.py b/openrag/core/config/test_indexation.py
new file mode 100644
index 000000000..1e6405ba8
--- /dev/null
+++ b/openrag/core/config/test_indexation.py
@@ -0,0 +1,31 @@
+"""Tests for indexation config — TranscriberConfig pipe-string parsing."""
+
+from __future__ import annotations
+
+from openrag.core.config.indexation import (
+ _DEFAULT_DIRECT_UPLOAD_SUFFIXES,
+ TranscriberConfig,
+)
+
+
+def test_transcriber_config_default_direct_upload_suffixes():
+ cfg = TranscriberConfig()
+ assert cfg.direct_upload_suffixes == set(_DEFAULT_DIRECT_UPLOAD_SUFFIXES)
+
+
+def test_transcriber_config_parses_pipe_delimited_string():
+ """The YAML default and TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES env var both
+ arrive as a pipe-delimited string. The validator must split + normalize
+ into a set of dot-prefixed lowercase suffixes."""
+ cfg = TranscriberConfig(direct_upload_suffixes=".wav|FLAC|mp3")
+ assert cfg.direct_upload_suffixes == {".wav", ".flac", ".mp3"}
+
+
+def test_transcriber_config_drops_empty_components():
+ cfg = TranscriberConfig(direct_upload_suffixes="|.wav||.mp3|")
+ assert cfg.direct_upload_suffixes == {".wav", ".mp3"}
+
+
+def test_transcriber_config_set_input_passes_through():
+ cfg = TranscriberConfig(direct_upload_suffixes={".wav", ".m4a"})
+ assert cfg.direct_upload_suffixes == {".wav", ".m4a"}
diff --git a/openrag/core/models/__init__.py b/openrag/core/models/__init__.py
index 2684036dd..9bc17ff38 100644
--- a/openrag/core/models/__init__.py
+++ b/openrag/core/models/__init__.py
@@ -6,7 +6,7 @@
from .conversation import Conversation, Message
from .document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock
from .prompt import Prompt, PromptType
-from .query import RetrievalQuery
+from .query import Query, RetrievalQuery, SearchQueries, TemporalPredicate
from .retrieval_response import RetrievalResponse
from .retrieval_result import RetrievalResult, ScoredChunk
from .user import ApiKey, OIDCSession, PartitionRole, TokenPayload, User, UserPartition
@@ -30,10 +30,13 @@
"ProcessedDocument",
"Prompt",
"PromptType",
+ "Query",
"RetrievalQuery",
"RetrievalResponse",
"RetrievalResult",
"ScoredChunk",
+ "SearchQueries",
+ "TemporalPredicate",
"TextBlock",
"TokenPayload",
"User",
diff --git a/openrag/core/models/chunk.py b/openrag/core/models/chunk.py
index 30a98f9c3..879a4f524 100644
--- a/openrag/core/models/chunk.py
+++ b/openrag/core/models/chunk.py
@@ -16,6 +16,26 @@ class ChunkType(str, Enum):
CONTEXTUALIZED = "contextualized"
+# Pre-Phase-5 chunkers stamped Document metadata with the raw MDElement
+# literal (`"image"`) for image elements. Deployments upgraded without
+# re-indexing have those values in Milvus; map them to the current enum
+# at read time so retrieval doesn't crash on legacy data.
+_CHUNK_TYPE_LEGACY_ALIASES = {"image": ChunkType.IMAGE_CAPTION}
+
+
+def _coerce_chunk_type(value: Any) -> ChunkType:
+ if isinstance(value, ChunkType):
+ return value
+ if value in _CHUNK_TYPE_LEGACY_ALIASES:
+ return _CHUNK_TYPE_LEGACY_ALIASES[value]
+ try:
+ return ChunkType(value)
+ except (ValueError, TypeError):
+ # Unknown value from upstream/legacy data — fall back to TEXT rather
+ # than crash the retrieval call.
+ return ChunkType.TEXT
+
+
class Chunk(BaseModel):
"""A chunk of text extracted from a document, optionally embedded."""
@@ -44,13 +64,19 @@ def from_langchain(cls, doc: Any) -> Chunk:
Import is deferred to method body so core/ stays pure at import time.
"""
metadata = dict(doc.metadata) if doc.metadata else {}
+ # Milvus assigns the primary key `_id` as INT64 (auto_id), so the value
+ # comes back as a Python int. Chunk.id is typed `str`, so coerce here
+ # rather than loosen the model — keeps the domain type strict while the
+ # store-specific shape is contained in the conversion boundary.
+ raw_id = metadata.pop("_id", None)
+ chunk_id = str(raw_id) if raw_id is not None else str(uuid.uuid4())
return cls(
- id=metadata.pop("_id", str(uuid.uuid4())),
+ id=chunk_id,
document_id=metadata.pop("file_id", ""),
text=doc.page_content,
partition=metadata.pop("partition", "default"),
page_number=metadata.pop("page", None),
- chunk_type=ChunkType(metadata.pop("chunk_type", "text")),
+ chunk_type=_coerce_chunk_type(metadata.pop("chunk_type", "text")),
metadata=metadata,
)
diff --git a/openrag/core/models/query.py b/openrag/core/models/query.py
index cbb33573c..a073b4332 100644
--- a/openrag/core/models/query.py
+++ b/openrag/core/models/query.py
@@ -2,10 +2,14 @@
from __future__ import annotations
-from typing import Any
+import logging
+from datetime import datetime
+from typing import Any, Literal
from pydantic import BaseModel, Field
+logger = logging.getLogger(__name__)
+
class RetrievalQuery(BaseModel):
"""A user query with retrieval parameters."""
@@ -21,3 +25,76 @@ class RetrievalQuery(BaseModel):
max_ancestor_depth: int | None = None
with_surrounding_chunks: bool = True
rerank: bool = True
+
+
+class TemporalPredicate(BaseModel):
+ """A single date constraint on a document's creation date.
+
+ Multiple predicates on the same ``Query`` are AND-combined. Closed
+ ranges (e.g. "last month") are encoded as two predicates, one per side.
+ """
+
+ field: Literal["created_at"] = Field(
+ default="created_at",
+ description="Document metadata field to filter on. Always `created_at` for now.",
+ )
+ operator: Literal[">", "<", ">=", "<="] = Field(
+ description="Comparison operator applied to the date field.",
+ )
+ value: str = Field(
+ description='ISO 8601 datetime with timezone, e.g. "2026-03-15T00:00:00+00:00".',
+ )
+
+
+class Query(BaseModel):
+ """A single vector-database search query plus optional temporal filters.
+
+ Two predicates yield an AND-range; an exclusion range (e.g. "last year
+ except March") is expressed as two separate ``Query`` objects.
+ """
+
+ query: str = Field(
+ description="A semantically enriched, descriptive query for vector similarity search.",
+ )
+ temporal_filters: list[TemporalPredicate] | None = Field(
+ default=None,
+ description="Date predicates on `created_at`, AND-combined.",
+ )
+
+ def to_milvus_filter(self) -> str | None:
+ """Render the AND-combined predicates as a Milvus filter expression.
+
+ Pydantic validates the field/operator types up front. The ``value``
+ field is parsed as ISO 8601 here defensively — predicates with an
+ unparseable value are dropped rather than crashing the search.
+ """
+ if not self.temporal_filters:
+ return None
+ parts: list[str] = []
+ for p in self.temporal_filters:
+ try:
+ datetime.fromisoformat(p.value)
+ except (TypeError, ValueError):
+ logger.warning(
+ "Dropping temporal predicate with non-ISO value: field=%s operator=%s value=%r",
+ p.field,
+ p.operator,
+ p.value,
+ )
+ continue
+ parts.append(f'{p.field} {p.operator} ISO "{p.value}"')
+ if not parts:
+ return None
+ return " and ".join(parts)
+
+ def __str__(self) -> str:
+ return f"Query: {self.query}, Filter: {self.to_milvus_filter()}"
+
+
+class SearchQueries(BaseModel):
+ """Collection of sub-queries produced by query decomposition."""
+
+ query_list: list[Query] = Field(..., description="Search sub-queries to retrieve relevant documents.")
+
+ def __str__(self) -> str:
+ return " --- ".join(str(q) for q in self.query_list)
diff --git a/openrag/core/models/test_chunk.py b/openrag/core/models/test_chunk.py
new file mode 100644
index 000000000..0fc01fb24
--- /dev/null
+++ b/openrag/core/models/test_chunk.py
@@ -0,0 +1,50 @@
+"""Tests for Chunk model — backward-compat coercions on from_langchain."""
+
+from __future__ import annotations
+
+from langchain_core.documents.base import Document
+
+from openrag.core.models.chunk import Chunk, ChunkType, _coerce_chunk_type
+
+
+def test_from_langchain_maps_legacy_image_chunk_type():
+ """Pre-Phase-5 chunkers stamped chunk_type='image' (raw MDElement
+ literal). Upgraded deployments still have those values in Milvus —
+ Chunk.from_langchain must not crash on them (ultrareview)."""
+ doc = Document(page_content="caption", metadata={"chunk_type": "image", "_id": "x", "file_id": "f1"})
+ chunk = Chunk.from_langchain(doc)
+ assert chunk.chunk_type == ChunkType.IMAGE_CAPTION
+
+
+def test_from_langchain_unknown_chunk_type_falls_back_to_text():
+ """Defensive: any historical value that isn't in the enum and isn't
+ in the legacy alias map should land on TEXT, not crash retrieval."""
+ doc = Document(page_content="x", metadata={"chunk_type": "unknown_legacy_value"})
+ chunk = Chunk.from_langchain(doc)
+ assert chunk.chunk_type == ChunkType.TEXT
+
+
+def test_from_langchain_accepts_canonical_values():
+ for value, expected in [
+ ("text", ChunkType.TEXT),
+ ("table", ChunkType.TABLE),
+ ("image_caption", ChunkType.IMAGE_CAPTION),
+ ("contextualized", ChunkType.CONTEXTUALIZED),
+ ]:
+ doc = Document(page_content="x", metadata={"chunk_type": value})
+ assert Chunk.from_langchain(doc).chunk_type == expected
+
+
+def test_coerce_chunk_type_passthrough_for_enum_input():
+ assert _coerce_chunk_type(ChunkType.TABLE) == ChunkType.TABLE
+
+
+def test_from_langchain_coerces_int_milvus_id_to_string():
+ """Milvus' `_id` primary key is INT64 (auto_id), so the value comes back
+ from the Ray actor as a Python int. Chunk.id is typed `str`; the
+ conversion boundary must coerce to avoid a ValidationError on every
+ retrieval call (CI api-tests regression)."""
+ doc = Document(page_content="hello", metadata={"_id": 466085833598567840, "file_id": "f1"})
+ chunk = Chunk.from_langchain(doc)
+ assert chunk.id == "466085833598567840"
+ assert isinstance(chunk.id, str)
diff --git a/openrag/core/prompts/__init__.py b/openrag/core/prompts/__init__.py
index e69de29bb..93c548dcc 100644
--- a/openrag/core/prompts/__init__.py
+++ b/openrag/core/prompts/__init__.py
@@ -0,0 +1,72 @@
+"""Prompt assembly helpers — pure string-formatting builders + disk loader."""
+
+from .chat_prompt_builder import (
+ EMPTY_CONTEXT_MESSAGE,
+ SOURCE_SEPARATOR,
+ WebSourceLike,
+ format_context,
+ format_web_context,
+ prepend_system_prompt,
+)
+from .contextualization_builder import (
+ BASE_CHUNK_FORMAT,
+ CHUNK_FORMAT,
+ wrap_chunk_with_context,
+)
+from .contextualization_builder import (
+ build_messages as build_contextualization_messages,
+)
+from .contextualization_builder import (
+ build_user_message as build_contextualization_user_message,
+)
+from .map_reduce_builder import (
+ SYSTEM_PROMPT_MAP,
+ USER_PROMPT_TEMPLATE,
+ build_map_messages,
+)
+from .query_rewriter import (
+ MULTI_QUERY_SEPARATOR,
+ build_hyde_prompt,
+ build_multi_query_prompt,
+ split_multi_query_response,
+)
+from .template_loader import load_template, load_template_by_key
+from .vlm_prompt_builder import (
+ IMAGE_DESCRIPTION_CLOSE,
+ IMAGE_DESCRIPTION_OPEN,
+ build_caption_messages,
+ wrap_caption,
+)
+
+__all__ = [
+ # template loader
+ "load_template",
+ "load_template_by_key",
+ # chat
+ "format_context",
+ "format_web_context",
+ "prepend_system_prompt",
+ "SOURCE_SEPARATOR",
+ "EMPTY_CONTEXT_MESSAGE",
+ "WebSourceLike",
+ # contextualization
+ "BASE_CHUNK_FORMAT",
+ "CHUNK_FORMAT",
+ "build_contextualization_messages",
+ "build_contextualization_user_message",
+ "wrap_chunk_with_context",
+ # query rewriter
+ "MULTI_QUERY_SEPARATOR",
+ "build_hyde_prompt",
+ "build_multi_query_prompt",
+ "split_multi_query_response",
+ # map-reduce
+ "build_map_messages",
+ "SYSTEM_PROMPT_MAP",
+ "USER_PROMPT_TEMPLATE",
+ # VLM
+ "build_caption_messages",
+ "wrap_caption",
+ "IMAGE_DESCRIPTION_OPEN",
+ "IMAGE_DESCRIPTION_CLOSE",
+]
diff --git a/openrag/core/prompts/chat_prompt_builder.py b/openrag/core/prompts/chat_prompt_builder.py
new file mode 100644
index 000000000..6d48fb42a
--- /dev/null
+++ b/openrag/core/prompts/chat_prompt_builder.py
@@ -0,0 +1,138 @@
+"""Chat-completion prompt builder.
+
+Pure helpers extracted from ``components/utils.py`` and ``components/pipeline.py``:
+
+* ``format_context`` — fit document snippets into a token budget,
+ numbering each as ``[Source N]``.
+* ``format_web_context`` — same, for web-search results, with continuous
+ numbering across RAG and web sources.
+* ``prepend_system_prompt`` — clone a message list and prepend a system
+ prompt rendered against ``context`` and
+ ``current_date``.
+* ``SOURCE_SEPARATOR`` — separator emitted between consecutive sources.
+
+Tokenizers are injected as ``Callable[[str], int]`` so this module stays pure
+(no LLM client, no LangChain).
+"""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Callable
+from typing import Protocol
+
+from openrag.core.utils.text import sanitize_text
+
+SOURCE_SEPARATOR = "-" * 10 + "\n\n"
+EMPTY_CONTEXT_MESSAGE = "No document found from the database"
+
+
+class WebSourceLike(Protocol):
+ """Minimal shape needed from a web-search result."""
+
+ title: str
+ url: str
+ snippet: str
+ content: str | None
+
+
+def format_context(
+ texts: list[str],
+ max_context_tokens: int,
+ length_function: Callable[[str], int],
+ *,
+ number_sources: bool = True,
+) -> tuple[str, list[int]]:
+ """Render ``texts`` as numbered ``[Source N]`` blocks within a token budget.
+
+ Args:
+ texts: Document texts (e.g. ``[d.page_content for d in docs]``).
+ max_context_tokens: Maximum total tokens for the context.
+ length_function: Token counter, e.g. ``llm.get_num_tokens``.
+ number_sources: If ``True``, prefix each block with ``[Source N]\\n``.
+
+ Returns:
+ ``(formatted_text, included_indices)`` — ``included_indices`` is the
+ positions in ``texts`` that fit within the budget; callers use it to
+ filter associated metadata down to the same set.
+ """
+ if not texts:
+ return EMPTY_CONTEXT_MESSAGE, []
+
+ reduced: list[str] = []
+ included: list[int] = []
+ total_tokens = 0
+
+ for i, text in enumerate(texts):
+ prefix = f"[Source {len(reduced) + 1}]\n" if number_sources else ""
+ n_tokens = length_function(text)
+ if prefix:
+ n_tokens += length_function(prefix)
+ if total_tokens + n_tokens > max_context_tokens:
+ break
+ reduced.append(f"{prefix}{text}")
+ included.append(i)
+ total_tokens += n_tokens
+
+ return SOURCE_SEPARATOR.join(reduced), included
+
+
+def format_web_context(
+ web_results: list[WebSourceLike],
+ length_function: Callable[[str], int],
+ *,
+ start_index: int = 1,
+ max_tokens: int = 2000,
+) -> tuple[str, list[int], int]:
+ """Render web results as numbered ``[Source N]`` blocks within a token budget.
+
+ Uses ``result.content`` when present, falling back to ``result.snippet``.
+
+ Args:
+ web_results: Web-search result objects (matching ``WebSourceLike``).
+ length_function: Token counter.
+ start_index: First source number — set to ``len(rag_sources) + 1`` so
+ web sources continue numbering after RAG sources.
+ max_tokens: Maximum total tokens for the web context.
+
+ Returns:
+ ``(formatted_text, source_numbers_used, total_tokens_used)``.
+ """
+ if not web_results:
+ return "", [], 0
+
+ parts: list[str] = []
+ source_numbers: list[int] = []
+ total_tokens = 0
+
+ for i, result in enumerate(web_results):
+ n = start_index + i
+ title = sanitize_text(result.title)
+ body_raw = result.content if result.content else result.snippet
+ body = sanitize_text(body_raw) if body_raw else ""
+ block = f"[Source {n}]\n{title}\n{body}"
+ block_tokens = length_function(block)
+ if total_tokens + block_tokens > max_tokens and parts:
+ break
+ parts.append(block)
+ source_numbers.append(n)
+ total_tokens += block_tokens
+
+ return SOURCE_SEPARATOR.join(parts), source_numbers, total_tokens
+
+
+def prepend_system_prompt(
+ messages: list[dict],
+ system_template: str,
+ *,
+ context: str,
+ current_date: str,
+) -> list[dict]:
+ """Return a deep-copied message list with a rendered system prompt prepended.
+
+ ``system_template`` must contain ``{context}`` and ``{current_date}``.
+ """
+ out = copy.deepcopy(messages)
+ rendered = system_template.format(context=context, current_date=current_date)
+ out.insert(0, {"role": "system", "content": rendered})
+ return out
diff --git a/openrag/core/prompts/contextualization_builder.py b/openrag/core/prompts/contextualization_builder.py
new file mode 100644
index 000000000..c5c45dccd
--- /dev/null
+++ b/openrag/core/prompts/contextualization_builder.py
@@ -0,0 +1,77 @@
+"""Chunk-contextualization prompt builder.
+
+Pure helpers extracted from ``components/indexer/chunker/chunker.py``. They
+produce the system+user message pair sent to the LLM when generating
+chunk-level context, and assemble the final wrapped chunk text used downstream.
+
+Format strings:
+ BASE_CHUNK_FORMAT — chunk wrapping when no LLM context is generated
+ CHUNK_FORMAT — chunk wrapping with leading [CONTEXT] block
+"""
+
+from __future__ import annotations
+
+BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]"
+CHUNK_FORMAT = "[CONTEXT]\n\n{chunk_context}\n\n" + BASE_CHUNK_FORMAT
+
+
+def build_user_message(
+ filename: str,
+ first_chunks_text: list[str],
+ prev_chunks_text: list[str],
+ current_chunk_text: str,
+ lang: str = "en",
+) -> str:
+ """Render the user-message body for a single chunk-contextualization call.
+
+ The system prompt is loaded from disk (``CHUNK_CONTEXTUALIZER_PROMPT``) and
+ paired with this user message by the caller.
+ """
+ first = "\n--\n".join(first_chunks_text)
+ previous = "\n--\n".join(prev_chunks_text)
+ return (
+ "\n"
+ " Here is the context to consider for generating the context:\n"
+ f" - Filename: {filename}\n"
+ " - First chunks:\n"
+ f" {first}\n\n"
+ " - Previous chunks:\n"
+ f" {previous}\n\n"
+ f" Here is the current chunk to contextualize strictly in this {lang} language:\n"
+ " - Current chunk:\n\n"
+ f" {current_chunk_text}\n "
+ )
+
+
+def build_messages(
+ system_prompt: str,
+ filename: str,
+ first_chunks_text: list[str],
+ prev_chunks_text: list[str],
+ current_chunk_text: str,
+ lang: str = "en",
+) -> list[dict[str, str]]:
+ """Build the system+user message list for a chunk-contextualization call."""
+ user = build_user_message(
+ filename=filename,
+ first_chunks_text=first_chunks_text,
+ prev_chunks_text=prev_chunks_text,
+ current_chunk_text=current_chunk_text,
+ lang=lang,
+ )
+ return [
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": user},
+ ]
+
+
+def wrap_chunk_with_context(content: str, filename: str, chunk_context: str = "") -> str:
+ """Wrap a chunk in the ``[CONTEXT] ... [CHUNK_START] ... [CHUNK_END]`` envelope.
+
+ If ``chunk_context`` is empty or whitespace-only, only the BASE_CHUNK_FORMAT
+ (no [CONTEXT] block) is used — preserves the legacy behavior for chunkers
+ that don't run contextualization.
+ """
+ if chunk_context and chunk_context.strip():
+ return CHUNK_FORMAT.format(content=content, chunk_context=chunk_context, filename=filename)
+ return BASE_CHUNK_FORMAT.format(content=content, filename=filename)
diff --git a/openrag/core/prompts/map_reduce_builder.py b/openrag/core/prompts/map_reduce_builder.py
new file mode 100644
index 000000000..632db5f8c
--- /dev/null
+++ b/openrag/core/prompts/map_reduce_builder.py
@@ -0,0 +1,38 @@
+"""Map-reduce prompt builder.
+
+The orchestrator (Phase 8) loops over chunks, calls the LLM with these
+messages, and reduces the structured outputs. The system + user-template
+strings live here so they're testable in isolation and can be evolved
+without touching pipeline code.
+"""
+
+from __future__ import annotations
+
+SYSTEM_PROMPT_MAP = """You are an AI assistant specialized in extracting and synthesizing relevant information from text.
+
+Your task:
+1. Analyze the provided text in relation to the user's question
+2. Extract only the essential information that directly addresses the query
+3. Preserve necessary context (Key words, project names or initiatives, dates, etc.) to maintain accuracy and clarity of the summary for it to be self-understandable
+
+Guidelines:
+- Present information clearly and concisely without unnecessary rephrasing or commentary
+- Focus on precision: include what matters, exclude what doesn't.
+- If a document does not have any relevant content with respect to the query, classify it as irrelevant without providing a `synthesis`.
+"""
+
+USER_PROMPT_TEMPLATE = """
+Here is a text:
+{content}
+
+From this document, identify and comprehensively summarize the information useful for answering the following question:
+{query}
+"""
+
+
+def build_map_messages(query: str, content: str) -> list[dict[str, str]]:
+ """Build the system+user message list for one map-step LLM call."""
+ return [
+ {"role": "system", "content": SYSTEM_PROMPT_MAP},
+ {"role": "user", "content": USER_PROMPT_TEMPLATE.format(query=query, content=content)},
+ ]
diff --git a/openrag/core/prompts/query_rewriter.py b/openrag/core/prompts/query_rewriter.py
new file mode 100644
index 000000000..a0b449ad7
--- /dev/null
+++ b/openrag/core/prompts/query_rewriter.py
@@ -0,0 +1,35 @@
+"""Query-rewriting prompt builders for HyDe and Multi-Query retrieval.
+
+Templates live on disk under ``prompts//`` and are loaded via
+``template_loader``. These functions are pure: they take a template string +
+substitution variables and return the formatted prompt.
+
+Template variables expected:
+ HyDe template: ``{question}``
+ Multi-query template: ``{query}``, ``{k_queries}``
+
+The multi-query helper also exposes the ``[SEP]`` separator used to split the
+LLM response into individual queries.
+"""
+
+from __future__ import annotations
+
+MULTI_QUERY_SEPARATOR = "[SEP]"
+
+
+def build_hyde_prompt(template: str, query: str) -> str:
+ """Format a HyDe prompt. ``template`` must contain ``{question}``."""
+ return template.format(question=query)
+
+
+def build_multi_query_prompt(template: str, query: str, k_queries: int) -> str:
+ """Format a multi-query prompt. ``template`` must contain ``{query}`` and ``{k_queries}``."""
+ return template.format(query=query, k_queries=k_queries)
+
+
+def split_multi_query_response(response: str, separator: str = MULTI_QUERY_SEPARATOR) -> list[str]:
+ """Split an LLM multi-query response into individual queries.
+
+ Drops empty entries and trims surrounding whitespace.
+ """
+ return [q.strip() for q in response.split(separator) if q.strip()]
diff --git a/openrag/core/prompts/template_loader.py b/openrag/core/prompts/template_loader.py
new file mode 100644
index 000000000..8e434be1c
--- /dev/null
+++ b/openrag/core/prompts/template_loader.py
@@ -0,0 +1,59 @@
+"""Disk-based prompt template loader.
+
+Pure I/O helper: given a directory and a filename, read and return the file
+contents as a string. Callers (typically the DI/composition layer) resolve
+the directory and the filename mapping from config; this function has no
+config dependency of its own.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def load_template(prompts_dir: str | Path, file_name: str) -> str:
+ """Read a prompt template from disk.
+
+ Args:
+ prompts_dir: Directory containing prompt template files.
+ file_name: Template filename (relative to ``prompts_dir``).
+
+ Returns:
+ The template contents as a string.
+
+ Raises:
+ FileNotFoundError: if the resolved path does not exist.
+ """
+ file_path = Path(prompts_dir) / file_name
+ if not file_path.exists():
+ raise FileNotFoundError(f"Prompt file not found: `{file_path}`")
+ return file_path.read_text(encoding="utf-8")
+
+
+def load_template_by_key(
+ prompts_dir: str | Path,
+ prompt_mapping: object,
+ prompt_key: str,
+) -> str:
+ """Read a prompt by logical key, looking up the filename on a mapping object.
+
+ The mapping object is typically the ``PromptsConfig`` Pydantic model with
+ attributes like ``sys_prompt``, ``hyde``, ``multi_query`` whose values are
+ template filenames.
+
+ Args:
+ prompts_dir: Directory containing prompt template files.
+ prompt_mapping: Object exposing prompt keys as attributes.
+ prompt_key: Attribute name on ``prompt_mapping`` (e.g. ``"hyde"``).
+
+ Returns:
+ The template contents as a string.
+
+ Raises:
+ ValueError: if ``prompt_key`` is not defined on ``prompt_mapping``.
+ FileNotFoundError: if the resolved path does not exist.
+ """
+ file_name = getattr(prompt_mapping, prompt_key, None)
+ if not file_name:
+ raise ValueError(f"No associated file name found for prompt: `{prompt_key}`")
+ return load_template(prompts_dir, file_name)
diff --git a/openrag/core/prompts/test_chat_prompt_builder.py b/openrag/core/prompts/test_chat_prompt_builder.py
new file mode 100644
index 000000000..7781103a8
--- /dev/null
+++ b/openrag/core/prompts/test_chat_prompt_builder.py
@@ -0,0 +1,110 @@
+"""Tests for chat_prompt_builder — these lock in the exact wire format."""
+
+from __future__ import annotations
+
+from openrag.core.prompts.chat_prompt_builder import (
+ EMPTY_CONTEXT_MESSAGE,
+ SOURCE_SEPARATOR,
+ format_context,
+ format_web_context,
+ prepend_system_prompt,
+)
+
+
+def _word_tokens(text: str) -> int:
+ """Simple deterministic token counter: 1 token per whitespace-delimited word."""
+ return len(text.split())
+
+
+def test_format_context_empty_returns_placeholder():
+ text, included = format_context([], max_context_tokens=100, length_function=_word_tokens)
+ assert text == EMPTY_CONTEXT_MESSAGE
+ assert included == []
+
+
+def test_format_context_numbers_sources_and_separates():
+ docs = ["alpha beta", "gamma delta epsilon"]
+ text, included = format_context(docs, max_context_tokens=100, length_function=_word_tokens)
+ assert "[Source 1]\nalpha beta" in text
+ assert "[Source 2]\ngamma delta epsilon" in text
+ assert SOURCE_SEPARATOR in text
+ assert included == [0, 1]
+
+
+def test_format_context_drops_to_fit_budget():
+ docs = ["one two", "three four five", "six"]
+ # _word_tokens("[Source 1]\n") = 1, doc1 = 2, prefix2 = 1, doc2 = 3 -> total 7
+ text, included = format_context(docs, max_context_tokens=4, length_function=_word_tokens)
+ assert "[Source 1]" in text
+ assert "[Source 2]" not in text
+ assert included == [0]
+
+
+def test_format_context_no_numbering():
+ docs = ["a", "b"]
+ text, included = format_context(docs, max_context_tokens=100, length_function=_word_tokens, number_sources=False)
+ assert "[Source" not in text
+ assert text == f"a{SOURCE_SEPARATOR}b"
+ assert included == [0, 1]
+
+
+class _FakeWeb:
+ def __init__(self, title: str, url: str, snippet: str, content: str | None = None):
+ self.title = title
+ self.url = url
+ self.snippet = snippet
+ self.content = content
+
+
+def test_format_web_context_uses_content_when_present():
+ results = [_FakeWeb("T1", "u1", "snip1", content="full body")]
+ text, nums, _ = format_web_context(results, length_function=_word_tokens, max_tokens=100)
+ assert "full body" in text
+ assert "snip1" not in text
+ assert nums == [1]
+
+
+def test_format_web_context_falls_back_to_snippet():
+ results = [_FakeWeb("T1", "u1", "snip1", content=None)]
+ text, _, _ = format_web_context(results, length_function=_word_tokens, max_tokens=100)
+ assert "snip1" in text
+
+
+def test_format_web_context_continues_numbering_with_start_index():
+ results = [_FakeWeb("T1", "u1", "snip1")]
+ text, nums, _ = format_web_context(results, length_function=_word_tokens, start_index=4, max_tokens=100)
+ assert "[Source 4]" in text
+ assert nums == [4]
+
+
+def test_prepend_system_prompt_does_not_mutate_input():
+ msgs = [{"role": "user", "content": "hi"}]
+ out = prepend_system_prompt(
+ msgs,
+ system_template="ctx={context} date={current_date}",
+ context="C",
+ current_date="2026-04-29",
+ )
+ assert msgs == [{"role": "user", "content": "hi"}]
+ assert out[0] == {"role": "system", "content": "ctx=C date=2026-04-29"}
+ assert out[1] == {"role": "user", "content": "hi"}
+
+
+def test_format_web_context_empty_returns_empty_tuple():
+ text, nums, total = format_web_context([], length_function=_word_tokens)
+ assert text == ""
+ assert nums == []
+ assert total == 0
+
+
+def test_format_web_context_drops_overflow_block_after_first_fits():
+ """If a later block would push past max_tokens we break — but only after
+ at least one block has been admitted (parts truthy guard)."""
+ results = [
+ _FakeWeb("T1", "u1", "short body"),
+ _FakeWeb("T2", "u2", "this snippet has many many many many many words that will overflow"),
+ ]
+ text, nums, _ = format_web_context(results, length_function=_word_tokens, max_tokens=10)
+ assert "[Source 1]" in text
+ assert "[Source 2]" not in text
+ assert nums == [1]
diff --git a/openrag/core/prompts/test_contextualization_builder.py b/openrag/core/prompts/test_contextualization_builder.py
new file mode 100644
index 000000000..0955a3528
--- /dev/null
+++ b/openrag/core/prompts/test_contextualization_builder.py
@@ -0,0 +1,82 @@
+"""Tests for the chunk-contextualization prompt builder."""
+
+from __future__ import annotations
+
+from openrag.core.prompts.contextualization_builder import (
+ BASE_CHUNK_FORMAT,
+ CHUNK_FORMAT,
+ build_messages,
+ build_user_message,
+ wrap_chunk_with_context,
+)
+
+
+def test_build_user_message_includes_filename_and_lang():
+ out = build_user_message(
+ filename="doc.pdf",
+ first_chunks_text=["intro chunk A", "intro chunk B"],
+ prev_chunks_text=["prev chunk"],
+ current_chunk_text="here is the current chunk",
+ lang="fr",
+ )
+ assert "doc.pdf" in out
+ assert "intro chunk A" in out
+ assert "intro chunk B" in out
+ assert "prev chunk" in out
+ assert "here is the current chunk" in out
+ assert "fr language" in out
+
+
+def test_build_user_message_handles_empty_history():
+ out = build_user_message(
+ filename="doc.pdf",
+ first_chunks_text=[],
+ prev_chunks_text=[],
+ current_chunk_text="solo chunk",
+ )
+ assert "solo chunk" in out
+ assert "en language" in out # default lang
+
+
+def test_build_messages_returns_system_then_user():
+ msgs = build_messages(
+ system_prompt="SYS",
+ filename="doc.pdf",
+ first_chunks_text=["a"],
+ prev_chunks_text=["b"],
+ current_chunk_text="c",
+ )
+ assert len(msgs) == 2
+ assert msgs[0] == {"role": "system", "content": "SYS"}
+ assert msgs[1]["role"] == "user"
+ assert "doc.pdf" in msgs[1]["content"]
+
+
+def test_wrap_chunk_with_context_uses_full_format_when_context_given():
+ out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context="ctx")
+ assert "[CONTEXT]" in out
+ assert "ctx" in out
+ assert "[CHUNK_START]" in out
+ assert "body" in out
+ assert "[CHUNK_END]" in out
+ # Sanity: result uses the documented CHUNK_FORMAT template.
+ assert out == CHUNK_FORMAT.format(content="body", chunk_context="ctx", filename="f.pdf")
+
+
+def test_wrap_chunk_with_context_uses_base_format_when_context_empty():
+ out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context="")
+ assert "[CONTEXT]" not in out
+ assert "[CHUNK_START]" in out
+ assert "body" in out
+ assert out == BASE_CHUNK_FORMAT.format(content="body", filename="f.pdf")
+
+
+def test_wrap_chunk_with_context_defaults_chunk_context_to_empty():
+ out = wrap_chunk_with_context(content="body", filename="f.pdf")
+ assert "[CONTEXT]" not in out
+
+
+def test_wrap_chunk_with_context_treats_whitespace_only_as_empty():
+ out = wrap_chunk_with_context(content="body", filename="f.pdf", chunk_context=" \n\t")
+ assert "[CONTEXT]" not in out
+ assert out == BASE_CHUNK_FORMAT.format(content="body", filename="f.pdf")
diff --git a/openrag/core/prompts/test_map_reduce_builder.py b/openrag/core/prompts/test_map_reduce_builder.py
new file mode 100644
index 000000000..a89a37b11
--- /dev/null
+++ b/openrag/core/prompts/test_map_reduce_builder.py
@@ -0,0 +1,27 @@
+"""Tests for the map-reduce prompt builder."""
+
+from __future__ import annotations
+
+from openrag.core.prompts.map_reduce_builder import (
+ SYSTEM_PROMPT_MAP,
+ USER_PROMPT_TEMPLATE,
+ build_map_messages,
+)
+
+
+def test_build_map_messages_returns_system_then_user_with_substitution():
+ msgs = build_map_messages(query="what is rag?", content="some doc body")
+ assert len(msgs) == 2
+ assert msgs[0] == {"role": "system", "content": SYSTEM_PROMPT_MAP}
+ assert msgs[1]["role"] == "user"
+ user_text = msgs[1]["content"]
+ assert "what is rag?" in user_text
+ assert "some doc body" in user_text
+
+
+def test_user_prompt_template_uses_named_placeholders():
+ """Lock the template's named placeholders so accidental positional refactors
+ don't silently change the wire format."""
+ rendered = USER_PROMPT_TEMPLATE.format(query="Q", content="C")
+ assert "Q" in rendered
+ assert "C" in rendered
diff --git a/openrag/core/prompts/test_query_rewriter.py b/openrag/core/prompts/test_query_rewriter.py
new file mode 100644
index 000000000..7a32a53d2
--- /dev/null
+++ b/openrag/core/prompts/test_query_rewriter.py
@@ -0,0 +1,40 @@
+"""Tests for the HyDe / multi-query prompt builders."""
+
+from __future__ import annotations
+
+from openrag.core.prompts.query_rewriter import (
+ MULTI_QUERY_SEPARATOR,
+ build_hyde_prompt,
+ build_multi_query_prompt,
+ split_multi_query_response,
+)
+
+
+def test_build_hyde_prompt_substitutes_question():
+ out = build_hyde_prompt("Q: {question}", "what is rag?")
+ assert out == "Q: what is rag?"
+
+
+def test_build_multi_query_prompt_substitutes_query_and_k():
+ out = build_multi_query_prompt("Generate {k_queries} variants of: {query}", "what is rag?", 5)
+ assert "5" in out
+ assert "what is rag?" in out
+
+
+def test_split_multi_query_response_splits_on_separator_and_trims():
+ raw = f" first query {MULTI_QUERY_SEPARATOR} second query {MULTI_QUERY_SEPARATOR}third"
+ assert split_multi_query_response(raw) == ["first query", "second query", "third"]
+
+
+def test_split_multi_query_response_drops_empty_entries():
+ raw = f" {MULTI_QUERY_SEPARATOR}only{MULTI_QUERY_SEPARATOR}{MULTI_QUERY_SEPARATOR} "
+ assert split_multi_query_response(raw) == ["only"]
+
+
+def test_split_multi_query_response_empty_string_returns_empty_list():
+ assert split_multi_query_response("") == []
+
+
+def test_split_multi_query_response_accepts_custom_separator():
+ raw = "a||b||c"
+ assert split_multi_query_response(raw, separator="||") == ["a", "b", "c"]
diff --git a/openrag/core/prompts/test_template_loader.py b/openrag/core/prompts/test_template_loader.py
new file mode 100644
index 000000000..e6ee5551c
--- /dev/null
+++ b/openrag/core/prompts/test_template_loader.py
@@ -0,0 +1,59 @@
+"""Tests for the disk-based prompt template loader."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from openrag.core.prompts.template_loader import load_template, load_template_by_key
+
+
+def test_load_template_reads_file_contents(tmp_path: Path):
+ target = tmp_path / "sys.txt"
+ target.write_text("hello {name}", encoding="utf-8")
+ assert load_template(tmp_path, "sys.txt") == "hello {name}"
+
+
+def test_load_template_accepts_str_path(tmp_path: Path):
+ target = tmp_path / "sys.txt"
+ target.write_text("hi", encoding="utf-8")
+ assert load_template(str(tmp_path), "sys.txt") == "hi"
+
+
+def test_load_template_raises_on_missing_file(tmp_path: Path):
+ with pytest.raises(FileNotFoundError, match="Prompt file not found"):
+ load_template(tmp_path, "nope.txt")
+
+
+class _Mapping:
+ """Stand-in for the PromptsConfig pydantic model."""
+
+ def __init__(self, **kwargs: str) -> None:
+ for k, v in kwargs.items():
+ setattr(self, k, v)
+
+
+def test_load_template_by_key_resolves_filename(tmp_path: Path):
+ (tmp_path / "hyde.txt").write_text("hyde body", encoding="utf-8")
+ mapping = _Mapping(hyde="hyde.txt")
+ assert load_template_by_key(tmp_path, mapping, "hyde") == "hyde body"
+
+
+def test_load_template_by_key_raises_when_attr_missing(tmp_path: Path):
+ mapping = _Mapping(hyde="hyde.txt")
+ with pytest.raises(ValueError, match="No associated file name"):
+ load_template_by_key(tmp_path, mapping, "multi_query")
+
+
+def test_load_template_by_key_raises_when_attr_falsy(tmp_path: Path):
+ """A mapping value of empty string / None should be treated as "not set"."""
+ mapping = _Mapping(multi_query="")
+ with pytest.raises(ValueError, match="No associated file name"):
+ load_template_by_key(tmp_path, mapping, "multi_query")
+
+
+def test_load_template_by_key_propagates_file_not_found(tmp_path: Path):
+ mapping = _Mapping(hyde="missing.txt")
+ with pytest.raises(FileNotFoundError):
+ load_template_by_key(tmp_path, mapping, "hyde")
diff --git a/openrag/core/prompts/test_vlm_prompt_builder.py b/openrag/core/prompts/test_vlm_prompt_builder.py
new file mode 100644
index 000000000..7fd6dc489
--- /dev/null
+++ b/openrag/core/prompts/test_vlm_prompt_builder.py
@@ -0,0 +1,44 @@
+"""Tests for the VLM (image-captioning) prompt builder."""
+
+from __future__ import annotations
+
+from openrag.core.prompts.vlm_prompt_builder import (
+ IMAGE_DESCRIPTION_CLOSE,
+ IMAGE_DESCRIPTION_OPEN,
+ build_caption_messages,
+ wrap_caption,
+)
+
+
+def test_build_caption_messages_shapes_for_openai_multimodal():
+ msgs = build_caption_messages(template="Describe this image.", image_url="https://example.com/x.png")
+ assert len(msgs) == 1
+ msg = msgs[0]
+ assert msg["role"] == "user"
+ parts = msg["content"]
+ assert {p["type"] for p in parts} == {"image_url", "text"}
+ image_part = next(p for p in parts if p["type"] == "image_url")
+ text_part = next(p for p in parts if p["type"] == "text")
+ assert image_part["image_url"] == {"url": "https://example.com/x.png"}
+ assert text_part["text"] == "Describe this image."
+
+
+def test_build_caption_messages_supports_data_uri():
+ msgs = build_caption_messages(template="caption", image_url="data:image/png;base64,abc")
+ image_part = next(p for p in msgs[0]["content"] if p["type"] == "image_url")
+ assert image_part["image_url"]["url"].startswith("data:image/png;base64,")
+
+
+def test_wrap_caption_uses_image_description_markers():
+ out = wrap_caption("a sunset over the sea")
+ assert out.startswith(IMAGE_DESCRIPTION_OPEN)
+ assert out.endswith(IMAGE_DESCRIPTION_CLOSE)
+ assert "a sunset over the sea" in out
+
+
+def test_wrap_caption_format_is_load_bearing_for_chunker():
+ """The chunker's image-element regex matches `...`,
+ so this exact pairing is part of the contract."""
+ out = wrap_caption("x")
+ assert "" in out
+ assert "" in out
diff --git a/openrag/core/prompts/vlm_prompt_builder.py b/openrag/core/prompts/vlm_prompt_builder.py
new file mode 100644
index 000000000..f966fe48a
--- /dev/null
+++ b/openrag/core/prompts/vlm_prompt_builder.py
@@ -0,0 +1,44 @@
+"""VLM (vision-language model) prompt builder.
+
+Pure helpers: given a captioning template and an image reference, produce a
+multimodal chat-message payload (OpenAI-style) and wrap captions in the
+```` markers downstream pipelines expect.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+IMAGE_DESCRIPTION_OPEN = ""
+IMAGE_DESCRIPTION_CLOSE = ""
+
+
+def build_caption_messages(template: str, image_url: str) -> list[dict[str, Any]]:
+ """Build a multimodal chat-message list for image captioning.
+
+ Args:
+ template: Image-captioning prompt text (no substitution required).
+ image_url: ``https://...`` URL or ``data:image/...;base64,...`` data URI.
+
+ Returns:
+ A single-message list shaped for OpenAI / vLLM chat completions:
+ ``[{"role": "user", "content": [{"type": "image_url", ...}, {"type": "text", ...}]}]``
+ """
+ return [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image_url", "image_url": {"url": image_url}},
+ {"type": "text", "text": template},
+ ],
+ }
+ ]
+
+
+def wrap_caption(caption: str) -> str:
+ """Wrap a raw caption in ```` markers.
+
+ Pipelines downstream (markdown image replacement, chunk parsing) look for
+ this exact marker, so the wrapping format is part of the contract.
+ """
+ return f"{IMAGE_DESCRIPTION_OPEN}\n\n{caption}\n\n{IMAGE_DESCRIPTION_CLOSE}"
diff --git a/openrag/core/retrieval/__init__.py b/openrag/core/retrieval/__init__.py
index e69de29bb..9a01aec35 100644
--- a/openrag/core/retrieval/__init__.py
+++ b/openrag/core/retrieval/__init__.py
@@ -0,0 +1,25 @@
+"""Retrieval domain logic: retriever strategies, RRF, and pipeline."""
+
+from .pipeline import RetrieverPipeline
+from .registry import retriever_registry
+from .retriever import (
+ BaseRetriever,
+ HyDeRetriever,
+ MultiQueryRetriever,
+ Retriever,
+ SingleRetriever,
+)
+from .rrf import rrf_reranking
+from .searcher import RetrievalSearcher
+
+__all__ = [
+ "Retriever",
+ "BaseRetriever",
+ "SingleRetriever",
+ "MultiQueryRetriever",
+ "HyDeRetriever",
+ "retriever_registry",
+ "RetrievalSearcher",
+ "RetrieverPipeline",
+ "rrf_reranking",
+]
diff --git a/openrag/core/retrieval/pipeline.py b/openrag/core/retrieval/pipeline.py
new file mode 100644
index 000000000..5f0ece5a7
--- /dev/null
+++ b/openrag/core/retrieval/pipeline.py
@@ -0,0 +1,152 @@
+"""Retrieval pipeline: per-query retrieval, optional temporal-filter fallback,
+optional reranking, optional related/ancestor expansion, and RRF fusion across
+sub-queries.
+
+Extracted from ``components/pipeline.py:RetrieverPipeline``. The legacy
+``RagPipeline`` (LLM-driven query generation, system-prompt assembly,
+streaming) lives in the orchestrator layer and is rebuilt in Phase 8.
+
+This pipeline depends only on core ABCs:
+
+ * ``Retriever`` — strategy that produces candidate chunks
+ * ``Reranker`` — optional cross-encoder reranker (per Phase 4 ABC:
+ returns ``[(idx, score), ...]`` over a list of texts)
+ * ``RetrievalSearcher`` is consumed by the retriever, not directly here.
+
+Config knobs are constructor arguments; there is no module-level config load.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+from typing import Any
+
+from openrag.core.models.chunk import Chunk
+from openrag.core.models.query import Query, SearchQueries
+from openrag.core.rerankers.reranker import Reranker
+from openrag.core.retrieval.retriever import Retriever
+from openrag.core.retrieval.rrf import rrf_reranking
+
+
+def _chunk_key(c: Chunk) -> Any:
+ """Identity key for fusion / dedup. Falls back to object id when missing."""
+ return c.id or id(c)
+
+
+async def _rerank_chunks(reranker: Reranker, query: str, chunks: list[Chunk]) -> list[Chunk]:
+ """Reorder chunks via the Reranker ABC.
+
+ The ABC scores text+query pairs and returns ``[(orig_index, score), ...]``;
+ we look up the original chunk for each ranked index. Items the reranker
+ drops are excluded.
+ """
+ if not chunks:
+ return chunks
+ ranking = await reranker.rerank(query=query, documents=[c.text for c in chunks], top_k=None)
+ return [chunks[idx] for idx, _ in ranking]
+
+
+class RetrieverPipeline:
+ """Orchestrates retrieval + reranking + expansion for a list of sub-queries.
+
+ Args:
+ retriever: Concrete retrieval strategy (Single / MultiQuery / HyDe).
+ reranker: Reranker implementation, or ``None`` to skip reranking.
+ reranker_top_k: When expansion is enabled, the top-K size used to
+ decide which results to expand.
+ allow_filterless_fallback: If a temporal filter wipes out all
+ candidates, retry once without it. When ``False``,
+ return zero docs rather than ones outside the
+ temporal range.
+ """
+
+ def __init__(
+ self,
+ retriever: Retriever,
+ reranker: Reranker | None = None,
+ reranker_top_k: int = 5,
+ allow_filterless_fallback: bool = True,
+ ) -> None:
+ self.retriever = retriever
+ self.reranker = reranker
+ self.reranker_top_k = reranker_top_k
+ self.allow_filterless_fallback = allow_filterless_fallback
+
+ @property
+ def reranker_enabled(self) -> bool:
+ return self.reranker is not None
+
+ @property
+ def expansion_enabled(self) -> bool:
+ # The retriever's BaseRetriever sets this; non-Base implementations
+ # may not. Treat absent attribute as no expansion.
+ return getattr(self.retriever, "expansion_enabled", False)
+
+ async def retrieve_docs(
+ self,
+ partition: list[str],
+ query: Query,
+ top_k: int | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ """Run a single ``Query`` through retrieval, expansion, and reranking."""
+ milvus_filter = query.to_milvus_filter()
+ chunks = await self.retriever.retrieve(
+ partition=partition,
+ query=query.query,
+ filter=milvus_filter,
+ filter_params=filter_params,
+ )
+
+ if not chunks and milvus_filter and self.allow_filterless_fallback:
+ # Temporal filter killed every candidate — retry without it so
+ # the user gets some results rather than none.
+ chunks = await self.retriever.retrieve(
+ partition=partition,
+ query=query.query,
+ filter=None,
+ filter_params=filter_params,
+ )
+
+ if not chunks:
+ return chunks
+
+ if self.reranker_enabled:
+ chunks = await _rerank_chunks(self.reranker, query.query, chunks)
+
+ if self.expansion_enabled:
+ limit = self.reranker_top_k if top_k is None else max(self.reranker_top_k, top_k)
+ head = copy.deepcopy(chunks[:limit])
+ expanded = await self.retriever.expand_search_results(results=head)
+ if len(expanded) > len(head):
+ chunks = expanded
+ if self.reranker_enabled:
+ chunks = await _rerank_chunks(self.reranker, query.query, chunks)
+
+ if top_k is not None:
+ chunks = chunks[:top_k]
+ return chunks
+
+ async def get_relevant_docs(
+ self,
+ partition: list[str],
+ search_queries: SearchQueries,
+ top_k: int | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ """Run every sub-query in parallel and fuse the per-query rankings via RRF."""
+ tasks = [
+ self.retrieve_docs(
+ partition=partition,
+ query=q,
+ top_k=top_k,
+ filter_params=filter_params,
+ )
+ for q in search_queries.query_list
+ ]
+ ranked_lists = await asyncio.gather(*tasks)
+ fused = rrf_reranking(ranked_lists, key_fn=_chunk_key)
+ if top_k is not None:
+ fused = fused[:top_k]
+ return fused
diff --git a/openrag/core/retrieval/registry.py b/openrag/core/retrieval/registry.py
new file mode 100644
index 000000000..11a92e626
--- /dev/null
+++ b/openrag/core/retrieval/registry.py
@@ -0,0 +1,12 @@
+"""Retriever strategy registry."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from openrag.core.utils.registry import Registry
+
+if TYPE_CHECKING:
+ from .retriever import Retriever
+
+retriever_registry: Registry[Retriever] = Registry("retriever")
diff --git a/openrag/core/retrieval/retriever.py b/openrag/core/retrieval/retriever.py
new file mode 100644
index 000000000..5bde40377
--- /dev/null
+++ b/openrag/core/retrieval/retriever.py
@@ -0,0 +1,274 @@
+"""Retriever strategies: Single, MultiQuery, HyDe.
+
+Rewritten from ``components/retriever.py``. Differences from the legacy:
+
+ * ``RetrievalSearcher`` (clean ABC) replaces ``get_vectordb()`` / Ray actor
+ direct access. The retriever has no Ray imports.
+ * ``LLM`` (clean ABC) replaces ``ChatOpenAI`` + LangChain chain assembly.
+ * Prompt templates are passed in as strings. The DI layer loads them
+ from disk via ``core/prompts/template_loader``.
+ * Returns are domain ``Chunk`` objects, not LangChain ``Document``.
+
+Concrete strategies register themselves with ``retriever_registry``
+(declared in :mod:`openrag.core.retrieval.registry`) via decorator at
+class-definition time, so the composition root can pick one by name
+(``single`` / ``multiQuery`` / ``hyde``) per the strategy doc's
+"every factory becomes a Registry" rule.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from itertools import chain as ichain
+from typing import Any
+
+from openrag.core.llm.llm import LLM
+from openrag.core.models.chunk import Chunk
+from openrag.core.prompts.query_rewriter import (
+ build_hyde_prompt,
+ build_multi_query_prompt,
+ split_multi_query_response,
+)
+from openrag.core.retrieval.registry import retriever_registry
+from openrag.core.retrieval.searcher import RetrievalSearcher
+
+
+class Retriever(ABC):
+ """Common surface for all retrieval strategies."""
+
+ @abstractmethod
+ async def retrieve(
+ self,
+ partition: list[str],
+ query: str,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ """Run the strategy and return scored chunks."""
+ ...
+
+ @abstractmethod
+ async def expand_search_results(self, results: list[Chunk]) -> list[Chunk]:
+ """Optionally enrich a result set with related/ancestor chunks."""
+ ...
+
+
+class BaseRetriever(Retriever):
+ """Single-query retriever — the building block for the others."""
+
+ def __init__(
+ self,
+ searcher: RetrievalSearcher,
+ top_k: int = 6,
+ similarity_threshold: float = 0.95,
+ with_surrounding_chunks: bool = True,
+ include_related: bool = False,
+ include_ancestors: bool = False,
+ related_limit: int = 10,
+ max_ancestor_depth: int | None = None,
+ **_: Any,
+ ) -> None:
+ self.searcher = searcher
+ self.top_k = top_k
+ self.similarity_threshold = similarity_threshold
+ self.with_surrounding_chunks = with_surrounding_chunks
+ self.include_related = include_related
+ self.include_ancestors = include_ancestors
+ self.related_limit = related_limit
+ self.max_ancestor_depth = max_ancestor_depth
+ self.expansion_enabled = include_related or include_ancestors
+
+ async def retrieve(
+ self,
+ partition: list[str],
+ query: str,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ return await self.searcher.search(
+ query=query,
+ partition=partition,
+ top_k=self.top_k,
+ filter=filter,
+ filter_params=filter_params,
+ similarity_threshold=self.similarity_threshold,
+ with_surrounding_chunks=self.with_surrounding_chunks,
+ )
+
+ async def expand_search_results(self, results: list[Chunk]) -> list[Chunk]:
+ return await _expand_with_related_chunks(
+ searcher=self.searcher,
+ results=results,
+ include_related=self.include_related,
+ include_ancestors=self.include_ancestors,
+ related_limit=self.related_limit,
+ max_ancestor_depth=self.max_ancestor_depth,
+ )
+
+
+@retriever_registry.register("single")
+class SingleRetriever(BaseRetriever):
+ """Default strategy — issues exactly one similarity search per query."""
+
+
+@retriever_registry.register("multiQuery")
+class MultiQueryRetriever(BaseRetriever):
+ """Generates K query variants via the LLM and unions their results."""
+
+ def __init__(
+ self,
+ searcher: RetrievalSearcher,
+ llm: LLM,
+ multi_query_template: str,
+ k_queries: int = 3,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(searcher=searcher, **kwargs)
+ if llm is None:
+ raise ValueError("llm must be provided for MultiQueryRetriever")
+ self.llm = llm
+ self.multi_query_template = multi_query_template
+ self.k_queries = k_queries
+
+ async def _generate_queries(self, query: str) -> list[str]:
+ prompt = build_multi_query_prompt(self.multi_query_template, query, self.k_queries)
+ response = await self.llm.chat([{"role": "user", "content": prompt}])
+ # Cap to k_queries — a non-compliant LLM response can otherwise fan
+ # out far more searches than configured.
+ queries = split_multi_query_response(response)[: self.k_queries]
+ return queries or [query]
+
+ async def retrieve(
+ self,
+ partition: list[str],
+ query: str,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ queries = await self._generate_queries(query)
+ return await self.searcher.multi_query_search(
+ queries=queries,
+ partition=partition,
+ top_k_per_query=self.top_k,
+ filter=filter,
+ filter_params=filter_params,
+ similarity_threshold=self.similarity_threshold,
+ with_surrounding_chunks=self.with_surrounding_chunks,
+ )
+
+
+@retriever_registry.register("hyde")
+class HyDeRetriever(BaseRetriever):
+ """Generates a hypothetical answer document and searches with it.
+
+ If ``combine`` is set, the original query is also issued and results
+ are unioned via the searcher's multi-query path.
+ """
+
+ def __init__(
+ self,
+ searcher: RetrievalSearcher,
+ llm: LLM,
+ hyde_template: str,
+ combine: bool = False,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(searcher=searcher, **kwargs)
+ if llm is None:
+ raise ValueError("llm must be provided for HyDeRetriever")
+ self.llm = llm
+ self.hyde_template = hyde_template
+ self.combine = combine
+
+ async def get_hyde(self, query: str) -> str:
+ prompt = build_hyde_prompt(self.hyde_template, query)
+ return await self.llm.chat([{"role": "user", "content": prompt}])
+
+ async def retrieve(
+ self,
+ partition: list[str],
+ query: str,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ ) -> list[Chunk]:
+ hyde = (await self.get_hyde(query)).strip()
+ if not hyde:
+ queries = [query]
+ else:
+ queries = [hyde, query] if self.combine else [hyde]
+ return await self.searcher.multi_query_search(
+ queries=queries,
+ partition=partition,
+ top_k_per_query=self.top_k,
+ filter=filter,
+ filter_params=filter_params,
+ similarity_threshold=self.similarity_threshold,
+ with_surrounding_chunks=self.with_surrounding_chunks,
+ )
+
+
+async def _expand_with_related_chunks(
+ searcher: RetrievalSearcher,
+ results: list[Chunk],
+ include_related: bool,
+ include_ancestors: bool,
+ related_limit: int = 10,
+ max_ancestor_depth: int | None = None,
+) -> list[Chunk]:
+ """Append related and/or ancestor chunks to a result set, deduplicated by id.
+
+ Failures on individual related/ancestor lookups are logged and treated
+ as empty results, matching legacy behavior so retrieval remains
+ resilient to per-document errors.
+ """
+ if not results or (not include_related and not include_ancestors):
+ return results
+
+ seen_ids = {c.id for c in results if c.id}
+ expanded: list[Chunk] = list(results)
+
+ relationship_ids: set[tuple[str, str]] = set()
+ file_infos: set[tuple[str, str]] = set()
+
+ for c in results:
+ if include_related:
+ rel_id = c.metadata.get("relationship_id")
+ if rel_id and c.partition:
+ relationship_ids.add((c.partition, rel_id))
+ if include_ancestors and c.partition and c.document_id:
+ file_infos.add((c.partition, c.document_id))
+
+ async def _safe_related(part: str, rel_id: str) -> list[Chunk]:
+ try:
+ return await searcher.get_related_chunks(partition=part, relationship_id=rel_id, limit=related_limit)
+ except Exception:
+ return []
+
+ async def _safe_ancestors(part: str, file_id: str) -> list[Chunk]:
+ try:
+ return await searcher.get_ancestor_chunks(
+ partition=part,
+ file_id=file_id,
+ limit=related_limit,
+ max_ancestor_depth=max_ancestor_depth,
+ )
+ except Exception:
+ return []
+
+ tasks: list[asyncio.Future] = []
+ if include_related:
+ tasks.extend(_safe_related(part, rid) for part, rid in relationship_ids)
+ if include_ancestors:
+ tasks.extend(_safe_ancestors(part, fid) for part, fid in file_infos if part and fid)
+
+ if tasks:
+ all_results = await asyncio.gather(*tasks)
+ for chunk in ichain.from_iterable(all_results):
+ if chunk.id and chunk.id in seen_ids:
+ continue
+ if chunk.id:
+ seen_ids.add(chunk.id)
+ expanded.append(chunk)
+
+ return expanded
diff --git a/openrag/core/retrieval/rrf.py b/openrag/core/retrieval/rrf.py
new file mode 100644
index 000000000..eadffdbe9
--- /dev/null
+++ b/openrag/core/retrieval/rrf.py
@@ -0,0 +1,66 @@
+"""Reciprocal Rank Fusion — pure math, no domain coupling.
+
+Combines multiple ranked lists into a single ranking by summing reciprocal
+ranks across lists. Items present in more lists, or higher-ranked in any
+list, sort to the top of the fused result.
+
+Formula:
+ score(item) = Σ_i 1 / (k + rank_i)
+
+with ``rank_i`` the 1-based rank of the item in list ``i``. Smaller ``k``
+amplifies the top of each list; ``k=60`` is the canonical default and
+balances rank sensitivity across lists.
+
+Identification of "the same item" is delegated to the caller via
+``key_fn`` — typically returning the chunk id, document id, or URL.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Hashable, Sequence
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def rrf_reranking(
+ ranked_lists: Sequence[Sequence[T]],
+ key_fn: Callable[[T], Hashable] | None = None,
+ k: int = 60,
+) -> list[T]:
+ """Fuse multiple ranked lists into one via Reciprocal Rank Fusion.
+
+ Args:
+ ranked_lists: Each inner sequence is a ranked list (best first).
+ key_fn: Returns the identity key for an item; items sharing a key
+ across lists have their RRF scores summed. Defaults to
+ ``id(item)`` (object identity), which prevents fusion across
+ lists for items lacking a logical id.
+ k: RRF dampening constant. ``60`` is canonical.
+
+ Returns:
+ A single ranked list, best first. Empty input -> empty list.
+ Single input list is returned as-is.
+
+ Raises:
+ ValueError: if ``k < 0`` (would produce a zero or negative
+ denominator at rank 1 or below and crash with ZeroDivisionError).
+ """
+ if k < 0:
+ raise ValueError(f"RRF k must be non-negative, got {k}")
+ if not ranked_lists:
+ return []
+ if len(ranked_lists) == 1:
+ return list(ranked_lists[0])
+
+ if key_fn is None:
+ key_fn = id # type: ignore[assignment]
+
+ fused: dict[Hashable, tuple[float, T]] = {}
+ for ranked in ranked_lists:
+ for rank, item in enumerate(ranked, start=1):
+ key = key_fn(item)
+ score, kept = fused.get(key, (0.0, item))
+ fused[key] = (score + 1.0 / (rank + k), kept)
+
+ return [item for _, item in sorted(fused.values(), key=lambda x: x[0], reverse=True)]
diff --git a/openrag/core/retrieval/searcher.py b/openrag/core/retrieval/searcher.py
new file mode 100644
index 000000000..fec714b3a
--- /dev/null
+++ b/openrag/core/retrieval/searcher.py
@@ -0,0 +1,81 @@
+"""Transitional port for chunk-level retrieval operations.
+
+The strict ``VectorStore`` ABC in ``core/vector_stores`` is intentionally
+narrow — ``search(embedding, top_k, ...)``. Phase 5 retrievers, however,
+still go through the legacy Milvus Ray actor which exposes higher-level
+operations:
+
+ * search by query string (embedding done internally, plus BM25)
+ * multi-query search (one round per query, dedup at the bottom)
+ * related-chunk lookup by ``relationship_id``
+ * ancestor lookup by ``file_id`` with depth bound
+
+Defining these on a dedicated port lets the retriever depend on a clean
+interface from day one, while a small shim in ``services/storage/``
+adapts the Ray actor to it. When the god object is decomposed in Phase 7
+this port either retires (operations move to ``VectorStore`` +
+``ChunkRepository``) or evolves into the shape MilvusVectorStore exposes
+directly.
+
+Returns are domain ``Chunk`` objects throughout — no LangChain types
+leak across this boundary.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from openrag.core.models.chunk import Chunk
+
+
+class RetrievalSearcher(ABC):
+ """Operations a retriever needs from the chunk store."""
+
+ @abstractmethod
+ async def search(
+ self,
+ query: str,
+ partition: list[str],
+ top_k: int,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ similarity_threshold: float = 0.0,
+ with_surrounding_chunks: bool = True,
+ ) -> list[Chunk]:
+ """Single-query similarity search."""
+ ...
+
+ @abstractmethod
+ async def multi_query_search(
+ self,
+ queries: list[str],
+ partition: list[str],
+ top_k_per_query: int,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ similarity_threshold: float = 0.0,
+ with_surrounding_chunks: bool = True,
+ ) -> list[Chunk]:
+ """Run one similarity search per query, return the merged result."""
+ ...
+
+ @abstractmethod
+ async def get_related_chunks(
+ self,
+ partition: str,
+ relationship_id: str,
+ limit: int,
+ ) -> list[Chunk]:
+ """Fetch other chunks belonging to the same relationship group."""
+ ...
+
+ @abstractmethod
+ async def get_ancestor_chunks(
+ self,
+ partition: str,
+ file_id: str,
+ limit: int,
+ max_ancestor_depth: int | None = None,
+ ) -> list[Chunk]:
+ """Walk parent links up the document tree from a file."""
+ ...
diff --git a/openrag/core/retrieval/test_pipeline.py b/openrag/core/retrieval/test_pipeline.py
new file mode 100644
index 000000000..89776e8c7
--- /dev/null
+++ b/openrag/core/retrieval/test_pipeline.py
@@ -0,0 +1,206 @@
+"""Tests for RetrieverPipeline using fake Retriever / Reranker."""
+
+from __future__ import annotations
+
+import pytest
+
+from openrag.core.models.chunk import Chunk
+from openrag.core.models.query import Query, SearchQueries, TemporalPredicate
+from openrag.core.retrieval.pipeline import RetrieverPipeline
+from openrag.core.retrieval.retriever import Retriever
+
+
+class FakeRetriever(Retriever):
+ """Plays back canned per-call results; records call kwargs."""
+
+ def __init__(self, expansion_enabled: bool = False) -> None:
+ self.calls: list[dict] = []
+ self.results_queue: list[list[Chunk]] = []
+ self.expand_input: list[Chunk] | None = None
+ self.expand_result: list[Chunk] | None = None
+ self.expansion_enabled = expansion_enabled
+
+ async def retrieve(self, partition, query, filter=None, filter_params=None):
+ self.calls.append({"partition": partition, "query": query, "filter": filter, "filter_params": filter_params})
+ if self.results_queue:
+ return self.results_queue.pop(0)
+ return []
+
+ async def expand_search_results(self, results):
+ self.expand_input = list(results)
+ return list(self.expand_result) if self.expand_result is not None else list(results)
+
+
+class FakeReranker:
+ """Reverses input ordering — easy to detect in assertions."""
+
+ def __init__(self) -> None:
+ self.calls: list[dict] = []
+
+ async def rerank(self, query, documents, top_k=None):
+ self.calls.append({"query": query, "documents": list(documents), "top_k": top_k})
+ # Reverse ranking, perfect score for the (now-)first item
+ return [(i, float(len(documents) - i)) for i in range(len(documents) - 1, -1, -1)]
+
+
+def _chunks(*ids: str) -> list[Chunk]:
+ return [Chunk(id=i, text=f"text-{i}", partition="p1") for i in ids]
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_no_filter_no_rerank_no_expand():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b", "c")]
+ p = RetrieverPipeline(retriever=r)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"))
+ assert [c.id for c in out] == ["a", "b", "c"]
+ assert r.calls[0]["filter"] is None
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_temporal_filter_passed_through():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a")]
+ p = RetrieverPipeline(retriever=r)
+ q = Query(
+ query="hi",
+ temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")],
+ )
+ await p.retrieve_docs(partition=["p1"], query=q)
+ assert "created_at" in r.calls[0]["filter"]
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_filterless_fallback_when_filter_returns_zero():
+ r = FakeRetriever()
+ r.results_queue = [[], _chunks("a")]
+ p = RetrieverPipeline(retriever=r, allow_filterless_fallback=True)
+ q = Query(
+ query="hi",
+ temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")],
+ )
+ out = await p.retrieve_docs(partition=["p1"], query=q)
+ assert [c.id for c in out] == ["a"]
+ assert r.calls[0]["filter"] is not None
+ assert r.calls[1]["filter"] is None
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_no_fallback_when_disabled():
+ r = FakeRetriever()
+ r.results_queue = [[]]
+ p = RetrieverPipeline(retriever=r, allow_filterless_fallback=False)
+ q = Query(
+ query="hi",
+ temporal_filters=[TemporalPredicate(operator=">=", value="2026-01-01T00:00:00+00:00")],
+ )
+ out = await p.retrieve_docs(partition=["p1"], query=q)
+ assert out == []
+ assert len(r.calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_runs_reranker_when_enabled():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b", "c")]
+ rer = FakeReranker()
+ p = RetrieverPipeline(retriever=r, reranker=rer)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"))
+ assert [c.id for c in out] == ["c", "b", "a"]
+ assert rer.calls[0]["query"] == "hi"
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_expansion_path_re_reranks():
+ r = FakeRetriever(expansion_enabled=True)
+ r.results_queue = [_chunks("a", "b")]
+ r.expand_result = _chunks("a", "b", "c")
+ rer = FakeReranker()
+ p = RetrieverPipeline(retriever=r, reranker=rer, reranker_top_k=2)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"))
+ # Two reranker invocations: pre-expansion (2 chunks), post-expansion (3 chunks)
+ assert len(rer.calls) == 2
+ assert len(rer.calls[0]["documents"]) == 2
+ assert len(rer.calls[1]["documents"]) == 3
+ assert {c.id for c in out} == {"a", "b", "c"}
+
+
+@pytest.mark.asyncio
+async def test_get_relevant_docs_runs_one_call_per_subquery_and_fuses():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b"), _chunks("b", "c")]
+ p = RetrieverPipeline(retriever=r)
+ sq = SearchQueries(query_list=[Query(query="q1"), Query(query="q2")])
+ out = await p.get_relevant_docs(partition=["p1"], search_queries=sq)
+ assert len(r.calls) == 2
+ assert {c.id for c in out} == {"a", "b", "c"}
+ # 'b' appears in both lists -> highest fused score
+ assert out[0].id == "b"
+
+
+@pytest.mark.asyncio
+async def test_get_relevant_docs_applies_top_k_cap():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b", "c")]
+ p = RetrieverPipeline(retriever=r)
+ sq = SearchQueries(query_list=[Query(query="q1")])
+ out = await p.get_relevant_docs(partition=["p1"], search_queries=sq, top_k=2)
+ assert len(out) == 2
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_expansion_no_new_chunks_skips_second_rerank():
+ r = FakeRetriever(expansion_enabled=True)
+ r.results_queue = [_chunks("a", "b")]
+ r.expand_result = _chunks("a", "b") # expansion returns same set
+ rer = FakeReranker()
+ p = RetrieverPipeline(retriever=r, reranker=rer, reranker_top_k=2)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"))
+ # Only the pre-expansion rerank fired.
+ assert len(rer.calls) == 1
+ assert {c.id for c in out} == {"a", "b"}
+
+
+@pytest.mark.asyncio
+async def test_rerank_chunks_short_circuits_on_empty_input():
+ """Direct cover of the early-return guard inside _rerank_chunks."""
+ from openrag.core.retrieval.pipeline import _rerank_chunks
+
+ rer = FakeReranker()
+ out = await _rerank_chunks(rer, "q", [])
+ assert out == []
+ assert rer.calls == []
+
+
+def test_pipeline_expansion_enabled_false_for_non_base_retriever():
+ """Retrievers without an expansion_enabled attr (e.g. custom impls) are
+ treated as non-expanding via getattr default."""
+
+ class MinimalRetriever(Retriever):
+ async def retrieve(self, partition, query, filter=None, filter_params=None):
+ return []
+
+ async def expand_search_results(self, results):
+ return results
+
+ p = RetrieverPipeline(retriever=MinimalRetriever())
+ assert p.expansion_enabled is False
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_caps_to_top_k():
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b", "c", "d")]
+ p = RetrieverPipeline(retriever=r)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"), top_k=2)
+ assert [c.id for c in out] == ["a", "b"]
+
+
+@pytest.mark.asyncio
+async def test_retrieve_docs_top_k_zero_returns_empty():
+ """top_k=0 must mean "zero results", not "treated as None" (the legacy bug)."""
+ r = FakeRetriever()
+ r.results_queue = [_chunks("a", "b")]
+ p = RetrieverPipeline(retriever=r)
+ out = await p.retrieve_docs(partition=["p1"], query=Query(query="hi"), top_k=0)
+ assert out == []
diff --git a/openrag/core/retrieval/test_retriever.py b/openrag/core/retrieval/test_retriever.py
new file mode 100644
index 000000000..e589b5c12
--- /dev/null
+++ b/openrag/core/retrieval/test_retriever.py
@@ -0,0 +1,275 @@
+"""Retriever strategy tests with fake searcher + LLM.
+
+These exercise the strategy logic without Ray, OpenAI, or LangChain — proving
+the new core/ retriever has clean dependencies.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from openrag.core.models.chunk import Chunk
+from openrag.core.retrieval.registry import retriever_registry
+from openrag.core.retrieval.retriever import (
+ HyDeRetriever,
+ MultiQueryRetriever,
+ SingleRetriever,
+)
+from openrag.core.retrieval.searcher import RetrievalSearcher
+
+
+class FakeSearcher(RetrievalSearcher):
+ """Records calls; returns canned chunks."""
+
+ def __init__(self) -> None:
+ self.search_calls: list[dict] = []
+ self.multi_calls: list[dict] = []
+ self.related_calls: list[dict] = []
+ self.ancestor_calls: list[dict] = []
+ self.search_result: list[Chunk] = []
+ self.multi_result: list[Chunk] = []
+ self.related_result: list[Chunk] = []
+ self.ancestor_result: list[Chunk] = []
+
+ async def search(self, **kwargs):
+ self.search_calls.append(kwargs)
+ return list(self.search_result)
+
+ async def multi_query_search(self, **kwargs):
+ self.multi_calls.append(kwargs)
+ return list(self.multi_result)
+
+ async def get_related_chunks(self, **kwargs):
+ self.related_calls.append(kwargs)
+ return list(self.related_result)
+
+ async def get_ancestor_chunks(self, **kwargs):
+ self.ancestor_calls.append(kwargs)
+ return list(self.ancestor_result)
+
+
+class FakeLLM:
+ def __init__(self, response: str) -> None:
+ self.response = response
+ self.chat_calls: list[list[dict]] = []
+
+ async def generate(self, prompt: str, **kwargs) -> str:
+ return self.response
+
+ async def chat(self, messages: list[dict], **kwargs) -> str:
+ self.chat_calls.append(messages)
+ return self.response
+
+
+def _chunk(idv: str, text: str = "x", document_id: str = "", partition: str = "p1") -> Chunk:
+ return Chunk(id=idv, text=text, document_id=document_id, partition=partition)
+
+
+def test_registry_has_three_strategies():
+ assert set(retriever_registry.list_registered()) == {"single", "multiQuery", "hyde"}
+
+
+@pytest.mark.asyncio
+async def test_single_retriever_passes_through_to_searcher():
+ s = FakeSearcher()
+ s.search_result = [_chunk("1"), _chunk("2")]
+ r = SingleRetriever(searcher=s, top_k=4, similarity_threshold=0.3, with_surrounding_chunks=False)
+ out = await r.retrieve(partition=["p1"], query="hello", filter="x>0", filter_params={"a": 1})
+ assert [c.id for c in out] == ["1", "2"]
+ assert s.search_calls == [
+ {
+ "query": "hello",
+ "partition": ["p1"],
+ "top_k": 4,
+ "filter": "x>0",
+ "filter_params": {"a": 1},
+ "similarity_threshold": 0.3,
+ "with_surrounding_chunks": False,
+ }
+ ]
+
+
+@pytest.mark.asyncio
+async def test_multi_query_retriever_splits_llm_response():
+ s = FakeSearcher()
+ s.multi_result = [_chunk("a")]
+ llm = FakeLLM(response="Q one[SEP]Q two[SEP]Q three")
+ r = MultiQueryRetriever(
+ searcher=s,
+ llm=llm,
+ multi_query_template="generate {k_queries} variants of: {query}",
+ k_queries=3,
+ top_k=5,
+ )
+ await r.retrieve(partition=["p1"], query="seed")
+ assert s.multi_calls[0]["queries"] == ["Q one", "Q two", "Q three"]
+ assert s.multi_calls[0]["top_k_per_query"] == 5
+
+
+@pytest.mark.asyncio
+async def test_multi_query_falls_back_to_seed_on_empty_response():
+ s = FakeSearcher()
+ llm = FakeLLM(response="")
+ r = MultiQueryRetriever(
+ searcher=s,
+ llm=llm,
+ multi_query_template="{query} {k_queries}",
+ k_queries=3,
+ )
+ await r.retrieve(partition=["p1"], query="seed")
+ assert s.multi_calls[0]["queries"] == ["seed"]
+
+
+@pytest.mark.asyncio
+async def test_hyde_retriever_uses_hyde_only_by_default():
+ s = FakeSearcher()
+ llm = FakeLLM(response="A hypothetical answer paragraph.")
+ r = HyDeRetriever(searcher=s, llm=llm, hyde_template="Answer: {question}")
+ await r.retrieve(partition=["p1"], query="real question")
+ assert s.multi_calls[0]["queries"] == ["A hypothetical answer paragraph."]
+
+
+@pytest.mark.asyncio
+async def test_hyde_retriever_combine_appends_original_query():
+ s = FakeSearcher()
+ llm = FakeLLM(response="hypothetical")
+ r = HyDeRetriever(searcher=s, llm=llm, hyde_template="Answer: {question}", combine=True)
+ await r.retrieve(partition=["p1"], query="real")
+ assert s.multi_calls[0]["queries"] == ["hypothetical", "real"]
+
+
+@pytest.mark.asyncio
+async def test_expansion_disabled_returns_unchanged():
+ s = FakeSearcher()
+ r = SingleRetriever(searcher=s)
+ initial = [_chunk("1")]
+ out = await r.expand_search_results(initial)
+ assert out is initial
+ assert not s.related_calls
+ assert not s.ancestor_calls
+
+
+@pytest.mark.asyncio
+async def test_expansion_with_related_dedupes_by_id():
+ s = FakeSearcher()
+ s.related_result = [_chunk("1"), _chunk("3")] # "1" already in results
+ r = SingleRetriever(searcher=s, include_related=True)
+ initial = [
+ Chunk(id="1", text="x", partition="p1", metadata={"relationship_id": "r1"}),
+ ]
+ out = await r.expand_search_results(initial)
+ assert [c.id for c in out] == ["1", "3"]
+
+
+@pytest.mark.asyncio
+async def test_expansion_with_ancestors_calls_searcher():
+ s = FakeSearcher()
+ s.ancestor_result = [_chunk("99", document_id="f1")]
+ r = SingleRetriever(searcher=s, include_ancestors=True, related_limit=20, max_ancestor_depth=2)
+ initial = [Chunk(id="1", text="x", partition="p1", document_id="f1")]
+ out = await r.expand_search_results(initial)
+ assert [c.id for c in out] == ["1", "99"]
+ assert s.ancestor_calls[0]["partition"] == "p1"
+ assert s.ancestor_calls[0]["file_id"] == "f1"
+ assert s.ancestor_calls[0]["limit"] == 20
+ assert s.ancestor_calls[0]["max_ancestor_depth"] == 2
+
+
+@pytest.mark.asyncio
+async def test_expansion_swallows_per_call_errors():
+ class BoomSearcher(FakeSearcher):
+ async def get_related_chunks(self, **kwargs):
+ raise RuntimeError("kaboom")
+
+ s = BoomSearcher()
+ r = SingleRetriever(searcher=s, include_related=True)
+ initial = [Chunk(id="1", text="x", partition="p1", metadata={"relationship_id": "r1"})]
+ out = await r.expand_search_results(initial)
+ assert [c.id for c in out] == ["1"]
+
+
+@pytest.mark.asyncio
+async def test_expansion_swallows_ancestor_errors():
+ class BoomSearcher(FakeSearcher):
+ async def get_ancestor_chunks(self, **kwargs):
+ raise RuntimeError("ancestor exploded")
+
+ s = BoomSearcher()
+ r = SingleRetriever(searcher=s, include_ancestors=True)
+ initial = [Chunk(id="1", text="x", partition="p1", document_id="f1")]
+ out = await r.expand_search_results(initial)
+ assert [c.id for c in out] == ["1"]
+
+
+def test_multi_query_retriever_rejects_missing_llm():
+ s = FakeSearcher()
+ with pytest.raises(ValueError, match="llm must be provided"):
+ MultiQueryRetriever(searcher=s, llm=None, multi_query_template="{query} {k_queries}")
+
+
+def test_hyde_retriever_rejects_missing_llm():
+ s = FakeSearcher()
+ with pytest.raises(ValueError, match="llm must be provided"):
+ HyDeRetriever(searcher=s, llm=None, hyde_template="{question}")
+
+
+@pytest.mark.asyncio
+async def test_multi_query_retriever_caps_response_to_k_queries():
+ """A non-compliant LLM that returns more variants than requested must
+ not fan out additional searches."""
+ s = FakeSearcher()
+ llm = FakeLLM(response="Q1[SEP]Q2[SEP]Q3[SEP]Q4[SEP]Q5")
+ r = MultiQueryRetriever(
+ searcher=s,
+ llm=llm,
+ multi_query_template="{query} {k_queries}",
+ k_queries=2,
+ )
+ await r.retrieve(partition=["p1"], query="seed")
+ assert s.multi_calls[0]["queries"] == ["Q1", "Q2"]
+
+
+@pytest.mark.asyncio
+async def test_hyde_retriever_falls_back_to_seed_on_blank_generation():
+ s = FakeSearcher()
+ llm = FakeLLM(response=" \n\t ")
+ r = HyDeRetriever(searcher=s, llm=llm, hyde_template="{question}")
+ await r.retrieve(partition=["p1"], query="seed")
+ assert s.multi_calls[0]["queries"] == ["seed"]
+
+
+@pytest.mark.asyncio
+async def test_hyde_retriever_falls_back_to_seed_when_combine_and_blank():
+ """Combine mode should also fall back to just [seed] on blank generation,
+ not [blank, seed]."""
+ s = FakeSearcher()
+ llm = FakeLLM(response="")
+ r = HyDeRetriever(searcher=s, llm=llm, hyde_template="{question}", combine=True)
+ await r.retrieve(partition=["p1"], query="seed")
+ assert s.multi_calls[0]["queries"] == ["seed"]
+
+
+@pytest.mark.asyncio
+async def test_expansion_dedupes_ancestor_fetches_per_file():
+ """Two chunks from the same (partition, document_id) must enqueue a
+ single ancestor fetch, not two."""
+
+ class CountingSearcher(FakeSearcher):
+ def __init__(self) -> None:
+ super().__init__()
+ self.ancestor_call_count = 0
+
+ async def get_ancestor_chunks(self, **kwargs):
+ self.ancestor_call_count += 1
+ return list(self.ancestor_result)
+
+ s = CountingSearcher()
+ s.ancestor_result = [_chunk("99", document_id="f1")]
+ r = SingleRetriever(searcher=s, include_ancestors=True)
+ initial = [
+ Chunk(id="1", text="x", partition="p1", document_id="f1"),
+ Chunk(id="2", text="y", partition="p1", document_id="f1"),
+ Chunk(id="3", text="z", partition="p1", document_id="f1"),
+ ]
+ await r.expand_search_results(initial)
+ assert s.ancestor_call_count == 1
diff --git a/openrag/core/retrieval/test_rrf.py b/openrag/core/retrieval/test_rrf.py
new file mode 100644
index 000000000..b84b0c2d3
--- /dev/null
+++ b/openrag/core/retrieval/test_rrf.py
@@ -0,0 +1,50 @@
+"""RRF unit tests — fusion semantics and edge cases."""
+
+from __future__ import annotations
+
+import pytest
+
+from openrag.core.retrieval.rrf import rrf_reranking
+
+
+def test_rrf_empty_returns_empty():
+ assert rrf_reranking([]) == []
+
+
+def test_rrf_single_list_returned_as_is():
+ items = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
+ assert rrf_reranking([items]) == items
+
+
+def test_rrf_fuses_overlapping_results():
+ # 'a' is rank 1 in list1 and rank 2 in list2 → top
+ # 'c' is rank 1 in list2 → second
+ # 'b' is rank 2 in list1 only
+ list1 = [{"id": "a"}, {"id": "b"}]
+ list2 = [{"id": "c"}, {"id": "a"}]
+ fused = rrf_reranking([list1, list2], key_fn=lambda x: x["id"])
+ ids = [item["id"] for item in fused]
+ assert ids[0] == "a"
+ assert set(ids) == {"a", "b", "c"}
+
+
+def test_rrf_without_key_fn_does_not_fuse():
+ list1 = [{"id": "a"}]
+ list2 = [{"id": "a"}] # different object, same logical id
+ fused = rrf_reranking([list1, list2])
+ # Object identity → two separate items in fused result
+ assert len(fused) == 2
+
+
+def test_rrf_smaller_k_emphasizes_top_ranks():
+ list1 = [{"id": "a"}, {"id": "b"}]
+ list2 = [{"id": "b"}, {"id": "a"}]
+ fused = rrf_reranking([list1, list2], key_fn=lambda x: x["id"], k=1)
+ # k=1: top-rank in any list dominates; with two top-1s for different items,
+ # both score the same — order is stable across implementations though
+ assert {item["id"] for item in fused} == {"a", "b"}
+
+
+def test_rrf_rejects_negative_k():
+ with pytest.raises(ValueError, match="non-negative"):
+ rrf_reranking([[{"id": "a"}], [{"id": "b"}]], key_fn=lambda x: x["id"], k=-1)
diff --git a/openrag/services/storage/milvus_ray_shim.py b/openrag/services/storage/milvus_ray_shim.py
new file mode 100644
index 000000000..0cdf40ee0
--- /dev/null
+++ b/openrag/services/storage/milvus_ray_shim.py
@@ -0,0 +1,157 @@
+"""Transitional ``RetrievalSearcher`` adapter wrapping the legacy Ray Milvus actor.
+
+The new core retriever (``core/retrieval/retriever.py``) talks to a clean
+ABC; this shim is what plugs the still-existing Ray god object into that
+ABC for the duration of Phase 5–6. Once Phase 7 decomposes the Vectordb
+actor into ``MilvusVectorStore`` + ``ChunkRepository`` this file goes away.
+
+Conversion: the Ray actor returns LangChain ``Document`` objects with
+metadata; we convert each one to a domain ``Chunk`` via
+``Chunk.from_langchain``. Ray actor calls are routed through
+``call_ray_actor_with_timeout`` so timeout/cancellation behavior matches
+the rest of the legacy app — bypassing this helper was the regression
+flagged in PR #352 (CodeRabbit).
+
+The Ray imports are deferred to method bodies so this module is importable
+in non-Ray contexts (tests, CLI tools).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from openrag.core.models.chunk import Chunk
+from openrag.core.retrieval.searcher import RetrievalSearcher
+
+# Match the legacy default in `components.pipeline` (config.ray.indexer.vectordb_timeout).
+# The composition root can override via ``MilvusRayShim(actor, timeout=...)``.
+DEFAULT_VECTORDB_TIMEOUT = 60.0
+
+
+def _to_chunks(docs: list[Any]) -> list[Chunk]:
+ """Convert LangChain Documents from the Ray actor into domain Chunks."""
+ return [Chunk.from_langchain(d) for d in docs]
+
+
+class MilvusRayShim(RetrievalSearcher):
+ """Adapter exposing the Vectordb Ray actor as a ``RetrievalSearcher``.
+
+ Args:
+ actor: Ray actor handle (typically ``ray.get_actor("Vectordb",
+ namespace="openrag")``). Accepts any object whose remote
+ methods match the legacy Vectordb actor — useful for tests.
+ timeout: Per-call timeout passed to ``call_ray_actor_with_timeout``.
+ """
+
+ def __init__(self, actor: Any, timeout: float = DEFAULT_VECTORDB_TIMEOUT) -> None:
+ self._actor = actor
+ self._timeout = timeout
+
+ async def _call(self, future: Any, task_description: str) -> Any:
+ # Deferred import: keeps this module importable without `ray` installed
+ # (legacy `components.ray_utils` pulls in ray at import time).
+ from openrag.components.ray_utils import call_ray_actor_with_timeout
+
+ return await call_ray_actor_with_timeout(
+ future=future,
+ timeout=self._timeout,
+ task_description=task_description,
+ )
+
+ async def search(
+ self,
+ query: str,
+ partition: list[str],
+ top_k: int,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ similarity_threshold: float = 0.0,
+ with_surrounding_chunks: bool = True,
+ ) -> list[Chunk]:
+ docs = await self._call(
+ self._actor.async_search.remote(
+ query=query,
+ partition=partition,
+ top_k=top_k,
+ filter=filter,
+ filter_params=filter_params,
+ similarity_threshold=similarity_threshold,
+ with_surrounding_chunks=with_surrounding_chunks,
+ ),
+ task_description=f"async_search(partition={partition})",
+ )
+ return _to_chunks(docs)
+
+ async def multi_query_search(
+ self,
+ queries: list[str],
+ partition: list[str],
+ top_k_per_query: int,
+ filter: str | None = None,
+ filter_params: dict | None = None,
+ similarity_threshold: float = 0.0,
+ with_surrounding_chunks: bool = True,
+ ) -> list[Chunk]:
+ docs = await self._call(
+ self._actor.async_multi_query_search.remote(
+ queries=queries,
+ partition=partition,
+ top_k_per_query=top_k_per_query,
+ filter=filter,
+ filter_params=filter_params,
+ similarity_threshold=similarity_threshold,
+ with_surrounding_chunks=with_surrounding_chunks,
+ ),
+ task_description=f"async_multi_query_search(n={len(queries)}, partition={partition})",
+ )
+ return _to_chunks(docs)
+
+ async def get_related_chunks(
+ self,
+ partition: str,
+ relationship_id: str,
+ limit: int,
+ ) -> list[Chunk]:
+ docs = await self._call(
+ self._actor.get_related_chunks.remote(
+ partition=partition,
+ relationship_id=relationship_id,
+ limit=limit,
+ ),
+ task_description=f"get_related_chunks(partition={partition}, rel={relationship_id})",
+ )
+ return _to_chunks(docs)
+
+ async def get_ancestor_chunks(
+ self,
+ partition: str,
+ file_id: str,
+ limit: int,
+ max_ancestor_depth: int | None = None,
+ ) -> list[Chunk]:
+ docs = await self._call(
+ self._actor.get_ancestor_chunks.remote(
+ partition=partition,
+ file_id=file_id,
+ limit=limit,
+ max_ancestor_depth=max_ancestor_depth,
+ ),
+ task_description=f"get_ancestor_chunks(partition={partition}, file={file_id})",
+ )
+ return _to_chunks(docs)
+
+
+def from_ray_namespace(
+ name: str = "Vectordb",
+ namespace: str = "openrag",
+ timeout: float = DEFAULT_VECTORDB_TIMEOUT,
+) -> MilvusRayShim:
+ """Look up the Vectordb Ray actor by name and wrap it.
+
+ Convenience for the composition root. The Ray import is deferred so
+ importing this module without Ray installed (e.g. in unit tests of
+ the retriever with a fake searcher) does not fail.
+ """
+ import ray
+
+ return MilvusRayShim(ray.get_actor(name, namespace=namespace), timeout=timeout)