import logging
from openjiuwen.core.foundation.tool.base import ToolCard
from openjiuwen.core.foundation.tool.function.function import LocalFunction
from openjiuwen_deepsearch.algorithm.research_collector.tool_log import tool_invoke_log_async
from openjiuwen_deepsearch.common.exception import CustomValueException
from openjiuwen_deepsearch.common.status_code import StatusCode
from openjiuwen_deepsearch.framework.openjiuwen.tools.search_api import (
LocalDatasetAPIWrapper,
NativeLocalSearchAPIWrapper,
load_external_search_tools,
)
from openjiuwen_deepsearch.utils.constants_utils.session_contextvars import local_search_context
from openjiuwen_deepsearch.utils.constants_utils.search_engine_constants import LocalSearch
from openjiuwen_deepsearch.utils.log_utils.log_manager import LogManager
logger = logging.getLogger(__name__)
local_search_mapping = {
LocalSearch.OPENAPI.value: LocalDatasetAPIWrapper,
LocalSearch.NATIVE.value: NativeLocalSearchAPIWrapper,
}
def update_local_search_mapping(func_path: str, func_name: str):
"""加载外部搜索工具，并更新本地搜索映射字典"""
engine_name, external_mapping = load_external_search_tools(func_path, func_name)
if external_mapping:
local_search_mapping["custom"] = external_mapping[engine_name]
return local_search_mapping
@tool_invoke_log_async
async def run_local_search(query: str, search_engine_name: str):
"""运行本地搜索"""
api_wrapper = local_search_context.get().get(search_engine_name)
if not api_wrapper:
raise CustomValueException(
StatusCode.LOCAL_SEARCH_INSTANCE_OBTAIN_ERROR.code,
StatusCode.LOCAL_SEARCH_INSTANCE_OBTAIN_ERROR.errmsg.format(name=search_engine_name)
)
try:
result = await api_wrapper.aresults(query)
except Exception as e:
if LogManager.is_sensitive():
logger.error(f"Error when run local search")
else:
logger.exception(f"Error when run local search: {e}")
return dict(search_engine=search_engine_name,
search_results=[],
error=f"Error when run local search {search_engine_name}: {repr(e)}")
return dict(search_engine=search_engine_name, search_results=result)
def create_local_search_tool():
"""获取本地搜索工具"""
card = ToolCard(
id="local_search_tool",
name="local_search_tool",
description="Use local search engine to get local dataset information.",
input_params={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query of current step."
},
"search_engine_name": {
"type": "string",
"description": "Name of the search engine to use."
}
},
"required": ["query", "search_engine_name"]
}
)
local_search_tool = LocalFunction(
card=card,
func=run_local_search
)
return local_search_tool