Merge pull request #1067 from jupyter-naas/refactor-models-loading

refactor(models): add include_models filter to selectively load models

Author: FlorentRvl

libs/naas-abi-cli/uv.lock

@@ -166,7 +166,9 @@ def on_load(self):
         # is additive.
         if self._engine.services.model_registry_available():
             ModuleModelLoader.load_models(
-                self.__class__, self._engine.services.model_registry
+                self.__class__,
+                self._engine.services.model_registry,
+                include_models=getattr(self._configuration, "include_models", None),
             )
 
     def on_initialized(self):

libs/naas-abi-core/naas_abi_core/module/Module.py

@@ -16,22 +16,79 @@
 import importlib
 import inspect
 import os
+import re
+from pathlib import Path
 from typing import TYPE_CHECKING
 
-from naas_abi_core.models.Model import Model, ModelDefinition
+from naas_abi_core.models.Model import CanonicalModelId, Model, ModelDefinition
 from naas_abi_core.module.ModuleUtils import find_class_module_root_path
 from naas_abi_core.utils.Logger import logger
 
 if TYPE_CHECKING:
     from naas_abi_core.services.model_registry.ModelRegistryPort import IModelRegistry
 
 
+# Matches ``CANONICAL_ID = CanonicalModelId.FOO`` /
+# ``CANONICAL_ID = CanonicalModelId["FOO"]`` / ``CANONICAL_ID = "foo"`` so a
+# model file's canonical id can be resolved *without importing it* (importing a
+# model file constructs its provider client, which can be slow).
+_CANONICAL_ID_ASSIGN_RE = re.compile(r"CANONICAL_ID\s*=")
+_CANONICAL_ID_VALUE_RE = re.compile(
+    r"CANONICAL_ID\s*=\s*(?:"
+    r"CanonicalModelId\.([A-Za-z0-9_]+)"  # enum-member access
+    r"|CanonicalModelId\[\s*[\"']([^\"']+)[\"']\s*\]"  # enum-subscript by name
+    r"|[\"']([^\"']+)[\"']"  # bare string literal
+    r")"
+)
+
+
+def _static_canonical_ids(source: str) -> tuple[set[str], bool]:
+    """Best-effort extraction of the canonical-id strings declared via
+    ``CANONICAL_ID = ...`` in a model source file, without importing it.
+
+    Returns ``(canonical_ids, fully_resolved)``. ``fully_resolved`` is ``False``
+    when at least one ``CANONICAL_ID`` assignment could not be resolved to a
+    concrete string — callers must not skip a file that isn't fully resolved.
+    """
+    ids: set[str] = set()
+    resolved = 0
+    for member, subscript, literal in _CANONICAL_ID_VALUE_RE.findall(source):
+        if member:
+            try:
+                ids.add(CanonicalModelId[member].value)
+                resolved += 1
+            except KeyError:
+                # Unknown member name — leave unresolved so we import to be safe.
+                pass
+        elif subscript:
+            ids.add(subscript)
+            resolved += 1
+        elif literal:
+            ids.add(literal)
+            resolved += 1
+    total = len(_CANONICAL_ID_ASSIGN_RE.findall(source))
+    return ids, resolved == total
+
+
 class ModuleModelLoader:
     @classmethod
-    def load_models(cls, class_: type, registry: "IModelRegistry") -> int:
+    def load_models(
+        cls,
+        class_: type,
+        registry: "IModelRegistry",
+        include_models: list[str] | None = None,
+    ) -> int:
         """Recursively walk ``<module_root>/models/`` and register every
         ``ModelDefinition`` subclass found, including those in provider
         subdirectories (e.g. ``models/openai/``, ``models/anthropic/``).
+
+        When ``include_models`` is provided, only models whose ``CANONICAL_ID``
+        is in that list are registered; every other model file is skipped
+        *without importing it* — importing a model file constructs its provider
+        client (e.g. a Bedrock client), which can add seconds per file. Files
+        whose canonical id cannot be resolved statically are imported anyway, so
+        the filter never drops a model it can't positively identify.
+
         Returns the count of successful registrations."""
         module_root_path = find_class_module_root_path(class_)
         models_path = module_root_path / "models"
