From b5825142d134d4b38e857325ba13896334abfad3 Mon Sep 17 00:00:00 2001 From: John Wang Date: Thu, 27 Jul 2023 17:39:36 +0800 Subject: [PATCH] feat: add current time tool in universal chat agent (#659) --- api/controllers/console/universal_chat/chat.py | 15 +++++++++++++++ api/core/agent/agent/openai_function_call.py | 12 ++---------- .../agent/agent/openai_multi_function_call.py | 11 ++--------- api/core/orchestrator_rule_parser.py | 15 ++++++++++++++- api/libs/helper.py | 6 ++++++ api/services/app_model_config_service.py | 2 +- 6 files changed, 40 insertions(+), 21 deletions(-) diff --git a/api/controllers/console/universal_chat/chat.py b/api/controllers/console/universal_chat/chat.py index faa009115f..2a95eb992b 100644 --- a/api/controllers/console/universal_chat/chat.py +++ b/api/controllers/console/universal_chat/chat.py @@ -43,6 +43,21 @@ class UniversalChatApi(UniversalChatResource): args['model_config']['model']['provider'] = llm_constant.models[args['model']] args['model_config']['agent_mode']['tools'] = args['tools'] + if not args['model_config']['agent_mode']['tools']: + args['model_config']['agent_mode']['tools'] = [ + { + "current_datetime": { + "enabled": True + } + } + ] + else: + args['model_config']['agent_mode']['tools'].append({ + "current_datetime": { + "enabled": True + } + }) + args['inputs'] = {} del args['model'] diff --git a/api/core/agent/agent/openai_function_call.py b/api/core/agent/agent/openai_function_call.py index 8e37d3414f..090d35d975 100644 --- a/api/core/agent/agent/openai_function_call.py +++ b/api/core/agent/agent/openai_function_call.py @@ -1,7 +1,5 @@ -from datetime import datetime from typing import List, Tuple, Any, Union, Sequence, Optional -import pytz from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent from langchain.agents.openai_functions_agent.base import _parse_ai_message, \ _format_intermediate_steps @@ -98,15 +96,9 @@ class AutoSummarizingOpenAIFunctionCallAgent(OpenAIFunctionsAgent, OpenAIFunctio @classmethod def get_system_message(cls): - # get current time - current_time = datetime.now() - current_timezone = pytz.timezone('UTC') - current_time = current_timezone.localize(current_time) - return SystemMessage(content="You are a helpful AI assistant.\n" - "Current time: {}\n" - "Respond directly if appropriate.".format( - current_time.strftime("%Y-%m-%d %H:%M:%S %Z%z"))) + "The current date or current time you know is wrong.\n" + "Respond directly if appropriate.") def return_stopped_response( self, diff --git a/api/core/agent/agent/openai_multi_function_call.py b/api/core/agent/agent/openai_multi_function_call.py index e00ebdf421..1524fc6975 100644 --- a/api/core/agent/agent/openai_multi_function_call.py +++ b/api/core/agent/agent/openai_multi_function_call.py @@ -1,7 +1,5 @@ -from datetime import datetime from typing import List, Tuple, Any, Union, Sequence, Optional -import pytz from langchain.agents import BaseMultiActionAgent from langchain.agents.openai_functions_multi_agent.base import OpenAIMultiFunctionsAgent, _format_intermediate_steps, \ _parse_ai_message @@ -99,11 +97,6 @@ class AutoSummarizingOpenMultiAIFunctionCallAgent(OpenAIMultiFunctionsAgent, Ope @classmethod def get_system_message(cls): # get current time - current_time = datetime.now() - current_timezone = pytz.timezone('UTC') - current_time = current_timezone.localize(current_time) - return SystemMessage(content="You are a helpful AI assistant.\n" - "Current time: {}\n" - "Respond directly if appropriate.".format( - current_time.strftime("%Y-%m-%d %H:%M:%S %Z%z"))) + "The current date or current time you know is wrong.\n" + "Respond directly if appropriate.") diff --git a/api/core/orchestrator_rule_parser.py b/api/core/orchestrator_rule_parser.py index 971f7ffed5..7890544f97 100644 --- a/api/core/orchestrator_rule_parser.py +++ b/api/core/orchestrator_rule_parser.py @@ -21,6 +21,7 @@ from core.tool.provider.serpapi_provider import SerpAPIToolProvider from core.tool.serpapi_wrapper import OptimizedSerpAPIWrapper, OptimizedSerpAPIInput from core.tool.web_reader_tool import WebReaderTool from extensions.ext_database import db +from libs import helper from models.dataset import Dataset, DatasetProcessRule from models.model import AppModelConfig @@ -167,6 +168,8 @@ class OrchestratorRuleParser: tool = self.to_google_search_tool() elif tool_type == "wikipedia": tool = self.to_wikipedia_tool() + elif tool_type == "current_datetime": + tool = self.to_current_datetime_tool() if tool: tool.callbacks.extend(callbacks) @@ -235,7 +238,7 @@ class OrchestratorRuleParser: name="google_search", description="A tool for performing a Google search and extracting snippets and webpages " "when you need to search for something you don't know or when your information " - "is not up to date." + "is not up to date. " "Input should be a search query.", func=OptimizedSerpAPIWrapper(**func_kwargs).run, args_schema=OptimizedSerpAPIInput, @@ -244,6 +247,16 @@ class OrchestratorRuleParser: return tool + def to_current_datetime_tool(self) -> Optional[BaseTool]: + tool = Tool( + name="current_datetime", + description="A tool when you want to get the current date or time. ", + func=helper.get_current_datetime, + callbacks=[DifyStdOutCallbackHandler()] + ) + + return tool + def to_wikipedia_tool(self) -> Optional[BaseTool]: class WikipediaInput(BaseModel): query: str = Field(..., description="search query.") diff --git a/api/libs/helper.py b/api/libs/helper.py index b306fee2a7..4be7b6e1eb 100644 --- a/api/libs/helper.py +++ b/api/libs/helper.py @@ -153,3 +153,9 @@ def get_remote_ip(request): def generate_text_hash(text: str) -> str: hash_text = str(text) + 'None' return sha256(hash_text.encode()).hexdigest() + + +def get_current_datetime(type: str) -> str: + # get current time + current_time = datetime.utcnow() + return current_time.strftime("%Y-%m-%d %H:%M:%S %Z%z") diff --git a/api/services/app_model_config_service.py b/api/services/app_model_config_service.py index a558dc75eb..a5af67e301 100644 --- a/api/services/app_model_config_service.py +++ b/api/services/app_model_config_service.py @@ -39,7 +39,7 @@ SUPPORT_AGENT_MODELS = [ "gpt-3.5-turbo-16k", ] -SUPPORT_TOOLS = ["dataset", "google_search", "web_reader", "wikipedia"] +SUPPORT_TOOLS = ["dataset", "google_search", "web_reader", "wikipedia", "current_datetime"] class AppModelConfigService: