Lazy load dspy to save memory on startup (baserow#4117)

Author: silvestrid

enterprise/backend/src/baserow_enterprise/assistant/adapter.py

@@ -1,14 +1,21 @@
-import dspy
-
 from .prompts import ASSISTANT_SYSTEM_PROMPT
 
 
-class ChatAdapter(dspy.ChatAdapter):
-    def format_field_description(self, signature: type[dspy.Signature]) -> str:
-        """
-        This is the first part of the prompt the LLM sees, so we prepend our custom
-        system prompt to it to give it the personality and context of Baserow.
-        """
+def get_chat_adapter():
+    import dspy  # local import to save memory when not used
+
+    class ChatAdapter(dspy.ChatAdapter):
+        def format_field_description(self, signature: type[dspy.Signature]) -> str:
+            """
+            This is the first part of the prompt the LLM sees, so we prepend our custom
+            system prompt to it to give it the personality and context of Baserow.
+            """
+
+            field_description = super().format_field_description(signature)
+            return (
+                ASSISTANT_SYSTEM_PROMPT
+                + "## TASK INSTRUCTIONS:\n\n"
+                + field_description
+            )
 
-        field_description = super().format_field_description(signature)
-        return ASSISTANT_SYSTEM_PROMPT + "## TASK INSTRUCTIONS:\n\n" + field_description
+    return ChatAdapter()

enterprise/backend/src/baserow_enterprise/assistant/assistant.py

@@ -3,19 +3,12 @@
 
 from django.conf import settings
 
-import dspy
-from dspy.primitives.prediction import Prediction
-from dspy.streaming import StreamListener, StreamResponse
-from dspy.utils.callback import BaseCallback
-from litellm import get_supported_openai_params
-
 from baserow.api.sessions import get_client_undo_redo_action_group_id
 from baserow_enterprise.assistant.exceptions import AssistantModelNotSupportedError
 from baserow_enterprise.assistant.tools.registries import assistant_tool_registry
 
-from .adapter import ChatAdapter
+from .adapter import get_chat_adapter
 from .models import AssistantChat, AssistantChatMessage, AssistantChatPrediction
-from .react import ReAct
 from .types import (
     AiMessage,
     AiMessageChunk,
@@ -28,90 +21,100 @@
 )
 
 
-class ChatSignature(dspy.Signature):
-    question: str = dspy.InputField()
-    history: dspy.History = dspy.InputField()
-    ui_context: UIContext | None = dspy.InputField(
-        default=None,
-        desc=(
-            "The frontend UI content the user is currently in. "
-            "Whenever make sense, use it to ground your answer."
-        ),
-    )
-    answer: str = dspy.OutputField()
-
-
 class AssistantMessagePair(TypedDict):
     question: str
     answer: str
 
 
-class AssistantCallbacks(BaseCallback):
-    def __init__(self):
-        self.tool_calls = {}
-        self.sources = []
+def get_assistant_callbacks():
+    from dspy.utils.callback import BaseCallback
 
-    def extend_sources(self, sources: list[str]) -> None:
-        """
-        Extends the current list of sources with new ones, avoiding duplicates.
+    class AssistantCallbacks(BaseCallback):
+        def __init__(self):
+            self.tool_calls = {}
+            self.sources = []
 
-        :param sources: The list of new source URLs to add.
-        :return: None
-        """
+        def extend_sources(self, sources: list[str]) -> None:
+            """
+            Extends the current list of sources with new ones, avoiding duplicates.
 
-        self.sources.extend([s for s in sources if s not in self.sources])
+            :param sources: The list of new source URLs to add.
+            :return: None
+            """
 
-    def on_tool_start(
-        self,
-        call_id: str,
-        instance: Any,
-        inputs: dict[str, Any],
-    ) -> None:
-        """
-        Called when a tool starts. It records the tool call and invokes the
-        corresponding tool's on_tool_start method if it exists.
+            self.sources.extend([s for s in sources if s not in self.sources])
 
-        :param call_id: The unique identifier of the tool call.
-        :param instance: The instance of the tool being called.
-        :param inputs: The inputs provided to the tool.
-        """
+        def on_tool_start(
+            self,
+            call_id: str,
+            instance: Any,
+            inputs: dict[str, Any],
+        ) -> None:
+            """
+            Called when a tool starts. It records the tool call and invokes the
+            corresponding tool's on_tool_start method if it exists.
 
-        try:
-            assistant_tool_registry.get(instance.name).on_tool_start(
-                call_id, instance, inputs
+            :param call_id: The unique identifier of the tool call.
+            :param instance: The instance of the tool being called.
+            :param inputs: The inputs provided to the tool.
+            """
+
+            try:
+                assistant_tool_registry.get(instance.name).on_tool_start(
+                    call_id, instance, inputs
+                )
+                self.tool_calls[call_id] = (instance, inputs)
+            except assistant_tool_registry.does_not_exist_exception_class:
+                pass
+
+        def on_tool_end(
+            self,
+            call_id: str,
+            outputs: dict[str, Any] | None,
+            exception: Exception | None = None,
+        ) -> None:
+            """
+            Called when a tool ends. It invokes the corresponding tool's on_tool_end
+            method if it exists and updates the sources if the tool produced any.
+
+            :param call_id: The unique identifier of the tool call.
+            :param outputs: The outputs returned by the tool, or None if there was an
+                exception.
+            :param exception: The exception raised by the tool, or None if it was
+                successful.
+            """
+
+            if call_id not in self.tool_calls:
+                return
+
+            instance, inputs = self.tool_calls.pop(call_id)
+            assistant_tool_registry.get(instance.name).on_tool_end(
+                call_id, instance, inputs, outputs, exception
             )
-            self.tool_calls[call_id] = (instance, inputs)
-        except assistant_tool_registry.does_not_exist_exception_class:
-            pass
 
-    def on_tool_end(
-        self,
-        call_id: str,
-        outputs: dict[str, Any] | None,
-        exception: Exception | None = None,
-    ) -> None:
-        """
-        Called when a tool ends. It invokes the corresponding tool's on_tool_end
-        method if it exists and updates the sources if the tool produced any.
-
-        :param call_id: The unique identifier of the tool call.
-        :param outputs: The outputs returned by the tool, or None if there was an
-            exception.
-        :param exception: The exception raised by the tool, or None if it was
-            successful.
-        """
+            # If the tool produced sources, add them to the overall list of sources.
+            if isinstance(outputs, dict) and "sources" in outputs:
+                self.extend_sources(outputs["sources"])
 
-        if call_id not in self.tool_calls:
-            return
+    return AssistantCallbacks()
 
-        instance, inputs = self.tool_calls.pop(call_id)
-        assistant_tool_registry.get(instance.name).on_tool_end(
-            call_id, instance, inputs, outputs, exception
+
+def get_chat_signature():
+    import dspy  # local import to save memory when not used
+
+    class ChatSignature(dspy.Signature):
+        question: str = dspy.InputField()
+        history: dspy.History = dspy.InputField()
+        ui_context: UIContext | None = dspy.InputField(
+            default=None,
+            desc=(
+                "The frontend UI content the user is currently in. "
+                "Whenever make sense, use it to ground your answer."
+            ),
         )
+        answer: str = dspy.OutputField()
 
-        # If the tool produced sources, add them to the overall list of sources.
-        if isinstance(outputs, dict) and "sources" in outputs:
-            self.extend_sources(outputs["sources"])
+    return ChatSignature
 
 
 class Assistant:
@@ -120,17 +123,27 @@ def __init__(self, chat: AssistantChat):
         self._user = chat.user
         self._workspace = chat.workspace
 
+        self._init_lm_client()
+        self._init_assistant()
+
+    def _init_lm_client(self):
+        import dspy  # local import to save memory when not used
+
         lm_model = settings.BASEROW_ENTERPRISE_ASSISTANT_LLM_MODEL
+
         self._lm_client = dspy.LM(
             model=lm_model,
             cache=not settings.DEBUG,
             max_retries=5,
         )
 
+    def _init_assistant(self):
+        from .react import ReAct  # local import to save memory when not used
+
         tools = assistant_tool_registry.list_all_usable_tools(
             self._user, self._workspace
         )
-        self._assistant = ReAct(ChatSignature, tools=tools)
+        self._assistant = ReAct(get_chat_signature(), tools=tools)
         self.history = None
 
     async def acreate_chat_message(
@@ -216,6 +229,8 @@ async def aload_chat_history(self, limit=20):
         :return: None
         """
 
+        import dspy  # local import to save memory when not used
+
         last_saved_messages: list[AssistantChatMessage] = [
             msg async for msg in self._chat.messages.order_by("-created_on")[:limit]
         ]
@@ -243,6 +258,9 @@ async def aload_chat_history(self, limit=20):
 
     @lru_cache(maxsize=1)
     def check_llm_ready_or_raise(self):
+        import dspy  # local import to save memory when not used
+        from litellm import get_supported_openai_params
+
         lm = self._lm_client
         params = get_supported_openai_params(lm.model)
         if params is None or "tools" not in params:
@@ -270,13 +288,17 @@ async def astream_messages(
         :return: An async generator that yields the response messages.
         """
 
-        callback_manager = AssistantCallbacks()
+        import dspy  # local import to save memory when not used
+        from dspy.primitives.prediction import Prediction
+        from dspy.streaming import StreamListener, StreamResponse
+
+        callback_manager = get_assistant_callbacks()
 
         with dspy.context(
             lm=self._lm_client,
             cache=not settings.DEBUG,
             callbacks=[*dspy.settings.config.callbacks, callback_manager],
-            adapter=ChatAdapter(),
+            adapter=get_chat_adapter(),
         ):
             if self.history is None:
                 await self.aload_chat_history()

enterprise/backend/src/baserow_enterprise/assistant/react.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Callable, Literal
+from typing import Any, Callable, Literal
 
 import dspy
 from dspy.adapters.types.tool import Tool
@@ -10,20 +10,14 @@
 
 from .types import ToolsUpgradeResponse
 
-if TYPE_CHECKING:
-    from dspy.signatures.signature import Signature
-
-
 # Variant of dspy.predict.react.ReAct that accepts a "meta-tool":
 # a callable that can produce tools at runtime (e.g. per-table schemas).
 # This lets a single ReAct instance handle many different table signatures
 # without creating a new agent for each request.
 
 
 class ReAct(Module):
-    def __init__(
-        self, signature: type["Signature"], tools: list[Callable], max_iters: int = 100
-    ):
+    def __init__(self, signature, tools: list[Callable], max_iters: int = 100):
         """
         ReAct stands for "Reasoning and Acting," a popular paradigm for building
         tool-using agents. In this approach, the language model is iteratively provided

enterprise/backend/src/baserow_enterprise/assistant/tools/database/tools.py

@@ -5,7 +5,6 @@
 from django.db import transaction
 from django.utils.translation import gettext as _
 
-import dspy
 from loguru import logger
 from pydantic import create_model
 
@@ -25,9 +24,9 @@
 from baserow_enterprise.assistant.tools.registries import AssistantToolType, ToolHelpers
 from baserow_enterprise.assistant.types import (
     TableNavigationType,
-    ToolSignature,
     ToolsUpgradeResponse,
     ViewNavigationType,
+    get_tool_signature,
 )
 
 from . import utils
@@ -284,6 +283,8 @@ def create_tables(
         - if add_sample_rows is True (default), add some example rows to each table
         """
 
+        import dspy  # local import to save memory when not used
+
         nonlocal user, workspace, tool_helpers
 
         if not tables:
@@ -354,7 +355,7 @@ def create_tables(
                     f"- Create 5 example rows for table_{created_table.id}. Fill every relationship with valid data when possible."
                 )
 
-            predictor = dspy.Predict(ToolSignature)
+            predictor = dspy.Predict(get_tool_signature())
             result = predictor(
                 question=("\n".join(instructions)),
                 tools=list(tools.values()),

enterprise/backend/src/baserow_enterprise/assistant/tools/database/utils.py

@@ -7,8 +7,6 @@
 from django.db.models import Q
 from django.utils.translation import gettext as _
 
-import dspy
-from dspy.adapters.types.tool import _resolve_json_schema_reference
 from pydantic import ConfigDict, Field, create_model
 
 from baserow.contrib.database.fields.actions import CreateFieldActionType
@@ -385,6 +383,9 @@ def get_view(user, view_id: int):
 def get_table_rows_tools(
     user, workspace: Workspace, tool_helpers: ToolHelpers, table: Table
 ):
+    import dspy  # local import to save memory when not used
+    from dspy.adapters.types.tool import _resolve_json_schema_reference
+
     row_model_for_create = get_create_row_model(table)
     row_model_for_update = get_update_row_model(table)
     row_model_for_response = create_model(

enterprise/backend/src/baserow_enterprise/assistant/tools/navigation/utils.py

@@ -1,6 +1,3 @@
-from dspy.dsp.utils.settings import settings
-from dspy.streaming.messages import sync_send_to_stream
-
 from baserow_enterprise.assistant.types import AiNavigationMessage, AnyNavigationType
 
 
@@ -13,6 +10,9 @@ def unsafe_navigate_to(location: AnyNavigationType) -> str:
     :param navigation_type: The type of navigation to perform.
     """
 
+    from dspy.dsp.utils.settings import settings
+    from dspy.streaming.messages import sync_send_to_stream
+
     stream = settings.send_stream
     if stream is not None:
         sync_send_to_stream(stream, AiNavigationMessage(location=location))