nvidia-nat-data-flywheel 1.3.0a20250828__py3-none-any.whl

nat/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA Agent Toolkit Subpackage
+This is a subpackage for NVIDIA Data Flywheel Blueprint integration for continuous model improvement.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nat/plugins/data_flywheel/observability/__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/plugins/data_flywheel/observability/exporter/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/plugins/data_flywheel/observability/exporter/dfw_elasticsearch_exporter.py

@@ -0,0 +1,74 @@
+# 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.
+
+from nat.builder.context import ContextState
+from nat.plugins.data_flywheel.observability.exporter.dfw_exporter import DFWExporter
+from nat.plugins.data_flywheel.observability.mixin.elasticsearch_mixin import ElasticsearchMixin
+from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch import ContractVersion
+
+
+class DFWElasticsearchExporter(ElasticsearchMixin, DFWExporter):
+    """Elasticsearch-specific Data Flywheel exporter."""
+
+    def __init__(self,
+                 context_state: ContextState | None = None,
+                 client_id: str = "default",
+                 contract_version: ContractVersion = ContractVersion.V1_1,
+                 batch_size: int = 100,
+                 flush_interval: float = 5.0,
+                 max_queue_size: int = 1000,
+                 drop_on_overflow: bool = False,
+                 shutdown_timeout: float = 10.0,
+                 **elasticsearch_kwargs):
+        """Initialize the Elasticsearch Data Flywheel exporter.
+
+        Args:
+            context_state: The context state to use for the exporter.
+            client_id: The client ID for the exporter.
+            contract_version: The contract version to use for the exporter.
+            batch_size: The batch size for exporting spans.
+            flush_interval: The flush interval in seconds for exporting spans.
+            max_queue_size: The maximum queue size for exporting spans.
+            drop_on_overflow: Whether to drop spans on overflow.
+            shutdown_timeout: The shutdown timeout in seconds.
+            **elasticsearch_kwargs: Additional arguments for ElasticsearchMixin:
+                - endpoint: The elasticsearch endpoint.
+                - index: The elasticsearch index name.
+                - elasticsearch_auth: The elasticsearch authentication credentials.
+                - headers: The elasticsearch headers.
+        """
+        # Initialize both mixins - ElasticsearchMixin expects elasticsearch_kwargs,
+        # DFWExporter expects the standard exporter parameters
+        self.contract_version = contract_version
+        super().__init__(export_contract=contract_version.get_contract_class(),
+                         context_state=context_state,
+                         batch_size=batch_size,
+                         flush_interval=flush_interval,
+                         max_queue_size=max_queue_size,
+                         drop_on_overflow=drop_on_overflow,
+                         shutdown_timeout=shutdown_timeout,
+                         client_id=client_id,
+                         **elasticsearch_kwargs)
+
+    async def export_processed(self, item: dict | list[dict]) -> None:
+        """Export processed DFW records to Elasticsearch.
+
+        Delegates to ElasticsearchMixin.export_processed() which handles
+        bulk operations for lists and single document indexing for individual records.
+
+        Args:
+            item (dict | list[dict]): Single dictionary or batch of dictionaries to export
+        """
+        await super().export_processed(item)

nat/plugins/data_flywheel/observability/exporter/dfw_exporter.py

