Bladeren bron

feat: add current time tool in universal chat agent (#659)

John Wang 1 jaar geleden
bovenliggende
commit
b5825142d1

+ 15 - 0
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']

+ 2 - 10
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,

+ 2 - 9
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.")

+ 14 - 1
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.")

+ 6 - 0
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")

+ 1 - 1
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: