deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -0,0 +1,922 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import base64
|
||||
import os
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from langchain_core.messages import AIMessageChunk, ToolMessage
|
||||
from langgraph.types import Command
|
||||
|
||||
from src.config.report_style import ReportStyle
|
||||
from src.server.app import _astream_workflow_generator, _make_event, app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
class TestMakeEvent:
|
||||
def test_make_event_with_content(self):
|
||||
event_type = "message_chunk"
|
||||
data = {"content": "Hello", "role": "assistant"}
|
||||
result = _make_event(event_type, data)
|
||||
expected = (
|
||||
'event: message_chunk\ndata: {"content": "Hello", "role": "assistant"}\n\n'
|
||||
)
|
||||
assert result == expected
|
||||
|
||||
def test_make_event_with_empty_content(self):
|
||||
event_type = "message_chunk"
|
||||
data = {"content": "", "role": "assistant"}
|
||||
result = _make_event(event_type, data)
|
||||
expected = 'event: message_chunk\ndata: {"role": "assistant"}\n\n'
|
||||
assert result == expected
|
||||
|
||||
def test_make_event_without_content(self):
|
||||
event_type = "tool_calls"
|
||||
data = {"role": "assistant", "tool_calls": []}
|
||||
result = _make_event(event_type, data)
|
||||
expected = (
|
||||
'event: tool_calls\ndata: {"role": "assistant", "tool_calls": []}\n\n'
|
||||
)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_astream_workflow_generator_preserves_clarification_history():
|
||||
messages = [
|
||||
{"role": "user", "content": "Research on renewable energy"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "What type of renewable energy would you like to know about?",
|
||||
},
|
||||
{"role": "user", "content": "Solar and wind energy"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please tell me the research dimensions you focus on, such as technological development or market applications.",
|
||||
},
|
||||
{"role": "user", "content": "Technological development"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please specify the time range you want to focus on, such as current status or future trends.",
|
||||
},
|
||||
{"role": "user", "content": "Current status and future trends"},
|
||||
]
|
||||
|
||||
captured_data = {}
|
||||
|
||||
def empty_async_iterator(*args, **kwargs):
|
||||
captured_data["workflow_input"] = args[1]
|
||||
captured_data["workflow_config"] = args[2]
|
||||
|
||||
class IteratorObject:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
return IteratorObject()
|
||||
|
||||
with (
|
||||
patch("src.server.app._process_initial_messages"),
|
||||
patch("src.server.app._stream_graph_events", side_effect=empty_async_iterator),
|
||||
):
|
||||
generator = _astream_workflow_generator(
|
||||
messages=messages,
|
||||
thread_id="clarification-thread",
|
||||
resources=[],
|
||||
max_plan_iterations=1,
|
||||
max_step_num=1,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=True,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=True,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await generator.__anext__()
|
||||
|
||||
workflow_input = captured_data["workflow_input"]
|
||||
assert workflow_input["clarification_history"] == [
|
||||
"Research on renewable energy",
|
||||
"Solar and wind energy",
|
||||
"Technological development",
|
||||
"Current status and future trends",
|
||||
]
|
||||
assert (
|
||||
workflow_input["clarified_research_topic"]
|
||||
== "Research on renewable energy - Solar and wind energy, Technological development, Current status and future trends"
|
||||
)
|
||||
|
||||
|
||||
class TestTTSEndpoint:
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VOLCENGINE_TTS_APPID": "test_app_id",
|
||||
"VOLCENGINE_TTS_ACCESS_TOKEN": "test_token",
|
||||
"VOLCENGINE_TTS_CLUSTER": "test_cluster",
|
||||
"VOLCENGINE_TTS_VOICE_TYPE": "test_voice",
|
||||
},
|
||||
)
|
||||
@patch("src.server.app.VolcengineTTS")
|
||||
def test_tts_success(self, mock_tts_class, client):
|
||||
mock_tts_instance = MagicMock()
|
||||
mock_tts_class.return_value = mock_tts_instance
|
||||
|
||||
# Mock successful TTS response
|
||||
audio_data_b64 = base64.b64encode(b"fake_audio_data").decode()
|
||||
mock_tts_instance.text_to_speech.return_value = {
|
||||
"success": True,
|
||||
"audio_data": audio_data_b64,
|
||||
}
|
||||
|
||||
request_data = {
|
||||
"text": "Hello world",
|
||||
"encoding": "mp3",
|
||||
"speed_ratio": 1.0,
|
||||
"volume_ratio": 1.0,
|
||||
"pitch_ratio": 1.0,
|
||||
"text_type": "plain",
|
||||
"with_frontend": True,
|
||||
"frontend_type": "unitTson",
|
||||
}
|
||||
|
||||
response = client.post("/api/tts", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "audio/mp3"
|
||||
assert b"fake_audio_data" in response.content
|
||||
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_tts_missing_app_id(self, client):
|
||||
request_data = {"text": "Hello world", "encoding": "mp3"}
|
||||
|
||||
response = client.post("/api/tts", json=request_data)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "VOLCENGINE_TTS_APPID is not set" in response.json()["detail"]
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"VOLCENGINE_TTS_APPID": "test_app_id", "VOLCENGINE_TTS_ACCESS_TOKEN": ""},
|
||||
)
|
||||
def test_tts_missing_access_token(self, client):
|
||||
request_data = {"text": "Hello world", "encoding": "mp3"}
|
||||
|
||||
response = client.post("/api/tts", json=request_data)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "VOLCENGINE_TTS_ACCESS_TOKEN is not set" in response.json()["detail"]
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VOLCENGINE_TTS_APPID": "test_app_id",
|
||||
"VOLCENGINE_TTS_ACCESS_TOKEN": "test_token",
|
||||
},
|
||||
)
|
||||
@patch("src.server.app.VolcengineTTS")
|
||||
def test_tts_api_error(self, mock_tts_class, client):
|
||||
mock_tts_instance = MagicMock()
|
||||
mock_tts_class.return_value = mock_tts_instance
|
||||
|
||||
# Mock TTS error response
|
||||
mock_tts_instance.text_to_speech.return_value = {
|
||||
"success": False,
|
||||
"error": "TTS API error",
|
||||
}
|
||||
|
||||
request_data = {"text": "Hello world", "encoding": "mp3"}
|
||||
|
||||
response = client.post("/api/tts", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert "Internal Server Error" in response.json()["detail"]
|
||||
|
||||
@pytest.mark.skip(reason="TTS server exception is catched")
|
||||
@patch("src.server.app.VolcengineTTS")
|
||||
def test_tts_api_exception(self, mock_tts_class, client):
|
||||
mock_tts_instance = MagicMock()
|
||||
mock_tts_class.return_value = mock_tts_instance
|
||||
|
||||
# Mock TTS error response
|
||||
mock_tts_instance.side_effect = Exception("TTS API error")
|
||||
|
||||
request_data = {"text": "Hello world", "encoding": "mp3"}
|
||||
|
||||
response = client.post("/api/tts", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert "Internal Server Error" in response.json()["detail"]
|
||||
|
||||
|
||||
class TestPodcastEndpoint:
|
||||
@patch("src.server.app.build_podcast_graph")
|
||||
def test_generate_podcast_success(self, mock_build_graph, client):
|
||||
mock_workflow = MagicMock()
|
||||
mock_build_graph.return_value = mock_workflow
|
||||
mock_workflow.invoke.return_value = {"output": b"fake_audio_data"}
|
||||
|
||||
request_data = {"content": "Test content for podcast"}
|
||||
|
||||
response = client.post("/api/podcast/generate", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "audio/mp3"
|
||||
assert response.content == b"fake_audio_data"
|
||||
|
||||
@patch("src.server.app.build_podcast_graph")
|
||||
def test_generate_podcast_error(self, mock_build_graph, client):
|
||||
mock_build_graph.side_effect = Exception("Podcast generation failed")
|
||||
|
||||
request_data = {"content": "Test content"}
|
||||
|
||||
response = client.post("/api/podcast/generate", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == "Internal Server Error"
|
||||
|
||||
|
||||
class TestPPTEndpoint:
|
||||
@patch("src.server.app.build_ppt_graph")
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=b"fake_ppt_data")
|
||||
def test_generate_ppt_success(self, mock_file, mock_build_graph, client):
|
||||
mock_workflow = MagicMock()
|
||||
mock_build_graph.return_value = mock_workflow
|
||||
mock_workflow.invoke.return_value = {
|
||||
"generated_file_path": "/fake/path/test.pptx"
|
||||
}
|
||||
|
||||
request_data = {"content": "Test content for PPT"}
|
||||
|
||||
response = client.post("/api/ppt/generate", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert (
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation"
|
||||
in response.headers["content-type"]
|
||||
)
|
||||
assert response.content == b"fake_ppt_data"
|
||||
|
||||
@patch("src.server.app.build_ppt_graph")
|
||||
def test_generate_ppt_error(self, mock_build_graph, client):
|
||||
mock_build_graph.side_effect = Exception("PPT generation failed")
|
||||
|
||||
request_data = {"content": "Test content"}
|
||||
|
||||
response = client.post("/api/ppt/generate", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == "Internal Server Error"
|
||||
|
||||
|
||||
class TestEnhancePromptEndpoint:
|
||||
@patch("src.server.app.build_prompt_enhancer_graph")
|
||||
def test_enhance_prompt_success(self, mock_build_graph, client):
|
||||
mock_workflow = MagicMock()
|
||||
mock_build_graph.return_value = mock_workflow
|
||||
mock_workflow.invoke.return_value = {"output": "Enhanced prompt"}
|
||||
|
||||
request_data = {
|
||||
"prompt": "Original prompt",
|
||||
"context": "Some context",
|
||||
"report_style": "academic",
|
||||
}
|
||||
|
||||
response = client.post("/api/prompt/enhance", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["result"] == "Enhanced prompt"
|
||||
|
||||
@patch("src.server.app.build_prompt_enhancer_graph")
|
||||
def test_enhance_prompt_with_different_styles(self, mock_build_graph, client):
|
||||
mock_workflow = MagicMock()
|
||||
mock_build_graph.return_value = mock_workflow
|
||||
mock_workflow.invoke.return_value = {"output": "Enhanced prompt"}
|
||||
|
||||
styles = [
|
||||
"ACADEMIC",
|
||||
"popular_science",
|
||||
"NEWS",
|
||||
"social_media",
|
||||
"invalid_style",
|
||||
]
|
||||
|
||||
for style in styles:
|
||||
request_data = {"prompt": "Test prompt", "report_style": style}
|
||||
|
||||
response = client.post("/api/prompt/enhance", json=request_data)
|
||||
assert response.status_code == 200
|
||||
|
||||
@patch("src.server.app.build_prompt_enhancer_graph")
|
||||
def test_enhance_prompt_error(self, mock_build_graph, client):
|
||||
mock_build_graph.side_effect = Exception("Enhancement failed")
|
||||
|
||||
request_data = {"prompt": "Test prompt"}
|
||||
|
||||
response = client.post("/api/prompt/enhance", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == "Internal Server Error"
|
||||
|
||||
|
||||
class TestMCPEndpoint:
|
||||
@patch("src.server.app.load_mcp_tools")
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"ENABLE_MCP_SERVER_CONFIGURATION": "true"},
|
||||
)
|
||||
def test_mcp_server_metadata_success(self, mock_load_tools, client):
|
||||
mock_load_tools.return_value = [
|
||||
{"name": "test_tool", "description": "Test tool"}
|
||||
]
|
||||
|
||||
request_data = {
|
||||
"transport": "stdio",
|
||||
"command": "test_command",
|
||||
"args": ["arg1", "arg2"],
|
||||
"env": {"ENV_VAR": "value"},
|
||||
}
|
||||
|
||||
response = client.post("/api/mcp/server/metadata", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
response_data = response.json()
|
||||
assert response_data["transport"] == "stdio"
|
||||
assert response_data["command"] == "test_command"
|
||||
assert len(response_data["tools"]) == 1
|
||||
|
||||
@patch("src.server.app.load_mcp_tools")
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"ENABLE_MCP_SERVER_CONFIGURATION": "true"},
|
||||
)
|
||||
def test_mcp_server_metadata_with_custom_timeout(self, mock_load_tools, client):
|
||||
mock_load_tools.return_value = []
|
||||
|
||||
request_data = {
|
||||
"transport": "stdio",
|
||||
"command": "test_command",
|
||||
"timeout_seconds": 600,
|
||||
}
|
||||
|
||||
response = client.post("/api/mcp/server/metadata", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
mock_load_tools.assert_called_once()
|
||||
|
||||
@patch("src.server.app.load_mcp_tools")
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"ENABLE_MCP_SERVER_CONFIGURATION": "true"},
|
||||
)
|
||||
def test_mcp_server_metadata_with_exception(self, mock_load_tools, client):
|
||||
mock_load_tools.side_effect = HTTPException(
|
||||
status_code=400, detail="MCP Server Error"
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"transport": "stdio",
|
||||
"command": "test_command",
|
||||
"args": ["arg1", "arg2"],
|
||||
"env": {"ENV_VAR": "value"},
|
||||
}
|
||||
|
||||
response = client.post("/api/mcp/server/metadata", json=request_data)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == "Internal Server Error"
|
||||
|
||||
@patch("src.server.app.load_mcp_tools")
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"ENABLE_MCP_SERVER_CONFIGURATION": ""},
|
||||
)
|
||||
def test_mcp_server_metadata_without_enable_configuration(
|
||||
self, mock_load_tools, client
|
||||
):
|
||||
request_data = {
|
||||
"transport": "stdio",
|
||||
"command": "test_command",
|
||||
"args": ["arg1", "arg2"],
|
||||
"env": {"ENV_VAR": "value"},
|
||||
}
|
||||
|
||||
response = client.post("/api/mcp/server/metadata", json=request_data)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert (
|
||||
response.json()["detail"]
|
||||
== "MCP server configuration is disabled. Set ENABLE_MCP_SERVER_CONFIGURATION=true to enable MCP features."
|
||||
)
|
||||
|
||||
|
||||
class TestRAGEndpoints:
|
||||
@patch("src.server.app.SELECTED_RAG_PROVIDER", "test_provider")
|
||||
def test_rag_config(self, client):
|
||||
response = client.get("/api/rag/config")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["provider"] == "test_provider"
|
||||
|
||||
@patch("src.server.app.build_retriever")
|
||||
def test_rag_resources_with_retriever(self, mock_build_retriever, client):
|
||||
mock_retriever = MagicMock()
|
||||
mock_retriever.list_resources.return_value = [
|
||||
{
|
||||
"uri": "test_uri",
|
||||
"title": "Test Resource",
|
||||
"description": "Test Description",
|
||||
}
|
||||
]
|
||||
mock_build_retriever.return_value = mock_retriever
|
||||
|
||||
response = client.get("/api/rag/resources?query=test")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()["resources"]) == 1
|
||||
|
||||
@patch("src.server.app.build_retriever")
|
||||
def test_rag_resources_without_retriever(self, mock_build_retriever, client):
|
||||
mock_build_retriever.return_value = None
|
||||
|
||||
response = client.get("/api/rag/resources")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["resources"] == []
|
||||
|
||||
|
||||
class TestChatStreamEndpoint:
|
||||
@patch("src.server.app.graph")
|
||||
def test_chat_stream_with_default_thread_id(self, mock_graph, client):
|
||||
# Mock the async stream
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1", "step1", {"test": "data"})
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
request_data = {
|
||||
"thread_id": "__default__",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"resources": [],
|
||||
"max_plan_iterations": 3,
|
||||
"max_step_num": 10,
|
||||
"max_search_results": 5,
|
||||
"auto_accepted_plan": True,
|
||||
"interrupt_feedback": "",
|
||||
"mcp_settings": {},
|
||||
"enable_background_investigation": False,
|
||||
"report_style": "academic",
|
||||
}
|
||||
|
||||
response = client.post("/api/chat/stream", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
|
||||
|
||||
@patch("src.server.app.graph")
|
||||
def test_chat_stream_with_mcp_settings(self, mock_graph, client):
|
||||
# Mock the async stream
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1", "step1", {"test": "data"})
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
request_data = {
|
||||
"thread_id": "__default__",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"resources": [],
|
||||
"max_plan_iterations": 3,
|
||||
"max_step_num": 10,
|
||||
"max_search_results": 5,
|
||||
"auto_accepted_plan": True,
|
||||
"interrupt_feedback": "",
|
||||
"mcp_settings": {
|
||||
"servers": {
|
||||
"mcp-github-trending": {
|
||||
"transport": "stdio",
|
||||
"command": "uvx",
|
||||
"args": ["mcp-github-trending"],
|
||||
"env": {"MCP_SERVER_ID": "mcp-github-trending"},
|
||||
"enabled_tools": ["get_github_trending_repositories"],
|
||||
"add_to_agents": ["researcher"],
|
||||
}
|
||||
}
|
||||
},
|
||||
"enable_background_investigation": False,
|
||||
"report_style": "academic",
|
||||
}
|
||||
|
||||
response = client.post("/api/chat/stream", json=request_data)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert (
|
||||
response.json()["detail"]
|
||||
== "MCP server configuration is disabled. Set ENABLE_MCP_SERVER_CONFIGURATION=true to enable MCP features."
|
||||
)
|
||||
|
||||
@patch("src.server.app.graph")
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{"ENABLE_MCP_SERVER_CONFIGURATION": "true"},
|
||||
)
|
||||
def test_chat_stream_with_mcp_settings_enabled(self, mock_graph, client):
|
||||
# Mock the async stream
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1", "step1", {"test": "data"})
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
request_data = {
|
||||
"thread_id": "__default__",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"resources": [],
|
||||
"max_plan_iterations": 3,
|
||||
"max_step_num": 10,
|
||||
"max_search_results": 5,
|
||||
"auto_accepted_plan": True,
|
||||
"interrupt_feedback": "",
|
||||
"mcp_settings": {
|
||||
"servers": {
|
||||
"mcp-github-trending": {
|
||||
"transport": "stdio",
|
||||
"command": "uvx",
|
||||
"args": ["mcp-github-trending"],
|
||||
"env": {"MCP_SERVER_ID": "mcp-github-trending"},
|
||||
"enabled_tools": ["get_github_trending_repositories"],
|
||||
"add_to_agents": ["researcher"],
|
||||
}
|
||||
}
|
||||
},
|
||||
"enable_background_investigation": False,
|
||||
"report_style": "academic",
|
||||
}
|
||||
|
||||
response = client.post("/api/chat/stream", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
|
||||
|
||||
|
||||
class TestAstreamWorkflowGenerator:
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_basic_flow(self, mock_graph):
|
||||
# Mock AI message chunk
|
||||
mock_message = AIMessageChunk(content="Hello world")
|
||||
mock_message.id = "msg_123"
|
||||
mock_message.response_metadata = {}
|
||||
mock_message.tool_calls = []
|
||||
mock_message.tool_call_chunks = []
|
||||
|
||||
# Mock the async stream - yield messages in the correct format
|
||||
async def mock_astream(*args, **kwargs):
|
||||
# Yield a tuple (message, metadata) instead of just [message]
|
||||
yield ("agent1:subagent", "messages", (mock_message, {}))
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
thread_id = "test_thread"
|
||||
resources = []
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=messages,
|
||||
thread_id=thread_id,
|
||||
resources=resources,
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: message_chunk" in events[0]
|
||||
assert "Hello world" in events[0]
|
||||
# Check for the actual agent name that appears in the output
|
||||
assert '"agent": "a"' in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_with_interrupt_feedback(self, mock_graph):
|
||||
# Mock the async stream
|
||||
async def mock_astream(*args, **kwargs):
|
||||
# Verify that Command is passed as input when interrupt_feedback is provided
|
||||
assert isinstance(args[0], Command)
|
||||
assert "[edit_plan] Hello" in args[0].resume
|
||||
yield ("agent1", "step1", {"test": "data"})
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=messages,
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=False,
|
||||
interrupt_feedback="edit_plan",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_interrupt_event(self, mock_graph):
|
||||
# Mock interrupt data
|
||||
mock_interrupt = MagicMock()
|
||||
mock_interrupt.ns = ["interrupt_id"]
|
||||
mock_interrupt.value = "Plan requires approval"
|
||||
|
||||
interrupt_data = {"__interrupt__": [mock_interrupt]}
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1", "step1", interrupt_data)
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=[],
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: interrupt" in events[0]
|
||||
assert "Plan requires approval" in events[0]
|
||||
assert "interrupt_id" in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_tool_message(self, mock_graph):
|
||||
# Mock tool message
|
||||
mock_tool_message = ToolMessage(content="Tool result", tool_call_id="tool_123")
|
||||
mock_tool_message.id = "msg_456"
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1:subagent", "step1", (mock_tool_message, {}))
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=[],
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: tool_call_result" in events[0]
|
||||
assert "Tool result" in events[0]
|
||||
assert "tool_123" in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_ai_message_with_tool_calls(
|
||||
self, mock_graph
|
||||
):
|
||||
# Mock AI message with tool calls
|
||||
mock_ai_message = AIMessageChunk(content="Making tool call")
|
||||
mock_ai_message.id = "msg_789"
|
||||
mock_ai_message.response_metadata = {"finish_reason": "tool_calls"}
|
||||
mock_ai_message.tool_calls = [{"name": "search", "args": {"query": "test"}}]
|
||||
mock_ai_message.tool_call_chunks = [{"name": "search"}]
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=[],
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: tool_calls" in events[0]
|
||||
assert "Making tool call" in events[0]
|
||||
assert "tool_calls" in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_ai_message_with_tool_call_chunks(
|
||||
self, mock_graph
|
||||
):
|
||||
# Mock AI message with only tool call chunks
|
||||
mock_ai_message = AIMessageChunk(content="Streaming tool call")
|
||||
mock_ai_message.id = "msg_101"
|
||||
mock_ai_message.response_metadata = {}
|
||||
mock_ai_message.tool_calls = []
|
||||
mock_ai_message.tool_call_chunks = [{"name": "search", "index": 0}]
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=[],
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: tool_call_chunks" in events[0]
|
||||
assert "Streaming tool call" in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_with_finish_reason(self, mock_graph):
|
||||
# Mock AI message with finish reason
|
||||
mock_ai_message = AIMessageChunk(content="Complete response")
|
||||
mock_ai_message.id = "msg_finish"
|
||||
mock_ai_message.response_metadata = {"finish_reason": "stop"}
|
||||
mock_ai_message.tool_calls = []
|
||||
mock_ai_message.tool_call_chunks = []
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
|
||||
|
||||
mock_graph.astream = mock_astream
|
||||
|
||||
generator = _astream_workflow_generator(
|
||||
messages=[],
|
||||
thread_id="test_thread",
|
||||
resources=[],
|
||||
max_plan_iterations=3,
|
||||
max_step_num=10,
|
||||
max_search_results=5,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="",
|
||||
mcp_settings={},
|
||||
enable_background_investigation=False,
|
||||
report_style=ReportStyle.ACADEMIC,
|
||||
enable_deep_thinking=False,
|
||||
enable_clarification=False,
|
||||
max_clarification_rounds=3,
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in generator:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) == 1
|
||||
assert "event: message_chunk" in events[0]
|
||||
assert "finish_reason" in events[0]
|
||||
assert "stop" in events[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.app.graph")
|
||||
async def test_astream_workflow_generator_config_passed_correctly(self, mock_graph):
|
||||
mock_ai_message = AIMessageChunk(content="Test")
|
||||
mock_ai_message.id = "test_id"
|
||||
mock_ai_message.response_metadata = {}
|
||||
mock_ai_message.tool_calls = []
|
||||
mock_ai_message.tool_call_chunks = []
|
||||
|
||||
async def verify_config(*args, **kwargs):
|
||||
config = kwargs.get("config", {})
|
||||
assert config["thread_id"] == "test_thread"
|
||||
assert config["max_plan_iterations"] == 5
|
||||
assert config["max_step_num"] == 20
|
||||
assert config["max_search_results"] == 10
|
||||
assert config["report_style"] == ReportStyle.NEWS.value
|
||||
yield ("agent1", "messages", [mock_ai_message])
|
||||
|
||||
|
||||
class TestGenerateProseEndpoint:
|
||||
@patch("src.server.app.build_prose_graph")
|
||||
def test_generate_prose_success(self, mock_build_graph, client):
|
||||
# Mock the workflow and its astream method
|
||||
mock_workflow = MagicMock()
|
||||
mock_build_graph.return_value = mock_workflow
|
||||
|
||||
class MockEvent:
|
||||
def __init__(self, content):
|
||||
self.content = content
|
||||
|
||||
async def mock_astream(*args, **kwargs):
|
||||
yield (None, [MockEvent("Generated prose 1")])
|
||||
yield (None, [MockEvent("Generated prose 2")])
|
||||
|
||||
mock_workflow.astream.return_value = mock_astream()
|
||||
request_data = {
|
||||
"prompt": "Write a story.",
|
||||
"option": "default",
|
||||
"command": "generate",
|
||||
}
|
||||
|
||||
response = client.post("/api/prose/generate", json=request_data)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
|
||||
# Read the streaming response content
|
||||
content = b"".join(response.iter_bytes())
|
||||
assert b"Generated prose 1" in content or b"Generated prose 2" in content
|
||||
|
||||
@patch("src.server.app.build_prose_graph")
|
||||
def test_generate_prose_error(self, mock_build_graph, client):
|
||||
mock_build_graph.side_effect = Exception("Prose generation failed")
|
||||
request_data = {
|
||||
"prompt": "Write a story.",
|
||||
"option": "default",
|
||||
"command": "generate",
|
||||
}
|
||||
response = client.post("/api/prose/generate", json=request_data)
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == "Internal Server Error"
|
||||
@@ -0,0 +1,168 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
|
||||
import src.server.mcp_utils as mcp_utils # Assuming mcp_utils is the module to test
|
||||
from src.config.report_style import ReportStyle
|
||||
from src.rag.retriever import Resource
|
||||
from src.server.chat_request import (
|
||||
ChatMessage,
|
||||
ChatRequest,
|
||||
ContentItem,
|
||||
EnhancePromptRequest,
|
||||
GeneratePodcastRequest,
|
||||
GeneratePPTRequest,
|
||||
GenerateProseRequest,
|
||||
TTSRequest,
|
||||
)
|
||||
|
||||
|
||||
def test_content_item_text_and_image():
|
||||
item_text = ContentItem(type="text", text="hello")
|
||||
assert item_text.type == "text"
|
||||
assert item_text.text == "hello"
|
||||
assert item_text.image_url is None
|
||||
|
||||
item_image = ContentItem(type="image", image_url="http://img.com/1.png")
|
||||
assert item_image.type == "image"
|
||||
assert item_image.text is None
|
||||
assert item_image.image_url == "http://img.com/1.png"
|
||||
|
||||
|
||||
def test_chat_message_with_string_content():
|
||||
msg = ChatMessage(role="user", content="Hello!")
|
||||
assert msg.role == "user"
|
||||
assert msg.content == "Hello!"
|
||||
|
||||
|
||||
def test_chat_message_with_content_items():
|
||||
items = [ContentItem(type="text", text="hi")]
|
||||
msg = ChatMessage(role="assistant", content=items)
|
||||
assert msg.role == "assistant"
|
||||
assert isinstance(msg.content, list)
|
||||
assert msg.content[0].type == "text"
|
||||
|
||||
|
||||
def test_chat_request_defaults():
|
||||
req = ChatRequest()
|
||||
assert req.messages == []
|
||||
assert req.resources == []
|
||||
assert req.debug is False
|
||||
assert req.thread_id == "__default__"
|
||||
assert req.max_plan_iterations == 1
|
||||
assert req.max_step_num == 3
|
||||
assert req.max_search_results == 3
|
||||
assert req.auto_accepted_plan is False
|
||||
assert req.interrupt_feedback is None
|
||||
assert req.mcp_settings is None
|
||||
assert req.enable_background_investigation is True
|
||||
assert req.report_style == ReportStyle.ACADEMIC
|
||||
|
||||
|
||||
def test_chat_request_with_values():
|
||||
resource = Resource(
|
||||
name="test", type="doc", uri="some-uri-value", title="some-title-value"
|
||||
)
|
||||
msg = ChatMessage(role="user", content="hi")
|
||||
req = ChatRequest(
|
||||
messages=[msg],
|
||||
resources=[resource],
|
||||
debug=True,
|
||||
thread_id="tid",
|
||||
max_plan_iterations=2,
|
||||
max_step_num=5,
|
||||
max_search_results=10,
|
||||
auto_accepted_plan=True,
|
||||
interrupt_feedback="stop",
|
||||
mcp_settings={"foo": "bar"},
|
||||
enable_background_investigation=False,
|
||||
report_style="academic",
|
||||
)
|
||||
assert req.messages[0].role == "user"
|
||||
assert req.debug is True
|
||||
assert req.thread_id == "tid"
|
||||
assert req.max_plan_iterations == 2
|
||||
assert req.max_step_num == 5
|
||||
assert req.max_search_results == 10
|
||||
assert req.auto_accepted_plan is True
|
||||
assert req.interrupt_feedback == "stop"
|
||||
assert req.mcp_settings == {"foo": "bar"}
|
||||
assert req.enable_background_investigation is False
|
||||
assert req.report_style == ReportStyle.ACADEMIC
|
||||
|
||||
|
||||
def test_tts_request_defaults():
|
||||
req = TTSRequest(text="hello")
|
||||
assert req.text == "hello"
|
||||
assert req.voice_type == "BV700_V2_streaming"
|
||||
assert req.encoding == "mp3"
|
||||
assert req.speed_ratio == 1.0
|
||||
assert req.volume_ratio == 1.0
|
||||
assert req.pitch_ratio == 1.0
|
||||
assert req.text_type == "plain"
|
||||
assert req.with_frontend == 1
|
||||
assert req.frontend_type == "unitTson"
|
||||
|
||||
|
||||
def test_generate_podcast_request():
|
||||
req = GeneratePodcastRequest(content="Podcast content")
|
||||
assert req.content == "Podcast content"
|
||||
|
||||
|
||||
def test_generate_ppt_request():
|
||||
req = GeneratePPTRequest(content="PPT content")
|
||||
assert req.content == "PPT content"
|
||||
|
||||
|
||||
def test_generate_prose_request():
|
||||
req = GenerateProseRequest(prompt="Write a poem", option="poet", command="rhyme")
|
||||
assert req.prompt == "Write a poem"
|
||||
assert req.option == "poet"
|
||||
assert req.command == "rhyme"
|
||||
|
||||
req2 = GenerateProseRequest(prompt="Write", option="short")
|
||||
assert req2.command == ""
|
||||
|
||||
|
||||
def test_enhance_prompt_request_defaults():
|
||||
req = EnhancePromptRequest(prompt="Improve this")
|
||||
assert req.prompt == "Improve this"
|
||||
assert req.context == ""
|
||||
assert req.report_style == "academic"
|
||||
|
||||
|
||||
def test_content_item_validation_error():
|
||||
with pytest.raises(ValidationError):
|
||||
ContentItem() # missing required 'type'
|
||||
|
||||
|
||||
def test_chat_message_validation_error():
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessage(role="user") # missing content
|
||||
|
||||
|
||||
def test_tts_request_validation_error():
|
||||
with pytest.raises(ValidationError):
|
||||
TTSRequest() # missing required 'text'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
|
||||
@patch("src.server.mcp_utils.StdioServerParameters")
|
||||
@patch("src.server.mcp_utils.stdio_client")
|
||||
async def test_load_mcp_tools_exception_handling(
|
||||
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
|
||||
): # Changed to async def
|
||||
mock_get_tools.side_effect = Exception("unexpected error")
|
||||
mock_StdioServerParameters.return_value = MagicMock()
|
||||
mock_stdio_client.return_value = MagicMock()
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mcp_utils.load_mcp_tools(server_type="stdio", command="foo") # Use await
|
||||
assert exc.value.status_code == 500
|
||||
assert "unexpected error" in exc.value.detail
|
||||
@@ -0,0 +1,74 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from src.server.mcp_request import MCPServerMetadataRequest, MCPServerMetadataResponse
|
||||
|
||||
|
||||
def test_mcp_server_metadata_request_required_fields():
|
||||
# 'transport' is required
|
||||
req = MCPServerMetadataRequest(transport="stdio")
|
||||
assert req.transport == "stdio"
|
||||
assert req.command is None
|
||||
assert req.args is None
|
||||
assert req.url is None
|
||||
assert req.env is None
|
||||
assert req.timeout_seconds is None
|
||||
|
||||
|
||||
def test_mcp_server_metadata_request_optional_fields():
|
||||
req = MCPServerMetadataRequest(
|
||||
transport="sse",
|
||||
command="run",
|
||||
args=["--foo", "bar"],
|
||||
url="http://localhost:8080",
|
||||
env={"FOO": "BAR"},
|
||||
timeout_seconds=30,
|
||||
)
|
||||
assert req.transport == "sse"
|
||||
assert req.command == "run"
|
||||
assert req.args == ["--foo", "bar"]
|
||||
assert req.url == "http://localhost:8080"
|
||||
assert req.env == {"FOO": "BAR"}
|
||||
assert req.timeout_seconds == 30
|
||||
|
||||
|
||||
def test_mcp_server_metadata_request_missing_transport():
|
||||
with pytest.raises(ValidationError):
|
||||
MCPServerMetadataRequest()
|
||||
|
||||
|
||||
def test_mcp_server_metadata_response_required_fields():
|
||||
resp = MCPServerMetadataResponse(transport="stdio")
|
||||
assert resp.transport == "stdio"
|
||||
assert resp.command is None
|
||||
assert resp.args is None
|
||||
assert resp.url is None
|
||||
assert resp.env is None
|
||||
assert resp.tools == []
|
||||
|
||||
|
||||
def test_mcp_server_metadata_response_optional_fields():
|
||||
resp = MCPServerMetadataResponse(
|
||||
transport="sse",
|
||||
command="run",
|
||||
args=["--foo", "bar"],
|
||||
url="http://localhost:8080",
|
||||
env={"FOO": "BAR"},
|
||||
tools=["tool1", "tool2"],
|
||||
)
|
||||
assert resp.transport == "sse"
|
||||
assert resp.command == "run"
|
||||
assert resp.args == ["--foo", "bar"]
|
||||
assert resp.url == "http://localhost:8080"
|
||||
assert resp.env == {"FOO": "BAR"}
|
||||
assert resp.tools == ["tool1", "tool2"]
|
||||
|
||||
|
||||
def test_mcp_server_metadata_response_tools_default_factory():
|
||||
resp1 = MCPServerMetadataResponse(transport="stdio")
|
||||
resp2 = MCPServerMetadataResponse(transport="stdio")
|
||||
resp1.tools.append("toolA")
|
||||
assert resp2.tools == [] # Should not share list between instances
|
||||
@@ -0,0 +1,132 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import src.server.mcp_utils as mcp_utils
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.mcp_utils.ClientSession")
|
||||
async def test__get_tools_from_client_session_success(mock_ClientSession):
|
||||
mock_read = AsyncMock()
|
||||
mock_write = AsyncMock()
|
||||
mock_callback = AsyncMock()
|
||||
mock_context_manager = AsyncMock()
|
||||
mock_context_manager.__aenter__.return_value = (
|
||||
mock_read,
|
||||
mock_write,
|
||||
mock_callback,
|
||||
)
|
||||
mock_context_manager.__aexit__.return_value = None
|
||||
|
||||
mock_session = AsyncMock()
|
||||
mock_session.__aenter__.return_value = mock_session
|
||||
mock_session.__aexit__.return_value = None
|
||||
mock_session.initialize = AsyncMock()
|
||||
mock_tools_obj = MagicMock()
|
||||
mock_tools_obj.tools = ["tool1", "tool2"]
|
||||
mock_session.list_tools = AsyncMock(return_value=mock_tools_obj)
|
||||
mock_ClientSession.return_value = mock_session
|
||||
|
||||
result = await mcp_utils._get_tools_from_client_session(
|
||||
mock_context_manager, timeout_seconds=5
|
||||
)
|
||||
assert result == ["tool1", "tool2"]
|
||||
mock_session.initialize.assert_awaited_once()
|
||||
mock_session.list_tools.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
|
||||
@patch("src.server.mcp_utils.StdioServerParameters")
|
||||
@patch("src.server.mcp_utils.stdio_client")
|
||||
async def test_load_mcp_tools_stdio_success(
|
||||
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
|
||||
):
|
||||
mock_get_tools.return_value = ["toolA"]
|
||||
params = MagicMock()
|
||||
mock_StdioServerParameters.return_value = params
|
||||
mock_client = MagicMock()
|
||||
mock_stdio_client.return_value = mock_client
|
||||
|
||||
result = await mcp_utils.load_mcp_tools(
|
||||
server_type="stdio",
|
||||
command="echo",
|
||||
args=["foo"],
|
||||
env={"FOO": "BAR"},
|
||||
timeout_seconds=3,
|
||||
)
|
||||
assert result == ["toolA"]
|
||||
mock_StdioServerParameters.assert_called_once_with(
|
||||
command="echo", args=["foo"], env={"FOO": "BAR"}
|
||||
)
|
||||
mock_stdio_client.assert_called_once_with(params)
|
||||
mock_get_tools.assert_awaited_once_with(mock_client, 3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_mcp_tools_stdio_missing_command():
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mcp_utils.load_mcp_tools(server_type="stdio")
|
||||
assert exc.value.status_code == 400
|
||||
assert "Command is required" in exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
|
||||
@patch("src.server.mcp_utils.sse_client")
|
||||
async def test_load_mcp_tools_sse_success(mock_sse_client, mock_get_tools):
|
||||
mock_get_tools.return_value = ["toolB"]
|
||||
mock_client = MagicMock()
|
||||
mock_sse_client.return_value = mock_client
|
||||
|
||||
result = await mcp_utils.load_mcp_tools(
|
||||
server_type="sse",
|
||||
url="http://localhost:1234",
|
||||
headers={"Authorization": "Bearer 1234567890"},
|
||||
timeout_seconds=7,
|
||||
)
|
||||
assert result == ["toolB"]
|
||||
mock_sse_client.assert_called_once_with(
|
||||
url="http://localhost:1234",
|
||||
headers={"Authorization": "Bearer 1234567890"},
|
||||
timeout=7,
|
||||
)
|
||||
mock_get_tools.assert_awaited_once_with(mock_client, 7)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_mcp_tools_sse_missing_url():
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mcp_utils.load_mcp_tools(server_type="sse")
|
||||
assert exc.value.status_code == 400
|
||||
assert "URL is required" in exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_mcp_tools_unsupported_type():
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mcp_utils.load_mcp_tools(server_type="unknown")
|
||||
assert exc.value.status_code == 400
|
||||
assert "Unsupported server type" in exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
|
||||
@patch("src.server.mcp_utils.StdioServerParameters")
|
||||
@patch("src.server.mcp_utils.stdio_client")
|
||||
async def test_load_mcp_tools_exception_handling(
|
||||
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
|
||||
):
|
||||
mock_get_tools.side_effect = Exception("unexpected error")
|
||||
mock_StdioServerParameters.return_value = MagicMock()
|
||||
mock_stdio_client.return_value = MagicMock()
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mcp_utils.load_mcp_tools(server_type="stdio", command="foo")
|
||||
assert exc.value.status_code == 500
|
||||
assert "unexpected error" in exc.value.detail
|
||||
@@ -0,0 +1,317 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Unit tests for tool call chunk processing.
|
||||
|
||||
Tests for the fix of issue #523: Tool name concatenation in consecutive tool calls.
|
||||
This ensures that tool call chunks are properly segregated by index to prevent
|
||||
tool names from being concatenated when multiple tool calls happen in sequence.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
# Import the functions to test
|
||||
# Note: We need to import from the app module
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Add src directory to path for imports
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../"))
|
||||
|
||||
from src.server.app import _process_tool_call_chunks, _validate_tool_call_chunks
|
||||
|
||||
|
||||
class TestProcessToolCallChunks:
|
||||
"""Test cases for _process_tool_call_chunks function."""
|
||||
|
||||
def test_empty_tool_call_chunks(self):
|
||||
"""Test processing empty tool call chunks."""
|
||||
result = _process_tool_call_chunks([])
|
||||
assert result == []
|
||||
|
||||
def test_single_tool_call_single_chunk(self):
|
||||
"""Test processing a single tool call with a single chunk."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{"query": "test"}', "id": "call_1", "index": 0}
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "web_search"
|
||||
assert result[0]["id"] == "call_1"
|
||||
assert result[0]["index"] == 0
|
||||
assert '"query": "test"' in result[0]["args"]
|
||||
|
||||
def test_consecutive_tool_calls_different_indices(self):
|
||||
"""Test that consecutive tool calls with different indices are not concatenated."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{"query": "test"}', "id": "call_1", "index": 0},
|
||||
{"name": "web_search", "args": '{"query": "test2"}', "id": "call_2", "index": 1},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0]["name"] == "web_search"
|
||||
assert result[0]["id"] == "call_1"
|
||||
assert result[0]["index"] == 0
|
||||
assert result[1]["name"] == "web_search"
|
||||
assert result[1]["id"] == "call_2"
|
||||
assert result[1]["index"] == 1
|
||||
# Verify names are NOT concatenated
|
||||
assert result[0]["name"] != "web_searchweb_search"
|
||||
assert result[1]["name"] != "web_searchweb_search"
|
||||
|
||||
def test_different_tools_different_indices(self):
|
||||
"""Test consecutive calls to different tools."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{"query": "test"}', "id": "call_1", "index": 0},
|
||||
{"name": "crawl_tool", "args": '{"url": "http://example.com"}', "id": "call_2", "index": 1},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0]["name"] == "web_search"
|
||||
assert result[1]["name"] == "crawl_tool"
|
||||
# Verify names are NOT concatenated (the issue bug scenario)
|
||||
assert "web_searchcrawl_tool" not in result[0]["name"]
|
||||
assert "web_searchcrawl_tool" not in result[1]["name"]
|
||||
|
||||
def test_streaming_chunks_same_index(self):
|
||||
"""Test streaming chunks with same index are properly accumulated."""
|
||||
chunks = [
|
||||
{"name": "web_", "args": '{"query"', "id": "call_1", "index": 0},
|
||||
{"name": "search", "args": ': "test"}', "id": "call_1", "index": 0},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 1
|
||||
# Name should NOT be concatenated when it's the same tool
|
||||
assert result[0]["name"] in ["web_", "search", "web_search"]
|
||||
assert result[0]["id"] == "call_1"
|
||||
# Args should be accumulated
|
||||
assert "query" in result[0]["args"]
|
||||
assert "test" in result[0]["args"]
|
||||
|
||||
def test_tool_call_index_collision_warning(self):
|
||||
"""Test that index collision with different names generates warning."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{}', "id": "call_1", "index": 0},
|
||||
{"name": "crawl_tool", "args": '{}', "id": "call_2", "index": 0},
|
||||
]
|
||||
|
||||
# This should trigger a warning
|
||||
with patch('src.server.app.logger') as mock_logger:
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
# Verify warning was logged
|
||||
mock_logger.warning.assert_called()
|
||||
call_args = mock_logger.warning.call_args[0][0]
|
||||
assert "Tool name mismatch detected" in call_args
|
||||
assert "web_search" in call_args
|
||||
assert "crawl_tool" in call_args
|
||||
|
||||
def test_chunks_without_explicit_index(self):
|
||||
"""Test handling chunks without explicit index (edge case)."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{}', "id": "call_1"} # No index
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["name"] == "web_search"
|
||||
|
||||
def test_chunk_sorting_by_index(self):
|
||||
"""Test that chunks are sorted by index in proper order."""
|
||||
chunks = [
|
||||
{"name": "tool_3", "args": '{}', "id": "call_3", "index": 2},
|
||||
{"name": "tool_1", "args": '{}', "id": "call_1", "index": 0},
|
||||
{"name": "tool_2", "args": '{}', "id": "call_2", "index": 1},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0]["index"] == 0
|
||||
assert result[1]["index"] == 1
|
||||
assert result[2]["index"] == 2
|
||||
|
||||
def test_args_accumulation(self):
|
||||
"""Test that arguments are properly accumulated for same index."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{"q', "id": "call_1", "index": 0},
|
||||
{"name": "web_search", "args": 'uery": "test"}', "id": "call_1", "index": 0},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 1
|
||||
# Sanitize removes json encoding, so just check it's accumulated
|
||||
assert len(result[0]["args"]) > 0
|
||||
|
||||
def test_preserve_tool_id(self):
|
||||
"""Test that tool IDs are preserved correctly."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{}', "id": "call_abc123", "index": 0},
|
||||
{"name": "web_search", "args": '{}', "id": "call_xyz789", "index": 1},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert result[0]["id"] == "call_abc123"
|
||||
assert result[1]["id"] == "call_xyz789"
|
||||
|
||||
def test_multiple_indices_detected(self):
|
||||
"""Test that multiple indices are properly detected and logged."""
|
||||
chunks = [
|
||||
{"name": "tool_a", "args": '{}', "id": "call_1", "index": 0},
|
||||
{"name": "tool_b", "args": '{}', "id": "call_2", "index": 1},
|
||||
{"name": "tool_c", "args": '{}', "id": "call_3", "index": 2},
|
||||
]
|
||||
|
||||
with patch('src.server.app.logger') as mock_logger:
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
# Should have debug logs for multiple indices
|
||||
debug_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
|
||||
# Check if any debug call mentions multiple indices
|
||||
multiple_indices_mentioned = any(
|
||||
"Multiple indices" in call for call in debug_calls
|
||||
)
|
||||
assert multiple_indices_mentioned or len(result) == 3
|
||||
|
||||
|
||||
class TestValidateToolCallChunks:
|
||||
"""Test cases for _validate_tool_call_chunks function."""
|
||||
|
||||
def test_validate_empty_chunks(self):
|
||||
"""Test validation of empty chunks."""
|
||||
# Should not raise any exception
|
||||
_validate_tool_call_chunks([])
|
||||
|
||||
def test_validate_logs_chunk_info(self):
|
||||
"""Test that validation logs chunk information."""
|
||||
chunks = [
|
||||
{"name": "web_search", "args": '{}', "id": "call_1", "index": 0},
|
||||
]
|
||||
|
||||
with patch('src.server.app.logger') as mock_logger:
|
||||
_validate_tool_call_chunks(chunks)
|
||||
|
||||
# Should have logged debug info
|
||||
assert mock_logger.debug.called
|
||||
|
||||
def test_validate_detects_multiple_indices(self):
|
||||
"""Test that validation detects multiple indices."""
|
||||
chunks = [
|
||||
{"name": "tool_1", "args": '{}', "id": "call_1", "index": 0},
|
||||
{"name": "tool_2", "args": '{}', "id": "call_2", "index": 1},
|
||||
]
|
||||
|
||||
with patch('src.server.app.logger') as mock_logger:
|
||||
_validate_tool_call_chunks(chunks)
|
||||
|
||||
# Should have logged about multiple indices
|
||||
debug_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
|
||||
multiple_indices_mentioned = any(
|
||||
"Multiple indices" in call for call in debug_calls
|
||||
)
|
||||
assert multiple_indices_mentioned
|
||||
|
||||
|
||||
class TestRealWorldScenarios:
|
||||
"""Test cases for real-world scenarios from issue #523."""
|
||||
|
||||
def test_issue_523_scenario_consecutive_web_search(self):
|
||||
"""
|
||||
Replicate issue #523: Consecutive web_search calls.
|
||||
Previously would result in "web_searchweb_search" error.
|
||||
"""
|
||||
# Simulate streaming chunks from two consecutive web_search calls
|
||||
chunks = [
|
||||
# First web_search call (index 0)
|
||||
{"name": "web_", "args": '{"query', "id": "call_1", "index": 0},
|
||||
{"name": "search", "args": '": "first query"}', "id": "call_1", "index": 0},
|
||||
# Second web_search call (index 1)
|
||||
{"name": "web_", "args": '{"query', "id": "call_2", "index": 1},
|
||||
{"name": "search", "args": '": "second query"}', "id": "call_2", "index": 1},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
# Should have 2 tool calls, not concatenated names
|
||||
assert len(result) >= 1 # At minimum should process without error
|
||||
|
||||
# Extract tool names from result
|
||||
tool_names = [chunk.get("name") for chunk in result]
|
||||
|
||||
# Verify "web_searchweb_search" error doesn't occur
|
||||
assert "web_searchweb_search" not in tool_names
|
||||
|
||||
# Both calls should have web_search (or parts of it)
|
||||
concatenated_names = "".join(tool_names)
|
||||
assert "web_search" in concatenated_names or "web_" in concatenated_names
|
||||
|
||||
def test_mixed_tools_consecutive_calls(self):
|
||||
"""Test realistic scenario with mixed tools in sequence."""
|
||||
chunks = [
|
||||
# web_search call
|
||||
{"name": "web_search", "args": '{"query": "python"}', "id": "1", "index": 0},
|
||||
# crawl_tool call
|
||||
{"name": "crawl_tool", "args": '{"url": "http://example.com"}', "id": "2", "index": 1},
|
||||
# Another web_search
|
||||
{"name": "web_search", "args": '{"query": "rust"}', "id": "3", "index": 2},
|
||||
]
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
assert len(result) == 3
|
||||
tool_names = [chunk.get("name") for chunk in result]
|
||||
|
||||
# No concatenation should occur
|
||||
assert "web_searchcrawl_tool" not in tool_names
|
||||
assert "crawl_toolweb_search" not in tool_names
|
||||
|
||||
def test_long_sequence_tool_calls(self):
|
||||
"""Test a long sequence of tool calls."""
|
||||
chunks = []
|
||||
for i in range(10):
|
||||
tool_name = "web_search" if i % 2 == 0 else "crawl_tool"
|
||||
chunks.append({
|
||||
"name": tool_name,
|
||||
"args": '{"query": "test"}' if tool_name == "web_search" else '{"url": "http://example.com"}',
|
||||
"id": f"call_{i}",
|
||||
"index": i
|
||||
})
|
||||
|
||||
result = _process_tool_call_chunks(chunks)
|
||||
|
||||
# Should process all 10 tool calls
|
||||
assert len(result) == 10
|
||||
|
||||
# Verify each tool call has correct name (not concatenated with other tool names)
|
||||
for i, chunk in enumerate(result):
|
||||
expected_name = "web_search" if i % 2 == 0 else "crawl_tool"
|
||||
actual_name = chunk.get("name", "")
|
||||
|
||||
# The actual name should be the expected name, not concatenated
|
||||
assert actual_name == expected_name, (
|
||||
f"Tool call {i} has name '{actual_name}', expected '{expected_name}'. "
|
||||
f"This indicates concatenation with adjacent tool call."
|
||||
)
|
||||
|
||||
# Verify IDs are correct
|
||||
assert chunk.get("id") == f"call_{i}"
|
||||
assert chunk.get("index") == i
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
Reference in New Issue
Block a user