TLK-2361 - Tools preamble extension (#876)

* TLK-2361 - Tools preamble extension

* TLK-2361 - Tools preamble extension

.env-template

@@ -15,6 +15,10 @@ PYTHON_INTERPRETER_URL=http://terrarium:8080
 TAVILY_API_KEY=<API_KEY_HERE>
 WOLFRAM_APP_ID=<APP_ID_HERE>
 
+# Use tools preamble to enhance the model with additional information for better tool usage.
+USE_TOOLS_PREAMBLE=True
+
+
 # CORAL WEB
 
 NEXT_PUBLIC_HAS_CUSTOM_LOGO=false

docs/config_details/config_description.md

@@ -25,6 +25,7 @@
   - tools - Tool configurations
      - python_interpreter - Python interpreter configurations
        - url - URL of the python interpreter tool
+     - use_tools_preamble - Use tools preamble - if set to true, the tools preamble will be used in the chat requests
   - feature_flags - Feature flags configurations
        - use_agents_view - Use agents view - if set to true, the frontend agents view will be available. 
          Please note that this setting is available only for the Coral web frontend. To change which frontend is used, set the context in the docker-compose file. 

src/backend/chat/custom/custom.py

@@ -6,14 +6,15 @@
 from backend.chat.custom.tool_calls import async_call_tools
 from backend.chat.custom.utils import get_deployment
 from backend.chat.enums import StreamEvent
+from backend.config import Settings
 from backend.config.tools import get_available_tools
 from backend.database_models.file import File
 from backend.model_deployments.base import BaseDeployment
 from backend.schemas.chat import ChatMessage, ChatRole, EventState
 from backend.schemas.cohere_chat import CohereChatRequest
 from backend.schemas.context import Context
 from backend.schemas.tool import Tool, ToolCategory
-from backend.services.chat import check_death_loop
+from backend.services.chat import check_death_loop, generate_tools_preamble
 from backend.services.file import get_file_service
 from backend.tools.utils.tools_checkers import tool_has_category
 
@@ -196,6 +197,9 @@ async def call_chat(
                 if tool.name not in file_reader_tools_names
             ]
 
+        if chat_request.tools and Settings().get("tools.use_tools_preamble"):
+            chat_request.preamble = generate_tools_preamble(chat_request)
+
         # Loop until there are no new tool calls
         for step in range(MAX_STEPS):
             logger.debug(

src/backend/config/configuration.template.yaml

@@ -37,6 +37,8 @@ tools:
   slack:
     user_scopes:
       - search:read
+  # To disable the use of the tools preamble, set it to false
+  use_tools_preamble: true
 feature_flags:
   # Experimental features
   use_agents_view: false

src/backend/config/settings.py

@@ -240,6 +240,10 @@ class ToolSettings(BaseSettings, BaseModel):
     gmail: Optional[GmailSettings] = Field(
         default=GmailSettings()
     )
+    use_tools_preamble: Optional[bool] = Field(
+        default=False,
+        validation_alias=AliasChoices("USE_TOOLS_PREAMBLE", "use_tools_preamble")
+    )
 
 
 class DatabaseSettings(BaseSettings, BaseModel):

src/backend/services/chat.py

@@ -52,6 +52,26 @@
 LOOKBACKS = [3, 5, 7]
 DEATHLOOP_SIMILARITY_THRESHOLDS = [0.5, 0.7, 0.9]
 
+
+def generate_tools_preamble(chat_request: CohereChatRequest) -> str:
+    available_tools = get_available_tools()
+    full_managed_tools = [
+        available_tools.get(tool.name)
+        for tool in chat_request.tools
+        if available_tools.get(tool.name)
+    ]
+    tools_preamble = ""
+    if full_managed_tools:
+        tools_preamble = " ".join(tool.implementation.TOOL_DEFAULT_PREAMBLE for tool in full_managed_tools if
+                                   tool.implementation.TOOL_DEFAULT_PREAMBLE)
+    passed_preamble = ""
+    if chat_request.preamble:
+        passed_preamble = chat_request.preamble.replace("## Task And Context", "")
+    tools_preamble = f"## Task And Context\n{tools_preamble}{passed_preamble}"
+
+    return tools_preamble
+
+
 def process_chat(
     session: DBSessionDep,
     chat_request: BaseChatRequest,
@@ -1017,15 +1037,15 @@ def check_death_loop(
         event_state.distances_actions.append(
             1
             - nltk.edit_distance(event_state.previous_action, action)
-            / max(len(event_state.previous_action), len(action))
+            / max(len(str(event_state.previous_action)), len(action))
         )
         check_similarity(event_state.distances_actions, ctx)
 
     if event_state.previous_plan:
         event_state.distances_plans.append(
             1
             - nltk.edit_distance(event_state.previous_plan, plan)
-            / max(len(event_state.previous_plan), len(plan))
+            / max(len(str(event_state.previous_plan)), len(plan))
         )
         check_similarity(event_state.distances_plans, ctx)
 

src/backend/tools/base.py

@@ -52,6 +52,32 @@ def __new__(cls, name, bases, class_dict):
         return super().__new__(cls, name, bases, class_dict)
 
 
+class BaseToolPreambleRegistry:
+    """
+    Registry to store default preamble values for tools.
+    Currently, it's implemented as a class with predefined values, but we can use set_preamble to set values.
+    """
+    _default_preambles = {
+        "toolkit_python_interpreter": "If you decide to use toolkit_python_interpreter tool and are going to plot something, try returning result as a png. Ensure that the generated code does not include any internet connection, and avoid using the module named requests. ",
+        "read_file": "When using the read_file tool, always ensure that the file parameter is prepared as a tuple in the format  (filename, file ID). The order of the tuple fields is critical. ",
+        "search_file":"When using the search_file tool, always ensure that the `files` parameter is prepared as a list of tuples in the format (filename, file ID). The order of the tuple fields is critical. "
+    }
+
+    @classmethod
+    def set_preamble(cls, tool_id, preamble):
+        """
+        Register a default preamble for a tool class. This method can be used to set a default preamble for a tool class.
+        """
+        cls._default_preambles[tool_id] = preamble
+
+    @classmethod
+    def get_preamble(cls, tool_class):
+        """
+        Retrieve the default preamble for a given tool class.
+        """
+        return cls._default_preambles.get(tool_class, None)
+
+
 class BaseTool(metaclass=ParametersValidationMeta):
     """
     Abstract base class for all Tools.
@@ -60,6 +86,15 @@ class BaseTool(metaclass=ParametersValidationMeta):
         ID (str): The name of the tool.
     """
     ID = None
+    TOOL_DEFAULT_PREAMBLE = None
+
+    def __init_subclass__(cls, **kwargs):
+        """
+        Automatically set TOOL_DEFAULT_PREAMBLE for child classes
+        using the registry class.
+        """
+        super().__init_subclass__(**kwargs)
+        cls.TOOL_DEFAULT_PREAMBLE = BaseToolPreambleRegistry.get_preamble(cls.ID)
 
     def __init__(self, *args, **kwargs):
         self._post_init_check()