Merge pull request #13 from bbednarski9/fix/issue-9-plugin-api-imports

justin-cechmanek · web-flow · commit 0833ba18f0ca · 2026-06-10T12:20:54.000-07:00
Route NAT plugin imports through public API facade
diff --git a/src/nat/plugins/redis/memory.py b/src/nat/plugins/redis/memory.py
@@ -13,13 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nat.builder.builder import Builder
-from nat.cli.register_workflow import register_memory
-from nat.data_models.common import OptionalSecretStr, get_secret_value
-from nat.data_models.component_ref import EmbedderRef
-from nat.data_models.memory import MemoryBaseConfig
 from pydantic import Field
 
+from nvidia_nat_redis._nat_api import (
+    Builder,
+    EmbedderRef,
+    LLMFrameworkEnum,
+    MemoryBaseConfig,
+    OptionalSecretStr,
+    get_secret_value,
+    register_memory,
+)
+
 
 class RedisMemoryClientConfig(MemoryBaseConfig, name="redis_memory"):
     host: str = Field(default="localhost", description="Redis server host")
@@ -35,8 +40,6 @@ class RedisMemoryClientConfig(MemoryBaseConfig, name="redis_memory"):
 @register_memory(config_type=RedisMemoryClientConfig)
 async def redis_memory_client(config: RedisMemoryClientConfig, builder: Builder):
 
-    from nat.builder.framework_enum import LLMFrameworkEnum
-
     import redis.asyncio as redis
 
     from .redis_editor import RedisEditor
diff --git a/src/nat/plugins/redis/object_store.py b/src/nat/plugins/redis/object_store.py
@@ -13,12 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nat.builder.builder import Builder
-from nat.cli.register_workflow import register_object_store
-from nat.data_models.common import OptionalSecretStr
-from nat.data_models.object_store import ObjectStoreBaseConfig
 from pydantic import Field, field_validator
 
