mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 02:13:47 +00:00
feat: per-source RAG configuration (retrieval strategies, chunking, exposure, prescreen)
Introduces a per-source config contract that makes RAG behavior strategy-dispatched instead of a single hardcoded path. Every source gains a validated JSONB config; an empty/absent config reproduces current behavior byte-for-byte, and the whole path is gated by PER_SOURCE_RETRIEVAL_ENABLED. Foundation: sources.config JSONB column + migration 0022_source_config; SourceConfig/ChunkingConfig/RetrievalConfig pydantic models (strict on write, lenient on read); ChunkerCreator and RetrieverCreator.register registries; config threaded through the upload routes, ingest/remote/connector workers, and reingest. Retrieval: a Dispatcher groups sources by retriever key (all-classic collapses to today's single ClassicRAG under one shared token budget; non-classic retrievers get their own instance), removing the previous single-global-retriever collapse in stream_processor. Per-source chunks, score_threshold (honored for pgvector/mongodb, safely ignored elsewhere), and rephrase_query toggle. New PATCH /api/sources/<id>/config with team-aware (effective_write_owner) authz and a requires_reingest signal. Chunking strategies: recursive, markdown, parent_child (selectable per source; re-ingest to apply). Search exposure: per-source prefetch vs agentic_tool for agentic/research agents. Map-reduce prescreen: optional LLM relevance pre-filter implemented as a composable post-retrieval stage that wraps any retriever. Backend and frontend (shared Retrieval options panel + edit modal) with tests; backend suite and frontend vitest green. Excludes the wiki and GraphRAG flagships.
This commit is contained in:
1 parent
cfda009be0
commit
f6400cd736
53 files changed
+4313
-74
No files matched your search
@@ -36,9 +36,11 @@ class TestChunkerInit:
|
||||
assert chunker.min_tokens == 50
|
||||
assert chunker.duplicate_headers is True
|
||||
|
||||
def test_invalid_strategy_raises(self):
|
||||
with pytest.raises(ValueError, match="Unsupported chunking strategy"):
|
||||
Chunker(chunking_strategy="unknown_strategy")
|
||||
def test_unknown_strategy_construction_no_longer_raises(self):
|
||||
# Strategy dispatch/whitelist moved to ChunkerCreator; the Chunker
|
||||
# constructor itself no longer rejects an unknown strategy string.
|
||||
chunker = Chunker(chunking_strategy="unknown_strategy")
|
||||
assert chunker.chunking_strategy == "unknown_strategy"
|
||||
|
||||
|
||||
# =====================================================================
|
||||
@@ -235,12 +237,14 @@ class TestChunkDispatcher:
|
||||
result = chunker.chunk([doc])
|
||||
assert len(result) == 1
|
||||
|
||||
def test_dispatch_unknown_raises(self):
|
||||
def test_chunk_runs_classic_regardless_of_strategy_attr(self):
|
||||
# Chunk() now always runs the classic implementation; strategy
|
||||
# selection happens at the ChunkerCreator level, not here.
|
||||
chunker = Chunker()
|
||||
chunker.chunking_strategy = "nonexistent"
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported chunking strategy"):
|
||||
chunker.chunk([Document(text="x", doc_id="d")])
|
||||
result = chunker.chunk([Document(text="x", doc_id="d")])
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Tests for ChunkerCreator registry + byte-identical classic_chunk parity."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.chunking import Chunker
|
||||
from application.parser.chunking_creator import ChunkerCreator
|
||||
from application.parser.schema.base import Document
|
||||
|
||||
|
||||
def _docs():
|
||||
"""Representative mix: a small doc, a large splittable doc, a tiny doc."""
|
||||
return [
|
||||
Document(text="A short paragraph of text.", doc_id="small"),
|
||||
Document(text="word " * 4000, doc_id="large"),
|
||||
Document(text="tiny", doc_id="below-min"),
|
||||
Document(text="header1\nheader2\nheader3\n" + "lorem " * 3000, doc_id="hdr"),
|
||||
]
|
||||
|
||||
|
||||
def _serialize(chunks):
|
||||
"""Stable, comparable view of a chunk list."""
|
||||
return [(c.doc_id, c.text, c.extra_info) for c in chunks]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestChunkerCreator:
|
||||
def test_classic_chunk_is_registered(self):
|
||||
assert "classic_chunk" in ChunkerCreator.chunkers
|
||||
assert ChunkerCreator.chunkers["classic_chunk"] is Chunker
|
||||
|
||||
def test_create_chunker_returns_chunker(self):
|
||||
chunker = ChunkerCreator.create_chunker(
|
||||
"classic_chunk", max_tokens=1250, min_tokens=150,
|
||||
)
|
||||
assert isinstance(chunker, Chunker)
|
||||
assert chunker.max_tokens == 1250
|
||||
assert chunker.min_tokens == 150
|
||||
|
||||
def test_create_chunker_unknown_strategy_raises(self):
|
||||
with pytest.raises(ValueError, match="No chunker class found"):
|
||||
ChunkerCreator.create_chunker("does_not_exist")
|
||||
|
||||
def test_register_is_idempotent_and_overrides(self):
|
||||
class Dummy:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
try:
|
||||
ChunkerCreator.register("dummy_strategy", Dummy)
|
||||
inst = ChunkerCreator.create_chunker("dummy_strategy", x=1)
|
||||
assert isinstance(inst, Dummy)
|
||||
assert inst.kwargs == {"x": 1}
|
||||
finally:
|
||||
ChunkerCreator.chunkers.pop("dummy_strategy", None)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestByteIdenticalParity:
|
||||
def test_creator_output_matches_direct_chunker(self):
|
||||
# The refactor must not change classic_chunk output: a chunker built
|
||||
# via ChunkerCreator must produce identical chunks to the old direct
|
||||
# ``Chunker(...)`` instantiation for the same input + params.
|
||||
params = dict(max_tokens=1250, min_tokens=150, duplicate_headers=False)
|
||||
|
||||
direct = Chunker(chunking_strategy="classic_chunk", **params)
|
||||
via_creator = ChunkerCreator.create_chunker(
|
||||
"classic_chunk", chunking_strategy="classic_chunk", **params,
|
||||
)
|
||||
|
||||
direct_out = direct.chunk(documents=_docs())
|
||||
creator_out = via_creator.chunk(documents=_docs())
|
||||
|
||||
assert _serialize(creator_out) == _serialize(direct_out)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Tests for the recursive / markdown / parent_child chunking strategies (D1)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.chunking import Chunker
|
||||
from application.parser.chunking_creator import ChunkerCreator
|
||||
from application.parser.chunking_strategies import (
|
||||
MarkdownChunker,
|
||||
ParentChildChunker,
|
||||
RecursiveChunker,
|
||||
)
|
||||
from application.parser.schema.base import Document
|
||||
from application.utils import get_encoding
|
||||
|
||||
|
||||
def _tok(text: str) -> int:
|
||||
return len(get_encoding().encode(text))
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRegistration:
|
||||
def test_strategies_registered(self):
|
||||
# create_chunker self-bootstraps the strategy module.
|
||||
ChunkerCreator.create_chunker("recursive")
|
||||
for key, cls in (
|
||||
("recursive", RecursiveChunker),
|
||||
("markdown", MarkdownChunker),
|
||||
("parent_child", ParentChildChunker),
|
||||
):
|
||||
assert ChunkerCreator.chunkers.get(key) is cls
|
||||
|
||||
def test_worker_kwargs_accepted(self):
|
||||
# The worker builds every strategy with the classic kwarg set.
|
||||
for strat in ("recursive", "markdown", "parent_child"):
|
||||
chunker = ChunkerCreator.create_chunker(
|
||||
strat,
|
||||
chunking_strategy=strat,
|
||||
max_tokens=200,
|
||||
min_tokens=20,
|
||||
duplicate_headers=False,
|
||||
)
|
||||
assert chunker.max_tokens == 200
|
||||
assert chunker.min_tokens == 20
|
||||
|
||||
def test_semantic_is_not_registered(self):
|
||||
# semantic is explicitly deferred.
|
||||
with pytest.raises(ValueError, match="No chunker class found"):
|
||||
ChunkerCreator.create_chunker("semantic")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRecursive:
|
||||
def test_caps_at_max_tokens(self):
|
||||
chunker = RecursiveChunker(max_tokens=40, min_tokens=5)
|
||||
docs = [Document(text="word " * 500, doc_id="d")]
|
||||
out = chunker.chunk(docs)
|
||||
assert len(out) > 1
|
||||
for c in out:
|
||||
assert _tok(c.text) <= 40
|
||||
assert c.extra_info["token_count"] == _tok(c.text)
|
||||
|
||||
def test_splits_on_separator_hierarchy(self):
|
||||
# Paragraph boundaries should drive the split before token slicing.
|
||||
text = "\n\n".join(["para " * 30 for _ in range(5)])
|
||||
chunker = RecursiveChunker(max_tokens=60, min_tokens=5)
|
||||
out = chunker.chunk([Document(text=text, doc_id="d")])
|
||||
assert len(out) >= 2
|
||||
for c in out:
|
||||
assert _tok(c.text) <= 60
|
||||
|
||||
def test_small_doc_single_chunk(self):
|
||||
chunker = RecursiveChunker(max_tokens=2000, min_tokens=1)
|
||||
out = chunker.chunk([Document(text="short text here", doc_id="d")])
|
||||
assert len(out) == 1
|
||||
assert out[0].text.strip() == "short text here"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestMarkdown:
|
||||
def test_splits_on_headings(self):
|
||||
text = "# A\nalpha\n\n## B\nbeta\n\n### C\ngamma"
|
||||
chunker = MarkdownChunker(max_tokens=2000, min_tokens=1)
|
||||
out = chunker.chunk([Document(text=text, doc_id="d")])
|
||||
# One section per heading.
|
||||
assert len(out) == 3
|
||||
assert out[0].text.startswith("# A")
|
||||
assert out[1].text.startswith("## B")
|
||||
|
||||
def test_oversized_section_token_capped(self):
|
||||
text = "# Big\n" + "word " * 400
|
||||
chunker = MarkdownChunker(max_tokens=50, min_tokens=5)
|
||||
out = chunker.chunk([Document(text=text, doc_id="d")])
|
||||
assert len(out) > 1
|
||||
for c in out:
|
||||
assert _tok(c.text) <= 50
|
||||
|
||||
def test_no_heading_falls_back_to_single_or_capped(self):
|
||||
chunker = MarkdownChunker(max_tokens=2000, min_tokens=1)
|
||||
out = chunker.chunk([Document(text="plain text no heading", doc_id="d")])
|
||||
assert len(out) == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestParentChild:
|
||||
def test_children_smaller_than_parent(self):
|
||||
chunker = ParentChildChunker(max_tokens=60, min_tokens=15)
|
||||
out = chunker.chunk([Document(text="alpha " * 200, doc_id="d")])
|
||||
assert len(out) > 1
|
||||
for c in out:
|
||||
assert _tok(c.text) <= 15
|
||||
assert _tok(c.extra_info["parent_text"]) <= 60
|
||||
assert _tok(c.text) <= _tok(c.extra_info["parent_text"])
|
||||
|
||||
def test_parent_text_reaches_vectorstore_metadata(self):
|
||||
chunker = ParentChildChunker(max_tokens=80, min_tokens=20)
|
||||
out = chunker.chunk([Document(text="beta " * 150, doc_id="d")])
|
||||
lc = out[0].to_langchain_format()
|
||||
# parent_text must survive the langchain conversion into metadata.
|
||||
assert "parent_text" in lc.metadata
|
||||
assert lc.metadata["parent_text"]
|
||||
assert lc.page_content == out[0].text
|
||||
|
||||
def test_child_size_defaults_when_min_zero(self):
|
||||
chunker = ParentChildChunker(max_tokens=200, min_tokens=0)
|
||||
out = chunker.chunk([Document(text="gamma " * 200, doc_id="d")])
|
||||
assert all("parent_text" in c.extra_info for c in out)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestClassicByteIdentical:
|
||||
def test_classic_chunk_unchanged(self):
|
||||
# The new strategies must not perturb the classic baseline.
|
||||
docs = [
|
||||
Document(text="A short paragraph.", doc_id="small"),
|
||||
Document(text="word " * 4000, doc_id="large"),
|
||||
]
|
||||
params = dict(max_tokens=1250, min_tokens=150, duplicate_headers=False)
|
||||
direct = Chunker(chunking_strategy="classic_chunk", **params).chunk(docs)
|
||||
via = ChunkerCreator.create_chunker("classic_chunk", **params).chunk(
|
||||
[
|
||||
Document(text="A short paragraph.", doc_id="small"),
|
||||
Document(text="word " * 4000, doc_id="large"),
|
||||
]
|
||||
)
|
||||
assert [(c.doc_id, c.text, c.extra_info) for c in via] == [
|
||||
(c.doc_id, c.text, c.extra_info) for c in direct
|
||||
]
|
||||
Reference in new issue
Block a user