nvidia-nat 1.2.0__py3-none-any.whl

nat/retriever/nemo_retriever/retriever.py

@@ -0,0 +1,190 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import os
+import typing
+from functools import partial
+from urllib.parse import urljoin
+
+import httpx
+from langchain_core.retrievers import BaseRetriever
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+
+from nat.retriever.interface import Retriever
+from nat.retriever.models import Document
+from nat.retriever.models import RetrieverError
+from nat.retriever.models import RetrieverOutput
+
+logger = logging.getLogger(__name__)
+
+
+class Collection(BaseModel):
+    id: str
+    name: str
+    meta: typing.Any
+    pipeline: str
+    created_at: str
+
+
+class RetrieverPayload(BaseModel):
+    query: str
+    top_k: int = Field(le=50, gt=0)
+
+
+class CollectionUnavailableError(RetrieverError):
+    pass
+
+
+class NemoRetriever(Retriever):
+    """
+    Client for retrieving document chunks from a Nemo Retriever service.
+    """
+
+    def __init__(self, uri: str | HttpUrl, timeout: int = 60, nvidia_api_key: str = None, **kwargs):
+
+        self.base_url = str(uri)
+        self.timeout = timeout
+        self._search_func = self._search
+        self.api_key = nvidia_api_key if nvidia_api_key else os.getenv('NVIDIA_API_KEY')
+        self._bound_params = []
+        if not self.api_key:
+            logger.warning("No API key was specified as part of configuration or as an environment variable.")
+
+    def bind(self, **kwargs) -> None:
+        """
+        Bind default values to the search method. Cannot bind the 'query' parameter.
+
+        Args:
+          kwargs (dict): Key value pairs corresponding to the default values of search parameters.
+        """
+        if "query" in kwargs:
+            kwargs = {k: v for k, v in kwargs.items() if k != "query"}
+        self._search_func = partial(self._search_func, **kwargs)
+        self._bound_params = list(kwargs.keys())
+        logger.debug("Binding paramaters for search function: %s", kwargs)
+
+    def get_unbound_params(self) -> list[str]:
+        """
+        Returns a list of unbound parameters which will need to be passed to the search function.
+        """
+        return [param for param in ["query", "collection_name", "top_k"] if param not in self._bound_params]
+
+    async def get_collections(self, client) -> list[Collection]:
+        """
+        Get a list of all available collections as pydantic `Collection` objects
+        """
+        collection_response = await client.get(urljoin(self.base_url, "/v1/collections"))
+        collection_response.raise_for_status()
+        if not collection_response or len(collection_response.json().get('collections', [])) == 0:
+            raise CollectionUnavailableError(f"No collections available at {self.base_url}")
+
+        collections = [
+            Collection.model_validate(collection) for collection in collection_response.json()["collections"]
+        ]
+
+        return collections
+
+    async def get_collection_by_name(self, collection_name, client) -> Collection:
+        """
+        Retrieve a collection using it's name. Will return the first collection found if the name is ambiguous.
+        """
+        collections = await self.get_collections(client)
+        if (collection := next((c for c in collections if c.name == collection_name), None)) is None:
+            raise CollectionUnavailableError(f"Collection {collection_name} not found")
+        return collection
+
+    async def search(self, query: str, **kwargs):
+        return await self._search_func(query=query, **kwargs)
+
+    async def _search(
+        self,
+        query: str,
+        collection_name: str,
+        top_k: str,
+        output_fields: list[str] = None,
+    ):
+        """
+        Retrieve document chunks from the configured Nemo Retriever Service.
+        """
+        output = []
+        try:
+            async with httpx.AsyncClient(headers={"Authorization": f"Bearer {self.api_key}"},
+                                         timeout=self.timeout) as client:
+                collection = await self.get_collection_by_name(collection_name, client)
+                url = urljoin(self.base_url, f"/v1/collections/{collection.id}/search")
+
+                payload = RetrieverPayload(query=query, top_k=top_k)
+                response = await client.post(url, content=json.dumps(payload.model_dump(mode="python")))
+
+                logger.debug("response.status_code=%s", response.status_code)
+
+                response.raise_for_status()
+                output = response.json().get("chunks")
+
+                # Handle output fields
+                output = [_flatten(chunk, output_fields) for chunk in output]
+
+                return _wrap_nemo_results(output=output, content_field="content")
+
+        except Exception as e:
+            logger.exception("Encountered an error when retrieving results from Nemo Retriever: %s", e)
+            raise CollectionUnavailableError(
+                f"Error when retrieving documents from {collection_name} for query '{query}'") from e
+
+
+def _wrap_nemo_results(output: list[dict], content_field: str):
+    return RetrieverOutput(results=[_wrap_nemo_single_results(o, content_field=content_field) for o in output])
+
+
+def _wrap_nemo_single_results(output: dict, content_field: str):
+    return Document(page_content=output[content_field],
+                    metadata={
+                        k: v
+                        for k, v in output.items() if k != content_field
+                    })
+
+
+def _flatten(obj: dict, output_fields: list[str]) -> list[str]:
+    base_fields = [
+        "format",
+        "id",
+    ]
+    if not output_fields:
+        output_fields = [
+            "format",
+            "id",
+        ]
+        output_fields.extend(list(obj["metadata"].keys()))
+    data = {"content": obj.get("content")}
+    for field in base_fields:
+        if field in output_fields:
+            data.update({field: obj[field]})
+
+    data.update({k: v for k, v in obj['metadata'].items() if k in output_fields})
+    return data
+
+
+class NemoLangchainRetriever(BaseRetriever, BaseModel):
+    client: NemoRetriever
+
+    def _get_relevant_documents(self, query, *, run_manager, **kwargs):
+        raise NotImplementedError
+
+    async def _aget_relevant_documents(self, query, *, run_manager, **kwargs):
+        return await self.client.search(query, **kwargs)

