fix: remove - in dataset retriever tool name (langgenius#3381)

gcc10055 · Apr 15, 2024 · f58c931 · f58c931
1 parent 61402a3
commit f58c931
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 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
@@ -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
@@ -46,7 +46,7 @@ def from_dataset(cls, dataset: Dataset, **kwargs):
 
         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
@@ -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')}>