diff --git a/src/config/__init__.py b/src/config/__init__.py index 54bf87c..104a580 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT +from .loader import load_yaml_config from .tools import SELECTED_SEARCH_ENGINE, SearchEngine from .questions import BUILT_IN_QUESTIONS, BUILT_IN_QUESTIONS_ZH_CN @@ -45,4 +46,5 @@ __all__ = [ "SearchEngine", "BUILT_IN_QUESTIONS", "BUILT_IN_QUESTIONS_ZH_CN", + load_yaml_config, ] diff --git a/src/crawler/crawler.py b/src/crawler/crawler.py index 9f3632d..01a2a6e 100644 --- a/src/crawler/crawler.py +++ b/src/crawler/crawler.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import sys from .article import Article from .jina_client import JinaClient diff --git a/src/llms/llm.py b/src/llms/llm.py index 15ba13b..aa78961 100644 --- a/src/llms/llm.py +++ b/src/llms/llm.py @@ -4,7 +4,6 @@ from pathlib import Path from typing import Any, Dict import os -import ssl import httpx from langchain_openai import ChatOpenAI diff --git a/src/prompt_enhancer/graph/enhancer_node.py b/src/prompt_enhancer/graph/enhancer_node.py index de7da03..8261d38 100644 --- a/src/prompt_enhancer/graph/enhancer_node.py +++ b/src/prompt_enhancer/graph/enhancer_node.py @@ -3,11 +3,11 @@ import logging -from langchain.schema import HumanMessage, SystemMessage +from langchain.schema import HumanMessage from src.config.agents import AGENT_LLM_MAP from src.llms.llm import get_llm_by_type -from src.prompts.template import env, apply_prompt_template +from src.prompts.template import apply_prompt_template from src.prompt_enhancer.graph.state import PromptEnhancerState logger = logging.getLogger(__name__) diff --git a/src/rag/vikingdb_knowledge_base.py b/src/rag/vikingdb_knowledge_base.py index 4f2672f..9c235e7 100644 --- a/src/rag/vikingdb_knowledge_base.py +++ b/src/rag/vikingdb_knowledge_base.py @@ -48,12 +48,12 @@ class VikingDBKnowledgeBaseProvider(Retriever): if params: for key in params: if ( - type(params[key]) == int - or type(params[key]) == float - or type(params[key]) == bool + type(params[key]) is int + or type(params[key]) is float + or type(params[key]) is bool ): params[key] = str(params[key]) - elif type(params[key]) == list: + elif type(params[key]) is list: if not doseq: params[key] = ",".join(params[key]) diff --git a/src/tools/search.py b/src/tools/search.py index 3f9f2bf..8852ea0 100644 --- a/src/tools/search.py +++ b/src/tools/search.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import json import logging import os from typing import List, Optional diff --git a/tests/integration/test_nodes.py b/tests/integration/test_nodes.py index 8f95e89..2eb3b71 100644 --- a/tests/integration/test_nodes.py +++ b/tests/integration/test_nodes.py @@ -1,9 +1,7 @@ from collections import namedtuple import json import pytest -import asyncio -import types -from unittest.mock import patch, MagicMock, AsyncMock +from unittest.mock import patch, MagicMock from src.graph.nodes import planner_node from src.graph.nodes import human_feedback_node from src.graph.nodes import coordinator_node @@ -39,7 +37,7 @@ def mock_state(): @pytest.fixture def mock_configurable(): mock = MagicMock() - mock.max_search_results = 5 + mock.max_search_results = 7 return mock @@ -668,8 +666,6 @@ def test_coordinator_node_tool_calls_exception_handling( patch_handoff_to_planner, patch_logger, ): - # tool_calls raises exception in processing - tool_calls = [{"name": "handoff_to_planner", "args": None}] with ( patch("src.graph.nodes.AGENT_LLM_MAP", {"coordinator": "basic"}), patch("src.graph.nodes.get_llm_by_type") as mock_get_llm, @@ -1037,24 +1033,6 @@ async def test_execute_agent_step_recursion_limit_env_negative( ) -@pytest.fixture -def mock_state_with_steps(mock_step, mock_completed_step): - # Simulate a plan with one completed and one unexecuted step - Plan = MagicMock() - Plan.steps = [mock_completed_step, mock_step] - return { - "current_plan": Plan, - "observations": ["obs1"], - "locale": "en-US", - "resources": [], - } - - -@pytest.fixture -def mock_config(): - return MagicMock() - - @pytest.fixture def mock_configurable_with_mcp(): mock = MagicMock() @@ -1297,27 +1275,6 @@ def mock_state_without_resources(): return {"other": "value"} -@pytest.fixture -def mock_config(): - return MagicMock() - - -@pytest.fixture -def mock_configurable(): - mock = MagicMock() - mock.max_search_results = 7 - return mock - - -@pytest.fixture -def patch_config_from_runnable_config(mock_configurable): - with patch( - "src.graph.nodes.Configuration.from_runnable_config", - return_value=mock_configurable, - ): - yield - - @pytest.fixture def patch_get_web_search_tool(): with patch("src.graph.nodes.get_web_search_tool") as mock: diff --git a/tests/test_state.py b/tests/test_state.py index c25b77b..598ba1c 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest import sys import os from typing import Annotated diff --git a/tests/unit/config/test_configuration.py b/tests/unit/config/test_configuration.py index a70aa79..b499b57 100644 --- a/tests/unit/config/test_configuration.py +++ b/tests/unit/config/test_configuration.py @@ -1,13 +1,8 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import os -import pytest import sys import types -from pathlib import Path -import builtins -import importlib from src.config.configuration import Configuration # Patch sys.path so relative import works diff --git a/tests/unit/config/test_loader.py b/tests/unit/config/test_loader.py index aaf75af..3d5538d 100644 --- a/tests/unit/config/test_loader.py +++ b/tests/unit/config/test_loader.py @@ -3,8 +3,6 @@ import os import tempfile -import yaml -import pytest from src.config.loader import load_yaml_config, process_dict, replace_env_vars diff --git a/tests/unit/crawler/test_article.py b/tests/unit/crawler/test_article.py index b2aa918..d9bd384 100644 --- a/tests/unit/crawler/test_article.py +++ b/tests/unit/crawler/test_article.py @@ -1,6 +1,5 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest from src.crawler.article import Article diff --git a/tests/unit/crawler/test_crawler_class.py b/tests/unit/crawler/test_crawler_class.py index eba2148..6f6429a 100644 --- a/tests/unit/crawler/test_crawler_class.py +++ b/tests/unit/crawler/test_crawler_class.py @@ -1,9 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest import src.crawler as crawler_module -from src.crawler import Crawler def test_crawler_sets_article_url(monkeypatch): diff --git a/tests/unit/llms/test_llm.py b/tests/unit/llms/test_llm.py index 2587080..dcdbc2f 100644 --- a/tests/unit/llms/test_llm.py +++ b/tests/unit/llms/test_llm.py @@ -1,8 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import os -import types import pytest from src.llms import llm diff --git a/tests/unit/prompt_enhancer/graph/test_builder.py b/tests/unit/prompt_enhancer/graph/test_builder.py index 575c270..f661f8b 100644 --- a/tests/unit/prompt_enhancer/graph/test_builder.py +++ b/tests/unit/prompt_enhancer/graph/test_builder.py @@ -6,7 +6,6 @@ from unittest.mock import patch, MagicMock from src.prompt_enhancer.graph.builder import build_graph from src.prompt_enhancer.graph.state import PromptEnhancerState -from src.config.report_style import ReportStyle class TestBuildGraph: @@ -48,7 +47,7 @@ class TestBuildGraph: mock_state_graph.return_value = mock_builder mock_builder.compile.return_value = mock_compiled_graph - result = build_graph() + build_graph() # Verify the correct node function was added mock_builder.add_node.assert_called_once_with("enhancer", mock_enhancer_node) diff --git a/tests/unit/prompt_enhancer/graph/test_state.py b/tests/unit/prompt_enhancer/graph/test_state.py index 8501882..3006b78 100644 --- a/tests/unit/prompt_enhancer/graph/test_state.py +++ b/tests/unit/prompt_enhancer/graph/test_state.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest from src.prompt_enhancer.graph.state import PromptEnhancerState from src.config.report_style import ReportStyle diff --git a/tests/unit/rag/test_ragflow.py b/tests/unit/rag/test_ragflow.py index 202cb5b..b5310ad 100644 --- a/tests/unit/rag/test_ragflow.py +++ b/tests/unit/rag/test_ragflow.py @@ -1,9 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import os import pytest -import requests from unittest.mock import patch, MagicMock from src.rag.ragflow import RAGFlowProvider, parse_uri @@ -144,30 +142,6 @@ def test_list_resources_success(mock_get, monkeypatch): assert resources[1].description == "desc2" -@patch("src.rag.ragflow.requests.get") -def test_list_resources_success(mock_get, monkeypatch): - monkeypatch.setenv("RAGFLOW_API_URL", "http://api") - monkeypatch.setenv("RAGFLOW_API_KEY", "key") - provider = RAGFlowProvider() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = { - "data": [ - {"id": "123", "name": "Dataset1", "description": "desc1"}, - {"id": "456", "name": "Dataset2", "description": "desc2"}, - ] - } - mock_get.return_value = mock_response - resources = provider.list_resources() - assert len(resources) == 2 - assert resources[0].uri == "rag://dataset/123" - assert resources[0].title == "Dataset1" - assert resources[0].description == "desc1" - assert resources[1].uri == "rag://dataset/456" - assert resources[1].title == "Dataset2" - assert resources[1].description == "desc2" - - @patch("src.rag.ragflow.requests.get") def test_list_resources_error(mock_get, monkeypatch): monkeypatch.setenv("RAGFLOW_API_URL", "http://api") diff --git a/tests/unit/rag/test_vikingdb_knowledge_base.py b/tests/unit/rag/test_vikingdb_knowledge_base.py index 22203e7..dbbea11 100644 --- a/tests/unit/rag/test_vikingdb_knowledge_base.py +++ b/tests/unit/rag/test_vikingdb_knowledge_base.py @@ -173,8 +173,8 @@ class TestVikingDBKnowledgeBaseProviderPrepareRequest: """Test basic request preparation""" with ( patch("src.rag.vikingdb_knowledge_base.Request") as mock_request, - patch("src.rag.vikingdb_knowledge_base.Credentials") as mock_credentials, - patch("src.rag.vikingdb_knowledge_base.SignerV4.sign") as mock_sign, + patch("src.rag.vikingdb_knowledge_base.Credentials") as _mock_credentials, + patch("src.rag.vikingdb_knowledge_base.SignerV4.sign") as _mock_sign, ): mock_req_instance = MagicMock() diff --git a/tests/unit/server/test_app.py b/tests/unit/server/test_app.py index 0016b50..f6eb503 100644 --- a/tests/unit/server/test_app.py +++ b/tests/unit/server/test_app.py @@ -2,31 +2,17 @@ # SPDX-License-Identifier: MIT import base64 -import json import os -from unittest.mock import AsyncMock, MagicMock, patch, mock_open -from uuid import uuid4 -from fastapi.responses import JSONResponse, StreamingResponse +from unittest.mock import MagicMock, patch, mock_open import pytest from fastapi.testclient import TestClient -from fastapi import HTTPException, logger +from fastapi import HTTPException from src.server.app import app, _make_event, _astream_workflow_generator -from src.server.mcp_request import MCPServerMetadataRequest -from src.server.rag_request import RAGResourceRequest from src.config.report_style import ReportStyle from langgraph.types import Command from langchain_core.messages import ToolMessage from langchain_core.messages import AIMessageChunk -from src.server.chat_request import ( - ChatRequest, - TTSRequest, - GeneratePodcastRequest, - GeneratePPTRequest, - GenerateProseRequest, - EnhancePromptRequest, -) - @pytest.fixture def client(): diff --git a/tests/unit/server/test_chat_request.py b/tests/unit/server/test_chat_request.py index 0556a7c..d3f5363 100644 --- a/tests/unit/server/test_chat_request.py +++ b/tests/unit/server/test_chat_request.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import asyncio # Ensure asyncio is imported import pytest from pydantic import ValidationError from src.config.report_style import ReportStyle diff --git a/tests/unit/tools/test_crawl.py b/tests/unit/tools/test_crawl.py index 405c44f..cdbe4cf 100644 --- a/tests/unit/tools/test_crawl.py +++ b/tests/unit/tools/test_crawl.py @@ -1,4 +1,3 @@ -import pytest from unittest.mock import Mock, patch from src.tools.crawl import crawl_tool diff --git a/tests/unit/tools/test_decorators.py b/tests/unit/tools/test_decorators.py index 4dbc2b5..5864dcb 100644 --- a/tests/unit/tools/test_decorators.py +++ b/tests/unit/tools/test_decorators.py @@ -1,10 +1,8 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest -import logging -from unittest.mock import Mock, call, patch, MagicMock -from src.tools.decorators import LoggedToolMixin, create_logged_tool +from unittest.mock import Mock, call, patch +from src.tools.decorators import create_logged_tool class MockBaseTool: @@ -58,7 +56,7 @@ class TestLoggedToolMixin: tool = LoggedTool() with patch("src.tools.decorators.logger.debug") as mock_debug: - result = tool._run("test_arg") + tool._run("test_arg") # Verify debug log was called with correct message mock_debug.assert_has_calls( diff --git a/tests/unit/tools/test_python_repl.py b/tests/unit/tools/test_python_repl.py index 963744e..728b544 100644 --- a/tests/unit/tools/test_python_repl.py +++ b/tests/unit/tools/test_python_repl.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: MIT import pytest -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import patch from src.tools.python_repl import python_repl_tool diff --git a/tests/unit/tools/test_search.py b/tests/unit/tools/test_search.py index cc914a2..1fb5c20 100644 --- a/tests/unit/tools/test_search.py +++ b/tests/unit/tools/test_search.py @@ -3,7 +3,7 @@ import os import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import patch from src.tools.search import get_web_search_tool from src.config import SearchEngine diff --git a/tests/unit/tools/test_tavily_search_api_wrapper.py b/tests/unit/tools/test_tavily_search_api_wrapper.py index 37d6242..eb60316 100644 --- a/tests/unit/tools/test_tavily_search_api_wrapper.py +++ b/tests/unit/tools/test_tavily_search_api_wrapper.py @@ -3,7 +3,6 @@ import json import pytest from unittest.mock import Mock, patch, AsyncMock, MagicMock -import aiohttp import requests from src.tools.tavily_search.tavily_search_api_wrapper import ( EnhancedTavilySearchAPIWrapper, diff --git a/tests/unit/tools/test_tavily_search_results_with_images.py b/tests/unit/tools/test_tavily_search_results_with_images.py index 963dbf1..4db69ee 100644 --- a/tests/unit/tools/test_tavily_search_results_with_images.py +++ b/tests/unit/tools/test_tavily_search_results_with_images.py @@ -4,7 +4,6 @@ import json import pytest from unittest.mock import Mock, patch, AsyncMock -from typing import Dict, Any from src.tools.tavily_search.tavily_search_results_with_images import ( TavilySearchResultsWithImages, ) diff --git a/tests/unit/tools/test_tools_retriever.py b/tests/unit/tools/test_tools_retriever.py index fa73b68..052504a 100644 --- a/tests/unit/tools/test_tools_retriever.py +++ b/tests/unit/tools/test_tools_retriever.py @@ -1,7 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import Mock, patch from langchain_core.callbacks import ( CallbackManagerForToolRun, AsyncCallbackManagerForToolRun, diff --git a/tests/unit/utils/test_json_utils.py b/tests/unit/utils/test_json_utils.py index 35f3de1..f90707a 100644 --- a/tests/unit/utils/test_json_utils.py +++ b/tests/unit/utils/test_json_utils.py @@ -1,9 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import pytest import json -from unittest.mock import patch from src.utils.json_utils import repair_json_output