diff --git a/llm/models.py b/llm/models.py index 892a4f2ba..02295e820 100644 --- a/llm/models.py +++ b/llm/models.py @@ -950,7 +950,9 @@ def log_to_db(self, db): "tool_call_id": tool_call.tool_call_id, } ) + tools_by_name = {tool.name: tool for tool in self.prompt.tools} for tool_result in self.prompt.tool_results: + matching_tool = tools_by_name.get(tool_result.name) instance_id = None if tool_result.instance: try: @@ -959,8 +961,15 @@ def log_to_db(self, db): db["tool_instances"] .insert( { - "plugin": tool.plugin, - "name": tool.name.split("_")[0], + "plugin": getattr( + matching_tool, "plugin", None + ) + or getattr(tool_result.instance, "plugin", None), + "name": ( + matching_tool.name.split("_")[0] + if matching_tool + else tool_result.instance.__class__.__name__ + ), "arguments": json.dumps( tool_result.instance._config ), diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 6777fd585..9193053e6 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -940,7 +940,7 @@ def test_toolbox_logging_async(logs_db, tmpdir): "name": "Memory_set", "output": "null", "instance": { - "name": "Filesystem", + "name": "Memory", "plugin": "ToolboxPlugin", "arguments": "{}", }, @@ -949,7 +949,7 @@ def test_toolbox_logging_async(logs_db, tmpdir): "name": "Memory_get", "output": "two", "instance": { - "name": "Filesystem", + "name": "Memory", "plugin": "ToolboxPlugin", "arguments": "{}", },