nat/retriever/register.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+import nat.retriever.milvus.register
+import nat.retriever.nemo_retriever.register

nat/runtime/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/runtime/loader.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import importlib.metadata
+import logging
+import time
+from contextlib import asynccontextmanager
+from enum import IntFlag
+from enum import auto
+from functools import lru_cache
+from functools import reduce
+
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.cli.type_registry import GlobalTypeRegistry
+from nat.data_models.config import Config
+from nat.runtime.session import SessionManager
+from nat.utils.data_models.schema_validator import validate_schema
+from nat.utils.debugging_utils import is_debugger_attached
+from nat.utils.io.yaml_tools import yaml_load
+from nat.utils.type_utils import StrPath
+
+logger = logging.getLogger(__name__)
+
+
+class PluginTypes(IntFlag):
+    COMPONENT = auto()
+    """
+    A plugin that is a component of the workflow. This includes tools, LLMs, retrievers, etc.
+    """
+    FRONT_END = auto()
+    """
+    A plugin that is a front end for the workflow. This includes FastAPI, Gradio, etc.
+    """
+    EVALUATOR = auto()
+    """
+    A plugin that is an evaluator for the workflow. This includes evaluators like RAGAS, SWE-bench, etc.
+    """
+    AUTHENTICATION = auto()
+    """
+    A plugin that is an API authentication provider for the workflow. This includes Oauth2, API Key, etc.
+    """
+    REGISTRY_HANDLER = auto()
+
+    # Convenience flag for groups of plugin types
+    CONFIG_OBJECT = COMPONENT | FRONT_END | EVALUATOR | AUTHENTICATION
+    """
+    Any plugin that can be specified in the NAT configuration file.
+    """
+    ALL = COMPONENT | FRONT_END | EVALUATOR | REGISTRY_HANDLER | AUTHENTICATION
+    """
+    All plugin types
+    """
+
+
+def load_config(config_file: StrPath) -> Config:
+    """
+    This is the primary entry point for loading a NAT configuration file. It ensures that all plugins are
+    loaded and then validates the configuration file against the Config schema.
+
+    Parameters
+    ----------
+    config_file : StrPath
+        The path to the configuration file
+
+    Returns
+    -------
+    Config
+        The validated Config object
+    """
+
+    # Ensure all of the plugins are loaded
+    discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+
+    config_yaml = yaml_load(config_file)
+
+    # Validate configuration adheres to NAT schemas
+    validated_nat_config = validate_schema(config_yaml, Config)
+
+    return validated_nat_config
+
+
+@asynccontextmanager
+async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
+    """
+    Load the NAT configuration file and create an Runner object. This is the primary entry point for running
+    NAT workflows.
+
+    Parameters
+    ----------
+    config_file : StrPath
+        The path to the configuration file
+    max_concurrency : int, optional
+        The maximum number of parallel workflow invocations to support. Specifying 0 or -1 will allow an unlimited
+        count, by default -1
+    """
+
+    # Load the config object
+    config = load_config(config_file)
+
+    # Must yield the workflow function otherwise it cleans up
+    async with WorkflowBuilder.from_config(config=config) as workflow:
+
+        yield SessionManager(workflow.build(), max_concurrency=max_concurrency)
+
+
+@lru_cache
+def discover_entrypoints(plugin_type: PluginTypes):
+    """
+    Discover all the requested plugin types which were registered via an entry point group and return them.
+    """
+
+    entry_points = importlib.metadata.entry_points()
+
+    plugin_groups = []
+
+    # Add the specified plugin type to the list of groups to load
+    # The aiq entrypoints are intentionally left in the list to maintain backwards compatibility.
+    if (plugin_type & PluginTypes.COMPONENT):
+        plugin_groups.extend(["aiq.plugins", "aiq.components", "nat.plugins", "nat.components"])
+    if (plugin_type & PluginTypes.FRONT_END):
+        plugin_groups.extend(["aiq.front_ends", "nat.front_ends"])
+    if (plugin_type & PluginTypes.REGISTRY_HANDLER):
+        plugin_groups.extend(["aiq.registry_handlers", "nat.registry_handlers"])
+    if (plugin_type & PluginTypes.EVALUATOR):
+        plugin_groups.extend(["aiq.evaluators", "nat.evaluators"])
+    if (plugin_type & PluginTypes.AUTHENTICATION):
+        plugin_groups.extend(["aiq.authentication_providers", "nat.authentication_providers"])
+
+    # Get the entry points for the specified groups
+    nat_plugins = reduce(lambda x, y: list(x) + list(y), [entry_points.select(group=y) for y in plugin_groups])
+
+    return nat_plugins
+
+
+@lru_cache
+def get_all_entrypoints_distro_mapping() -> dict[str, str]:
+    """
+    Get the mapping of all NAT entry points to their distribution names.
+    """
+
+    mapping = {}
+    nat_entrypoints = discover_entrypoints(PluginTypes.ALL)
+    for ep in nat_entrypoints:
+        ep_module_parts = ep.module.split(".")
+        current_parts = []
+        for part in ep_module_parts:
+            current_parts.append(part)
+            module_prefix = ".".join(current_parts)
+            mapping[module_prefix] = ep.dist.name
+
+    return mapping
+
+
+def discover_and_register_plugins(plugin_type: PluginTypes):
+    """
+    Discover all the requested plugin types which were registered via an entry point group and register them into the
+    GlobalTypeRegistry.
+    """
+
+    # Get the entry points for the specified groups
+    nat_plugins = discover_entrypoints(plugin_type)
+
+    count = 0
+
+    # Pause registration hooks for performance. This is useful when loading a large number of plugins.
+    with GlobalTypeRegistry.get().pause_registration_changed_hooks():
+
+        for entry_point in nat_plugins:
+            try:
+                logger.debug("Loading module '%s' from entry point '%s'...", entry_point.module, entry_point.name)
+
+                start_time = time.time()
+
+                entry_point.load()
+
+                elapsed_time = (time.time() - start_time) * 1000
+
+                logger.debug("Loading module '%s' from entry point '%s'...Complete (%f ms)",
+                             entry_point.module,
+                             entry_point.name,
+                             elapsed_time)
+
+                # Log a warning if the plugin took a long time to load. This can be useful for debugging slow imports.
+                # The threshold is 300 ms if no plugins have been loaded yet, and 100 ms otherwise. Triple the threshold
+                # if a debugger is attached.
+                if (elapsed_time > (300.0 if count == 0 else 150.0) * (3 if is_debugger_attached() else 1)):
+                    logger.debug(
+                        "Loading module '%s' from entry point '%s' took a long time (%f ms). "
+                        "Ensure all imports are inside your registered functions.",
+                        entry_point.module,
+                        entry_point.name,
+                        elapsed_time)
+
+            except ImportError:
+                logger.warning("Failed to import plugin '%s'", entry_point.name, exc_info=True)
+                # Optionally, you can mark the plugin as unavailable or take other actions
+
+            except Exception:
+                logger.exception("An error occurred while loading plugin '%s': {e}", entry_point.name, exc_info=True)
+
+            finally:
+                count += 1
+
+
+# Compatibility alias
+get_all_aiq_entrypoints_distro_mapping = get_all_entrypoints_distro_mapping

