From 44448ba68dabdfbacab92a1e7046789eb90ccccf Mon Sep 17 00:00:00 2001
From: Yeuoly <45712896+Yeuoly@users.noreply.github.com>
Date: Fri, 12 Apr 2024 11:12:52 +0800
Subject: [PATCH] fix: remove - in dataset retriever tool name (#3381)

---
 .../tool/dataset_retriever/dataset_multi_retriever_tool.py  | 4 ++--
 .../tools/tool/dataset_retriever/dataset_retriever_tool.py  | 2 +-
 web/app/components/app/chat/thought/tool.tsx                | 6 +++---
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py
index d9934acff9..e977772fd1 100644
--- a/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py
+++ b/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py
@@ -31,7 +31,7 @@ class DatasetMultiRetrieverToolInput(BaseModel):
 
 class DatasetMultiRetrieverTool(BaseTool):
     """Tool for querying multi dataset."""
-    name: str = "dataset-"
+    name: str = "dataset_"
     args_schema: type[BaseModel] = DatasetMultiRetrieverToolInput
     description: str = "dataset multi retriever and rerank. "
     tenant_id: str
@@ -47,7 +47,7 @@ class DatasetMultiRetrieverTool(BaseTool):
     @classmethod
     def from_dataset(cls, dataset_ids: list[str], tenant_id: str, **kwargs):
         return cls(
-            name=f'dataset-{tenant_id}',
+            name=f"dataset_{tenant_id.replace('-', '_')}",
             tenant_id=tenant_id,
             dataset_ids=dataset_ids,
             **kwargs
diff --git a/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py
index 13331d981b..b99e392de0 100644
--- a/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py
+++ b/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py
@@ -46,7 +46,7 @@ class DatasetRetrieverTool(BaseTool):
 
         description = description.replace('\n', '').replace('\r', '')
         return cls(
-            name=f'dataset-{dataset.id}',
+            name=f"dataset_{dataset.id.replace('-', '_')}",
             tenant_id=dataset.tenant_id,
             dataset_id=dataset.id,
             description=description,
diff --git a/web/app/components/app/chat/thought/tool.tsx b/web/app/components/app/chat/thought/tool.tsx
index 322df438b4..324fcde7bc 100644
--- a/web/app/components/app/chat/thought/tool.tsx
+++ b/web/app/components/app/chat/thought/tool.tsx
@@ -19,7 +19,7 @@ type Props = {
 }
 
 const getIcon = (toolName: string, allToolIcons: Record<string, string | Emoji>) => {
-  if (toolName.startsWith('dataset-'))
+  if (toolName.startsWith('dataset_'))
     return <DataSetIcon className='shrink-0'></DataSetIcon>
   const icon = allToolIcons[toolName]
   if (!icon)
@@ -50,9 +50,9 @@ const Tool: FC<Props> = ({
 }) => {
   const { t } = useTranslation()
   const { name, input, isFinished, output } = payload
-  const toolName = name.startsWith('dataset-') ? t('dataset.knowledge') : name
+  const toolName = name.startsWith('dataset_') ? t('dataset.knowledge') : name
   const [isShowDetail, setIsShowDetail] = useState(false)
-  const icon = getIcon(toolName, allToolIcons) as any
+  const icon = getIcon(name, allToolIcons) as any
   return (
     <div>
       <div className={cn(!isShowDetail && 'shadow-sm', !isShowDetail && 'inline-block', 'max-w-full overflow-x-auto bg-white rounded-md')}>