@@ -40,6 +97,7 @@ def load_models(cls, class_: type, registry: "IModelRegistry") -> int:
             return 0
 
         logger.debug(f"Loading models from {models_path}")
+        include_set = set(include_models) if include_models is not None else None
         top_package = class_.__module__.split(".")[0]
         registered = 0
 
@@ -64,6 +122,25 @@ def load_models(cls, class_: type, registry: "IModelRegistry") -> int:
                         f"{class_.__module__}.models.{subpackage}.{file[:-3]}"
                     )
 
+                # Honour include_models without paying the import cost: if the
+                # file's canonical id(s) can be resolved statically and none are
+                # in the include set, skip before importing.
+                if include_set is not None:
+                    try:
+                        source = (Path(dirpath) / file).read_text(encoding="utf-8")
+                    except OSError:
+                        source = None
+                    if source is not None:
+                        ids, fully_resolved = _static_canonical_ids(source)
+                        if fully_resolved and ids and ids.isdisjoint(include_set):
+                            logger.debug(
+                                "ModuleModelLoader: skipping %s — canonical id(s) "
+                                "%s not in include_models.",
+                                model_module_path,
+                                sorted(ids),
+                            )
+                            continue
+
                 try:
                     model_module = importlib.import_module(model_module_path)
                 except Exception as exc:

libs/naas-abi-core/naas_abi_core/module/ModuleModelLoader.py

@@ -47,6 +47,8 @@ class Configuration(ModuleConfiguration):
             region_name: "us-east-1"
             validate_on_load: true
             validation_model_id: null  # optional: invoke this model to prove access
+            include_models:            # optional: register only these canonical ids
+              - "gpt-oss-120b"
         """
 
         aws_access_key_id: Optional[str] = None
@@ -55,6 +57,13 @@ class Configuration(ModuleConfiguration):
         region_name: str = "us-east-1"
         datastore_path: str = "bedrock"
 
+        # When set, only models whose CANONICAL_ID is in this list are
+        # registered on load; every other models/*.py file is skipped *without
+        # importing it*, which avoids constructing the (slow) Bedrock client for
+        # models you don't use. Values are canonical model ids, e.g.
+        # ["gpt-oss-120b"]. When None (default), all discovered models load.
+        include_models: Optional[list[str]] = None
+
         # When true (default), the module verifies on load that:
         #   1. boto3 can resolve AWS credentials,
         #   2. sts:GetCallerIdentity succeeds,

libs/naas-abi-marketplace/naas_abi_marketplace/ai/bedrock/__init__.py

@@ -2,6 +2,9 @@
 
 from typing import Optional
 
+from langchain_core.language_models import BaseChatModel
+
+from naas_abi_core.models.Model import ChatModel
 from naas_abi_core.services.agent.IntentAgent import (
     AgentConfiguration,
     AgentSharedState,
@@ -86,6 +89,7 @@ def New(
             if cls.MODEL_ID and "/" in cls.MODEL_ID
             else cls.MODEL_ID
         )
+        chat_model: BaseChatModel | ChatModel
         if lookup_id in registry.list_canonical_ids():
             chat_model = registry.get_chat_model(lookup_id, provider="openrouter")
         else:

libs/naas-abi-marketplace/naas_abi_marketplace/applications/openrouter/agents/OpenRouterAgent.py

@@ -218,6 +218,8 @@ def _public_modules_url(path: str) -> str:
     from naas_abi import ABIModule
 
     public_api_host = ABIModule.get_instance().configuration.global_config.public_api_host
+    if not public_api_host.startswith("https://"):
+        public_api_host = f"https://{public_api_host}"
     return f"{public_api_host}/modules/{path.lstrip('/')}"