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.
1615 lines
56 KiB
Python
1615 lines
56 KiB
Python
"""Comprehensive tests for docsgpt/agents/research_agent.py
|
|
|
|
Covers: CitationManager, ResearchAgent (init, budget, timeout, phases:
|
|
clarification, planning, research step, synthesis, _extract_text,
|
|
JSON parsing, tool setup, is_follow_up).
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from docsgpt.agents.research_agent import (
|
|
COMPLEXITY_CAPS,
|
|
CitationManager,
|
|
ResearchAgent,
|
|
DEFAULT_MAX_STEPS,
|
|
DEFAULT_MAX_SUB_ITERATIONS,
|
|
DEFAULT_TIMEOUT_SECONDS,
|
|
DEFAULT_TOKEN_BUDGET,
|
|
DEFAULT_PARALLEL_WORKERS,
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# CitationManager
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCitationManager:
|
|
|
|
def test_add_returns_citation_number(self):
|
|
cm = CitationManager()
|
|
num = cm.add({"source": "s1", "title": "T1"})
|
|
assert num == 1
|
|
|
|
def test_add_deduplicates(self):
|
|
cm = CitationManager()
|
|
n1 = cm.add({"source": "s1", "title": "T1"})
|
|
n2 = cm.add({"source": "s1", "title": "T1"})
|
|
assert n1 == n2
|
|
assert len(cm.citations) == 1
|
|
|
|
def test_add_different_sources(self):
|
|
cm = CitationManager()
|
|
n1 = cm.add({"source": "s1", "title": "T1"})
|
|
n2 = cm.add({"source": "s2", "title": "T2"})
|
|
assert n1 != n2
|
|
assert len(cm.citations) == 2
|
|
|
|
def test_add_same_source_different_title(self):
|
|
cm = CitationManager()
|
|
n1 = cm.add({"source": "s1", "title": "T1"})
|
|
n2 = cm.add({"source": "s1", "title": "T2"})
|
|
assert n1 != n2
|
|
|
|
def test_add_docs_returns_mapping(self):
|
|
cm = CitationManager()
|
|
docs = [
|
|
{"source": "s1", "title": "Doc A"},
|
|
{"source": "s2", "title": "Doc B"},
|
|
]
|
|
text = cm.add_docs(docs)
|
|
assert "[1] Doc A" in text
|
|
assert "[2] Doc B" in text
|
|
|
|
def test_add_docs_deduplication(self):
|
|
cm = CitationManager()
|
|
docs = [
|
|
{"source": "s1", "title": "Doc A"},
|
|
{"source": "s1", "title": "Doc A"},
|
|
]
|
|
text = cm.add_docs(docs)
|
|
assert text.count("[1]") == 2
|
|
|
|
def test_format_references(self):
|
|
cm = CitationManager()
|
|
cm.add({
|
|
"source": "http://example.com",
|
|
"title": "Example",
|
|
"filename": "ex.md",
|
|
})
|
|
refs = cm.format_references()
|
|
assert "[1]" in refs
|
|
assert "ex.md" in refs
|
|
assert "http://example.com" in refs
|
|
|
|
def test_format_references_uses_title_when_no_filename(self):
|
|
cm = CitationManager()
|
|
cm.add({"source": "http://example.com", "title": "My Title"})
|
|
refs = cm.format_references()
|
|
assert "My Title" in refs
|
|
|
|
def test_format_references_empty(self):
|
|
cm = CitationManager()
|
|
assert "No sources" in cm.format_references()
|
|
|
|
def test_get_all_docs(self):
|
|
cm = CitationManager()
|
|
cm.add({"source": "s1", "title": "T1"})
|
|
cm.add({"source": "s2", "title": "T2"})
|
|
docs = cm.get_all_docs()
|
|
assert len(docs) == 2
|
|
|
|
def test_format_references_sorted(self):
|
|
cm = CitationManager()
|
|
cm.add({"source": "s1", "title": "A"})
|
|
cm.add({"source": "s2", "title": "B"})
|
|
cm.add({"source": "s3", "title": "C"})
|
|
refs = cm.format_references()
|
|
lines = refs.strip().split("\n")
|
|
assert lines[0].startswith("[1]")
|
|
assert lines[1].startswith("[2]")
|
|
assert lines[2].startswith("[3]")
|
|
|
|
|
|
# =====================================================================
|
|
# ResearchAgent Init & Constants
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentInit:
|
|
|
|
def test_initialization(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(**agent_base_params)
|
|
assert isinstance(agent, ResearchAgent)
|
|
assert agent.max_steps == DEFAULT_MAX_STEPS
|
|
assert agent.timeout_seconds == DEFAULT_TIMEOUT_SECONDS
|
|
assert agent.token_budget == DEFAULT_TOKEN_BUDGET
|
|
assert agent.max_sub_iterations == DEFAULT_MAX_SUB_ITERATIONS
|
|
assert agent.parallel_workers == DEFAULT_PARALLEL_WORKERS
|
|
assert agent.retriever_config == {}
|
|
|
|
def test_custom_budget(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(
|
|
max_steps=3,
|
|
timeout_seconds=60,
|
|
token_budget=50_000,
|
|
max_sub_iterations=2,
|
|
parallel_workers=1,
|
|
**agent_base_params,
|
|
)
|
|
assert agent.max_steps == 3
|
|
assert agent.timeout_seconds == 60
|
|
assert agent.token_budget == 50_000
|
|
assert agent.max_sub_iterations == 2
|
|
assert agent.parallel_workers == 1
|
|
|
|
def test_with_retriever_config(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
rc = {"source": {"active_docs": ["abc"]}}
|
|
agent = ResearchAgent(retriever_config=rc, **agent_base_params)
|
|
assert agent.retriever_config == rc
|
|
|
|
def test_constants(self):
|
|
assert DEFAULT_MAX_STEPS == 6
|
|
assert DEFAULT_MAX_SUB_ITERATIONS == 5
|
|
assert DEFAULT_TIMEOUT_SECONDS == 300
|
|
assert DEFAULT_TOKEN_BUDGET == 100_000
|
|
assert DEFAULT_PARALLEL_WORKERS == 3
|
|
|
|
def test_complexity_caps(self):
|
|
assert COMPLEXITY_CAPS["simple"] == 2
|
|
assert COMPLEXITY_CAPS["moderate"] == 4
|
|
assert COMPLEXITY_CAPS["complex"] == 6
|
|
|
|
|
|
# =====================================================================
|
|
# Budget & Timeout
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentBudget:
|
|
|
|
def _make_agent(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, **kwargs
|
|
):
|
|
return ResearchAgent(**kwargs, **agent_base_params)
|
|
|
|
def test_timeout_detection(
|
|
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,
|
|
timeout_seconds=0,
|
|
)
|
|
agent._start_time = time.monotonic() - 1
|
|
assert agent._is_timed_out() is True
|
|
|
|
def test_not_timed_out(
|
|
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,
|
|
timeout_seconds=300,
|
|
)
|
|
agent._start_time = time.monotonic()
|
|
assert agent._is_timed_out() is False
|
|
|
|
def test_token_budget_tracking(
|
|
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,
|
|
token_budget=1000,
|
|
)
|
|
agent._track_tokens(500)
|
|
assert agent._budget_remaining() == 500
|
|
assert agent._is_over_budget() is False
|
|
|
|
agent._track_tokens(500)
|
|
assert agent._budget_remaining() == 0
|
|
assert agent._is_over_budget() is True
|
|
|
|
def test_over_budget_returns_zero_remaining(
|
|
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,
|
|
token_budget=100,
|
|
)
|
|
agent._track_tokens(200)
|
|
assert agent._budget_remaining() == 0
|
|
|
|
def test_snapshot_llm_tokens_returns_delta(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
agent = self._make_agent(
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
)
|
|
mock_llm.token_usage = {"prompt_tokens": 100, "generated_tokens": 50}
|
|
|
|
delta1 = agent._snapshot_llm_tokens()
|
|
assert delta1 == 150
|
|
|
|
mock_llm.token_usage = {"prompt_tokens": 200, "generated_tokens": 100}
|
|
delta2 = agent._snapshot_llm_tokens()
|
|
assert delta2 == 150
|
|
|
|
def test_elapsed(
|
|
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,
|
|
)
|
|
agent._start_time = time.monotonic() - 1.5
|
|
elapsed = agent._elapsed()
|
|
assert elapsed >= 1.0
|
|
|
|
|
|
# =====================================================================
|
|
# Clarification Phase
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentClarification:
|
|
|
|
def test_is_follow_up_no_history(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(**agent_base_params)
|
|
assert agent._is_follow_up() is False
|
|
|
|
def test_is_follow_up_with_clarification_metadata(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["chat_history"] = [
|
|
{
|
|
"prompt": "What?",
|
|
"response": "Clarify",
|
|
"metadata": {"is_clarification": True},
|
|
},
|
|
]
|
|
agent = ResearchAgent(**agent_base_params)
|
|
assert agent._is_follow_up() is True
|
|
|
|
def test_is_follow_up_without_metadata(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["chat_history"] = [
|
|
{"prompt": "What?", "response": "Normal answer"},
|
|
]
|
|
agent = ResearchAgent(**agent_base_params)
|
|
assert agent._is_follow_up() is False
|
|
|
|
def test_is_follow_up_empty_metadata(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent_base_params["chat_history"] = [
|
|
{"prompt": "What?", "response": "X", "metadata": {}},
|
|
]
|
|
agent = ResearchAgent(**agent_base_params)
|
|
assert agent._is_follow_up() is False
|
|
|
|
def test_clarification_returns_none_on_no_clarification_needed(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
response = Mock()
|
|
response.choices = [Mock()]
|
|
response.choices[0].message = Mock()
|
|
response.choices[0].message.content = json.dumps(
|
|
{"needs_clarification": False, "reason": "Clear enough"}
|
|
)
|
|
mock_llm.gen = Mock(return_value=response)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("What is Python?")
|
|
assert result is None
|
|
|
|
def test_clarification_returns_questions(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
clarification_json = json.dumps({
|
|
"needs_clarification": True,
|
|
"questions": ["Which version?", "What context?"],
|
|
})
|
|
mock_llm.gen = Mock(return_value=clarification_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("Tell me about it")
|
|
assert result is not None
|
|
assert "Which version?" in result
|
|
assert "What context?" in result
|
|
assert "1." in result
|
|
assert "2." in result
|
|
|
|
def test_clarification_limits_questions_to_three(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
clarification_json = json.dumps({
|
|
"needs_clarification": True,
|
|
"questions": ["q1", "q2", "q3", "q4", "q5"],
|
|
})
|
|
mock_llm.gen = Mock(return_value=clarification_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("complex question")
|
|
# Should only show 3 questions
|
|
assert "3." in result
|
|
assert "4." not in result
|
|
|
|
def test_clarification_returns_none_on_empty_questions(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
clarification_json = json.dumps({
|
|
"needs_clarification": True,
|
|
"questions": [],
|
|
})
|
|
mock_llm.gen = Mock(return_value=clarification_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("question")
|
|
assert result is None
|
|
|
|
def test_clarification_returns_none_on_llm_error(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm.gen = Mock(side_effect=Exception("LLM error"))
|
|
mock_llm.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("question")
|
|
assert result is None
|
|
|
|
|
|
# =====================================================================
|
|
# Planning Phase
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentPlanning:
|
|
|
|
def test_planning_returns_steps_and_complexity(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
plan_json = json.dumps({
|
|
"complexity": "moderate",
|
|
"steps": [
|
|
{"query": "sub-question 1", "rationale": "reason 1"},
|
|
{"query": "sub-question 2", "rationale": "reason 2"},
|
|
],
|
|
})
|
|
mock_llm.gen = Mock(return_value=plan_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("Compare A and B")
|
|
|
|
assert complexity == "moderate"
|
|
assert len(steps) == 2
|
|
assert steps[0]["query"] == "sub-question 1"
|
|
|
|
def test_planning_caps_steps_by_complexity(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
plan_json = json.dumps({
|
|
"complexity": "simple",
|
|
"steps": [{"query": f"q{i}", "rationale": f"r{i}"} for i in range(10)],
|
|
})
|
|
response = Mock()
|
|
response.choices = [Mock()]
|
|
response.choices[0].message = Mock()
|
|
response.choices[0].message.content = plan_json
|
|
mock_llm.gen = Mock(return_value=response)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("Simple question")
|
|
|
|
assert complexity == "simple"
|
|
assert len(steps) <= 2
|
|
|
|
def test_planning_caps_steps_for_complex(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
plan_json = json.dumps({
|
|
"complexity": "complex",
|
|
"steps": [{"query": f"q{i}", "rationale": f"r{i}"} for i in range(10)],
|
|
})
|
|
mock_llm.gen = Mock(return_value=plan_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("Complex analysis")
|
|
|
|
assert complexity == "complex"
|
|
assert len(steps) <= 6
|
|
|
|
def test_planning_fallback_on_error(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
mock_llm.gen = Mock(side_effect=Exception("LLM down"))
|
|
mock_llm.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("Anything")
|
|
|
|
assert complexity == "simple"
|
|
assert len(steps) == 1
|
|
assert steps[0]["query"] == "Anything"
|
|
|
|
def test_planning_list_response(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
plan_json = json.dumps([
|
|
{"query": "q1", "rationale": "r1"},
|
|
{"query": "q2", "rationale": "r2"},
|
|
])
|
|
mock_llm.gen = Mock(return_value=plan_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("question")
|
|
|
|
assert complexity == "moderate"
|
|
assert len(steps) == 2
|
|
|
|
|
|
# =====================================================================
|
|
# Extract Text
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentExtractText:
|
|
|
|
def _make_agent(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
return ResearchAgent(**agent_base_params)
|
|
|
|
def test_extract_from_string(
|
|
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._extract_text("hello") == "hello"
|
|
|
|
def test_extract_from_openai_response(
|
|
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
|
|
)
|
|
response = Mock()
|
|
response.choices = [Mock()]
|
|
response.choices[0].message = Mock()
|
|
response.choices[0].message.content = "OpenAI content"
|
|
response.message = None
|
|
response.content = None
|
|
assert agent._extract_text(response) == "OpenAI content"
|
|
|
|
def test_extract_from_anthropic_response(
|
|
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
|
|
)
|
|
text_block = Mock()
|
|
text_block.text = "Anthropic content"
|
|
response = Mock()
|
|
response.content = [text_block]
|
|
response.message = None
|
|
response.choices = None
|
|
assert agent._extract_text(response) == "Anthropic content"
|
|
|
|
def test_extract_from_message_content(
|
|
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
|
|
)
|
|
response = Mock()
|
|
response.message = Mock()
|
|
response.message.content = "From message"
|
|
assert agent._extract_text(response) == "From message"
|
|
|
|
def test_extract_from_none(
|
|
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._extract_text(None) == ""
|
|
|
|
|
|
# =====================================================================
|
|
# Parse JSON
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentParseJson:
|
|
|
|
def _make_agent(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
return ResearchAgent(**agent_base_params)
|
|
|
|
def test_parse_plan_direct_json(
|
|
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
|
|
)
|
|
text = '{"steps": [{"query": "q1"}], "complexity": "simple"}'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, dict)
|
|
assert len(result["steps"]) == 1
|
|
|
|
def test_parse_plan_list(
|
|
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
|
|
)
|
|
text = '[{"query": "q1"}]'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
|
|
def test_parse_plan_from_code_fence(
|
|
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
|
|
)
|
|
text = 'Here is the plan:\n```json\n{"steps": [{"query": "q1"}]}\n```'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, dict)
|
|
|
|
def test_parse_plan_from_plain_code_fence(
|
|
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
|
|
)
|
|
text = 'Result:\n```\n{"steps": [{"query": "q1"}]}\n```'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, dict)
|
|
|
|
def test_parse_plan_embedded_json_object(
|
|
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
|
|
)
|
|
text = 'Here is the plan: {"steps": [{"query": "q1"}]} end.'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, dict)
|
|
|
|
def test_parse_plan_invalid_returns_empty(
|
|
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
|
|
)
|
|
result = agent._parse_plan_json("not json at all")
|
|
assert result == []
|
|
|
|
def test_parse_clarification_json(
|
|
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
|
|
)
|
|
text = '{"needs_clarification": false, "reason": "clear"}'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result["needs_clarification"] is False
|
|
|
|
def test_parse_clarification_json_from_code_fence(
|
|
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
|
|
)
|
|
text = '```json\n{"needs_clarification": true, "questions": ["q1"]}\n```'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result["needs_clarification"] is True
|
|
|
|
def test_parse_clarification_embedded_json(
|
|
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
|
|
)
|
|
text = 'Here: {"needs_clarification": true, "questions": ["q1"]} done.'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result["needs_clarification"] is True
|
|
|
|
def test_parse_clarification_json_invalid(
|
|
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
|
|
)
|
|
result = agent._parse_clarification_json("not json")
|
|
assert result is None
|
|
|
|
|
|
# =====================================================================
|
|
# Tool Setup
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentToolSetup:
|
|
|
|
def test_setup_tools_includes_think_and_internal(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(
|
|
retriever_config={
|
|
"source": {"active_docs": ["abc"]},
|
|
"retriever_name": "classic",
|
|
},
|
|
**agent_base_params,
|
|
)
|
|
|
|
with patch.object(
|
|
agent.tool_executor, "get_tools", return_value={}
|
|
), patch(
|
|
"docsgpt.agents.research_agent.add_internal_search_tool"
|
|
) as mock_add:
|
|
tools = agent._setup_tools()
|
|
mock_add.assert_called_once()
|
|
assert "think" in tools
|
|
|
|
def test_setup_tools_no_retriever_config(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
with patch.object(
|
|
agent.tool_executor, "get_tools", return_value={}
|
|
), patch(
|
|
"docsgpt.agents.research_agent.add_internal_search_tool"
|
|
) as mock_add:
|
|
tools = agent._setup_tools()
|
|
mock_add.assert_called_once()
|
|
assert "think" in tools
|
|
|
|
|
|
# =====================================================================
|
|
# Collect Step Sources
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCollectStepSources:
|
|
|
|
def test_collects_from_internal_search_tool(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
mock_tool = Mock()
|
|
mock_tool.retrieved_docs = [
|
|
{"source": "s1", "title": "T1"},
|
|
{"source": "s2", "title": "T2"},
|
|
]
|
|
|
|
cache_key = f"internal_search:internal:{agent.user or ''}"
|
|
agent.tool_executor._loaded_tools[cache_key] = mock_tool
|
|
|
|
agent._collect_step_sources()
|
|
|
|
assert len(agent.citations.citations) == 2
|
|
|
|
def test_no_tool_no_error(
|
|
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
|
|
):
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._collect_step_sources()
|
|
assert len(agent.citations.citations) == 0
|
|
|
|
|
|
# =====================================================================
|
|
# _gen_inner (full orchestration tests)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGenInner:
|
|
|
|
def test_gen_inner_clarification_path(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""When clarification is needed, _gen_inner yields clarification output and returns."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
with patch.object(agent, "_is_follow_up", return_value=False), \
|
|
patch.object(agent, "_clarification_phase", return_value="Please clarify:\n1. Which version?"), \
|
|
patch.object(agent, "_setup_tools", return_value={}):
|
|
events = list(agent._gen_inner("ambiguous question", log_context))
|
|
|
|
# Should have: metadata, answer, sources, tool_calls
|
|
meta_events = [e for e in events if isinstance(e, dict) and "metadata" in e]
|
|
assert len(meta_events) == 1
|
|
assert meta_events[0]["metadata"]["is_clarification"] is True
|
|
|
|
answer_events = [e for e in events if isinstance(e, dict) and "answer" in e]
|
|
assert len(answer_events) == 1
|
|
assert "Please clarify" in answer_events[0]["answer"]
|
|
|
|
source_events = [e for e in events if isinstance(e, dict) and "sources" in e]
|
|
assert len(source_events) == 1
|
|
assert source_events[0]["sources"] == []
|
|
|
|
def test_gen_inner_skips_clarification_on_follow_up(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""When user is responding to clarification, skip clarification phase."""
|
|
agent_base_params["chat_history"] = [
|
|
{"prompt": "What?", "response": "clarify", "metadata": {"is_clarification": True}},
|
|
]
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
plan_steps = [{"query": "test query", "rationale": "direct"}]
|
|
|
|
with patch.object(agent, "_setup_tools", return_value={}), \
|
|
patch.object(agent, "_planning_phase", return_value=(plan_steps, "simple")), \
|
|
patch.object(agent, "_research_step", return_value="findings here"), \
|
|
patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "result"}])), \
|
|
patch.object(agent, "_get_truncated_tool_calls", return_value=[]):
|
|
events = list(agent._gen_inner("Python 3.10", log_context))
|
|
|
|
# Should NOT have clarification metadata
|
|
meta_events = [e for e in events if isinstance(e, dict) and e.get("metadata", {}).get("is_clarification")]
|
|
assert len(meta_events) == 0
|
|
|
|
# Should have planning event
|
|
plan_events = [e for e in events if isinstance(e, dict) and e.get("type") == "research_plan"]
|
|
assert len(plan_events) == 1
|
|
|
|
def test_gen_inner_empty_plan_fallback(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""When planning returns no steps, _gen_inner uses a fallback single step."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
with patch.object(agent, "_setup_tools", return_value={}), \
|
|
patch.object(agent, "_is_follow_up", return_value=True), \
|
|
patch.object(agent, "_planning_phase", return_value=([], "moderate")), \
|
|
patch.object(agent, "_research_step", return_value="direct findings"), \
|
|
patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "done"}])), \
|
|
patch.object(agent, "_get_truncated_tool_calls", return_value=[]):
|
|
events = list(agent._gen_inner("What is X?", log_context))
|
|
|
|
plan_events = [e for e in events if isinstance(e, dict) and e.get("type") == "research_plan"]
|
|
assert len(plan_events) == 1
|
|
# Fallback plan should have one step with the original query
|
|
assert plan_events[0]["data"]["steps"][0]["query"] == "What is X?"
|
|
assert plan_events[0]["data"]["complexity"] == "simple"
|
|
|
|
def test_gen_inner_timeout_during_research(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Timeout during research steps stops early and proceeds to synthesis."""
|
|
agent = ResearchAgent(timeout_seconds=0, **agent_base_params)
|
|
|
|
plan_steps = [
|
|
{"query": "step1", "rationale": "r1"},
|
|
{"query": "step2", "rationale": "r2"},
|
|
]
|
|
|
|
with patch.object(agent, "_setup_tools", return_value={}), \
|
|
patch.object(agent, "_is_follow_up", return_value=True), \
|
|
patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")):
|
|
# Set start time in the past to trigger timeout
|
|
agent._start_time = time.monotonic() - 1
|
|
|
|
with patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "partial"}])), \
|
|
patch.object(agent, "_get_truncated_tool_calls", return_value=[]):
|
|
events = list(agent._gen_inner("question", log_context))
|
|
|
|
# No research progress events with status "researching" expected (timed out before any step)
|
|
researching = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "researching"
|
|
]
|
|
assert len(researching) == 0
|
|
|
|
# Should still have synthesis event
|
|
synth = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "synthesizing"
|
|
]
|
|
assert len(synth) == 1
|
|
|
|
def test_gen_inner_budget_exhausted_during_research(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Token budget exhaustion during research stops early."""
|
|
agent = ResearchAgent(token_budget=10, **agent_base_params)
|
|
|
|
plan_steps = [
|
|
{"query": "step1", "rationale": "r1"},
|
|
{"query": "step2", "rationale": "r2"},
|
|
]
|
|
|
|
with patch.object(agent, "_setup_tools", return_value={}), \
|
|
patch.object(agent, "_is_follow_up", return_value=True), \
|
|
patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")):
|
|
agent._start_time = time.monotonic()
|
|
agent._tokens_used = 100 # Over budget
|
|
|
|
with patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "partial"}])), \
|
|
patch.object(agent, "_get_truncated_tool_calls", return_value=[]):
|
|
events = list(agent._gen_inner("question", log_context))
|
|
|
|
researching = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "researching"
|
|
]
|
|
assert len(researching) == 0
|
|
|
|
def test_gen_inner_full_flow(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Full flow: plan, research multiple steps, synthesize."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
plan_steps = [
|
|
{"query": "step1", "rationale": "r1"},
|
|
{"query": "step2", "rationale": "r2"},
|
|
]
|
|
|
|
with patch.object(agent, "_setup_tools", return_value={}), \
|
|
patch.object(agent, "_is_follow_up", return_value=True), \
|
|
patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")), \
|
|
patch.object(agent, "_research_step", side_effect=["report1", "report2"]), \
|
|
patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "final report"}])), \
|
|
patch.object(agent, "_get_truncated_tool_calls", return_value=[{"tool": "search"}]):
|
|
events = list(agent._gen_inner("Compare A and B", log_context))
|
|
|
|
# Planning event
|
|
plan_events = [e for e in events if isinstance(e, dict) and e.get("type") == "research_plan"]
|
|
assert len(plan_events) == 1
|
|
|
|
# Research progress events: 2 researching + 2 complete
|
|
researching = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "researching"
|
|
]
|
|
assert len(researching) == 2
|
|
|
|
complete = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "complete"
|
|
]
|
|
assert len(complete) == 2
|
|
|
|
# Synthesis event
|
|
synth = [
|
|
e for e in events
|
|
if isinstance(e, dict) and e.get("type") == "research_progress"
|
|
and e.get("data", {}).get("status") == "synthesizing"
|
|
]
|
|
assert len(synth) == 1
|
|
|
|
# Sources and tool_calls events
|
|
source_events = [e for e in events if isinstance(e, dict) and "sources" in e]
|
|
assert len(source_events) == 1
|
|
|
|
tc_events = [e for e in events if isinstance(e, dict) and "tool_calls" in e]
|
|
assert len(tc_events) == 1
|
|
|
|
|
|
# =====================================================================
|
|
# _synthesis_phase
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSynthesisPhase:
|
|
|
|
def test_synthesis_phase_builds_correct_prompt(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Synthesis phase constructs prompt from plan and findings."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
agent.citations.add({"source": "s1", "title": "T1", "filename": "f1.md"})
|
|
|
|
plan = [
|
|
{"query": "q1", "rationale": "reason1"},
|
|
{"query": "q2", "rationale": "reason2"},
|
|
]
|
|
reports = [
|
|
{"step": plan[0], "content": "Found X"},
|
|
{"step": plan[1], "content": "Found Y"},
|
|
]
|
|
|
|
mock_llm.gen_stream = Mock(return_value=iter(["chunk1", "chunk2"]))
|
|
|
|
with patch.object(agent, "_handle_response", return_value=iter([
|
|
{"answer": "Synthesized report"},
|
|
])):
|
|
events = list(agent._synthesis_phase(
|
|
"test question", plan, reports, {}, log_context
|
|
))
|
|
|
|
answer_events = [e for e in events if isinstance(e, dict) and "answer" in e]
|
|
assert len(answer_events) == 1
|
|
|
|
# Verify gen_stream was called
|
|
mock_llm.gen_stream.assert_called_once()
|
|
call_kwargs = mock_llm.gen_stream.call_args
|
|
messages = call_kwargs[1]["messages"] if "messages" in call_kwargs[1] else call_kwargs[0][1] if len(call_kwargs[0]) > 1 else None
|
|
if messages is None:
|
|
messages = call_kwargs.kwargs.get("messages", call_kwargs.args[1] if len(call_kwargs.args) > 1 else [])
|
|
|
|
def test_synthesis_phase_with_empty_reports(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
log_context,
|
|
):
|
|
"""Synthesis handles empty reports."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
|
|
mock_llm.gen_stream = Mock(return_value=iter([]))
|
|
|
|
with patch.object(agent, "_handle_response", return_value=iter([
|
|
{"answer": "No findings available."},
|
|
])):
|
|
events = list(agent._synthesis_phase(
|
|
"test question", [], [], {}, log_context
|
|
))
|
|
|
|
answer_events = [e for e in events if isinstance(e, dict) and "answer" in e]
|
|
assert len(answer_events) == 1
|
|
|
|
|
|
# =====================================================================
|
|
# _research_step and _research_step_with_executor
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchStep:
|
|
|
|
def test_research_step_no_tool_call(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""LLM returns direct answer without tool calls."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
# LLM returns a direct response
|
|
mock_response = Mock()
|
|
mock_llm.gen = Mock(return_value=mock_response)
|
|
|
|
from docsgpt.llm.handlers.base import LLMResponse
|
|
parsed = LLMResponse(
|
|
content="Direct answer to the question",
|
|
tool_calls=[],
|
|
finish_reason="stop",
|
|
raw_response=mock_response,
|
|
)
|
|
mock_llm_handler.parse_response = Mock(return_value=parsed)
|
|
|
|
report = agent._research_step("What is Python?", {})
|
|
assert report == "Direct answer to the question"
|
|
|
|
def test_research_step_with_tool_calls(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""LLM makes a tool call, then returns final answer."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
mock_response1 = Mock()
|
|
mock_response2 = Mock()
|
|
mock_llm.gen = Mock(side_effect=[mock_response1, mock_response2])
|
|
|
|
from docsgpt.llm.handlers.base import LLMResponse, ToolCall
|
|
|
|
tool_call = ToolCall(id="tc1", name="internal__search", arguments={"query": "python"})
|
|
parsed_with_tool = LLMResponse(
|
|
content="",
|
|
tool_calls=[tool_call],
|
|
finish_reason="tool_calls",
|
|
raw_response=mock_response1,
|
|
)
|
|
parsed_final = LLMResponse(
|
|
content="Python is a programming language.",
|
|
tool_calls=[],
|
|
finish_reason="stop",
|
|
raw_response=mock_response2,
|
|
)
|
|
mock_llm_handler.parse_response = Mock(side_effect=[parsed_with_tool, parsed_final])
|
|
|
|
# Mock tool execution
|
|
with patch.object(agent, "_execute_step_tools_with_refinement",
|
|
return_value=([], False)):
|
|
report = agent._research_step("What is Python?", {})
|
|
assert report == "Python is a programming language."
|
|
|
|
def test_research_step_timeout_mid_iteration(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Research step times out and returns summary."""
|
|
agent = ResearchAgent(timeout_seconds=0, **agent_base_params)
|
|
agent._start_time = time.monotonic() - 1 # Already timed out
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
# Summary response when max iterations hit
|
|
mock_llm.gen = Mock(return_value="Summary of findings")
|
|
|
|
report = agent._research_step("query", {})
|
|
assert "Summary" in report or "completed" in report
|
|
|
|
def test_research_step_budget_exhausted(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Research step hits token budget and returns summary."""
|
|
agent = ResearchAgent(token_budget=10, **agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
agent._tokens_used = 100 # Over budget
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
mock_llm.gen = Mock(return_value="Budget summary")
|
|
|
|
report = agent._research_step("query", {})
|
|
assert "Budget summary" in report or "completed" in report
|
|
|
|
def test_research_step_llm_error(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Research step handles LLM error gracefully."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
mock_llm.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
|
|
# First gen call fails
|
|
mock_llm.gen = Mock(side_effect=[
|
|
Exception("LLM error"),
|
|
"Fallback summary",
|
|
])
|
|
|
|
report = agent._research_step("query", {})
|
|
assert "completed" in report or "Fallback" in report
|
|
|
|
def test_research_step_max_iterations_summary(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""After max iterations, research step asks for summary."""
|
|
agent = ResearchAgent(max_sub_iterations=1, **agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
from docsgpt.llm.handlers.base import LLMResponse, ToolCall
|
|
|
|
tool_call = ToolCall(id="tc1", name="internal__search", arguments={"query": "test"})
|
|
|
|
mock_response1 = Mock()
|
|
parsed_with_tool = LLMResponse(
|
|
content="",
|
|
tool_calls=[tool_call],
|
|
finish_reason="tool_calls",
|
|
raw_response=mock_response1,
|
|
)
|
|
mock_llm_handler.parse_response = Mock(return_value=parsed_with_tool)
|
|
|
|
# First gen returns tool call, second gen (summary request) returns text
|
|
mock_llm.gen = Mock(side_effect=[mock_response1, "Final summary after max iters"])
|
|
|
|
with patch.object(agent, "_execute_step_tools_with_refinement",
|
|
return_value=([], False)):
|
|
report = agent._research_step("query", {})
|
|
|
|
assert "Final summary" in report
|
|
|
|
def test_research_step_summary_fails_gracefully(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""When summary LLM call fails, returns fallback text."""
|
|
agent = ResearchAgent(max_sub_iterations=0, **agent_base_params)
|
|
agent._start_time = time.monotonic()
|
|
mock_llm.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
|
|
# Summary call fails
|
|
mock_llm.gen = Mock(side_effect=Exception("gen failed"))
|
|
|
|
report = agent._research_step("query", {})
|
|
assert report == "Research step completed."
|
|
|
|
|
|
# =====================================================================
|
|
# _execute_step_tools_with_refinement
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExecuteStepToolsWithRefinement:
|
|
|
|
def test_basic_tool_execution(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Tool execution appends messages correctly."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
from docsgpt.llm.handlers.base import ToolCall
|
|
|
|
call = ToolCall(id="tc1", name="internal__search", arguments={"query": "test"})
|
|
|
|
def fake_execute(tools_dict, tc, llm_class):
|
|
gen_result = ("Search result text", "tc1")
|
|
return gen_result
|
|
yield # noqa: E501 - makes it a generator
|
|
|
|
# Build a proper generator mock
|
|
def gen_execute(tools_dict, tc, llm_class):
|
|
yield {"type": "tool_call", "data": {"action_name": "search", "status": "pending"}}
|
|
return ("Search result text", "tc1")
|
|
|
|
agent.tool_executor.execute = gen_execute
|
|
mock_llm_handler.create_tool_message = Mock(
|
|
return_value={"role": "tool", "content": "Search result text"}
|
|
)
|
|
|
|
messages = [{"role": "user", "content": "query"}]
|
|
result_msgs, was_empty = agent._execute_step_tools_with_refinement(
|
|
[call], {}, messages, agent.tool_executor, False
|
|
)
|
|
|
|
assert len(result_msgs) > 1
|
|
assert any(m.get("role") == "assistant" for m in result_msgs)
|
|
assert any(m.get("role") == "tool" for m in result_msgs)
|
|
|
|
def test_empty_search_result_refinement(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""When search returns empty twice, adds refinement hint."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
from docsgpt.llm.handlers.base import ToolCall
|
|
|
|
call = ToolCall(id="tc1", name="internal__search", arguments={"query": "test"})
|
|
|
|
def gen_execute(tools_dict, tc, llm_class):
|
|
yield {"type": "tool_call", "data": {"action_name": "search", "status": "pending"}}
|
|
return ("No documents found for the query", "tc1")
|
|
|
|
agent.tool_executor.execute = gen_execute
|
|
mock_llm_handler.create_tool_message = Mock(
|
|
return_value={"role": "tool", "content": "No documents found"}
|
|
)
|
|
|
|
messages = [{"role": "user", "content": "query"}]
|
|
# First call with last_search_empty=True to trigger refinement
|
|
result_msgs, was_empty = agent._execute_step_tools_with_refinement(
|
|
[call], {}, messages, agent.tool_executor, True
|
|
)
|
|
|
|
assert was_empty is True
|
|
|
|
def test_non_search_tool_no_refinement(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_handler,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Non-search tools don't trigger empty search logic."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
|
|
from docsgpt.llm.handlers.base import ToolCall
|
|
|
|
call = ToolCall(id="tc1", name="think__think", arguments={"thought": "hmm"})
|
|
|
|
def gen_execute(tools_dict, tc, llm_class):
|
|
yield {"type": "tool_call", "data": {"action_name": "think", "status": "pending"}}
|
|
return ("Thought processed", "tc1")
|
|
|
|
agent.tool_executor.execute = gen_execute
|
|
mock_llm_handler.create_tool_message = Mock(
|
|
return_value={"role": "tool", "content": "Thought processed"}
|
|
)
|
|
|
|
messages = [{"role": "user", "content": "query"}]
|
|
result_msgs, was_empty = agent._execute_step_tools_with_refinement(
|
|
[call], {}, messages, agent.tool_executor, False
|
|
)
|
|
|
|
assert was_empty is False
|
|
|
|
|
|
# =====================================================================
|
|
# _planning_phase extended (edge cases in JSON parsing)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestPlanningPhaseExtended:
|
|
|
|
def test_planning_unknown_complexity_uses_default_cap(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Unknown complexity level uses max_steps as cap."""
|
|
plan_json = json.dumps({
|
|
"complexity": "extreme",
|
|
"steps": [{"query": f"q{i}", "rationale": f"r{i}"} for i in range(10)],
|
|
})
|
|
mock_llm.gen = Mock(return_value=plan_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
steps, complexity = agent._planning_phase("Hard question")
|
|
|
|
assert complexity == "extreme"
|
|
assert len(steps) <= agent.max_steps
|
|
|
|
def test_parse_plan_json_dict_without_steps_key(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""JSON dict without 'steps' key is not treated as a plan."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
# Returns empty list since it's a dict but no 'steps'
|
|
result = agent._parse_plan_json('{"complexity": "simple"}')
|
|
assert result == []
|
|
|
|
def test_parse_plan_json_code_fence_with_list(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""JSON list inside code fence is parsed correctly."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = 'Plan:\n```json\n[{"query": "q1", "rationale": "r1"}]\n```'
|
|
result = agent._parse_plan_json(text)
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
|
|
|
|
# =====================================================================
|
|
# Additional coverage: lines 326, 328, 335-336, 346-352, 360
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestClarificationPhaseAdditional:
|
|
|
|
def test_clarification_returns_formatted_questions(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover lines 326, 328, 335-336: clarification with questions."""
|
|
clarification_json = json.dumps({
|
|
"needs_clarification": True,
|
|
"questions": ["What version?", "Which platform?", "What scope?"],
|
|
})
|
|
mock_llm.gen = Mock(return_value=clarification_json)
|
|
mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5}
|
|
|
|
agent = ResearchAgent(**agent_base_params)
|
|
result = agent._clarification_phase("Tell me about it")
|
|
|
|
assert result is not None
|
|
assert "1." in result
|
|
assert "2." in result
|
|
assert "3." in result
|
|
assert "clarify" in result.lower()
|
|
|
|
def test_parse_clarification_json_code_fence_invalid(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover lines 346-352: invalid JSON inside code fence falls through."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = '```json\nnot valid json\n```'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result is None
|
|
|
|
def test_parse_clarification_json_embedded_invalid(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 360: embedded JSON with invalid content."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = 'Before {invalid json} after'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result is None
|
|
|
|
def test_parse_clarification_code_fence_no_closing(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 358: code fence without closing marker."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = '```json\n{"needs_clarification": true, "questions": ["q1"]}'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result is not None
|
|
assert result["needs_clarification"] is True
|
|
|
|
def test_parse_plan_json_embedded_dict_without_steps(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 463: embedded dict without 'steps' key."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = 'Here is a plan: {"key": "value"} done.'
|
|
result = agent._parse_plan_json(text)
|
|
assert result == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Coverage — additional uncovered lines: 326, 328, 335-336, 360
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResearchAgentClarificationCoverage:
|
|
|
|
def test_clarification_no_needs_clarification(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 326: data has needs_clarification=False returns None."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
# Mock _generate_response to return valid JSON without clarification
|
|
agent._generate_response = lambda *a, **kw: None
|
|
agent._extract_text = lambda r: '{"needs_clarification": false}'
|
|
agent._snapshot_llm_tokens = lambda: {}
|
|
agent._track_tokens = lambda t: None
|
|
|
|
result = agent._clarification_phase("test query")
|
|
assert result is None
|
|
|
|
def test_clarification_with_questions(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover lines 328, 335-336: questions returned as formatted response."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._generate_response = lambda *a, **kw: None
|
|
agent._extract_text = lambda r: '{"needs_clarification": true, "questions": ["What scope?", "What depth?"]}'
|
|
agent._snapshot_llm_tokens = lambda: {}
|
|
agent._track_tokens = lambda t: None
|
|
|
|
result = agent._clarification_phase("test query")
|
|
assert result is not None
|
|
assert "What scope?" in result
|
|
assert "What depth?" in result
|
|
assert "Before I begin" in result
|
|
|
|
def test_clarification_empty_questions_returns_none(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 328: needs_clarification=True but empty questions."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
agent._generate_response = lambda *a, **kw: None
|
|
agent._extract_text = lambda r: '{"needs_clarification": true, "questions": []}'
|
|
agent._snapshot_llm_tokens = lambda: {}
|
|
agent._track_tokens = lambda t: None
|
|
|
|
result = agent._clarification_phase("test query")
|
|
assert result is None
|
|
|
|
def test_parse_clarification_json_with_code_fence_json(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 360: JSON in code fence marker parsed."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = '```json\n{"needs_clarification": true, "questions": ["q1"]}\n```'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result is not None
|
|
assert result["needs_clarification"] is True
|
|
|
|
def test_parse_clarification_json_embedded_object(
|
|
self,
|
|
agent_base_params,
|
|
mock_llm_creator,
|
|
mock_llm_handler_creator,
|
|
):
|
|
"""Cover line 360+: JSON object embedded in text."""
|
|
agent = ResearchAgent(**agent_base_params)
|
|
text = 'Here is my response: {"needs_clarification": false} end.'
|
|
result = agent._parse_clarification_json(text)
|
|
assert result == {"needs_clarification": False}
|