mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 20:12:55 +00:00
arc53-machine
029189fe0c
The answer routes refuse such a request with 401, but only after building the agent, so a public agent's prompt tools were pre-fetched, and their actions run, for nobody. Anonymous chat without an agent key is not supported, so the processor now stops before any setup and the route answers 401 as before.
184 lines
7.5 KiB
Python
184 lines
7.5 KiB
Python
"""Tests for per-source search exposure partitioning (E1 / D11)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from docsgpt.api.answer.services.stream_processor import StreamProcessor
|
|
from docsgpt.storage.db.source_config import RetrievalConfig
|
|
|
|
|
|
def _processor() -> StreamProcessor:
|
|
"""A bare StreamProcessor with only the fields the exposure helpers touch.
|
|
|
|
Built via ``__new__`` to skip the DB-heavy ``__init__``; the exposure
|
|
methods read ``all_sources`` / ``source`` / ``agent_config`` only.
|
|
"""
|
|
sp = StreamProcessor.__new__(StreamProcessor)
|
|
sp.all_sources = []
|
|
sp.source = {}
|
|
sp.agent_config = {}
|
|
sp.retriever_config = {"retriever_name": "classic", "chunks": 2, "doc_token_limit": 50000}
|
|
sp.data = {}
|
|
sp.decoded_token = {"sub": "u"}
|
|
return sp
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExposurePartition:
|
|
def test_no_config_all_default_to_prefetch(self):
|
|
sp = _processor()
|
|
sp.all_sources = [{"id": "a", "retrieval": None}, {"id": "b", "retrieval": None}]
|
|
prefetch, agentic = sp._exposure_partition()
|
|
assert [e["id"] for e in prefetch] == ["a", "b"]
|
|
assert agentic == []
|
|
|
|
def test_mixed_partition(self):
|
|
sp = _processor()
|
|
sp.all_sources = [
|
|
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
|
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
|
{"id": "c", "retrieval": RetrievalConfig()}, # default prefetch
|
|
]
|
|
prefetch, agentic = sp._exposure_partition()
|
|
assert sorted(e["id"] for e in prefetch) == ["a", "c"]
|
|
assert [e["id"] for e in agentic] == ["b"]
|
|
|
|
def test_exposure_of_dict_and_model_and_missing(self):
|
|
sp = _processor()
|
|
assert sp._exposure_of(RetrievalConfig(exposure="agentic_tool")) == "agentic_tool"
|
|
assert sp._exposure_of({"exposure": "agentic_tool"}) == "agentic_tool"
|
|
assert sp._exposure_of(None) == "prefetch"
|
|
assert sp._exposure_of({}) == "prefetch"
|
|
|
|
def test_source_for_docs(self):
|
|
sp = _processor()
|
|
assert sp._source_for_docs(["a", "b"]) == {"active_docs": ["a", "b"]}
|
|
assert sp._source_for_docs([]) == {}
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBuildAgentExposure:
|
|
def _agentic_processor(self):
|
|
sp = _processor()
|
|
sp.agent_config = {"agent_type": "agentic"}
|
|
sp.initialize = MagicMock()
|
|
sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
|
|
sp.pre_fetch_tools = MagicMock(return_value={"tool": {}})
|
|
sp.create_agent = MagicMock(return_value="AGENT")
|
|
return sp
|
|
|
|
def test_all_prefetch_is_today_no_partition(self):
|
|
# No agentic_tool source → today's behavior: no pre-fetch, no
|
|
# agentic_sources passed (tool exposes all sources).
|
|
sp = self._agentic_processor()
|
|
sp.all_sources = [
|
|
{"id": "a", "retrieval": RetrievalConfig()},
|
|
{"id": "b", "retrieval": None},
|
|
]
|
|
result = sp.build_agent("q")
|
|
assert result == "AGENT"
|
|
sp.pre_fetch_docs.assert_not_called()
|
|
_, kwargs = sp.create_agent.call_args
|
|
assert "agentic_sources" not in kwargs
|
|
|
|
def test_mixed_prefetches_and_scopes_tool(self):
|
|
sp = self._agentic_processor()
|
|
sp.all_sources = [
|
|
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
|
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
|
]
|
|
sp.build_agent("q")
|
|
# Pre-fetch ran scoped to the prefetch subset.
|
|
sp.pre_fetch_docs.assert_called_once()
|
|
_, pf_kwargs = sp.pre_fetch_docs.call_args
|
|
assert pf_kwargs.get("exposure") == "prefetch"
|
|
# create_agent received only the agentic_tool subset as the tool sources.
|
|
_, kwargs = sp.create_agent.call_args
|
|
assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]
|
|
|
|
def _classic_processor(self):
|
|
sp = _processor()
|
|
sp.agent_config = {"agent_type": "classic"}
|
|
sp.initialize = MagicMock()
|
|
sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
|
|
sp.pre_fetch_tools = MagicMock(return_value=None)
|
|
sp.create_agent = MagicMock(return_value="AGENT")
|
|
return sp
|
|
|
|
def test_classic_default_prefetches_all_no_partition(self):
|
|
# Default classic (all prefetch / no config): byte-identical to today —
|
|
# unscoped pre-fetch and no agentic_sources, so no search tool is added.
|
|
sp = self._classic_processor()
|
|
sp.all_sources = [
|
|
{"id": "a", "retrieval": RetrievalConfig()},
|
|
{"id": "b", "retrieval": None},
|
|
]
|
|
result = sp.build_agent("q")
|
|
assert result == "AGENT"
|
|
sp.pre_fetch_docs.assert_called_once_with("q")
|
|
_, kwargs = sp.create_agent.call_args
|
|
assert "agentic_sources" not in kwargs
|
|
|
|
def test_classic_no_per_source_detail_is_today(self):
|
|
# Single-source / no-config requests carry no per-source detail; classic
|
|
# must fall back to the unscoped pre-fetch (no exposure scoping).
|
|
sp = self._classic_processor()
|
|
sp.all_sources = []
|
|
sp.source = {"active_docs": ["a"]}
|
|
sp.build_agent("q")
|
|
sp.pre_fetch_docs.assert_called_once_with("q")
|
|
_, kwargs = sp.create_agent.call_args
|
|
assert "agentic_sources" not in kwargs
|
|
|
|
def test_classic_mixed_prefetches_and_scopes_tool(self):
|
|
# Classic with an agentic_tool source now pre-fetches only the prefetch
|
|
# subset and exposes the agentic_tool subset via the search tool.
|
|
sp = self._classic_processor()
|
|
sp.all_sources = [
|
|
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
|
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
|
]
|
|
sp.build_agent("q")
|
|
sp.pre_fetch_docs.assert_called_once()
|
|
_, pf_kwargs = sp.pre_fetch_docs.call_args
|
|
assert pf_kwargs.get("exposure") == "prefetch"
|
|
_, kwargs = sp.create_agent.call_args
|
|
assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestNonAgentSourceConfig:
|
|
"""The non-agent chat path must also load per-source config so exposure
|
|
(and other per-source overrides) are honored, not just the agent path."""
|
|
|
|
def test_active_docs_loads_per_source_retrieval_config(self):
|
|
sp = StreamProcessor.__new__(StreamProcessor)
|
|
sp._agent_data = None
|
|
sp.data = {"active_docs": "s1"}
|
|
sp.initial_user_id = "user1"
|
|
sp.source = {}
|
|
sp.all_sources = []
|
|
fake_source = {
|
|
"config": {"retrieval": {"exposure": "agentic_tool", "chunks": 7}}
|
|
}
|
|
with patch(
|
|
"docsgpt.api.answer.services.stream_processor.db_readonly"
|
|
), patch(
|
|
"docsgpt.api.answer.services.stream_processor.can_access",
|
|
return_value=True,
|
|
), patch(
|
|
"docsgpt.api.answer.services.stream_processor.SourcesRepository"
|
|
) as repo:
|
|
# Read unscoped after the access check, so a team grantee gets the
|
|
# source's real config instead of silently falling back to defaults.
|
|
repo.return_value.get_by_id.return_value = fake_source
|
|
sp._configure_source()
|
|
assert sp.source == {"active_docs": "s1"}
|
|
assert len(sp.all_sources) == 1
|
|
assert sp.all_sources[0]["id"] == "s1"
|
|
assert sp.all_sources[0]["retrieval"].exposure == "agentic_tool"
|
|
assert sp.all_sources[0]["retrieval"].chunks == 7
|