mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-18 14:13:10 +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
@@ -872,3 +872,160 @@ class TestDirectoryStructure:
|
||||
data = response.json
|
||||
assert data["provider"] == "gdrive"
|
||||
assert data["base_path"] == "/data/nested"
|
||||
|
||||
|
||||
class TestSourceConfigResource:
|
||||
def test_returns_401_unauthenticated(self, app):
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/sources/x/config", method="PATCH", json={}
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = None
|
||||
response = SourceConfigResource().patch("x")
|
||||
assert response.status_code == 401
|
||||
|
||||
def test_invalid_config_rejected(self, app, pg_conn):
|
||||
# Strict-on-write: an unknown field fails validation → 400.
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
|
||||
user = "u-cfg-bad"
|
||||
src = _seed_source(pg_conn, user, name="cfg-src", type="file")
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/sources/{src['id']}/config",
|
||||
method="PATCH",
|
||||
json={"retrieval": {"bogus_field": 1}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = SourceConfigResource().patch(str(src["id"]))
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_owner_updates_retrieval_no_reingest(self, app, pg_conn):
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
user = "u-cfg-owner"
|
||||
src = _seed_source(pg_conn, user, name="cfg-live", type="file")
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/sources/{src['id']}/config",
|
||||
method="PATCH",
|
||||
json={"retrieval": {"chunks": 7, "rephrase_query": False}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = SourceConfigResource().patch(str(src["id"]))
|
||||
|
||||
assert response.status_code == 200
|
||||
# Retrieval-only change takes effect live, no re-ingest needed.
|
||||
assert response.json["requires_reingest"] is False
|
||||
got = SourcesRepository(pg_conn).get_any(str(src["id"]), user)
|
||||
assert got["config"]["retrieval"]["chunks"] == 7
|
||||
assert got["config"]["retrieval"]["rephrase_query"] is False
|
||||
|
||||
def test_chunking_change_requires_reingest(self, app, pg_conn):
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
|
||||
user = "u-cfg-chunk"
|
||||
src = _seed_source(pg_conn, user, name="cfg-chunk", type="file")
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/sources/{src['id']}/config",
|
||||
method="PATCH",
|
||||
json={"chunking": {"max_tokens": 800}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = SourceConfigResource().patch(str(src["id"]))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json["requires_reingest"] is True
|
||||
|
||||
def test_viewer_rejected(self, app, pg_conn):
|
||||
# A team VIEWER (not editor) cannot edit config → 403.
|
||||
import uuid
|
||||
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
from application.storage.db.repositories.team_members import (
|
||||
TeamMembersRepository,
|
||||
)
|
||||
from application.storage.db.repositories.team_resource_grants import (
|
||||
TeamResourceGrantsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.teams import TeamsRepository
|
||||
|
||||
owner = "alice-cfg"
|
||||
viewer = "bob-cfg-viewer"
|
||||
src = _seed_source(pg_conn, owner, name="shared-cfg", type="file")
|
||||
sid = str(src["id"])
|
||||
team = TeamsRepository(pg_conn).create(
|
||||
"AcmeCfg", f"acmecfg-{uuid.uuid4().hex[:8]}", owner
|
||||
)
|
||||
TeamMembersRepository(pg_conn).add_member(
|
||||
team["id"], viewer, role="team_member"
|
||||
)
|
||||
TeamResourceGrantsRepository(pg_conn).grant(
|
||||
team["id"], "source", sid, owner_id=owner, granted_by=owner,
|
||||
access_level="viewer",
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/sources/{sid}/config",
|
||||
method="PATCH",
|
||||
json={"retrieval": {"chunks": 5}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": viewer}
|
||||
response = SourceConfigResource().patch(sid)
|
||||
assert response.status_code == 403
|
||||
# The write must NOT have landed.
|
||||
got = SourcesRepository(pg_conn).get_any(sid, owner)
|
||||
assert got["config"] == {}
|
||||
|
||||
def test_team_editor_writes_under_owner(self, app, pg_conn):
|
||||
# A team EDITOR can edit; the write lands under the OWNER's id.
|
||||
import uuid
|
||||
|
||||
from application.api.user.sources.routes import SourceConfigResource
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
from application.storage.db.repositories.team_members import (
|
||||
TeamMembersRepository,
|
||||
)
|
||||
from application.storage.db.repositories.team_resource_grants import (
|
||||
TeamResourceGrantsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.teams import TeamsRepository
|
||||
|
||||
owner = "alice-cfg-edit"
|
||||
editor = "bob-cfg-editor"
|
||||
src = _seed_source(pg_conn, owner, name="shared-cfg-edit", type="file")
|
||||
sid = str(src["id"])
|
||||
team = TeamsRepository(pg_conn).create(
|
||||
"AcmeEdit", f"acmeedit-{uuid.uuid4().hex[:8]}", owner
|
||||
)
|
||||
TeamMembersRepository(pg_conn).add_member(
|
||||
team["id"], editor, role="team_member"
|
||||
)
|
||||
TeamResourceGrantsRepository(pg_conn).grant(
|
||||
team["id"], "source", sid, owner_id=owner, granted_by=owner,
|
||||
access_level="editor",
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/sources/{sid}/config",
|
||||
method="PATCH",
|
||||
json={"retrieval": {"chunks": 9}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": editor}
|
||||
response = SourceConfigResource().patch(sid)
|
||||
|
||||
assert response.status_code == 200
|
||||
# The write landed under the owner id (not the editor's), so the
|
||||
# owner-scoped read sees it.
|
||||
got = SourcesRepository(pg_conn).get_any(sid, owner)
|
||||
assert got["config"]["retrieval"]["chunks"] == 9
|
||||
@@ -33,7 +33,7 @@ class TestIngestTask:
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
|
||||
file_name_map=None, idempotency_key=None, source_id=None,
|
||||
file_name_map=None, config=None, idempotency_key=None, source_id=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -50,7 +50,8 @@ class TestIngestTask:
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
|
||||
file_name_map=name_map, idempotency_key=None, source_id=None,
|
||||
file_name_map=name_map, config=None, idempotency_key=None,
|
||||
source_id=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -66,7 +67,7 @@ class TestIngestRemoteTask:
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, {"url": "http://x"}, "job1", "user1", "web",
|
||||
idempotency_key=None, source_id=None,
|
||||
config=None, idempotency_key=None, source_id=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -168,6 +169,7 @@ class TestIngestConnectorTask:
|
||||
operation_mode="upload",
|
||||
doc_id=None,
|
||||
sync_frequency="never",
|
||||
config=None,
|
||||
idempotency_key=None,
|
||||
source_id=None,
|
||||
)
|
||||
@@ -207,6 +209,7 @@ class TestIngestConnectorTask:
|
||||
operation_mode="sync",
|
||||
doc_id="doc1",
|
||||
sync_frequency="daily",
|
||||
config=None,
|
||||
idempotency_key=None,
|
||||
source_id=None,
|
||||
)
|
||||
@@ -600,7 +603,7 @@ class TestIngestIdempotency:
|
||||
worker_calls = []
|
||||
|
||||
def _fake_worker(self, directory, formats, job_name, file_path,
|
||||
filename, user, file_name_map=None,
|
||||
filename, user, file_name_map=None, config=None,
|
||||
idempotency_key=None, source_id=None):
|
||||
worker_calls.append(filename)
|
||||
return {"status": "ok", "directory": directory}
|
||||
|
||||
Reference in new issue
Block a user