wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
1615 lines
56 KiB
Python
1615 lines
56 KiB
Python
"""Comprehensive tests for application/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 application.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(
|
|
"application.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(
|
|
"application.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 application.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 application.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 application.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 application.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 application.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 application.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}
|