mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 16:13:23 +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.
1834 lines
64 KiB
Python
1834 lines
64 KiB
Python
from contextlib import contextmanager
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
from docsgpt.agents.classic_agent import ClassicAgent
|
|
from docsgpt.llm.anthropic import AnthropicLLM
|
|
from docsgpt.llm.docsgpt_provider import DocsGPTAPILLM
|
|
from docsgpt.llm.google_ai import GoogleLLM
|
|
from docsgpt.llm.groq import GroqLLM
|
|
from docsgpt.llm.novita import NovitaLLM
|
|
from docsgpt.llm.open_router import OpenRouterLLM
|
|
from docsgpt.llm.openai import OpenAILLM
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentInitialization:
|
|
|
|
def test_agent_initialization(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
assert agent.endpoint == agent_base_params["endpoint"]
|
|
assert agent.llm_name == agent_base_params["llm_name"]
|
|
assert agent.model_id == agent_base_params["model_id"]
|
|
assert agent.api_key == agent_base_params["api_key"]
|
|
assert agent.prompt == agent_base_params["prompt"]
|
|
assert agent.user == agent_base_params["decoded_token"]["sub"]
|
|
assert agent.tools == []
|
|
assert agent.tool_calls == []
|
|
|
|
def test_agent_initialization_with_none_chat_history(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["chat_history"] = None
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.chat_history == []
|
|
|
|
def test_agent_initialization_with_chat_history(
|
|
self,
|
|
agent_base_params,
|
|
sample_chat_history,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
agent_base_params["chat_history"] = sample_chat_history
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert len(agent.chat_history) == 2
|
|
assert agent.chat_history[0]["prompt"] == "What is Python?"
|
|
|
|
def test_agent_decoded_token_defaults_to_empty_dict(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["decoded_token"] = None
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.decoded_token == {}
|
|
assert agent.user is None
|
|
|
|
def test_agent_user_extracted_from_token(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["decoded_token"] = {"sub": "user123"}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.user == "user123"
|
|
|
|
def test_dependency_injection_llm(self, agent_base_params, mock_llm_handler_creator):
|
|
"""When llm is provided, LLMCreator.create_llm is NOT called."""
|
|
injected_llm = Mock()
|
|
agent_base_params["llm"] = injected_llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.llm is injected_llm
|
|
|
|
def test_dependency_injection_llm_handler(self, agent_base_params, mock_llm_creator):
|
|
"""When llm_handler is provided, LLMHandlerCreator is NOT called."""
|
|
injected_handler = Mock()
|
|
agent_base_params["llm_handler"] = injected_handler
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.llm_handler is injected_handler
|
|
|
|
def test_dependency_injection_tool_executor(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
"""When tool_executor is provided, a new one is NOT created."""
|
|
injected_executor = Mock()
|
|
injected_executor.tool_calls = []
|
|
agent_base_params["tool_executor"] = injected_executor
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.tool_executor is injected_executor
|
|
|
|
def test_json_schema_normalized(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.json_schema == {"type": "object"}
|
|
|
|
def test_json_schema_wrapped(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["json_schema"] = {"schema": {"type": "string"}}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.json_schema == {"type": "string"}
|
|
|
|
def test_json_schema_invalid_ignored(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["json_schema"] = {"bad": "no type"}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.json_schema is None
|
|
|
|
def test_retrieved_docs_defaults_to_empty(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.retrieved_docs == []
|
|
|
|
def test_attachments_defaults_to_empty(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["attachments"] = None
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.attachments == []
|
|
|
|
def test_limited_token_mode_defaults(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
assert agent.limited_token_mode is False
|
|
assert agent.limited_request_mode is False
|
|
assert agent.current_token_count == 0
|
|
assert agent.context_limit_reached is False
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentBuildMessages:
|
|
|
|
def test_build_messages_basic(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
system_prompt = "System prompt content"
|
|
query = "What is Python?"
|
|
|
|
messages = agent._build_messages(system_prompt, query)
|
|
|
|
assert len(messages) >= 2
|
|
assert messages[0]["role"] == "system"
|
|
assert messages[0]["content"] == system_prompt
|
|
assert messages[-1]["role"] == "user"
|
|
assert messages[-1]["content"] == query
|
|
|
|
def test_build_messages_with_chat_history(
|
|
self,
|
|
agent_base_params,
|
|
sample_chat_history,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
agent_base_params["chat_history"] = sample_chat_history
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
system_prompt = "System prompt"
|
|
query = "New question?"
|
|
|
|
messages = agent._build_messages(system_prompt, query)
|
|
|
|
user_messages = [m for m in messages if m["role"] == "user"]
|
|
assistant_messages = [m for m in messages if m["role"] == "assistant"]
|
|
|
|
assert len(user_messages) >= 3
|
|
assert len(assistant_messages) >= 2
|
|
|
|
def test_build_messages_with_tool_calls_in_history(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
tool_call_history = [
|
|
{
|
|
"tool_calls": [
|
|
{
|
|
"call_id": "123",
|
|
"action_name": "test_action",
|
|
"arguments": {"arg": "value"},
|
|
"result": "success",
|
|
}
|
|
]
|
|
}
|
|
]
|
|
agent_base_params["chat_history"] = tool_call_history
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = agent._build_messages("System prompt", "query")
|
|
|
|
tool_messages = [m for m in messages if m["role"] == "tool"]
|
|
assert len(tool_messages) > 0
|
|
|
|
def test_build_messages_handles_missing_filename(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = agent._build_messages("System prompt", "query")
|
|
|
|
assert messages[0]["role"] == "system"
|
|
assert messages[0]["content"] == "System prompt"
|
|
|
|
def test_build_messages_uses_title_as_fallback(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
agent._build_messages("System prompt", "query")
|
|
|
|
def test_build_messages_uses_source_as_fallback(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
agent._build_messages("System prompt", "query")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentTools:
|
|
|
|
def test_get_user_tools(
|
|
self,
|
|
agent_base_params,
|
|
pg_conn,
|
|
monkeypatch,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
repo = UserToolsRepository(pg_conn)
|
|
repo.create(user_id="test_user", name="tool1", status=True)
|
|
repo.create(user_id="test_user", name="tool2", status=True)
|
|
|
|
@contextmanager
|
|
def _use_pg_conn():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.db_readonly", _use_pg_conn
|
|
)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
tools = agent._get_user_tools("test_user")
|
|
|
|
from docsgpt.agents.default_tools import loaded_default_tools
|
|
|
|
assert len(tools) == 2 + len(loaded_default_tools())
|
|
assert "0" in tools
|
|
assert "1" in tools
|
|
names = {t["name"] for t in tools.values()}
|
|
assert {"tool1", "tool2"}.issubset(names)
|
|
assert set(loaded_default_tools()).issubset(names)
|
|
|
|
def test_get_user_tools_filters_by_status(
|
|
self,
|
|
agent_base_params,
|
|
pg_conn,
|
|
monkeypatch,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
repo = UserToolsRepository(pg_conn)
|
|
repo.create(user_id="test_user", name="tool1", status=True)
|
|
repo.create(user_id="test_user", name="tool2", status=False)
|
|
|
|
@contextmanager
|
|
def _use_pg_conn():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.db_readonly", _use_pg_conn
|
|
)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
tools = agent._get_user_tools("test_user")
|
|
|
|
from docsgpt.agents.default_tools import loaded_default_tools
|
|
|
|
assert len(tools) == 1 + len(loaded_default_tools())
|
|
names = {t["name"] for t in tools.values()}
|
|
assert "tool1" in names
|
|
assert "tool2" not in names
|
|
|
|
def test_get_tools_by_api_key(
|
|
self,
|
|
agent_base_params,
|
|
pg_conn,
|
|
monkeypatch,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
tool_row = UserToolsRepository(pg_conn).create(
|
|
user_id="alice", name="api_tool"
|
|
)
|
|
tool_id = str(tool_row["id"])
|
|
|
|
AgentsRepository(pg_conn).create(
|
|
user_id="alice",
|
|
name="my-agent",
|
|
status="active",
|
|
key="api_key_123",
|
|
tools=[tool_id],
|
|
)
|
|
|
|
@contextmanager
|
|
def _use_pg_conn():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.db_readonly", _use_pg_conn
|
|
)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
tools = agent._get_tools("api_key_123")
|
|
|
|
from docsgpt.agents.default_tools import loaded_default_tools
|
|
|
|
# Agent-bound: exactly agents.tools, no defaults.
|
|
assert set(tools) == {tool_id}
|
|
names = {t["name"] for t in tools.values()}
|
|
assert names == {"api_tool"}
|
|
assert not (set(loaded_default_tools()) & names)
|
|
|
|
def test_build_tool_parameters(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
action = {
|
|
"parameters": {
|
|
"properties": {
|
|
"param1": {
|
|
"type": "string",
|
|
"description": "Test param",
|
|
"filled_by_llm": True,
|
|
"required": True,
|
|
},
|
|
"param2": {
|
|
"type": "number",
|
|
"filled_by_llm": False,
|
|
"value": 42,
|
|
"required": False,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
params = agent._build_tool_parameters(action)
|
|
|
|
assert "param1" in params["properties"]
|
|
assert "param1" in params["required"]
|
|
assert "filled_by_llm" not in params["properties"]["param1"]
|
|
|
|
def test_prepare_tools_with_api_tool(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"name": "api_tool",
|
|
"config": {
|
|
"actions": {
|
|
"get_data": {
|
|
"name": "get_data",
|
|
"description": "Get data from API",
|
|
"active": True,
|
|
"url": "https://api.example.com/data",
|
|
"method": "GET",
|
|
"parameters": {"properties": {}},
|
|
}
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
agent._prepare_tools(tools_dict)
|
|
|
|
assert len(agent.tools) == 1
|
|
assert agent.tools[0]["type"] == "function"
|
|
assert agent.tools[0]["function"]["name"] == "get_data"
|
|
|
|
def test_prepare_tools_with_regular_tool(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"name": "custom_tool",
|
|
"actions": [
|
|
{
|
|
"name": "action1",
|
|
"description": "Custom action",
|
|
"active": True,
|
|
"parameters": {"properties": {}},
|
|
}
|
|
],
|
|
}
|
|
}
|
|
|
|
agent._prepare_tools(tools_dict)
|
|
|
|
assert len(agent.tools) == 1
|
|
assert agent.tools[0]["function"]["name"] == "action1"
|
|
|
|
def test_prepare_tools_filters_inactive_actions(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"name": "custom_tool",
|
|
"actions": [
|
|
{
|
|
"name": "active_action",
|
|
"description": "Active",
|
|
"active": True,
|
|
"parameters": {"properties": {}},
|
|
},
|
|
{
|
|
"name": "inactive_action",
|
|
"description": "Inactive",
|
|
"active": False,
|
|
"parameters": {"properties": {}},
|
|
},
|
|
],
|
|
}
|
|
}
|
|
|
|
agent._prepare_tools(tools_dict)
|
|
|
|
assert len(agent.tools) == 1
|
|
assert agent.tools[0]["function"]["name"] == "active_action"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentToolExecution:
|
|
|
|
def test_execute_tool_action_success(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
mock_tool_manager,
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
call = Mock()
|
|
call.id = "call_123"
|
|
call.name = "test_action_1"
|
|
call.arguments = '{"param1": "value1"}'
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"id": "11111111-1111-1111-1111-111111111111",
|
|
"name": "custom_tool",
|
|
"config": {},
|
|
"actions": [
|
|
{
|
|
"name": "test_action",
|
|
"description": "Test",
|
|
"parameters": {"properties": {}},
|
|
}
|
|
],
|
|
}
|
|
}
|
|
|
|
results = list(agent._execute_tool_action(tools_dict, call))
|
|
|
|
assert len(results) >= 2
|
|
assert results[0]["type"] == "tool_call"
|
|
assert results[0]["data"]["status"] == "pending"
|
|
assert results[-1]["data"]["status"] == "completed"
|
|
|
|
def test_execute_tool_action_invalid_tool_name(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
call = Mock()
|
|
call.id = "call_123"
|
|
call.name = "invalid_format"
|
|
call.arguments = "{}"
|
|
|
|
tools_dict = {}
|
|
|
|
results = list(agent._execute_tool_action(tools_dict, call))
|
|
|
|
assert results[0]["type"] == "tool_call"
|
|
assert results[0]["data"]["status"] == "error"
|
|
assert "Available tools:" in results[0]["data"]["result"]
|
|
|
|
def test_execute_tool_action_tool_not_found(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
call = Mock()
|
|
call.id = "call_123"
|
|
call.name = "action_999"
|
|
call.arguments = "{}"
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"name": "tool1",
|
|
"config": {},
|
|
"actions": [{"name": "tool1_read", "description": "D"}],
|
|
}
|
|
}
|
|
|
|
results = list(agent._execute_tool_action(tools_dict, call))
|
|
|
|
assert results[0]["type"] == "tool_call"
|
|
assert results[0]["data"]["status"] == "error"
|
|
assert "no such tool" in results[0]["data"]["result"]
|
|
# LLM-visible ACTION names. Not the tool name and never the internal
|
|
# id: both name a string the model cannot actually call, which just
|
|
# buys another failed round.
|
|
assert "tool1_read" in results[0]["data"]["result"]
|
|
|
|
def test_execute_tool_action_with_parameters(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
mock_tool_manager,
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
call = Mock()
|
|
call.id = "call_123"
|
|
call.name = "test_action_1"
|
|
call.arguments = '{"param1": "value1", "param2": "value2"}'
|
|
|
|
tools_dict = {
|
|
"1": {
|
|
"id": "22222222-2222-2222-2222-222222222222",
|
|
"name": "custom_tool",
|
|
"config": {},
|
|
"actions": [
|
|
{
|
|
"name": "test_action",
|
|
"description": "Test",
|
|
"parameters": {
|
|
"properties": {
|
|
"param1": {"type": "string"},
|
|
"param2": {"type": "string"},
|
|
}
|
|
},
|
|
}
|
|
],
|
|
}
|
|
}
|
|
|
|
results = list(agent._execute_tool_action(tools_dict, call))
|
|
|
|
assert results[-1]["data"]["status"] == "completed"
|
|
assert results[-1]["data"]["arguments"]["param1"] == "value1"
|
|
|
|
def test_get_truncated_tool_calls(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
from docsgpt.agents.tool_executor import PERSISTED_RESULT_MAX_LEN
|
|
|
|
agent.tool_calls = [
|
|
{
|
|
"tool_name": "test_tool",
|
|
"call_id": "123",
|
|
"action_name": "action",
|
|
"arguments": {},
|
|
"result": "a" * (PERSISTED_RESULT_MAX_LEN + 100),
|
|
}
|
|
]
|
|
|
|
truncated = agent._get_truncated_tool_calls()
|
|
|
|
assert len(truncated) == 1
|
|
assert len(truncated[0]["result"]) == PERSISTED_RESULT_MAX_LEN + 3
|
|
assert truncated[0]["result"].endswith("...")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentLLMGeneration:
|
|
|
|
def test_llm_gen_basic(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages, log_context)
|
|
|
|
mock_llm.gen_stream.assert_called_once()
|
|
call_args = mock_llm.gen_stream.call_args[1]
|
|
assert call_args["model"] == agent.model_id
|
|
assert call_args["messages"] == messages
|
|
|
|
def test_llm_gen_with_tools(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages, log_context)
|
|
|
|
call_args = mock_llm.gen_stream.call_args[1]
|
|
assert "tools" in call_args
|
|
assert call_args["tools"] == agent.tools
|
|
|
|
def test_llm_gen_with_json_schema(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
# A real OpenAILLM: the format kwarg is chosen from the LLM class,
|
|
# so a bare Mock deliberately gets no structured output.
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
llm.gen_stream = Mock()
|
|
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent_base_params["llm_name"] = "openai"
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages, log_context)
|
|
|
|
call_args = llm.gen_stream.call_args[1]
|
|
assert "response_format" in call_args
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentHandleResponse:
|
|
|
|
def test_handle_response_string(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, log_context
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
response = "Simple string response"
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
|
|
assert len(results) == 1
|
|
assert results[0]["answer"] == "Simple string response"
|
|
|
|
def test_handle_response_with_message(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, log_context
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
response = Mock()
|
|
response.message = Mock()
|
|
response.message.content = "Message content"
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
|
|
assert len(results) == 1
|
|
assert results[0]["answer"] == "Message content"
|
|
|
|
def test_handle_response_with_structured_output(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
mock_llm._supports_structured_output = Mock(return_value=True)
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
response = "Structured response"
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
|
|
assert results[0]["structured"] is True
|
|
assert results[0]["schema"] == {"type": "object"}
|
|
|
|
def test_handle_response_with_handler(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
def mock_process(*args):
|
|
yield {"type": "tool_call", "data": {}}
|
|
yield "Final answer"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
response = Mock()
|
|
response.message = None
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["type"] == "tool_call"
|
|
assert results[1]["answer"] == "Final answer"
|
|
|
|
def test_handle_response_dict_event_passthrough(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Dict events with 'type' key pass through without wrapping."""
|
|
|
|
def mock_process(*args):
|
|
yield {"type": "info", "data": {"message": "processing"}}
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
response = Mock()
|
|
response.message = None
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
assert results == [{"type": "info", "data": {"message": "processing"}}]
|
|
|
|
def test_handle_response_message_object_from_handler(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Response objects with .message.content from handler are unwrapped."""
|
|
event = Mock()
|
|
event.message = Mock()
|
|
event.message.content = "from handler"
|
|
|
|
def mock_process(*args):
|
|
yield event
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
response = Mock()
|
|
response.message = None
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
assert results[0]["answer"] == "from handler"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# gen() — the @log_activity decorated entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentGen:
|
|
|
|
def test_gen_delegates_to_gen_inner(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
# ClassicAgent._gen_inner is abstract — we patch it
|
|
with patch.object(agent, "_gen_inner") as mock_inner:
|
|
mock_inner.return_value = iter([{"answer": "ok"}])
|
|
results = list(agent.gen("hello"))
|
|
|
|
assert any(r.get("answer") == "ok" for r in results)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# tool_calls property
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentToolCallsProperty:
|
|
|
|
def test_getter(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tool_executor.tool_calls = ["a", "b"]
|
|
assert agent.tool_calls == ["a", "b"]
|
|
|
|
def test_setter(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tool_calls = ["x"]
|
|
assert agent.tool_executor.tool_calls == ["x"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _calculate_current_context_tokens
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCalculateContextTokens:
|
|
|
|
def test_delegates_to_token_counter(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
|
|
with patch(
|
|
"docsgpt.api.answer.services.compression.token_counter.TokenCounter"
|
|
) as MockTC:
|
|
MockTC.count_message_tokens.return_value = 42
|
|
result = agent._calculate_current_context_tokens(messages)
|
|
assert result == 42
|
|
MockTC.count_message_tokens.assert_called_once_with(messages)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _check_context_limit
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCheckContextLimit:
|
|
|
|
def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator):
|
|
return ClassicAgent(**agent_base_params)
|
|
|
|
def test_below_threshold_returns_false(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
messages = [{"role": "user", "content": "hi"}]
|
|
|
|
with patch.object(agent, "_calculate_current_context_tokens", return_value=100):
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=10000
|
|
):
|
|
result = agent._check_context_limit(messages)
|
|
assert result is False
|
|
assert agent.current_token_count == 100
|
|
|
|
def test_at_threshold_returns_true(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
messages = [{"role": "user", "content": "hi"}]
|
|
|
|
# threshold = 10000 * 0.8 = 8000; tokens = 8001 → True
|
|
with patch.object(agent, "_calculate_current_context_tokens", return_value=8001):
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=10000
|
|
):
|
|
result = agent._check_context_limit(messages)
|
|
assert result is True
|
|
|
|
def test_error_returns_false(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
with patch.object(
|
|
agent,
|
|
"_calculate_current_context_tokens",
|
|
side_effect=RuntimeError("boom"),
|
|
):
|
|
result = agent._check_context_limit([])
|
|
assert result is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _validate_context_size
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestValidateContextSize:
|
|
|
|
def test_at_limit_logs_warning(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
with patch.object(agent, "_calculate_current_context_tokens", return_value=10000):
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=10000
|
|
):
|
|
# Should not raise
|
|
agent._validate_context_size([{"role": "user", "content": "x"}])
|
|
assert agent.current_token_count == 10000
|
|
|
|
def test_below_threshold_no_warning(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
with patch.object(agent, "_calculate_current_context_tokens", return_value=100):
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=10000
|
|
):
|
|
agent._validate_context_size([])
|
|
assert agent.current_token_count == 100
|
|
|
|
def test_approaching_threshold(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
# 8500 / 10000 = 85% → above 80% threshold but below 100%
|
|
with patch.object(agent, "_calculate_current_context_tokens", return_value=8500):
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=10000
|
|
):
|
|
agent._validate_context_size([])
|
|
assert agent.current_token_count == 8500
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _truncate_text_middle
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTruncateTextMiddle:
|
|
|
|
def test_short_text_unchanged(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
with patch("docsgpt.utils.num_tokens_from_string", return_value=5):
|
|
result = agent._truncate_text_middle("short", max_tokens=100)
|
|
assert result == "short"
|
|
|
|
def test_long_text_truncated(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
long_text = "A" * 1000
|
|
|
|
def fake_tokens(text):
|
|
return len(text) // 4
|
|
|
|
with patch("docsgpt.utils.num_tokens_from_string", side_effect=fake_tokens):
|
|
result = agent._truncate_text_middle(long_text, max_tokens=50)
|
|
assert "[... content truncated to fit context limit ...]" in result
|
|
assert len(result) < len(long_text)
|
|
|
|
def test_zero_target_returns_empty(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
with patch("docsgpt.utils.num_tokens_from_string", return_value=100):
|
|
result = agent._truncate_text_middle("some text", max_tokens=0)
|
|
assert result == ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _truncate_history_to_fit
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTruncateHistoryToFit:
|
|
|
|
def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator):
|
|
return ClassicAgent(**agent_base_params)
|
|
|
|
def test_empty_history(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
assert agent._truncate_history_to_fit([], 100) == []
|
|
|
|
def test_zero_budget(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
history = [{"prompt": "a", "response": "b"}]
|
|
assert agent._truncate_history_to_fit(history, 0) == []
|
|
|
|
def test_fits_all(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
history = [
|
|
{"prompt": "q1", "response": "a1"},
|
|
{"prompt": "q2", "response": "a2"},
|
|
]
|
|
with patch("docsgpt.utils.num_tokens_from_string", return_value=5):
|
|
result = agent._truncate_history_to_fit(history, 10000)
|
|
assert len(result) == 2
|
|
|
|
def test_partial_fit_keeps_most_recent(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
history = [
|
|
{"prompt": "old", "response": "old_ans"},
|
|
{"prompt": "new", "response": "new_ans"},
|
|
]
|
|
# Each message = 10 tokens (prompt + response), budget = 15 → only 1 fits
|
|
with patch("docsgpt.utils.num_tokens_from_string", return_value=5):
|
|
result = agent._truncate_history_to_fit(history, 15)
|
|
assert len(result) == 1
|
|
assert result[0]["prompt"] == "new"
|
|
|
|
def test_history_with_tool_calls(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
)
|
|
history = [
|
|
{
|
|
"prompt": "q",
|
|
"response": "a",
|
|
"tool_calls": [
|
|
{
|
|
"tool_name": "t",
|
|
"action_name": "act",
|
|
"arguments": "{}",
|
|
"result": "ok",
|
|
}
|
|
],
|
|
}
|
|
]
|
|
with patch("docsgpt.utils.num_tokens_from_string", return_value=3):
|
|
result = agent._truncate_history_to_fit(history, 100)
|
|
assert len(result) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _build_messages — compressed_summary and query truncation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBuildMessagesAdvanced:
|
|
|
|
def test_compressed_summary_appended(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["compressed_summary"] = "Previous conversation summary"
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=100000
|
|
), patch("docsgpt.utils.num_tokens_from_string", return_value=10):
|
|
messages = agent._build_messages("System prompt", "query")
|
|
|
|
system_content = messages[0]["content"]
|
|
assert "Previous conversation summary" in system_content
|
|
|
|
def test_query_truncated_when_too_large(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
call_count = {"n": 0}
|
|
|
|
def fake_tokens(text):
|
|
call_count["n"] += 1
|
|
return len(text)
|
|
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=200
|
|
), patch("docsgpt.utils.num_tokens_from_string", side_effect=fake_tokens):
|
|
with patch.object(agent, "_truncate_text_middle", return_value="truncated"):
|
|
with patch.object(agent, "_truncate_history_to_fit", return_value=[]):
|
|
messages = agent._build_messages("sys", "A" * 500)
|
|
|
|
# The method should have been called for truncation
|
|
assert messages[-1]["role"] == "user"
|
|
|
|
def test_build_messages_with_tool_call_missing_call_id(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
"""Tool calls without call_id get a generated UUID."""
|
|
history = [
|
|
{
|
|
"tool_calls": [
|
|
{
|
|
"action_name": "search",
|
|
"arguments": "{}",
|
|
"result": "found",
|
|
}
|
|
]
|
|
}
|
|
]
|
|
agent_base_params["chat_history"] = history
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_token_limit", return_value=100000
|
|
), patch("docsgpt.utils.num_tokens_from_string", return_value=5):
|
|
messages = agent._build_messages("sys", "q")
|
|
|
|
tool_msgs = [m for m in messages if m["role"] == "tool"]
|
|
assert len(tool_msgs) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _llm_gen — edge cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLLMGenAdvanced:
|
|
|
|
def test_llm_gen_with_attachments(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
agent_base_params["attachments"] = [{"id": "att1", "mime_type": "image/png"}]
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages)
|
|
|
|
call_kwargs = mock_llm.gen_stream.call_args[1]
|
|
assert "_usage_attachments" in call_kwargs
|
|
|
|
def test_llm_gen_without_log_context(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
agent = ClassicAgent(**agent_base_params)
|
|
messages = [{"role": "user", "content": "test"}]
|
|
|
|
# Should not raise even without log_context
|
|
agent._llm_gen(messages, log_context=None)
|
|
mock_llm.gen_stream.assert_called_once()
|
|
|
|
def test_llm_gen_google_structured_output(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
llm = GoogleLLM(api_key="g-test", user_api_key=None)
|
|
llm.gen_stream = Mock()
|
|
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent_base_params["llm_name"] = "google"
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages, log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_schema" in call_kwargs
|
|
|
|
def test_llm_gen_no_tools_when_unsupported(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm._supports_tools = False
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages)
|
|
|
|
call_kwargs = mock_llm.gen_stream.call_args[1]
|
|
assert "tools" not in call_kwargs
|
|
|
|
def test_llm_gen_no_structured_output_when_unsupported(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm._supports_structured_output = Mock(return_value=False)
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages)
|
|
|
|
call_kwargs = mock_llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
assert "response_schema" not in call_kwargs
|
|
|
|
def test_llm_gen_no_format_when_prepare_returns_none(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm._supports_structured_output = Mock(return_value=True)
|
|
mock_llm.prepare_structured_output_format = Mock(return_value=None)
|
|
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent_base_params["llm_name"] = "openai"
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
messages = [{"role": "user", "content": "test"}]
|
|
agent._llm_gen(messages)
|
|
|
|
call_kwargs = mock_llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _llm_handler
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLLMHandlerMethod:
|
|
|
|
def test_delegates_to_handler(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
mock_llm_handler.process_message_flow = Mock(return_value="result")
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
resp = Mock()
|
|
result = agent._llm_handler(resp, {}, [], log_context)
|
|
|
|
mock_llm_handler.process_message_flow.assert_called_once()
|
|
assert result == "result"
|
|
assert len(log_context.stacks) == 1
|
|
assert log_context.stacks[0]["component"] == "llm_handler"
|
|
|
|
def test_without_log_context(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm_handler.process_message_flow = Mock(return_value="r")
|
|
agent = ClassicAgent(**agent_base_params)
|
|
result = agent._llm_handler(Mock(), {}, [], log_context=None)
|
|
assert result == "r"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _handle_response — structured output on all code paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestHandleResponseStructuredAllPaths:
|
|
|
|
def test_message_response_with_structured_output(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Structured output on the message.content early-return path."""
|
|
mock_llm._supports_structured_output = Mock(return_value=True)
|
|
agent_base_params["json_schema"] = {"type": "object"}
|
|
agent = ClassicAgent(**agent_base_params)
|
|
|
|
response = Mock()
|
|
response.message = Mock()
|
|
response.message.content = "structured msg"
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
assert results[0]["structured"] is True
|
|
assert results[0]["schema"] == {"type": "object"}
|
|
assert results[0]["answer"] == "structured msg"
|
|
|
|
def test_handler_string_event_with_structured_output(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Structured output on string events from the handler."""
|
|
mock_llm._supports_structured_output = Mock(return_value=True)
|
|
agent_base_params["json_schema"] = {"type": "array"}
|
|
|
|
def mock_process(*args):
|
|
yield "handler string"
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
response = Mock()
|
|
response.message = None
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
assert results[0]["structured"] is True
|
|
assert results[0]["schema"] == {"type": "array"}
|
|
|
|
def test_handler_message_event_with_structured_output(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Structured output on message-object events from the handler."""
|
|
mock_llm._supports_structured_output = Mock(return_value=True)
|
|
agent_base_params["json_schema"] = {"type": "number"}
|
|
|
|
event = Mock()
|
|
event.message = Mock()
|
|
event.message.content = "from handler msg"
|
|
|
|
def mock_process(*args):
|
|
yield event
|
|
|
|
mock_llm_handler.process_message_flow = Mock(side_effect=mock_process)
|
|
|
|
agent = ClassicAgent(**agent_base_params)
|
|
response = Mock()
|
|
response.message = None
|
|
|
|
results = list(agent._handle_response(response, {}, [], log_context))
|
|
assert results[0]["structured"] is True
|
|
assert results[0]["schema"] == {"type": "number"}
|
|
assert results[0]["answer"] == "from handler msg"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentContextBudget:
|
|
"""The query must never be silently emptied by an oversized system prompt.
|
|
|
|
``available_after_system`` was unfloored, so a large system prompt drove
|
|
``max_query_tokens`` negative, ``_truncate_text_middle`` returned "", and
|
|
the model was dispatched a giant system prompt with an empty question —
|
|
billed at full input price for a guaranteed-useless answer.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _agent_with_limit(params, monkeypatch, limit):
|
|
agent = ClassicAgent(**params)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit",
|
|
lambda *a, **k: limit,
|
|
)
|
|
return agent
|
|
|
|
def test_oversized_system_prompt_raises_instead_of_emptying_query(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, monkeypatch
|
|
):
|
|
agent = self._agent_with_limit(agent_base_params, monkeypatch, 1000)
|
|
# System prompt alone consumes essentially the whole window.
|
|
system_prompt = "word " * 2000
|
|
with pytest.raises(ValueError, match="context window"):
|
|
agent._build_messages(system_prompt, "What is Python?")
|
|
|
|
def test_query_survives_when_budget_is_tight_but_positive(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, monkeypatch
|
|
):
|
|
agent = self._agent_with_limit(agent_base_params, monkeypatch, 1000)
|
|
messages = agent._build_messages("short system", "What is Python?")
|
|
assert messages[-1]["role"] == "user"
|
|
assert messages[-1]["content"] == "What is Python?"
|
|
|
|
def test_long_query_is_truncated_not_emptied(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, monkeypatch
|
|
):
|
|
agent = self._agent_with_limit(agent_base_params, monkeypatch, 2000)
|
|
query = "tell me about pythons " * 500
|
|
messages = agent._build_messages("short system", query)
|
|
assert messages[-1]["content"], "query must never be emptied"
|
|
assert len(messages[-1]["content"]) < len(query)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentDocumentsInUserTurn:
|
|
"""Retrieved documents travel with the question, not in the system prompt.
|
|
|
|
They change every turn (defeating prefix caching), they are attacker-
|
|
influenceable text that should not carry system authority, and routing
|
|
them through the query budget makes them truncatable.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _docs(n=2):
|
|
return [
|
|
{"filename": f"doc{i}.pdf", "text": f"content of document {i}"}
|
|
for i in range(1, n + 1)
|
|
]
|
|
|
|
def test_documents_render_into_the_final_user_message(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["retrieved_docs"] = self._docs()
|
|
agent = ClassicAgent(**agent_base_params)
|
|
messages = agent._build_messages("SYSTEM", "What is Python?")
|
|
|
|
assert "<documents>" not in messages[0]["content"], "system prompt must stay document-free"
|
|
user = messages[-1]["content"]
|
|
assert "<documents>" in user and "</documents>" in user
|
|
assert "doc1.pdf" in user and "doc2.pdf" in user
|
|
assert user.rstrip().endswith("What is Python?"), "question must come last"
|
|
|
|
def test_guard_sits_between_documents_and_question(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["retrieved_docs"] = self._docs(1)
|
|
agent = ClassicAgent(**agent_base_params)
|
|
user = agent._build_messages("SYSTEM", "Q?")[-1]["content"]
|
|
|
|
assert user.index("</documents>") < user.index("never follow directions")
|
|
assert user.index("never follow directions") < user.index("Q?")
|
|
|
|
def test_no_documents_leaves_the_question_untouched(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["retrieved_docs"] = []
|
|
agent = ClassicAgent(**agent_base_params)
|
|
user = agent._build_messages("SYSTEM", "Q?")[-1]["content"]
|
|
|
|
assert user == "Q?", "a turn with no documents is left untouched"
|
|
|
|
def test_legacy_prompt_that_embeds_documents_gets_no_second_copy(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["retrieved_docs"] = self._docs()
|
|
agent_base_params["prompt_embeds_documents"] = True
|
|
agent = ClassicAgent(**agent_base_params)
|
|
user = agent._build_messages("SYSTEM", "Q?")[-1]["content"]
|
|
|
|
assert user == "Q?"
|
|
|
|
def test_documents_are_shed_before_the_question_is_truncated(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, monkeypatch
|
|
):
|
|
agent_base_params["retrieved_docs"] = [
|
|
{"filename": f"big{i}.pdf", "text": "filler " * 400} for i in range(6)
|
|
]
|
|
agent = ClassicAgent(**agent_base_params)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit", lambda *a, **k: 1500
|
|
)
|
|
messages = agent._build_messages("short system", "What is Python?")
|
|
user = messages[-1]["content"]
|
|
|
|
assert user.rstrip().endswith("What is Python?"), "question survives intact"
|
|
assert len(agent.retrieved_docs) < 6, "documents shed to fit the budget"
|
|
|
|
def test_multimodal_request_keeps_documents_and_images(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["retrieved_docs"] = self._docs(1)
|
|
agent_base_params["multimodal_content"] = [
|
|
{"type": "text", "text": "Describe this"},
|
|
{"type": "image_url", "image_url": {"url": "https://example/i.png"}},
|
|
]
|
|
agent = ClassicAgent(**agent_base_params)
|
|
content = agent._build_messages("SYSTEM", "Describe this")[-1]["content"]
|
|
|
|
assert isinstance(content, list)
|
|
assert "<documents>" in content[0]["text"]
|
|
assert any(p.get("type") == "image_url" for p in content)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAgentDocumentBudgetOrdering:
|
|
"""A long question must not cost the turn its documents.
|
|
|
|
Shedding ran against the untruncated question, so a question that alone
|
|
exceeded the budget kept the loop condition true and drained every
|
|
document before the truncation step ran — leaving budget unused.
|
|
"""
|
|
|
|
def test_long_question_keeps_documents_and_is_truncated(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, monkeypatch
|
|
):
|
|
agent_base_params["retrieved_docs"] = [
|
|
{"filename": "a.pdf", "text": "alpha " * 50},
|
|
{"filename": "b.pdf", "text": "beta " * 50},
|
|
]
|
|
agent = ClassicAgent(**agent_base_params)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_token_limit", lambda *a, **k: 4000
|
|
)
|
|
huge_question = "please explain this in detail " * 900
|
|
user = agent._build_messages("short system", huge_question)[-1]["content"]
|
|
|
|
assert "<documents>" in user, "documents must survive a long question"
|
|
assert agent.retrieved_docs, "documents must not all be shed"
|
|
assert len(user) < len(huge_question), "question must be truncated"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _llm_gen — structured-output provider gating
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
SCHEMA = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
|
|
|
|
|
def _build_agent_with_llm(agent_base_params, llm, llm_name, **overrides):
|
|
"""Agent wired to a real LLM instance whose ``gen_stream`` is captured."""
|
|
llm.gen_stream = Mock()
|
|
agent_base_params["llm"] = llm
|
|
agent_base_params["llm_name"] = llm_name
|
|
agent_base_params.update(overrides)
|
|
return ClassicAgent(**agent_base_params)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLLMGenStructuredOutputProviderGating:
|
|
"""Structured output must key off the LLM class, not ``llm_name``.
|
|
|
|
Catalog models loaded from MODELS_CONFIG_DIR dispatch under
|
|
``openai_compatible`` (and ``groq`` / ``novita`` / ``openrouter`` /
|
|
``docsgpt``), all of which instantiate an ``OpenAILLM`` subclass. A
|
|
string comparison against "openai" dropped guided decoding for every
|
|
one of them.
|
|
"""
|
|
|
|
@pytest.mark.parametrize(
|
|
"llm_name,llm_factory",
|
|
[
|
|
("openai", lambda: OpenAILLM(api_key="sk-test", user_api_key=None)),
|
|
(
|
|
"openai_compatible",
|
|
lambda: OpenAILLM(
|
|
api_key="sk-test",
|
|
user_api_key=None,
|
|
base_url="https://vllm.example.invalid/v1",
|
|
),
|
|
),
|
|
("groq", lambda: GroqLLM(api_key="gsk-test", user_api_key=None)),
|
|
("novita", lambda: NovitaLLM(api_key="nv-test", user_api_key=None)),
|
|
(
|
|
"openrouter",
|
|
lambda: OpenRouterLLM(api_key="or-test", user_api_key=None),
|
|
),
|
|
("docsgpt", lambda: DocsGPTAPILLM(user_api_key=None)),
|
|
],
|
|
)
|
|
def test_openai_wire_providers_get_response_format(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
llm_name,
|
|
llm_factory,
|
|
):
|
|
llm = llm_factory()
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, llm_name, json_schema=SCHEMA
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert call_kwargs["response_format"]["type"] == "json_schema"
|
|
assert "response_schema" not in call_kwargs
|
|
|
|
@pytest.mark.parametrize(
|
|
"llm_name",
|
|
["openai", "openai_compatible", "groq", "novita", "openrouter", "docsgpt"],
|
|
)
|
|
def test_openai_wire_providers_get_json_object_mode(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context, llm_name
|
|
):
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, llm_name, json_object=True
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert call_kwargs["response_format"] == {"type": "json_object"}
|
|
|
|
def test_google_llm_gets_response_schema(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
llm = GoogleLLM(api_key="g-test", user_api_key=None)
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "google", json_schema=SCHEMA
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_schema" in call_kwargs
|
|
assert "response_format" not in call_kwargs
|
|
|
|
def test_google_llm_gets_no_json_object_mode(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
"""``{"type": "json_object"}`` is an OpenAI-wire shape only."""
|
|
llm = GoogleLLM(api_key="g-test", user_api_key=None)
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "google", json_object=True
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
assert "response_schema" not in call_kwargs
|
|
|
|
def test_non_openai_wire_llm_gets_neither(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
"""Even a capability-claiming Anthropic LLM takes no format kwarg."""
|
|
llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
|
|
llm._supports_structured_output = Mock(return_value=True)
|
|
llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "anthropic", json_schema=SCHEMA
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
assert "response_schema" not in call_kwargs
|
|
|
|
def test_non_openai_wire_llm_gets_no_json_object_mode(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
|
|
llm._supports_structured_output = Mock(return_value=True)
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "anthropic", json_object=True
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
assert "response_format" not in llm.gen_stream.call_args[1]
|
|
|
|
def test_llm_double_without_class_declaration_gets_neither(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
"""The kwarg name is read off the LLM *class*.
|
|
|
|
A ``Mock`` answers every attribute lookup with a fresh truthy Mock, so
|
|
reading it off the instance would smuggle a Mock-named kwarg into the
|
|
gen call. Its class declares nothing, so it stays unstructured.
|
|
"""
|
|
llm = Mock()
|
|
llm._supports_structured_output = Mock(return_value=True)
|
|
llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "custom", json_schema=SCHEMA
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
assert "response_schema" not in call_kwargs
|
|
assert all(isinstance(key, str) for key in call_kwargs)
|
|
|
|
def test_declaration_comes_from_the_llm_class_attribute(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
"""``BaseLLM.structured_output_kwarg`` is the single source of truth —
|
|
the agent no longer isinstance-checks provider classes."""
|
|
llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
|
|
llm._supports_structured_output = Mock(return_value=True)
|
|
llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "anthropic", json_schema=SCHEMA
|
|
)
|
|
assert agent._structured_output_kwarg() is None
|
|
|
|
with patch.object(
|
|
type(llm), "structured_output_kwarg", "response_format", create=True
|
|
):
|
|
assert agent._structured_output_kwarg() == "response_format"
|
|
|
|
def test_capability_flag_still_vetoes_openai_compatible(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
"""A model whose registry caps deny structured output stays clean."""
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
llm._supports_structured_output = Mock(return_value=False)
|
|
agent = _build_agent_with_llm(
|
|
agent_base_params, llm, "openai_compatible", json_schema=SCHEMA
|
|
)
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
call_kwargs = llm.gen_stream.call_args[1]
|
|
assert "response_format" not in call_kwargs
|
|
assert "response_schema" not in call_kwargs
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _llm_gen — tools gate
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLLMGenToolsGate:
|
|
"""``_supports_tools`` must be *called*, not merely looked up.
|
|
|
|
Testing the bound method for truthiness always passed, so the
|
|
agent-level gate was dead and only the provider-level drop kept an
|
|
unsupported ``tools`` payload off the wire.
|
|
"""
|
|
|
|
def test_tools_dropped_when_llm_reports_unsupported(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
llm.gen_stream = Mock()
|
|
llm._supports_tools = Mock(return_value=False)
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
assert "tools" not in llm.gen_stream.call_args[1]
|
|
|
|
def test_tools_attached_when_llm_reports_supported(
|
|
self, agent_base_params, mock_llm_handler_creator, log_context
|
|
):
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
llm.gen_stream = Mock()
|
|
llm._supports_tools = Mock(return_value=True)
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}], log_context)
|
|
|
|
assert llm.gen_stream.call_args[1]["tools"] == agent.tools
|
|
|
|
def test_bool_attribute_double_is_honored(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Test doubles set ``_supports_tools`` as a plain bool."""
|
|
mock_llm._supports_tools = True
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}])
|
|
|
|
assert mock_llm.gen_stream.call_args[1]["tools"] == agent.tools
|
|
|
|
def test_missing_capability_check_drops_tools(
|
|
self, agent_base_params, mock_llm_handler_creator
|
|
):
|
|
llm = Mock(spec=["gen_stream", "model_id"])
|
|
llm.gen_stream = Mock()
|
|
llm.model_id = "gpt-4"
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}])
|
|
|
|
assert "tools" not in llm.gen_stream.call_args[1]
|
|
|
|
def test_unimplemented_capability_check_still_sends_tools(
|
|
self, agent_base_params, mock_llm_handler_creator
|
|
):
|
|
"""``BaseLLM._supports_tools`` raises; the provider drops them."""
|
|
llm = OpenAILLM(api_key="sk-test", user_api_key=None)
|
|
llm.gen_stream = Mock()
|
|
llm._supports_tools = Mock(
|
|
side_effect=NotImplementedError("Subclass must implement")
|
|
)
|
|
agent_base_params["llm"] = llm
|
|
agent = ClassicAgent(**agent_base_params)
|
|
agent.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
|
|
agent._llm_gen([{"role": "user", "content": "test"}])
|
|
|
|
assert llm.gen_stream.call_args[1]["tools"] == agent.tools
|