nat/runtime/runner.py

@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import typing
+from enum import Enum
+
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.builder.function import Function
+from nat.data_models.invocation_node import InvocationNode
+from nat.observability.exporter_manager import ExporterManager
+from nat.utils.reactive.subject import Subject
+
+logger = logging.getLogger(__name__)
+
+
+class UserManagerBase:
+    pass
+
+
+class RunnerState(Enum):
+    UNINITIALIZED = 0
+    INITIALIZED = 1
+    RUNNING = 2
+    COMPLETED = 3
+    FAILED = 4
+
+
+_T = typing.TypeVar("_T")
+
+
+class Runner:
+
+    def __init__(self,
+                 input_message: typing.Any,
+                 entry_fn: Function,
+                 context_state: ContextState,
+                 exporter_manager: ExporterManager):
+        """
+        The Runner class is used to run a workflow. It handles converting input and output data types and running the
+        workflow with the specified concurrency.
+
+        Parameters
+        ----------
+        input_message : typing.Any
+            The input message to the workflow
+        entry_fn : Function
+            The entry function to the workflow
+        context_state : ContextState
+            The context state to use
+        exporter_manager : ExporterManager
+            The exporter manager to use
+        """
+
+        if (entry_fn is None):
+            raise ValueError("entry_fn cannot be None")
+
+        self._entry_fn = entry_fn
+        self._context_state = context_state
+        self._context = Context(self._context_state)
+
+        self._state = RunnerState.UNINITIALIZED
+
+        self._input_message_token = None
+
+        # Before we start, we need to convert the input message to the workflow input type
+        self._input_message = input_message
+
+        self._exporter_manager = exporter_manager
+
+    @property
+    def context(self) -> Context:
+        return self._context
+
+    def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
+        return self._entry_fn.convert(value, to_type)
+
+    async def __aenter__(self):
+
+        # Set the input message on the context
+        self._input_message_token = self._context_state.input_message.set(self._input_message)
+
+        # Create reactive event stream
+        self._context_state.event_stream.set(Subject())
+        self._context_state.active_function.set(InvocationNode(
+            function_name="root",
+            function_id="root",
+        ))
+
+        if (self._state == RunnerState.UNINITIALIZED):
+            self._state = RunnerState.INITIALIZED
+        else:
+            raise ValueError("Cannot enter the context more than once")
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+
+        if (self._input_message_token is None):
+            raise ValueError("Cannot exit the context without entering it")
+
+        self._context_state.input_message.reset(self._input_message_token)
+
+        if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
+            raise ValueError("Cannot exit the context without completing the workflow")
+
+    @typing.overload
+    async def result(self) -> typing.Any:
+        ...
+
+    @typing.overload
+    async def result(self, to_type: type[_T]) -> _T:
+        ...
+
+    async def result(self, to_type: type | None = None):
+
+        if (self._state != RunnerState.INITIALIZED):
+            raise ValueError("Cannot run the workflow without entering the context")
+
+        try:
+            self._state = RunnerState.RUNNING
+
+            if (not self._entry_fn.has_single_output):
+                raise ValueError("Workflow does not support single output")
+
+            async with self._exporter_manager.start(context_state=self._context_state):
+                # Run the workflow
+                result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)
+
+                # Close the intermediate stream
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
+
+            self._state = RunnerState.COMPLETED
+
+            return result
+        except Exception as e:
+            logger.exception("Error running workflow: %s", e)
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
+            self._state = RunnerState.FAILED
+
+            raise
+
+    async def result_stream(self, to_type: type | None = None):
+
+        if (self._state != RunnerState.INITIALIZED):
+            raise ValueError("Cannot run the workflow without entering the context")
+
+        try:
+            self._state = RunnerState.RUNNING
+
+            if (not self._entry_fn.has_streaming_output):
+                raise ValueError("Workflow does not support streaming output")
+
+            # Run the workflow
+            async with self._exporter_manager.start(context_state=self._context_state):
+                async for m in self._entry_fn.astream(self._input_message, to_type=to_type):
+                    yield m
+
+                self._state = RunnerState.COMPLETED
+
+                # Close the intermediate stream
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
+
+        except Exception as e:
+            logger.exception("Error running workflow: %s", e)
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
+            self._state = RunnerState.FAILED
+
+            raise
+
+
+# Compatibility aliases with previous releases
+AIQRunnerState = RunnerState
+AIQRunner = Runner