Drop @register_assistant decorator

django_ai_assistant/__init__.py

@@ -2,7 +2,6 @@
 
 from django_ai_assistant.helpers.assistants import (  # noqa
     AIAssistant,
-    register_assistant,
 )
 from django_ai_assistant.langchain.tools import (  # noqa
     BaseModel,

django_ai_assistant/helpers/assistants.py

@@ -55,22 +55,9 @@ class AIAssistant(abc.ABC):  # noqa: F821
     _init_kwargs: dict[str, Any]
     _method_tools: Sequence[BaseTool]
 
-    def __init__(self, *, user=None, request=None, view=None, **kwargs):
-        if not hasattr(self, "id"):
-            raise AIAssistantMisconfiguredError(
-                f"Assistant id is not defined at {self.__class__.__name__}"
-            )
-        if self.id is None:
-            raise AIAssistantMisconfiguredError(
-                f"Assistant id is None at {self.__class__.__name__}"
-            )
-        if not re.match(r"^[a-zA-Z0-9_-]+$", self.id):
-            # id should match the pattern '^[a-zA-Z0-9_-]+$ to support as_tool in OpenAI
-            raise AIAssistantMisconfiguredError(
-                f"Assistant id '{self.id}' does not match the pattern '^[a-zA-Z0-9_-]+$'"
-                f"at {self.__class__.__name__}"
-            )
+    _registry: ClassVar[dict[str, type["AIAssistant"]]] = {}
 
+    def __init__(self, *, user=None, request=None, view=None, **kwargs):
         self._user = user
         self._request = request
         self._view = view
@@ -80,6 +67,33 @@ def __init__(self, *, user=None, request=None, view=None, **kwargs):
 
         self._set_method_tools()
 
+    def __init_subclass__(cls, **kwargs):
+        """
+        Called when a class is subclassed from AIAssistant.
+
+        This method is automatically invoked when a new subclass of AIAssistant
+        is created. It allows AIAssistant to perform additional setup or configuration
+        for the subclass, such as registering the subclass in a registry.
+
+        Args:
+            cls (type): The newly created subclass.
+            **kwargs: Additional keyword arguments passed during subclass creation.
+        """
+        super().__init_subclass__(**kwargs)
+
+        if not hasattr(cls, "id"):
+            raise AIAssistantMisconfiguredError(f"Assistant id is not defined at {cls.__name__}")
+        if cls.id is None:
+            raise AIAssistantMisconfiguredError(f"Assistant id is None at {cls.__name__}")
+        if not re.match(r"^[a-zA-Z0-9_-]+$", cls.id):
+            # id should match the pattern '^[a-zA-Z0-9_-]+$ to support as_tool in OpenAI
+            raise AIAssistantMisconfiguredError(
+                f"Assistant id '{cls.id}' does not match the pattern '^[a-zA-Z0-9_-]+$'"
+                f"at {cls.__name__}"
+            )
+
+        cls._registry[cls.id] = cls
+
     def _set_method_tools(self):
         # Find tool methods (decorated with `@method_tool` from django_ai_assistant/tools.py):
         members = inspect.getmembers(
@@ -113,6 +127,20 @@ def _set_method_tools(self):
 
         self._method_tools = tools
 
+    @classmethod
+    def get_cls_registry(cls) -> dict[str, type["AIAssistant"]]:
+        """Get the registry of AIAssistant classes."""
+        return cls._registry
+
+    @classmethod
+    def get_cls(cls, assistant_id: str) -> type["AIAssistant"]:
+        """Get the AIAssistant class for the given assistant ID."""
+        return cls.get_cls_registry()[assistant_id]
+
+    @classmethod
+    def clear_cls_registry(cls: type["AIAssistant"]) -> None:
+        cls._registry.clear()
+
     def get_name(self):
         return self.name
 
@@ -306,11 +334,3 @@ def as_tool(self, description) -> BaseTool:
             name=self.id,
             description=description,
         )
-
-
-ASSISTANT_CLS_REGISTRY: dict[str, type[AIAssistant]] = {}
-
-
-def register_assistant(cls: type[AIAssistant]):
-    ASSISTANT_CLS_REGISTRY[cls.id] = cls
-    return cls

django_ai_assistant/helpers/use_cases.py

@@ -8,7 +8,7 @@
     AIAssistantNotDefinedError,
     AIUserNotAllowedError,
 )
-from django_ai_assistant.helpers.assistants import ASSISTANT_CLS_REGISTRY
+from django_ai_assistant.helpers.assistants import AIAssistant
 from django_ai_assistant.langchain.chat_message_histories import DjangoChatMessageHistory
 from django_ai_assistant.models import Message, Thread
 from django_ai_assistant.permissions import (
@@ -21,14 +21,14 @@
 )
 
 