@@ -0,0 +1,99 @@
+# 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 logging
+from abc import abstractmethod
+
+from pydantic import BaseModel
+
+from nat.builder.context import ContextState
+from nat.data_models.span import Span
+from nat.observability.exporter.span_exporter import SpanExporter
+from nat.observability.processor.batching_processor import BatchingProcessor
+from nat.observability.processor.falsy_batch_filter_processor import DictBatchFilterProcessor
+from nat.observability.processor.processor_factory import processor_factory_from_type
+from nat.observability.processor.processor_factory import processor_factory_to_type
+from nat.plugins.data_flywheel.observability.processor import DFWToDictProcessor
+from nat.plugins.data_flywheel.observability.processor import SpanToDFWRecordProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class DictBatchingProcessor(BatchingProcessor[dict]):
+    """Processor that batches dictionary objects for bulk operations.
+
+    Specializes BatchingProcessor with explicit dict typing to support
+    bulk export operations to sinks.
+    """
+    pass
+
+
+class DFWExporter(SpanExporter[Span, dict]):
+    """Abstract base class for Data Flywheel exporters."""
+
+    def __init__(self,
+                 export_contract: type[BaseModel],
+                 context_state: ContextState | None = None,
+                 batch_size: int = 100,
+                 flush_interval: float = 5.0,
+                 max_queue_size: int = 1000,
+                 drop_on_overflow: bool = False,
+                 shutdown_timeout: float = 10.0,
+                 client_id: str = "default"):
+        """Initialize the Data Flywheel exporter.
+
+        Args:
+            export_contract: The Pydantic model type for the export contract.
+            context_state: The context state to use for the exporter.
+            batch_size: The batch size for exporting spans.
+            flush_interval: The flush interval in seconds for exporting spans.
+            max_queue_size: The maximum queue size for exporting spans.
+            drop_on_overflow: Whether to drop spans on overflow.
+            shutdown_timeout: The shutdown timeout in seconds.
+            client_id: The client ID for the exporter.
+        """
+        super().__init__(context_state)
+
+        # Store the contract for property access
+        self._export_contract = export_contract
+
+        # Define the processor chain
+        ConcreteSpanToDFWRecordProcessor = processor_factory_to_type(SpanToDFWRecordProcessor, export_contract)
+        ConcreteDFWToDictProcessor = processor_factory_from_type(DFWToDictProcessor, export_contract)
+        self.add_processor(ConcreteSpanToDFWRecordProcessor(client_id=client_id))  # type: ignore
+        self.add_processor(ConcreteDFWToDictProcessor())
+        self.add_processor(
+            DictBatchingProcessor(batch_size=batch_size,
+                                  flush_interval=flush_interval,
+                                  max_queue_size=max_queue_size,
+                                  drop_on_overflow=drop_on_overflow,
+                                  shutdown_timeout=shutdown_timeout))
+        self.add_processor(DictBatchFilterProcessor())
+
+    @property
+    def export_contract(self) -> type[BaseModel]:
+        """The export contract used for processing spans before converting to dict.
+
+        This type defines the structure of records that spans are converted to
+        before being serialized to dictionaries for export.
+
+        Returns:
+            type[BaseModel]: The Pydantic model type for the export contract.
+        """
+        return self._export_contract
+
+    @abstractmethod
+    async def export_processed(self, item: dict | list[dict]) -> None:
+        pass

nat/plugins/data_flywheel/observability/mixin/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/plugins/data_flywheel/observability/mixin/elasticsearch_mixin.py

@@ -0,0 +1,75 @@
+# 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
+
+from elasticsearch import AsyncElasticsearch
+
+logger = logging.getLogger(__name__)
+
+
+class ElasticsearchMixin:
+    """Mixin for elasticsearch exporters.
+
+    This mixin provides elasticsearch-specific functionality for SpanExporter exporters.
+    It handles elasticsearch-specific resource tagging and uses the AsyncElasticsearch client.
+    """
+
+    def __init__(self,
+                 *args,
+                 endpoint: str,
+                 index: str,
+                 elasticsearch_auth: tuple[str, str],
+                 headers: dict[str, str] | None = None,
+                 **kwargs):
+        """Initialize the elasticsearch exporter.
+
+        Args:
+            endpoint (str): The elasticsearch endpoint.
+            index (str): The elasticsearch index.
+            elasticsearch_auth (tuple[str, str]): The elasticsearch authentication credentials.
+            headers (dict[str, str] | None): The elasticsearch headers.
+        """
+        if headers is None:
+            headers = {"Accept": "application/vnd.elasticsearch+json; compatible-with=8"}
+
+        self._elastic_client = AsyncElasticsearch(endpoint, basic_auth=elasticsearch_auth, headers=headers)
+        self._index = index
+        super().__init__(*args, **kwargs)
+
+    async def export_processed(self, item: dict | list[dict]) -> None:
+        """Export a batch of spans.
+
+        Args:
+            item (dict | list[dict]): Dictionary or list of dictionaries to export to Elasticsearch.
+        """
+        if isinstance(item, list):
+            if not item:  # Empty list
+                return
+            if not all(isinstance(doc, dict) for doc in item):
+                raise ValueError("All items in list must be dictionaries")
+
+            # Format for bulk operations: each document needs an action/metadata line
+            bulk_operations = []
+            for doc in item:
+                bulk_operations.append({"index": {"_index": self._index}})  # action/metadata with index
+                bulk_operations.append(doc)  # document
+
+            await self._elastic_client.bulk(operations=bulk_operations)
+        elif isinstance(item, dict):
+            # Single document export
+            await self._elastic_client.index(index=self._index, document=item)
+        else:
+            raise ValueError(f"Invalid item type: {type(item)}. Expected dict or list[dict]")

nat/plugins/data_flywheel/observability/processor/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+
+# DFW Record Processors
+from .dfw_record_processor import DFWToDictProcessor
+from .dfw_record_processor import SpanToDFWRecordProcessor
+
+# Trace Source Registry
+from .trace_conversion.trace_adapter_registry import TraceAdapterRegistry
+
+__all__ = [
+    "SpanToDFWRecordProcessor",  # DFW Record Processors
+    "DFWToDictProcessor",
+    "TraceAdapterRegistry",  # Trace Source Registry
+]

