diff --git a/src/python/txtai/agent/tool/factory.py b/src/python/txtai/agent/tool/factory.py index 3ef493561..3eb570ff8 100644 --- a/src/python/txtai/agent/tool/factory.py +++ b/src/python/txtai/agent/tool/factory.py @@ -95,11 +95,13 @@ def fromdocs(target, config): Tool """ - name = target.__class__.__name__ if hasattr(target, "__call__") else target.__name__ - target = target.__call__ if hasattr(target, "__call__") else target + # Get function name and target - use target if it's a function or method, else use target.__call__ + name = target.__name__ if isinstance(target, (FunctionType, MethodType)) or not hasattr(target, "__call__") else target.__class__.__name__ + target = target if isinstance(target, (FunctionType, MethodType)) or not hasattr(target, "__call__") else target.__call__ - doc = inspect.getdoc(target).strip() - description, parameters, _ = chat_template_utils.parse_google_format_docstring(doc) + # Extract target documentation + doc = inspect.getdoc(target) + description, parameters, _ = chat_template_utils.parse_google_format_docstring(doc.strip()) if doc else (None, {}, None) # Get list of required parameters signature = inspect.signature(target) diff --git a/test/python/testagent.py b/test/python/testagent.py index 8085e14ff..872585fc8 100644 --- a/test/python/testagent.py +++ b/test/python/testagent.py @@ -69,20 +69,24 @@ def __call__(self, iso): today = {"name": "today", "description": "Gets the current date and time", "target": DateTime()} - def current() -> str: + def current(iso) -> str: """ Gets the current date and time + Args: + iso: date will be converted to iso format if True + Returns: current date and time """ - return datetime.today().isoformat() + return datetime.today().isoformat() if iso else datetime.today() agent = Agent(tools=[today, current, "websearch"], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_iterations=1) self.assertIsNotNone(agent) self.assertIsInstance(agent.tools["today"](True), str) + self.assertIsInstance(agent.tools["current"](True), str) def testToolsEmbeddings(self): """