diff --git a/python/packages/devui/agent_framework_devui/_discovery.py b/python/packages/devui/agent_framework_devui/_discovery.py index 70447f2624a..bb20f80e312 100644 --- a/python/packages/devui/agent_framework_devui/_discovery.py +++ b/python/packages/devui/agent_framework_devui/_discovery.py @@ -20,6 +20,19 @@ logger = logging.getLogger(__name__) +def _global_names(scope: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef) -> set[str]: + """Return the names a function or class body declares ``global``, not counting nested scopes.""" + names: set[str] = set() + pending: list[ast.AST] = list(scope.body) + while pending: + node = pending.pop() + if isinstance(node, ast.Global): + names.update(node.names) + elif not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + pending.extend(ast.iter_child_nodes(node)) + return names + + class EntityDiscovery: """Discovery for Agent Framework entities - agents and workflows.""" @@ -555,12 +568,28 @@ def _has_entity_exports(self, file_path: Path) -> bool: source = file_path.read_text(encoding="utf-8") tree = ast.parse(source, filename=str(file_path)) - # Look for module-level assignments of 'agent' or 'workflow' - for node in ast.walk(tree): + # Look for assignments of 'agent' or 'workflow' that create module attributes. In a + # function or class body that is only the case for names the scope declares global; + # otherwise they are locals or class attributes. + exports = {"agent", "workflow"} + pending: list[tuple[ast.AST, set[str]]] = [(node, exports) for node in tree.body] + while pending: + node, names = pending.pop() + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + scope_names = exports & _global_names(node) + pending.extend((child, scope_names) for child in node.body) + continue + pending.extend((child, names) for child in ast.iter_child_nodes(node)) if isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id in ("agent", "workflow"): - return True + targets = node.targets + elif isinstance(node, ast.AnnAssign) and node.value is not None: + # Annotated assignment, e.g. `agent: Agent = Agent(...)` + targets = [node.target] + else: + continue + for target in targets: + if isinstance(target, ast.Name) and target.id in names: + return True except Exception as e: logger.debug(f"Could not parse {file_path} for entity exports: {e}") return False diff --git a/python/packages/devui/tests/devui/test_discovery.py b/python/packages/devui/tests/devui/test_discovery.py index f7cd22e7be5..e36659cb608 100644 --- a/python/packages/devui/tests/devui/test_discovery.py +++ b/python/packages/devui/tests/devui/test_discovery.py @@ -56,6 +56,27 @@ async def test_empty_directory(): assert len(entities) == 0 +async def test_discover_file_entity_with_annotated_export(): + """A file exporting `agent: Agent = ...` is discovered like `agent = ...`.""" + with tempfile.TemporaryDirectory() as temp_dir: + (Path(temp_dir) / "typed_agent.py").write_text("agent: object = object()\n") + (Path(temp_dir) / "declared_only.py").write_text("agent: object\n") + (Path(temp_dir) / "conditional_agent.py").write_text("if True:\n agent: object = object()\n") + (Path(temp_dir) / "local_only.py").write_text( + "def build():\n agent: object = object()\n return agent\n\n\n" + "class Holder:\n workflow = object()\n\n\n" + "def other():\n global name\n agent = object()\n" + ) + (Path(temp_dir) / "global_factory.py").write_text( + "def build():\n global agent\n if True:\n agent = object()\n\n\nbuild()\n" + ) + + discovery = EntityDiscovery(temp_dir) + entities = await discovery.discover_entities() + + assert sorted(entity.id for entity in entities) == ["conditional_agent", "global_factory", "typed_agent"] + + async def test_discovery_accepts_agents_with_only_run(): """Test that discovery accepts agents with only run() method.