From 6d8853b7c70e9ef64784c98a3a4e7936575e8e73 Mon Sep 17 00:00:00 2001 From: DanielWalnut <45447813+hetaoBackend@users.noreply.github.com> Date: Tue, 22 Jul 2025 14:49:04 +0800 Subject: [PATCH] refine the research prompt (#460) --- .../tavily_search_results_with_images.py | 11 ++- .../test_tavily_search_results_with_images.py | 68 +------------------ 2 files changed, 12 insertions(+), 67 deletions(-) diff --git a/src/tools/tavily_search/tavily_search_results_with_images.py b/src/tools/tavily_search/tavily_search_results_with_images.py index e18680d..edd4a07 100644 --- a/src/tools/tavily_search/tavily_search_results_with_images.py +++ b/src/tools/tavily_search/tavily_search_results_with_images.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT +import logging import json from typing import Dict, List, Optional, Tuple, Union @@ -15,6 +16,8 @@ from src.tools.tavily_search.tavily_search_api_wrapper import ( EnhancedTavilySearchAPIWrapper, ) +logger = logging.getLogger(__name__) + class TavilySearchResultsWithImages(TavilySearchResults): # type: ignore[override, override] """Tool that queries the Tavily Search API and gets back json. @@ -123,7 +126,9 @@ class TavilySearchResultsWithImages(TavilySearchResults): # type: ignore[overri except Exception as e: return repr(e), {} cleaned_results = self.api_wrapper.clean_results_with_images(raw_results) - print("sync", json.dumps(cleaned_results, indent=2, ensure_ascii=False)) + logger.debug( + "sync: %s", json.dumps(cleaned_results, indent=2, ensure_ascii=False) + ) return cleaned_results, raw_results async def _arun( @@ -147,5 +152,7 @@ class TavilySearchResultsWithImages(TavilySearchResults): # type: ignore[overri except Exception as e: return repr(e), {} cleaned_results = self.api_wrapper.clean_results_with_images(raw_results) - print("async", json.dumps(cleaned_results, indent=2, ensure_ascii=False)) + logger.debug( + "async: %s", json.dumps(cleaned_results, indent=2, ensure_ascii=False) + ) return cleaned_results, raw_results 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 4db69ee..562089c 100644 --- a/tests/unit/tools/test_tavily_search_results_with_images.py +++ b/tests/unit/tools/test_tavily_search_results_with_images.py @@ -1,9 +1,8 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -import json import pytest -from unittest.mock import Mock, patch, AsyncMock +from unittest.mock import Mock, AsyncMock from src.tools.tavily_search.tavily_search_results_with_images import ( TavilySearchResultsWithImages, ) @@ -77,10 +76,8 @@ class TestTavilySearchResultsWithImages: assert tool.max_results == 10 assert tool.include_image_descriptions is True - @patch("builtins.print") def test_run_success( self, - mock_print, search_tool, mock_api_wrapper, sample_raw_results, @@ -110,10 +107,8 @@ class TestTavilySearchResultsWithImages: mock_api_wrapper.clean_results_with_images.assert_called_once_with( sample_raw_results ) - mock_print.assert_called_once() - @patch("builtins.print") - def test_run_exception(self, mock_print, search_tool, mock_api_wrapper): + def test_run_exception(self, search_tool, mock_api_wrapper): """Test synchronous run with exception.""" mock_api_wrapper.raw_results.side_effect = Exception("API Error") @@ -124,10 +119,8 @@ class TestTavilySearchResultsWithImages: mock_api_wrapper.clean_results_with_images.assert_not_called() @pytest.mark.asyncio - @patch("builtins.print") async def test_arun_success( self, - mock_print, search_tool, mock_api_wrapper, sample_raw_results, @@ -157,11 +150,9 @@ class TestTavilySearchResultsWithImages: mock_api_wrapper.clean_results_with_images.assert_called_once_with( sample_raw_results ) - mock_print.assert_called_once() @pytest.mark.asyncio - @patch("builtins.print") - async def test_arun_exception(self, mock_print, search_tool, mock_api_wrapper): + async def test_arun_exception(self, search_tool, mock_api_wrapper): """Test asynchronous run with exception.""" mock_api_wrapper.raw_results_async = AsyncMock( side_effect=Exception("Async API Error") @@ -173,10 +164,8 @@ class TestTavilySearchResultsWithImages: assert raw == {} mock_api_wrapper.clean_results_with_images.assert_not_called() - @patch("builtins.print") def test_run_with_run_manager( self, - mock_print, search_tool, mock_api_wrapper, sample_raw_results, @@ -193,10 +182,8 @@ class TestTavilySearchResultsWithImages: assert raw == sample_raw_results @pytest.mark.asyncio - @patch("builtins.print") async def test_arun_with_run_manager( self, - mock_print, search_tool, mock_api_wrapper, sample_raw_results, @@ -213,52 +200,3 @@ class TestTavilySearchResultsWithImages: assert result == sample_cleaned_results assert raw == sample_raw_results - - @patch("builtins.print") - def test_print_output_format( - self, - mock_print, - search_tool, - mock_api_wrapper, - sample_raw_results, - sample_cleaned_results, - ): - """Test that print outputs correctly formatted JSON.""" - mock_api_wrapper.raw_results.return_value = sample_raw_results - mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results - - search_tool._run("test query") - - # Verify print was called with expected format - call_args = mock_print.call_args[0] - assert call_args[0] == "sync" - assert isinstance(call_args[1], str) # Should be JSON string - - # Verify it's valid JSON - json_data = json.loads(call_args[1]) - assert json_data == sample_cleaned_results - - @pytest.mark.asyncio - @patch("builtins.print") - async def test_async_print_output_format( - self, - mock_print, - search_tool, - mock_api_wrapper, - sample_raw_results, - sample_cleaned_results, - ): - """Test that async print outputs correctly formatted JSON.""" - mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results) - mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results - - await search_tool._arun("test query") - - # Verify print was called with expected format - call_args = mock_print.call_args[0] - assert call_args[0] == "async" - assert isinstance(call_args[1], str) # Should be JSON string - - # Verify it's valid JSON - json_data = json.loads(call_args[1]) - assert json_data == sample_cleaned_results