Split up display name and name for Tools (#241)

* CHange * change
cohere-ai · Jun 20, 2024 · ed8eb80 · ed8eb80
1 parent 457df1d
commit ed8eb80
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 12 deletions.
diff --git a/src/backend/config/tools.py b/src/backend/config/tools.py
@@ -26,17 +26,18 @@
 
 
 class ToolName(StrEnum):
-    Wiki_Retriever_LangChain = "Wikipedia"
+    Wiki_Retriever_LangChain = "wikipedia"
     Search_File = "search_file"
     Read_File = "read_document"
-    Python_Interpreter = "Python_Interpreter"
-    Calculator = "Calculator"
-    Tavily_Internet_Search = "Internet_Search"
+    Python_Interpreter = "python_interpreter"
+    Calculator = "calculator"
+    Tavily_Internet_Search = "internet_search"
 
 
 ALL_TOOLS = {
     ToolName.Wiki_Retriever_LangChain: ManagedTool(
         name=ToolName.Wiki_Retriever_LangChain,
+        display_name="Wikipedia",
         implementation=LangChainWikiRetriever,
         parameter_definitions={
             "query": {
@@ -54,6 +55,7 @@ class ToolName(StrEnum):
     ),
     ToolName.Search_File: ManagedTool(
         name=ToolName.Search_File,
+        display_name="Search File",
         implementation=SearchFileTool,
         parameter_definitions={
             "search_query": {
@@ -75,6 +77,7 @@ class ToolName(StrEnum):
     ),
     ToolName.Read_File: ManagedTool(
         name=ToolName.Read_File,
+        display_name="Read Document",
         implementation=ReadFileTool,
         parameter_definitions={
             "filename": {
@@ -91,6 +94,7 @@ class ToolName(StrEnum):
     ),
     ToolName.Python_Interpreter: ManagedTool(
         name=ToolName.Python_Interpreter,
+        display_name="Python Interpreter",
         implementation=PythonInterpreter,
         parameter_definitions={
             "code": {
@@ -107,6 +111,7 @@ class ToolName(StrEnum):
     ),
     ToolName.Calculator: ManagedTool(
         name=ToolName.Calculator,
+        display_name="Calculator",
         implementation=Calculator,
         parameter_definitions={
             "code": {
@@ -123,6 +128,7 @@ class ToolName(StrEnum):
     ),
     ToolName.Tavily_Internet_Search: ManagedTool(
         name=ToolName.Tavily_Internet_Search,
+        display_name="Web Search",
         implementation=TavilyInternetSearch,
         parameter_definitions={
             "query": {

diff --git a/src/backend/schemas/tool.py b/src/backend/schemas/tool.py
@@ -16,6 +16,7 @@ class ToolInput(BaseModel):
 
 class Tool(BaseModel):
     name: str
+    display_name: str = ""
     description: Optional[str] = ""
     parameter_definitions: Optional[dict] = {}
 

diff --git a/src/community/config/tools.py b/src/community/config/tools.py
@@ -13,17 +13,18 @@
 
 
 class CommunityToolName(StrEnum):
-    Arxiv = "Arxiv"
-    Connector = "Connector"
-    Pub_Med = "Pub Med"
-    File_Upload_LlamaIndex = "File Reader - LlamaIndex"
-    Wolfram_Alpha = "Wolfram_Alpha"
-    ClinicalTrials = "ClinicalTrials"
+    Arxiv = "arxiv"
+    Connector = "example_connector"
+    Pub_Med = "pub_med"
+    File_Upload_LlamaIndex = "file_reader_llamaindex"
+    Wolfram_Alpha = "wolfram_alpha"
+    ClinicalTrials = "clinical_trials"
 
 
 COMMUNITY_TOOLS = {
     CommunityToolName.Arxiv: ManagedTool(
         name=CommunityToolName.Arxiv,
+        display_name="Arxiv",
         implementation=ArxivRetriever,
         parameter_definitions={
             "query": {
@@ -40,6 +41,7 @@ class CommunityToolName(StrEnum):
     ),
     CommunityToolName.Connector: ManagedTool(
         name=CommunityToolName.Connector,
+        display_name="Example Connector",
         implementation=ConnectorRetriever,
         is_visible=True,
         is_available=ConnectorRetriever.is_available(),
@@ -49,6 +51,7 @@ class CommunityToolName(StrEnum):
     ),
     CommunityToolName.Pub_Med: ManagedTool(
         name=CommunityToolName.Pub_Med,
+        display_name="PubMed",
         implementation=PubMedRetriever,
         parameter_definitions={
             "query": {
@@ -65,6 +68,7 @@ class CommunityToolName(StrEnum):
     ),
     CommunityToolName.File_Upload_LlamaIndex: ManagedTool(
         name=CommunityToolName.File_Upload_LlamaIndex,
+        display_name="File Reader",
         implementation=LlamaIndexUploadPDFRetriever,
         is_visible=True,
         is_available=LlamaIndexUploadPDFRetriever.is_available(),
@@ -74,6 +78,7 @@ class CommunityToolName(StrEnum):
     ),
     CommunityToolName.Wolfram_Alpha: ManagedTool(
         name=CommunityToolName.Wolfram_Alpha,
+        display_name="Wolfram Alpha",
         implementation=WolframAlpha,
         is_visible=False,
         is_available=WolframAlpha.is_available(),
@@ -83,6 +88,7 @@ class CommunityToolName(StrEnum):
     ),
     CommunityToolName.ClinicalTrials: ManagedTool(
         name=CommunityToolName.ClinicalTrials,
+        display_name="Clinical Trials",
         implementation=ClinicalTrials,
         is_visible=True,
         is_available=ClinicalTrials.is_available(),

diff --git a/src/interfaces/coral_web/src/cohere-client/generated/models/ManagedTool.ts b/src/interfaces/coral_web/src/cohere-client/generated/models/ManagedTool.ts
@@ -9,6 +9,7 @@ import type { Category } from './Category';
 
 export type ManagedTool = {
   name: string;
+  display_name?: string;
   description?: string | null;
   parameter_definitions?: Record<string, any> | null;
   kwargs?: Record<string, any>;

diff --git a/src/interfaces/coral_web/src/cohere-client/generated/models/Tool.ts b/src/interfaces/coral_web/src/cohere-client/generated/models/Tool.ts
@@ -4,6 +4,7 @@
 /* eslint-disable */
 export type Tool = {
   name: string;
+  display_name?: string;
   description?: string | null;
   parameter_definitions?: Record<string, any> | null;
 };
diff --git a/src/interfaces/coral_web/src/components/Configuration/Tools.tsx b/src/interfaces/coral_web/src/components/Configuration/Tools.tsx
@@ -95,7 +95,7 @@ const ToolSection = () => {
             />
           </div>
           <div className="flex flex-col gap-y-5">
-            {tools.map(({ name, is_available, description, error_message }) => {
+            {tools.map(({ name, display_name, is_available, description, error_message }) => {
               const enabledTool = enabledTools.find(
                 (enabledTool) => enabledTool.name.toLocaleLowerCase() === name.toLocaleLowerCase()
               );
@@ -109,7 +109,7 @@ const ToolSection = () => {
                     onChange={(e) => {
                       handleToolToggle(name, e.target.checked);
                     }}
-                    label={name}
+                    label={display_name}
                     name={name}
                     theme="secondary"
                     dataTestId={`checkbox-tool-${name}`}