+from nvidia_nat_redis._nat_api import Builder, ObjectStoreBaseConfig, OptionalSecretStr, register_object_store
+
 
 class RedisObjectStoreClientConfig(ObjectStoreBaseConfig, name="redis"):
     """
diff --git a/src/nat/plugins/redis/redis_editor.py b/src/nat/plugins/redis/redis_editor.py
@@ -16,12 +16,12 @@
 import logging
 import secrets
 
-from nat.memory.interfaces import MemoryEditor
-from nat.memory.models import MemoryItem
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import VectorQuery
 from redisvl.query.filter import Tag
 
+from nvidia_nat_redis._nat_api import MemoryEditor, MemoryItem
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/src/nat/plugins/redis/redis_object_store.py b/src/nat/plugins/redis/redis_object_store.py
@@ -15,12 +15,8 @@
 
 import logging
 
-from nat.data_models.object_store import KeyAlreadyExistsError, NoSuchKeyError
-from nat.object_store.interfaces import ObjectStore
-from nat.object_store.models import ObjectStoreItem
-from nat.utils.type_utils import override
-
 import redis.asyncio as redis
+from nvidia_nat_redis._nat_api import KeyAlreadyExistsError, NoSuchKeyError, ObjectStore, ObjectStoreItem
 
 logger = logging.getLogger(__name__)
 
@@ -88,7 +84,6 @@ async def __aexit__(self, exc_type, exc_value, traceback) -> None:
     def _make_key(self, key: str) -> str:
         return f"nat/object_store/{self._bucket_name}/{key}"
 
-    @override
     async def put_object(self, key: str, item: ObjectStoreItem):
 
         if not self._client:
@@ -103,7 +98,6 @@ async def put_object(self, key: str, item: ObjectStoreItem):
                 key=key, additional_message=f"Redis bucket {self._bucket_name} already has key {key}"
             )
 
-    @override
     async def upsert_object(self, key: str, item: ObjectStoreItem):
 
         if not self._client:
@@ -113,7 +107,6 @@ async def upsert_object(self, key: str, item: ObjectStoreItem):
         item_json = item.model_dump_json()
         await self._client.set(full_key, item_json, ex=self._ttl)
 
-    @override
     async def get_object(self, key: str) -> ObjectStoreItem:
 
         if not self._client:
@@ -127,7 +120,6 @@ async def get_object(self, key: str) -> ObjectStoreItem:
             )
         return ObjectStoreItem.model_validate_json(data)
 
-    @override
     async def delete_object(self, key: str):
 
         if not self._client:
diff --git a/src/nvidia_nat_redis/_nat_api.py b/src/nvidia_nat_redis/_nat_api.py
@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, Redis
+# SPDX-License-Identifier: Apache-2.0
+
+"""NAT public plugin API imports with compatibility for older supported NAT releases.
+
+NAT's third-party plugin guidance asks runtime plugin code to import plugin-authoring
+symbols from ``nat.plugin_api``. The current supported range also includes older NAT
+releases that do not ship that facade, so this module is the only compatibility
+fallback to implementation imports.
+"""
+
+from __future__ import annotations
+
+try:
+    from nat.plugin_api import (
+        Builder,
+        EmbedderRef,
+        Function,
+        FunctionBaseConfig,
+        FunctionInfo,
+        FunctionRef,
+        KeyAlreadyExistsError,
+        LLMFrameworkEnum,
+        MemoryBaseConfig,
+        MemoryEditor,
+        MemoryItem,
+        MemoryRef,
+        NoSuchKeyError,
+        ObjectStore,
+        ObjectStoreBaseConfig,
+        ObjectStoreItem,
+        OptionalSecretStr,
+        get_secret_value,
+        register_function,
+        register_memory,
+        register_object_store,
+    )
+
+    USING_PLUGIN_API = True
+except ImportError:  # pragma: no cover - exercised only on older or partial NAT releases
+    from nat.builder.builder import Builder
+    from nat.builder.framework_enum import LLMFrameworkEnum
+    from nat.builder.function import Function
+    from nat.builder.function_info import FunctionInfo
+    from nat.cli.register_workflow import register_function, register_memory, register_object_store
+    from nat.data_models.common import OptionalSecretStr, get_secret_value
+    from nat.data_models.component_ref import EmbedderRef, FunctionRef, MemoryRef
+    from nat.data_models.function import FunctionBaseConfig
+    from nat.data_models.memory import MemoryBaseConfig
+    from nat.data_models.object_store import KeyAlreadyExistsError, NoSuchKeyError, ObjectStoreBaseConfig
+    from nat.memory.interfaces import MemoryEditor
+    from nat.memory.models import MemoryItem
+    from nat.object_store.interfaces import ObjectStore
+    from nat.object_store.models import ObjectStoreItem
+
+    USING_PLUGIN_API = False
+
+
+__all__ = [
+    "Builder",
+    "EmbedderRef",
+    "Function",
+    "FunctionBaseConfig",
+    "FunctionInfo",
+    "FunctionRef",
+    "KeyAlreadyExistsError",
+    "LLMFrameworkEnum",
+    "MemoryBaseConfig",
+    "MemoryEditor",
+    "MemoryItem",
+    "MemoryRef",
+    "NoSuchKeyError",
+    "ObjectStore",
+    "ObjectStoreBaseConfig",
+    "ObjectStoreItem",
+    "OptionalSecretStr",
+    "USING_PLUGIN_API",
+    "get_secret_value",
+    "register_function",
+    "register_memory",
+    "register_object_store",
+]
diff --git a/src/nvidia_nat_redis/redis_agent_memory/auto_memory/config.py b/src/nvidia_nat_redis/redis_agent_memory/auto_memory/config.py
@@ -7,10 +7,10 @@
 
 from typing import TYPE_CHECKING, Any, Literal
 
-from nat.data_models.component_ref import FunctionRef, MemoryRef
-from nat.data_models.function import FunctionBaseConfig
 from pydantic import BaseModel, Field, PositiveInt, field_validator
 
+from nvidia_nat_redis._nat_api import FunctionBaseConfig, FunctionRef, MemoryRef
+
 from .._text import normalize_optional_string
 
 if TYPE_CHECKING:
diff --git a/src/nvidia_nat_redis/redis_agent_memory/auto_memory/register.py b/src/nvidia_nat_redis/redis_agent_memory/auto_memory/register.py
@@ -8,12 +8,14 @@
 from collections.abc import AsyncGenerator
 
 from agent_memory_client import create_memory_client
-from nat.builder.builder import Builder
-from nat.builder.function_info import FunctionInfo
-from nat.cli.register_workflow import register_function
+
+# Chat request models and retry patching are runtime helpers not exported by
+# nat.plugin_api in the currently supported NAT versions.
 from nat.data_models.api_server import ChatRequest, ChatRequestOrMessage, ChatResponse
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 
+from nvidia_nat_redis._nat_api import Builder, FunctionInfo, register_function
+
 from ..memory import RedisAgentMemoryBackendConfig
 from .config import RedisAgentMemoryAutoMemoryConfig
 from .service import RedisAgentMemoryAutoMemoryService
diff --git a/src/nvidia_nat_redis/redis_agent_memory/auto_memory/service.py b/src/nvidia_nat_redis/redis_agent_memory/auto_memory/service.py
@@ -9,8 +9,9 @@
 from datetime import UTC, datetime
 from typing import Any
 
+# Context, chat request models, and type conversion are runtime helpers not
+# exported by nat.plugin_api in the currently supported NAT versions.
 from nat.builder.context import Context
-from nat.builder.function import Function
 from nat.data_models.api_server import (
     ChatRequest,
     ChatRequestOrMessage,
@@ -20,6 +21,8 @@
 )
 from nat.utils.type_converter import GlobalTypeConverter
 
+from nvidia_nat_redis._nat_api import Function
+
 from .._text import message_content_to_text, normalize_optional_string
 from .config import RedisAgentMemoryAutoMemoryConfig
 
diff --git a/src/nvidia_nat_redis/redis_agent_memory/editor.py b/src/nvidia_nat_redis/redis_agent_memory/editor.py
@@ -11,9 +11,11 @@
 from typing import TYPE_CHECKING, Any
 
 from agent_memory_client import MemoryAPIClient
+
+# Runtime request context is not exported by nat.plugin_api.
 from nat.builder.context import Context
-from nat.memory.interfaces import MemoryEditor
-from nat.memory.models import MemoryItem
+
+from nvidia_nat_redis._nat_api import MemoryEditor, MemoryItem
 
 from ._text import message_content_to_text
 
diff --git a/src/nvidia_nat_redis/redis_agent_memory/memory.py b/src/nvidia_nat_redis/redis_agent_memory/memory.py
@@ -8,14 +8,15 @@
 from collections.abc import AsyncGenerator
 
 from agent_memory_client import create_memory_client
-from nat.builder.builder import Builder
-from nat.cli.register_workflow import register_memory
-from nat.data_models.memory import MemoryBaseConfig
+
+# RetryMixin and patch_with_retry are retry implementation helpers not exported
+# by nat.plugin_api in the currently supported NAT versions.
 from nat.data_models.retry_mixin import RetryMixin
-from nat.memory.interfaces import MemoryEditor
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 from pydantic import Field
 
+from nvidia_nat_redis._nat_api import Builder, MemoryBaseConfig, MemoryEditor, register_memory
+
 from .editor import RedisAgentMemoryEditor
 
 
diff --git a/tests/test_nat_api_imports.py b/tests/test_nat_api_imports.py
@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, Redis
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import ast
+import builtins
+import importlib
+from pathlib import Path
+
+PROHIBITED_DIRECT_IMPORTS = {
+    "nat.builder.builder",
+    "nat.builder.framework_enum",
+    "nat.builder.function",
+    "nat.builder.function_info",
+    "nat.cli.register_workflow",
+    "nat.data_models.common",
+    "nat.data_models.component_ref",
+    "nat.data_models.function",
+    "nat.data_models.memory",
+    "nat.data_models.object_store",
+    "nat.memory.interfaces",
+    "nat.memory.models",
+    "nat.object_store.interfaces",
+    "nat.object_store.models",
+}
+
+
+def _prohibited_import_violations(source: str, filename: str) -> list[str]:
+    tree = ast.parse(source, filename=filename)
+    violations: list[str] = []
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom) and node.module:
+            candidate_modules = {node.module}
+            candidate_modules.update(f"{node.module}.{alias.name}" for alias in node.names)
+            for module_name in candidate_modules:
+                if module_name in PROHIBITED_DIRECT_IMPORTS:
+                    violations.append(f"{filename}:{node.lineno} imports {module_name}")
+                    break
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.name in PROHIBITED_DIRECT_IMPORTS:
+                    violations.append(f"{filename}:{node.lineno} imports {alias.name}")
+                    break
+
+    return violations
+
+
+def test_runtime_plugin_code_uses_nat_api_facade_for_public_symbols() -> None:
+    """Public plugin-authoring imports should flow through nvidia_nat_redis._nat_api."""
+    root = Path(__file__).resolve().parents[1]
+    source_files = sorted((root / "src").rglob("*.py"))
+    violations: list[str] = []
+
+    for source_file in source_files:
+        if source_file.name == "_nat_api.py":
+            continue
+
+        violations.extend(
+            _prohibited_import_violations(source_file.read_text(encoding="utf-8"), str(source_file.relative_to(root)))
+        )
+
+    assert not violations, "\n".join(violations)
+
+
+def test_import_guard_catches_prohibited_import_forms() -> None:
+    source = """
+from nat.builder.builder import Builder
+from nat.builder import builder
+import nat.builder.builder
+import nat.plugins.redis.register
+"""
+
+    assert _prohibited_import_violations(source, "sample.py") == [
+        "sample.py:2 imports nat.builder.builder",
+        "sample.py:3 imports nat.builder.builder",
+        "sample.py:4 imports nat.builder.builder",
+    ]
+
+
+def test_nat_api_facade_exposes_plugin_authoring_symbols() -> None:
+    from nvidia_nat_redis import _nat_api
+
+    for name in (
+        "Builder",
+        "FunctionBaseConfig",
+        "FunctionInfo",
+        "MemoryBaseConfig",
+        "MemoryEditor",
+        "ObjectStoreBaseConfig",
+        "register_function",
+        "register_memory",
+        "register_object_store",
+    ):
+        assert hasattr(_nat_api, name)
+
+
+def test_nat_api_facade_falls_back_when_plugin_api_is_partial(monkeypatch) -> None:
+    import nvidia_nat_redis._nat_api as nat_api
+
+    real_import = builtins.__import__
+
+    def import_with_partial_plugin_api(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "nat.plugin_api":
+            raise ImportError("cannot import name 'MemoryRef' from 'nat.plugin_api'")
+        return real_import(name, globals, locals, fromlist, level)
+
+    with monkeypatch.context() as context:
+        context.setattr(builtins, "__import__", import_with_partial_plugin_api)
+        reloaded_api = importlib.reload(nat_api)
+
+    try:
+        assert reloaded_api.USING_PLUGIN_API is False
+        assert hasattr(reloaded_api, "Builder")
+        assert hasattr(reloaded_api, "register_memory")
+    finally:
+        importlib.reload(nat_api)
