fix: dataset desc (#1045)

This commit is contained in:
takatost
2023-08-29 09:07:27 +08:00
committed by GitHub
parent a55ba6e614
commit 7b3314c5db
5 changed files with 17 additions and 18 deletions

View File

@@ -1,5 +1,6 @@
import json
import logging
from json import JSONDecodeError
from typing import Any, Dict, List, Union, Optional
@@ -44,10 +45,15 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
input_str: str,
**kwargs: Any,
) -> None:
# tool_name = serialized.get('name')
input_dict = json.loads(input_str.replace("'", "\""))
dataset_id = input_dict.get('dataset_id')
query = input_dict.get('query')
tool_name: str = serialized.get('name')
dataset_id = tool_name.removeprefix('dataset-')
try:
input_dict = json.loads(input_str.replace("'", "\""))
query = input_dict.get('query')
except JSONDecodeError:
query = input_str
self.conversation_message_task.on_dataset_query_end(DatasetQueryObj(dataset_id=dataset_id, query=query))
def on_tool_end(