From 2a32152e40636d8d4c32b95dcaf31efe5a396758 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 22 Oct 2024 16:34:38 -0400 Subject: [PATCH] switch to content_and_artifact, add output_format --- .../tools/ddg_search/tool.py | 37 ++++++++++++++----- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/libs/community/langchain_community/tools/ddg_search/tool.py b/libs/community/langchain_community/tools/ddg_search/tool.py index 54b93b46d348e..c92a7925e32d5 100644 --- a/libs/community/langchain_community/tools/ddg_search/tool.py +++ b/libs/community/langchain_community/tools/ddg_search/tool.py @@ -1,7 +1,8 @@ """Tool for the DuckDuckGo search API.""" +import json import warnings -from typing import Any, List, Optional, Type +from typing import Any, List, Literal, Optional, Type, Union from langchain_core.callbacks import CallbackManagerForToolRun from langchain_core.tools import BaseTool @@ -74,14 +75,13 @@ def _run( class DuckDuckGoSearchResults(BaseTool): - """Tool that queries the DuckDuckGo search API and gets back List of Dictionary.""" + """Tool that queries the DuckDuckGo search API and returns the results in `output_format`.""" name: str = "duckduckgo_results_json" description: str = ( "A wrapper around Duck Duck Go Search. " "Useful for when you need to answer questions about current events. " - "Input should be a search query. Output is a List of Dictionaries of the " - "query results" + "Input should be a search query." ) max_results: int = Field(alias="num_results", default=4) api_wrapper: DuckDuckGoSearchAPIWrapper = Field( @@ -93,25 +93,44 @@ class DuckDuckGoSearchResults(BaseTool): """Which keys from each result to include. If None all keys are included.""" results_separator: str = ", " """Character for separating results.""" + output_format: Literal["string", "json", "list"] = "string" + """Output format of the search results. + + - 'string': Return a concatenated string of the search results. + - 'json': Return a JSON string of the search results. + - 'list': Return a list of dictionaries of the search results. + """ + response_format: Literal["content_and_artifact"] = "content_and_artifact" def _run( self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None, - ) -> list[dict]: + ) -> tuple[Union[List[dict], str], List[dict]]: """Use the tool.""" - res = self.api_wrapper.results(query, self.max_results, source=self.backend) - + raw_results = self.api_wrapper.results( + query, self.max_results, source=self.backend + ) results = [ { k: v for k, v in d.items() if not self.keys_to_include or k in self.keys_to_include } - for d in res + for d in raw_results ] - return results + if self.output_format == "list": + return results, raw_results + elif self.output_format == "json": + return json.dumps(results), raw_results + elif self.output_format == "string": + res_strs = [", ".join([f"{k}: {v}" for k, v in d.items()]) for d in results] + return self.results_separator.join(res_strs), raw_results + else: + raise ValueError( + f"Invalid output_format: {self.output_format}. Needs to be one of 'string', 'json', 'list'." + ) def DuckDuckGoSearchTool(*args: Any, **kwargs: Any) -> DuckDuckGoSearchRun: