mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 14:17:08 +00:00
feat: default tools (#2485)
* feat: default tools * fix: tests * feat scheduler * fix: scheduler UI * Agent switch breadcrumbs * fix: minor fixes to scheduler * fix: tests
This commit is contained in:
107 files changed
+14212
-458
No files matched your search
@@ -429,7 +429,7 @@ class TestConfigureRetriever:
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid_search"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
|
||||
def test_request_overrides_agent(self):
|
||||
def test_agent_wins_over_request_on_agent_bound(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
@@ -438,9 +438,33 @@ class TestConfigureRetriever:
|
||||
)
|
||||
sp._agent_data = {"retriever": "hybrid_search", "chunks": 5}
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid_search"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
|
||||
def test_body_wins_on_agentless(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor(
|
||||
{"retriever": "duckdb", "chunks": 7}, {"sub": "u"},
|
||||
)
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "duckdb"
|
||||
assert sp.retriever_config["chunks"] == 7
|
||||
|
||||
def test_agent_bound_drops_body_chunks_and_retriever(self):
|
||||
# Missing agent values fall back to system defaults, not body's.
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor(
|
||||
{"retriever": "duckdb", "chunks": 7}, {"sub": "u"},
|
||||
)
|
||||
sp._agent_data = {}
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "classic"
|
||||
assert sp.retriever_config["chunks"] == 2
|
||||
|
||||
def test_invalid_agent_chunks_falls_back(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
@@ -569,10 +593,11 @@ class TestPreFetchTools:
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_no_user_tools_returns_none(self, pg_conn):
|
||||
def test_no_template_skips_default_tool_prefetch(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({}, {"sub": "no-tools-user"})
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
@@ -580,3 +605,347 @@ class TestPreFetchTools:
|
||||
):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_no_template_skips_only_default_rows_not_explicit(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.user_tools import (
|
||||
UserToolsRepository,
|
||||
)
|
||||
|
||||
UserToolsRepository(pg_conn).create(
|
||||
user_id="u-explicit-prefetch", name="read_webpage", status=True
|
||||
)
|
||||
sp = StreamProcessor({}, {"sub": "u-explicit-prefetch"})
|
||||
fetched = []
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
fetched.append(tool_doc)
|
||||
return {"ok": True}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert "read_webpage" in got
|
||||
assert all(not d.get("default") for d in fetched)
|
||||
assert any(d.get("name") == "read_webpage" for d in fetched)
|
||||
|
||||
def test_default_tool_prefetched_when_template_references_it(
|
||||
self, pg_conn
|
||||
):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({}, {"sub": "u-tpl-default"})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
fetched = []
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
fetched.append(tool_doc)
|
||||
return {"ok": True}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert any(
|
||||
d.get("name") == "read_webpage" and d.get("default")
|
||||
for d in fetched
|
||||
)
|
||||
# Defaults are reachable by synthetic id only — not by name.
|
||||
assert default_tool_id("read_webpage") in got
|
||||
|
||||
def test_agent_bound_invocation_omits_default_tool_prefetch(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({"agent_id": "agent-xyz"}, {"sub": "u-ag"})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_template_name_key_favors_explicit_over_default(self, pg_conn):
|
||||
"""An explicit row and the synthesized default of the same name
|
||||
coexist: name key stays on the explicit, default reachable by
|
||||
synthetic id only."""
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.user_tools import (
|
||||
UserToolsRepository,
|
||||
)
|
||||
|
||||
user = "u-collision"
|
||||
explicit = UserToolsRepository(pg_conn).create(
|
||||
user_id=user, name="read_webpage", status=True,
|
||||
)
|
||||
explicit_id = str(explicit["id"])
|
||||
default_id = default_tool_id("read_webpage")
|
||||
|
||||
sp = StreamProcessor({}, {"sub": user})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
return {
|
||||
"is_default": bool(tool_doc.get("default")),
|
||||
"id": str(tool_doc.get("_id") or tool_doc.get("id")),
|
||||
}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert got["read_webpage"]["is_default"] is False
|
||||
assert got["read_webpage"]["id"] == explicit_id
|
||||
assert got[explicit_id]["is_default"] is False
|
||||
assert got[default_id]["is_default"] is True
|
||||
|
||||
|
||||
class TestValidateAndSetModelAgentAuthority:
|
||||
"""Agent-bound chats: agent's ``default_model_id`` is authoritative."""
|
||||
|
||||
def test_agent_bound_ignores_body_model_id(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "caller"})
|
||||
sp._agent_data = {"user": "owner"}
|
||||
sp.agent_config = {
|
||||
"default_model_id": "agent-model",
|
||||
"user_id": "owner",
|
||||
}
|
||||
captured: list = []
|
||||
|
||||
def _fake_validate(model_id, user_id=None):
|
||||
captured.append((model_id, user_id))
|
||||
return True
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
side_effect=_fake_validate,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="global-default",
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "agent-model"
|
||||
# Resolved under the agent owner, not the caller.
|
||||
assert sp.model_user_id == "owner"
|
||||
assert ("agent-model", "owner") in captured
|
||||
|
||||
def test_agent_bound_no_default_falls_back_to_system(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "u"})
|
||||
sp._agent_data = {"user": "u"}
|
||||
sp.agent_config = {"default_model_id": "", "user_id": "u"}
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="global-default",
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "global-default"
|
||||
assert sp.model_user_id is None
|
||||
|
||||
def test_agentless_body_model_still_wins(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "u"})
|
||||
sp._agent_data = None
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=True,
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "body-model"
|
||||
assert sp.model_user_id == "u"
|
||||
|
||||
|
||||
class TestGetDataFromApiKeySourceUnion:
|
||||
"""`_get_data_from_api_key`: primary ∪ extras, deduplicated, primary first."""
|
||||
|
||||
def _make_sp(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
return StreamProcessor({}, {"sub": "u"})
|
||||
|
||||
def test_union_primary_and_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-both"
|
||||
sources_repo = SourcesRepository(pg_conn)
|
||||
primary = sources_repo.create(name="primary", user_id=owner)
|
||||
extra1 = sources_repo.create(name="extra1", user_id=owner)
|
||||
extra2 = sources_repo.create(name="extra2", user_id=owner)
|
||||
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
owner, "agent-merge", "published",
|
||||
key="merge-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(extra1["id"]), str(extra2["id"])],
|
||||
retriever="hybrid",
|
||||
chunks=5,
|
||||
)
|
||||
assert agent is not None
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("merge-key")
|
||||
ids = [s["id"] for s in data["sources"]]
|
||||
assert ids == [
|
||||
str(primary["id"]),
|
||||
str(extra1["id"]),
|
||||
str(extra2["id"]),
|
||||
]
|
||||
assert data["source"] == str(primary["id"])
|
||||
|
||||
def test_only_primary(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-primary-only"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="primary", user_id=owner,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "primary-only", "published",
|
||||
key="primary-only-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("primary-only-key")
|
||||
assert [s["id"] for s in data["sources"]] == [str(primary["id"])]
|
||||
assert data["source"] == str(primary["id"])
|
||||
|
||||
def test_only_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-extras-only"
|
||||
e1 = SourcesRepository(pg_conn).create(name="e1", user_id=owner)
|
||||
e2 = SourcesRepository(pg_conn).create(name="e2", user_id=owner)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "extras-only", "published",
|
||||
key="extras-only-key",
|
||||
extra_source_ids=[str(e1["id"]), str(e2["id"])],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("extras-only-key")
|
||||
assert [s["id"] for s in data["sources"]] == [
|
||||
str(e1["id"]), str(e2["id"]),
|
||||
]
|
||||
assert data["source"] is None
|
||||
|
||||
def test_dedupe_primary_repeated_in_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-dedupe"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="dup-primary", user_id=owner,
|
||||
)
|
||||
extra = SourcesRepository(pg_conn).create(
|
||||
name="dup-extra", user_id=owner,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "dedupe", "published",
|
||||
key="dedupe-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(primary["id"]), str(extra["id"])],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("dedupe-key")
|
||||
ids = [s["id"] for s in data["sources"]]
|
||||
assert ids == [str(primary["id"]), str(extra["id"])]
|
||||
|
||||
|
||||
class TestAgentBoundFieldsAuthoritative:
|
||||
"""End-to-end regression: agent's source/model/chunks/retriever win."""
|
||||
|
||||
def test_agent_values_win_over_body(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-regr-agent-authority"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="primary", user_id=owner,
|
||||
)
|
||||
extra = SourcesRepository(pg_conn).create(
|
||||
name="extra", user_id=owner,
|
||||
)
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "authoritative", "published",
|
||||
key="auth-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(extra["id"])],
|
||||
default_model_id="model-A",
|
||||
retriever="hybrid",
|
||||
chunks=5,
|
||||
)
|
||||
|
||||
# Body sends different values for every field; all must be ignored.
|
||||
body = {
|
||||
"api_key": "auth-key",
|
||||
"model_id": "body-model-Z",
|
||||
"retriever": "duckdb",
|
||||
"chunks": 99,
|
||||
"active_docs": "body-source-id",
|
||||
}
|
||||
sp = StreamProcessor(body, {"sub": owner})
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="system-default",
|
||||
):
|
||||
sp._configure_agent()
|
||||
sp._validate_and_set_model()
|
||||
sp._configure_source()
|
||||
sp._configure_retriever()
|
||||
|
||||
assert sp.model_id == "model-A"
|
||||
assert sp.model_user_id == owner
|
||||
assert sp.agent_config["default_model_id"] == "model-A"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid"
|
||||
assert sp.source == {
|
||||
"active_docs": [str(primary["id"]), str(extra["id"])],
|
||||
}
|
||||
Reference in new issue
Block a user