-def get_assistant_cls(
+def get_cls(
     assistant_id: str,
     user: Any,
     request: HttpRequest | None = None,
 ):
-    if assistant_id not in ASSISTANT_CLS_REGISTRY:
+    if assistant_id not in AIAssistant.get_cls_registry():
         raise AIAssistantNotDefinedError(f"Assistant with id={assistant_id} not found")
-    assistant_cls = ASSISTANT_CLS_REGISTRY[assistant_id]
+    assistant_cls = AIAssistant.get_cls(assistant_id)
     if not can_run_assistant(
         assistant_cls=assistant_cls,
         user=user,
@@ -43,7 +43,7 @@ def get_single_assistant_info(
     user: Any,
     request: HttpRequest | None = None,
 ):
-    assistant_cls = get_assistant_cls(assistant_id, user, request)
+    assistant_cls = get_cls(assistant_id, user, request)
 
     return {
         "id": assistant_id,
@@ -56,8 +56,8 @@ def get_assistants_info(
     request: HttpRequest | None = None,
 ):
     return [
-        get_assistant_cls(assistant_id=assistant_id, user=user, request=request)
-        for assistant_id in ASSISTANT_CLS_REGISTRY.keys()
+        get_cls(assistant_id=assistant_id, user=user, request=request)
+        for assistant_id in AIAssistant.get_cls_registry().keys()
     ]
 
 
@@ -68,7 +68,7 @@ def create_message(
     content: Any,
     request: HttpRequest | None = None,
 ):
-    assistant_cls = get_assistant_cls(assistant_id, user, request)
+    assistant_cls = get_cls(assistant_id, user, request)
 
     if not can_create_message(thread=thread, user=user, request=request):
         raise AIUserNotAllowedError("User is not allowed to create messages in this thread")

docs/tutorial.md

@@ -29,13 +29,13 @@ This is possible by defining "tools" the AI can use. These tools are methods in
 To create an AI Assistant, you need to:
 
 1. Create an `ai_assistants.py` file;
-2. Define a class that inherits from `AIAssistant` with the decorator `@register_assistant` over it;
+2. Define a class that inherits from `AIAssistant`;
 3. Provide an `id`, a `name`, some `instructions` for the LLM (a system prompt), and a `model` name:
 
 ```python title="myapp/ai_assistants.py"
-from django_ai_assistant import AIAssistant, register_assistant
+from django_ai_assistant import AIAssistant
+
 
-@register_assistant
 class WeatherAIAssistant(AIAssistant):
     id = "weather_assistant"
     name = "Weather Assistant"
@@ -52,10 +52,10 @@ Use the `@method_tool` decorator to define a tool method in the AI Assistant:
 
 ```{.python title="myapp/ai_assistants.py" hl_lines="15-22"}
 from django.utils import timezone
-from django_ai_assistant import AIAssistant, method_tool, register_assistant
+from django_ai_assistant import AIAssistant, method_tool
 import json
 
-@register_assistant
+
 class WeatherAIAssistant(AIAssistant):
     id = "weather_assistant"
     name = "Weather Assistant"
@@ -95,9 +95,9 @@ AI: The weather in NYC is sunny with a temperature of 25°C.
 You have access to the current request user in tools:
 
 ```{.python title="myapp/ai_assistants.py" hl_lines=13}
-from django_ai_assistant import AIAssistant, method_tool, register_assistant
+from django_ai_assistant import AIAssistant, method_tool
+
 
-@register_assistant
 class PersonalAIAssistant(AIAssistant):
     id = "personal_assistant"
     name = "Personal Assistant"
@@ -113,10 +113,10 @@ class PersonalAIAssistant(AIAssistant):
 You can also add any Django logic to tools, such as querying the database:
 
 ```{.python title="myapp/ai_assistants.py" hl_lines=14-16}
-from django_ai_assistant import AIAssistant, method_tool, register_assistant
+from django_ai_assistant import AIAssistant, method_tool
 import json
 
-@register_assistant
+
 class IssueManagementAIAssistant(AIAssistant):
     id = "issue_mgmt_assistant"
     name = "Issue Management Assistant"
@@ -154,10 +154,10 @@ Then, set the `TAVILY_API_KEY` environment variable. You'll need to sign up at [
 Finally, add the tool to your AI Assistant class by overriding the `get_tools` method:
 
 ```{.python title="myapp/ai_assistants.py"  hl_lines="2 20"}
-from django_ai_assistant import AIAssistant, register_assistant
+from django_ai_assistant import AIAssistant
 from langchain_community.tools.tavily_search import TavilySearchResults
 
-@register_assistant
+
 class MovieSearchAIAssistant(AIAssistant):
     id = "movie_search_assistant"  # noqa: A003
     instructions = (
@@ -191,6 +191,7 @@ You can manually call an AI Assistant from anywhere in your Django application:
 ```python
 from myapp.ai_assistants import WeatherAIAssistant
 
+
 assistant = WeatherAIAssistant()
 output = assistant.run("What's the weather in New York City?")
 assert output == "The weather in NYC is sunny with a temperature of 25°C."
@@ -209,10 +210,11 @@ and automatically retrieved then passed to the LLM when calling the AI Assistant
 
 To create a `Thread`, you can use a helper from the `django_ai_assistant.use_cases` module. For example:
 
-```{.python hl_lines="4 8"}
+```{.python hl_lines="5 9"}
 from django_ai_assistant.use_cases import create_thread, get_thread_messages
 from myapp.ai_assistants import WeatherAIAssistant
 
+
 thread = create_thread(name="Weather Chat", user=some_user)
 assistant = WeatherAIAssistant()
 assistant.run("What's the weather in New York City?", thread_id=thread.id)
@@ -231,6 +233,7 @@ such as a React application or a mobile app. Add the following to your Django pr
 ```python title="myproject/urls.py"
 from django.urls import include, path
 
+
 urlpatterns = [
     path("ai-assistant/", include("django_ai_assistant.urls")),
     ...
@@ -267,6 +270,7 @@ Thread permission signatures look like this:
 from django_ai_assistant.models import Thread
 from django.http import HttpRequest
 
+
 def check_custom_thread_permission(
         thread: Thread,
         user: Any,
@@ -280,6 +284,7 @@ While Message permission signatures look like this:
 from django_ai_assistant.models import Thread, Message
 from django.http import HttpRequest
 
+
 def check_custom_message_permission(
         message: Message,
         thread: Thread,
@@ -296,10 +301,10 @@ By default the supported models are OpenAI ones,
 but you can use [any chat model from Langchain that supports Tool Calling](https://python.langchain.com/v0.2/docs/integrations/chat/#advanced-features) by overriding `get_llm`:
 
 ```python title="myapp/ai_assistants.py"
-from django_ai_assistant import AIAssistant, register_assistant
+from django_ai_assistant import AIAssistant
 from langchain_anthropic import ChatAnthropic
 
-@register_assistant
+
 class WeatherAIAssistant(AIAssistant):
     id = "weather_assistant"
     name = "Weather Assistant"
@@ -324,16 +329,15 @@ class WeatherAIAssistant(AIAssistant):
 One AI Assistant can call another AI Assistant as a tool. This is useful for composing complex AI Assistants.
 Use the `as_tool` method for that:
 
-```{.python title="myapp/ai_assistants.py"  hl_lines="15 17"}
-@register_assistant
+```{.python title="myapp/ai_assistants.py"  hl_lines="14 16"}
 class SimpleAssistant(AIAssistant):
     ...
 
-@register_assistant
+
 class AnotherSimpleAssistant(AIAssistant):
     ...
 
-@register_assistant
+
 class ComplexAssistant(AIAssistant):
     ...
 
@@ -366,9 +370,9 @@ For this to work, your must do the following in your AI Assistant:
 For example:
 
 ```{.python title="myapp/ai_assistants.py"  hl_lines="12 16 18"}
-from django_ai_assistant import AIAssistant, register_assistant
+from django_ai_assistant import AIAssistant
+
 
-@register_assistant
 class DocsAssistant(AIAssistant):
     id = "docs_assistant"  # noqa: A003
     name = "Docs Assistant"

example/movies/ai_assistants.py

@@ -10,7 +10,7 @@
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_core.tools import BaseTool
 
-from django_ai_assistant import AIAssistant, method_tool, register_assistant
+from django_ai_assistant import AIAssistant, method_tool
 from movies.models import MovieBacklogItem
 
 
@@ -53,7 +53,6 @@ def run_as_tool(self, message: str, **kwargs):
         return super().run_as_tool(message, **kwargs)
 
 
-@register_assistant
 class MovieRecommendationAIAssistant(AIAssistant):
     id = "movie_recommendation_assistant"  # noqa: A003
     instructions = (

example/rag/ai_assistants.py

@@ -2,11 +2,10 @@
 from langchain_core.retrievers import BaseRetriever
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-from django_ai_assistant import AIAssistant, register_assistant
+from django_ai_assistant import AIAssistant
 from rag.models import DjangoDocPage
 
 
-@register_assistant
 class DjangoDocsAssistant(AIAssistant):
     id = "django_docs_assistant"  # noqa: A003
     name = "Django Docs Assistant"

example/weather/ai_assistants.py

@@ -3,14 +3,13 @@
 
 import requests
 
-from django_ai_assistant import AIAssistant, BaseModel, Field, method_tool, register_assistant
+from django_ai_assistant import AIAssistant, BaseModel, Field, method_tool
 
 
 BASE_URL = "https://api.weatherapi.com/v1/"
 TIMEOUT = 10
 
 
-@register_assistant
 class WeatherAIAssistant(AIAssistant):
     id = "weather_assistant"  # noqa: A003
     name = "Weather Assistant"