Merge pull request #152 from bioimage-io/fix-tool-info

Fix tool info

bioimageio_chatbot/chatbot.py

@@ -88,6 +88,7 @@ async def respond_to_user(
         extensions_by_id = {ext.id: ext for ext in builtin_extensions}
         extensions_by_name = {ext.name: ext for ext in builtin_extensions}
         extensions_by_tool_name = {}
+        tool_info = {}
 
         tools = []
         tool_prompts = {}
@@ -115,6 +116,11 @@ async def respond_to_user(
                 assert len(extension.description) <= max_length, f"Extension tool prompt is too long: {extension.description}"
                 tool_prompts[create_tool_name(extension.id) + "*"] = extension.description.replace("\n", ";")[:max_length]
             extensions_by_tool_name.update({t.__name__: extension for t in ts})
+
+            if extension.info:
+                for t in ts:
+                    tool_info[t.__name__] = extension.info.get(t.__tool_id__)
+
             tools += ts
             if extension.get_state:
                 # the state of the extension is a dictionary with keys and values for the states of the extension
@@ -152,24 +158,13 @@ class ThoughtsSchema(BaseModel):
             tool_usage_prompt=tool_usage_prompt,
         )
         result_steps = metadata["steps"]
-
-        # Create a map to store the mapping of result step names to their corresponding info
-        info_map = {}
-        # Iterate through each extension in the dictionary
-        for ext_id, ext in extensions_by_id.items():
-            # Iterate through the info dictionary of each extension
-            for tool_key, tool_info in ext.info.items():
-                # remove all special characters and convert to lowercase, and add `ext_id` to the beginning
-                step_name = f"{ext_id}-{tool_key}".replace(" ", "").replace("-", "").replace("_","").lower()
-                # Map the step name to the tool info
-                info_map[step_name] = tool_info
 
         for idx, step_list in enumerate(result_steps):
             for step in step_list:
                 # Get the step name
-                step_name = step['name'].lower()
+                step_name = step['name']
                 # Get the corresponding info for the step
-                step_info = info_map.get(step_name)
+                step_info = tool_info.get(step_name)
                 if step_info:
                     # Add the info to the step details
                     step["info"] = step_info

bioimageio_chatbot/chatbot_extensions/__init__.py

@@ -61,12 +61,14 @@ async def extension_to_tools(extension: ChatbotExtension):
             assert k in extension.tools, f"Tool `{k}` not found in extension `{extension.id}`."
             ext_tool = extension.tools[k]
             tool = tool_factory(extension.id, k, ext_tool, schemas[k])
+            ext_tool.__tool_id__ = k
             tools.append(tool)
     else:
         tools = []
         for k in extension.tools:
             ext_tool = extension.tools[k]
             ext_tool.__name__ = create_tool_name(extension.id, k)
+            ext_tool.__tool_id__ = k
             tools.append(ext_tool)
 
     return tools

pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "bioimageio-chatbot"
-version = "0.2.10"
+version = "0.2.11"
 readme = "README.md"
 description = "Your Personal Assistant in Computational BioImaging."
 dependencies = [

requirements.txt

@@ -4,7 +4,7 @@ requests
 pypdf
 pillow
 matplotlib
-hypha>=0.15.50
+hypha==0.15.53
 tqdm
 aiofiles
 serpapi