nat/plugins/data_flywheel/observability/processor/dfw_record_processor.py

@@ -0,0 +1,86 @@
+# 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
+from typing import TypeVar
+from typing import cast
+
+from pydantic import BaseModel
+
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.span import Span
+from nat.observability.mixin.type_introspection_mixin import TypeIntrospectionMixin
+from nat.observability.processor.processor import Processor
+from nat.plugins.data_flywheel.observability.processor.trace_conversion import span_to_dfw_record
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+DFWRecordT = TypeVar("DFWRecordT", bound=BaseModel)
+
+
+class DFWToDictProcessor(Processor[DFWRecordT, dict]):
+    """Processor that converts a Data Flywheel record to a dictionary.
+
+    Serializes Pydantic DFW record models to dictionaries using model_dump_json()
+    for consistent field aliasing and proper JSON serialization.
+    """
+
+    @override
+    async def process(self, item: DFWRecordT | None) -> dict:
+        """Convert a DFW record to a dictionary.
+
+        Args:
+            item (DFWRecordT | None): The DFW record to convert.
+
+        Returns:
+            dict: The converted dictionary.
+        """
+        if item is None:
+            logger.debug("Cannot process 'None' item, returning empty dict")
+            return {}
+
+        return json.loads(item.model_dump_json(by_alias=True))
+
+
+class SpanToDFWRecordProcessor(Processor[Span, DFWRecordT], TypeIntrospectionMixin):
+    """Processor that converts a Span to a Data Flywheel record.
+
+    Extracts trace data from spans and uses the trace adapter registry to convert
+    it to the target DFW record format.
+    """
+
+    def __init__(self, client_id: str):
+        self._client_id = client_id
+
+    @override
+    async def process(self, item: Span) -> DFWRecordT | None:
+        """Convert a Span to a DFW record.
+
+        Args:
+            item (Span): The Span to convert.
+
+        Returns:
+            DFWRecordT | None: The converted DFW record.
+        """
+
+        match item.attributes.get("nat.event_type"):
+            case IntermediateStepType.LLM_START:
+                dfw_record = span_to_dfw_record(span=item, to_type=self.output_type, client_id=self._client_id)
+                return cast(DFWRecordT | None, dfw_record)
+            case _:
+                logger.debug("Unsupported event type: '%s'", item.attributes.get("nat.event_type"))
+                return None

nat/plugins/data_flywheel/observability/processor/trace_conversion/__init__.py

@@ -0,0 +1,30 @@
+# 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 .span_extractor import extract_timestamp
+from .span_extractor import extract_token_usage
+from .span_extractor import extract_usage_info
+from .span_to_dfw_record import span_to_dfw_record
+from .trace_adapter_registry import TraceAdapterRegistry
+from .trace_adapter_registry import register_adapter
+
+__all__ = [
+    "extract_timestamp",
+    "extract_usage_info",
+    "extract_token_usage",
+    "span_to_dfw_record",
+    "register_adapter",
+    "TraceAdapterRegistry",
+]

nat/plugins/data_flywheel/observability/processor/trace_conversion/adapter/__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/plugins/data_flywheel/observability/processor/trace_conversion/adapter/elasticsearch/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/plugins/data_flywheel/observability/processor/trace_conversion/adapter/elasticsearch/nim_converter.py

@@ -0,0 +1,44 @@
+# 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 logging
+
+# pylint: disable=line-too-long
+# flake8: noqa
+from nat.plugins.data_flywheel.observability.processor.trace_conversion.adapter.elasticsearch.openai_converter import \
+    convert_langchain_openai
+from nat.plugins.data_flywheel.observability.processor.trace_conversion.trace_adapter_registry import \
+    register_adapter
+from nat.plugins.data_flywheel.observability.schema.provider.nim_trace_source import \
+    NIMTraceSource
+from nat.plugins.data_flywheel.observability.schema.sink.elasticsearch.dfw_es_record import \
+    DFWESRecord
+from nat.plugins.data_flywheel.observability.schema.trace_container import \
+    TraceContainer
+
+logger = logging.getLogger(__name__)
+
+
+@register_adapter(trace_source_model=NIMTraceSource)
+def convert_langchain_nim(trace_source: TraceContainer) -> DFWESRecord:
+    """Convert a LangChain Nim trace source to a DFWESRecord.
+
+    Args:
+        trace_source (TraceContainer): The trace source to convert
+
+    Returns:
+        DFWESRecord: The converted DFW record
+    """
+    return convert_langchain_openai(trace_source)