mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 22:13:08 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
376 lines
12 KiB
Python
376 lines
12 KiB
Python
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from docsgpt.agents.classic_agent import ClassicAgent
|
|
from docsgpt.agents.tools.internal_search import INTERNAL_TOOL_ID
|
|
|
|
|
|
@pytest.fixture
|
|
def _no_tools(monkeypatch):
|
|
"""Stub ToolExecutor.get_tools to avoid DB hits for most tests."""
|
|
|
|
def _fake_get_tools(self):
|
|
return {}
|
|
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolExecutor.get_tools", _fake_get_tools
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def _no_dir_structure(monkeypatch):
|
|
"""Stub the DB-backed directory-structure lookup used by the search tool."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tools.internal_search.sources_have_directory_structure",
|
|
lambda source: False,
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestClassicAgent:
|
|
|
|
def test_classic_agent_initialization(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
assert isinstance(agent, ClassicAgent)
|
|
assert agent.endpoint == agent_base_params["endpoint"]
|
|
assert agent.llm_name == agent_base_params["llm_name"]
|
|
|
|
def test_gen_inner_basic_flow(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
log_context,
|
|
):
|
|
def mock_gen_stream(*args, **kwargs):
|
|
yield "Answer chunk 1"
|
|
yield "Answer chunk 2"
|
|
|
|
mock_llm.gen_stream = Mock(return_value=mock_gen_stream())
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed answer"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
results = list(agent._gen_inner("Test query", log_context))
|
|
|
|
assert len(results) >= 2
|
|
sources = [r for r in results if "sources" in r]
|
|
tool_calls = [r for r in results if "tool_calls" in r]
|
|
|
|
assert len(sources) == 1
|
|
assert len(tool_calls) == 1
|
|
|
|
def test_gen_inner_retrieves_documents(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
log_context,
|
|
):
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
list(agent._gen_inner("Test query", log_context))
|
|
|
|
def test_gen_inner_uses_user_tools(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
monkeypatch,
|
|
log_context,
|
|
):
|
|
# Inject a fake user tool dict via get_tools rather than touching DB.
|
|
fake_tools = {
|
|
"t1": {
|
|
"id": "t1",
|
|
"name": "test_tool",
|
|
"config": {},
|
|
"actions": [
|
|
{
|
|
"name": "do_thing",
|
|
"description": "",
|
|
"active": True,
|
|
"parameters": {"properties": {}},
|
|
}
|
|
],
|
|
}
|
|
}
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolExecutor.get_tools",
|
|
lambda self: fake_tools,
|
|
)
|
|
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
list(agent._gen_inner("Test query", log_context))
|
|
|
|
# After _prepare_tools, the fake user tool should have become an llm tool schema.
|
|
assert any(
|
|
t["function"]["name"] == "do_thing" for t in agent.tools
|
|
)
|
|
|
|
def test_gen_inner_builds_correct_messages(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
log_context,
|
|
):
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
list(agent._gen_inner("Test query", log_context))
|
|
|
|
call_kwargs = mock_llm.gen_stream.call_args[1]
|
|
messages = call_kwargs["messages"]
|
|
|
|
assert len(messages) >= 2
|
|
assert messages[0]["role"] == "system"
|
|
assert messages[-1]["role"] == "user"
|
|
assert messages[-1]["content"] == "Test query"
|
|
|
|
def test_gen_inner_logs_tool_calls(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
log_context,
|
|
):
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tool_calls = [{"tool": "test", "result": "success"}]
|
|
|
|
list(agent._gen_inner("Test query", log_context))
|
|
|
|
agent_logs = [s for s in log_context.stacks if s["component"] == "agent"]
|
|
assert len(agent_logs) == 1
|
|
assert "tool_calls" in agent_logs[0]["data"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestClassicAgentSearchExposure:
|
|
"""ClassicAgent honors per-source exposure via an internal_search tool."""
|
|
|
|
def test_default_registers_no_internal_search(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
log_context,
|
|
):
|
|
# No retriever_config (every source at default prefetch) → no search
|
|
# tool is added and pre-fetched sources are preserved (unchanged).
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.retriever_config == {}
|
|
|
|
captured = {}
|
|
original_prepare = agent._prepare_tools
|
|
|
|
def _spy(tools_dict):
|
|
captured["tools_dict"] = tools_dict
|
|
return original_prepare(tools_dict)
|
|
|
|
agent._prepare_tools = _spy
|
|
prefetched = [{"text": "pre", "title": "Pre", "source": "s"}]
|
|
agent.retrieved_docs = list(prefetched)
|
|
|
|
list(agent._gen_inner("q", log_context))
|
|
|
|
assert INTERNAL_TOOL_ID not in captured["tools_dict"]
|
|
# Pre-fetched sources untouched (no _collect_internal_sources overwrite).
|
|
assert agent.retrieved_docs == prefetched
|
|
|
|
def test_with_agentic_tool_source_registers_internal_search(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
_no_dir_structure,
|
|
log_context,
|
|
):
|
|
# A retriever_config (agentic_tool subset) → the internal_search tool is
|
|
# registered so the LLM can search that subset on demand.
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
retriever_config = {
|
|
"source": {"active_docs": ["b"]},
|
|
"retriever_name": "classic",
|
|
"chunks": 2,
|
|
"sources": [{"id": "b", "retrieval": None}],
|
|
}
|
|
agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params)
|
|
|
|
captured = {}
|
|
original_prepare = agent._prepare_tools
|
|
|
|
def _spy(tools_dict):
|
|
captured["tools_dict"] = tools_dict
|
|
return original_prepare(tools_dict)
|
|
|
|
agent._prepare_tools = _spy
|
|
|
|
list(agent._gen_inner("q", log_context))
|
|
|
|
assert INTERNAL_TOOL_ID in captured["tools_dict"]
|
|
assert captured["tools_dict"][INTERNAL_TOOL_ID]["name"] == "internal_search"
|
|
|
|
def test_collect_internal_sources_surfaces_tool_docs(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
retriever_config = {"source": {"active_docs": ["b"]}}
|
|
agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params)
|
|
|
|
mock_tool = Mock()
|
|
mock_tool.retrieved_docs = [{"text": "Found", "title": "Doc", "source": "t"}]
|
|
cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{agent.user or ''}"
|
|
agent.tool_executor._loaded_tools[cache_key] = mock_tool
|
|
|
|
agent._collect_internal_sources()
|
|
assert [d["title"] for d in agent.retrieved_docs] == ["Doc"]
|
|
|
|
def test_collect_internal_sources_merges_with_prefetched(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
# Mixed exposure: a pre-fetched source's docs must survive alongside the
|
|
# tool-retrieved docs (not be overwritten), so both are cited.
|
|
retriever_config = {"source": {"active_docs": ["b"]}}
|
|
agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params)
|
|
agent.retrieved_docs = [
|
|
{"text": "Prefetched", "title": "Prefetched Doc", "source": "a"}
|
|
]
|
|
mock_tool = Mock()
|
|
mock_tool.retrieved_docs = [
|
|
{"text": "Found", "title": "Tool Doc", "source": "b"}
|
|
]
|
|
cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{agent.user or ''}"
|
|
agent.tool_executor._loaded_tools[cache_key] = mock_tool
|
|
|
|
agent._collect_internal_sources()
|
|
assert [d["title"] for d in agent.retrieved_docs] == [
|
|
"Prefetched Doc",
|
|
"Tool Doc",
|
|
]
|
|
|
|
def test_collect_internal_sources_dedupes(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
retriever_config = {"source": {"active_docs": ["b"]}}
|
|
agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params)
|
|
dup = {"text": "X", "title": "Dup", "source": "b"}
|
|
agent.retrieved_docs = [dup]
|
|
mock_tool = Mock()
|
|
mock_tool.retrieved_docs = [dict(dup)]
|
|
cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{agent.user or ''}"
|
|
agent.tool_executor._loaded_tools[cache_key] = mock_tool
|
|
|
|
agent._collect_internal_sources()
|
|
assert [d["title"] for d in agent.retrieved_docs] == ["Dup"]
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestClassicAgentIntegration:
|
|
|
|
def test_gen_method_with_logging(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
):
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
results = list(agent.gen("Test query"))
|
|
|
|
assert len(results) >= 1
|
|
|
|
def test_gen_method_decorator_applied(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
_no_tools,
|
|
):
|
|
mock_llm.gen_stream = Mock(return_value=iter(["Answer"]))
|
|
|
|
def mock_handler(*args, **kwargs):
|
|
yield "Processed"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
assert hasattr(agent.gen, "__wrapped__")
|