Merge pull request #2 from wagtail/fix/registry-register-type

Improve DX for `AgentRegistry.register()`

Author: tomusher

src/django_ai_core/contrib/agents/base.py

@@ -1,14 +1,16 @@
 import inspect
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Annotated, get_args, get_origin
+from typing import Annotated, Callable, TypeVar, get_args, get_origin
 
 from django.core.exceptions import ValidationError
 from django.core.validators import validate_slug
 
 from .permissions import BasePermission
 from .views import AgentExecutionView
 
+AgentT = TypeVar("AgentT", bound="Agent")
+
 
 @dataclass
 class AgentParameter:
@@ -79,15 +81,20 @@ class AgentRegistry:
     def __init__(self):
         self._agents: dict[str, type[Agent]] = {}
 
-    def register(self):
-        """Decorator to register an agent."""
-
-        def decorator(cls: type[Agent]) -> type[Agent]:
-            agent_slug = cls.slug
-            self._agents[agent_slug] = cls
-            return cls
-
-        return decorator
+    def register(
+        self, cls: type[AgentT] | None = None
+    ) -> type[AgentT] | Callable[[type[AgentT]], type[AgentT]]:
+        def decorator(agent_cls: type[AgentT]) -> type[AgentT]:
+            agent_slug = agent_cls.slug
+            self._agents[agent_slug] = agent_cls
+            return agent_cls
+
+        if cls is None:
+            # Called with parentheses: @registry.register()
+            return decorator
+        else:
+            # Called without parentheses: @registry.register
+            return decorator(cls)
 
     def get(self, slug: str) -> type[Agent]:
         if slug not in self._agents:

tests/unit/contrib/agents/test_agent_registry.py

@@ -31,6 +31,16 @@ def test_agent_registry_register_uses_agent_slug():
     assert registry._agents["test-one"] is TestAgentOne
 
 
+def test_agent_registry_register_without_parentheses():
+    """Test that register can be called without parentheses."""
+    registry = AgentRegistry()
+
+    decorated = registry.register(TestAgentOne)
+    assert decorated is TestAgentOne
+    assert "test-one" in registry._agents
+    assert registry._agents["test-one"] is TestAgentOne
+
+
 def test_agent_registry_register_multiple_agents():
     """Test registering multiple agents with the registry."""
     registry = AgentRegistry()