nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/llm/litellm_llm.py

@@ -0,0 +1,80 @@
+# 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 collections.abc import AsyncIterator
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.common import OptionalSecretStr
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
+from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+
+
+class LiteLlmModelConfig(
+        LLMBaseConfig,
+        OptimizableMixin,
+        RetryMixin,
+        ThinkingMixin,
+        name="litellm",
+):
+    """A LiteLlm provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: OptionalSecretStr = Field(default=None, description="API key to interact with hosted model.")
+    base_url: str | None = Field(default=None,
+                                 description="Base url to the hosted model.",
+                                 validation_alias=AliasChoices("base_url", "api_base"),
+                                 serialization_alias="api_base")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The LiteLlm hosted model name.")
+    seed: int | None = Field(default=None, description="Random seed to set for generation.")
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        description="Sampling temperature to control randomness in the output.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))
+    top_p: float | None = OptimizableField(default=None,
+                                           ge=0.0,
+                                           le=1.0,
+                                           description="Top-p for distribution sampling.",
+                                           space=SearchSpace(high=1.0, low=0.5, step=0.1))
+
+
+@register_llm_provider(config_type=LiteLlmModelConfig)
+async def litellm_model(
+    config: LiteLlmModelConfig,
+    _builder: Builder,
+) -> AsyncIterator[LLMProviderInfo]:
+    """Litellm model provider.
+
+    Args:
+        config (LiteLlmModelConfig): The LiteLlm model configuration.
+        _builder (Builder): The NAT builder instance.
+
+    Returns:
+        AsyncIterator[LLMProviderInfo]: An async iterator that yields an LLMProviderInfo object.
+    """
+    yield LLMProviderInfo(config=config, description="A LiteLlm model for use with an LLM client.")

nat/llm/nim_llm.py

@@ -21,24 +21,38 @@ from pydantic import PositiveInt
 from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.common import OptionalSecretStr
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
-from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
-from nat.data_models.top_p_mixin import TopPMixin
 
 
-class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, OptimizableMixin, ThinkingMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
-    api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
+    api_key: OptionalSecretStr = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
     base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The model name for the hosted NIM.")
-    max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The model name for the hosted NIM.")
+    max_tokens: PositiveInt = OptimizableField(default=300,
+                                               description="Maximum number of tokens to generate.",
+                                               space=SearchSpace(high=2176, low=128, step=512))
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        description="Sampling temperature to control randomness in the output.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))
+    top_p: float | None = OptimizableField(default=None,
+                                           ge=0.0,
+                                           le=1.0,
+                                           description="Top-p for distribution sampling.",
+                                           space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
 
 @register_llm_provider(config_type=NIMModelConfig)

nat/llm/openai_llm.py

@@ -20,25 +20,37 @@ from pydantic import Field
 from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.common import OptionalSecretStr
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
-from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
-from nat.data_models.top_p_mixin import TopPMixin
 
 
-class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, OptimizableMixin, ThinkingMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
-    api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
+    api_key: OptionalSecretStr = Field(default=None, description="OpenAI API key to interact with hosted model.")
     base_url: str | None = Field(default=None, description="Base url to the hosted model.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The OpenAI hosted model name.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The OpenAI hosted model name.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
     max_retries: int = Field(default=10, description="The max number of retries for the request.")
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        description="Sampling temperature to control randomness in the output.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))
+    top_p: float | None = OptimizableField(default=None,
+                                           ge=0.0,
+                                           le=1.0,
+                                           description="Top-p for distribution sampling.",
+                                           space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
 
 @register_llm_provider(config_type=OpenAIModelConfig)

nat/llm/register.py

@@ -15,9 +15,13 @@
 
 # flake8: noqa
 # isort:skip_file
+"""Register LLM providers via import side effects.
 
+This module is imported by the NeMo Agent Toolkit runtime to ensure providers are registered and discoverable.
+"""
 # Import any providers which need to be automatically registered here
 from . import aws_bedrock_llm
 from . import azure_openai_llm
+from . import litellm_llm
 from . import nim_llm
 from . import openai_llm

nat/llm/utils/thinking.py

@@ -19,10 +19,10 @@ import logging
 import types
 from abc import abstractmethod
 from collections.abc import AsyncGenerator
+from collections.abc import Callable
 from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Any
-from typing import Callable
 from typing import TypeVar
 
 ModelType = TypeVar("ModelType")

nat/observability/exporter/base_exporter.py

@@ -372,7 +372,7 @@ class BaseExporter(Exporter):
         try:
             # Wait for all tasks to complete with a timeout
             await asyncio.wait_for(asyncio.gather(*self._tasks, return_exceptions=True), timeout=timeout)
-        except asyncio.TimeoutError:
+        except TimeoutError:
             logger.warning("%s: Some tasks did not complete within %s seconds", self.name, timeout)
         except Exception as e:
             logger.exception("%s: Error while waiting for tasks: %s", self.name, e)

nat/observability/exporter/processing_exporter.py

@@ -53,6 +53,8 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
     - Configurable None filtering: processors returning None can drop items from pipeline
     - Automatic type validation before export
     """
+    # All ProcessingExporter instances automatically use this for signature checking
+    _signature_method = '_process_pipeline'
 
     def __init__(self, context_state: ContextState | None = None, drop_nones: bool = True):
         """Initialize the processing exporter.
@@ -294,8 +296,6 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             self._check_processor_compatibility(predecessor,
                                                 processor,
                                                 "predecessor",
-                                                predecessor.output_class,
-                                                processor.input_class,
                                                 str(predecessor.output_type),
                                                 str(processor.input_type))
 
@@ -304,8 +304,6 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             self._check_processor_compatibility(processor,
                                                 successor,
                                                 "successor",
-                                                processor.output_class,
-                                                successor.input_class,
                                                 str(processor.output_type),
                                                 str(successor.input_type))
 
@@ -313,34 +311,22 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
                                        source_processor: Processor,
                                        target_processor: Processor,
                                        relationship: str,
-                                       source_class: type,
-                                       target_class: type,
                                        source_type: str,
                                        target_type: str) -> None:
-        """Check type compatibility between two processors.
+        """Check type compatibility between two processors using Pydantic validation.
 
         Args:
             source_processor (Processor): The processor providing output
             target_processor (Processor): The processor receiving input
             relationship (str): Description of relationship ("predecessor" or "successor")
-            source_class (type): The output class of source processor
-            target_class (type): The input class of target processor
             source_type (str): String representation of source type
             target_type (str): String representation of target type
         """
-        try:
-            if not issubclass(source_class, target_class):
-                raise ValueError(f"Processor {target_processor.__class__.__name__} input type {target_type} "
-                                 f"is not compatible with {relationship} {source_processor.__class__.__name__} "
-                                 f"output type {source_type}")
-        except TypeError:
-            logger.warning(
-                "Cannot use issubclass() for type compatibility check between "
-                "%s (%s) and %s (%s). Skipping compatibility check.",
-                source_processor.__class__.__name__,
-                source_type,
-                target_processor.__class__.__name__,
-                target_type)
+        # Use Pydantic-based type compatibility checking
+        if not source_processor.is_output_compatible_with(target_processor.input_type):
+            raise ValueError(f"Processor {target_processor.__class__.__name__} input type {target_type} "
+                             f"is not compatible with {relationship} {source_processor.__class__.__name__} "
+                             f"output type {source_type}")
 
     async def _pre_start(self) -> None:
 
@@ -350,36 +336,21 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             last_processor = self._processors[-1]
 
             # validate that the first processor's input type is compatible with the exporter's input type
-            try:
-                if not issubclass(self.input_class, first_processor.input_class):
-                    raise ValueError(f"Processor {first_processor.__class__.__name__} input type "
-                                     f"{first_processor.input_type} is not compatible with the "
-                                     f"{self.input_type} input type")
-            except TypeError as e:
-                # Handle cases where classes are generic types that can't be used with issubclass
-                logger.warning(
-                    "Cannot validate type compatibility between %s (%s) "
-                    "and exporter (%s): %s. Skipping validation.",
-                    first_processor.__class__.__name__,
-                    first_processor.input_type,
-                    self.input_type,
-                    e)
-
+            if not first_processor.is_compatible_with_input(self.input_type):
+                logger.error("First processor %s input=%s incompatible with exporter input=%s",
+                             first_processor.__class__.__name__,
+                             first_processor.input_type,
+                             self.input_type)
+                raise ValueError("First processor incompatible with exporter input")
             # Validate that the last processor's output type is compatible with the exporter's output type
-            try:
-                if not DecomposedType.is_type_compatible(last_processor.output_type, self.output_type):
-                    raise ValueError(f"Processor {last_processor.__class__.__name__} output type "
-                                     f"{last_processor.output_type} is not compatible with the "
-                                     f"{self.output_type} output type")
-            except TypeError as e:
-                # Handle cases where classes are generic types that can't be used with issubclass
-                logger.warning(
-                    "Cannot validate type compatibility between %s (%s) "
-                    "and exporter (%s): %s. Skipping validation.",
-                    last_processor.__class__.__name__,
-                    last_processor.output_type,
-                    self.output_type,
-                    e)
+            # Use DecomposedType.is_type_compatible for the final export stage to allow batch compatibility
+            # This enables BatchingProcessor[T] -> Exporter[T] patterns where the exporter handles both T and list[T]
+            if not DecomposedType.is_type_compatible(last_processor.output_type, self.output_type):
+                logger.error("Last processor %s output=%s incompatible with exporter output=%s",
+                             last_processor.__class__.__name__,
+                             last_processor.output_type,
+                             self.output_type)
+                raise ValueError("Last processor incompatible with exporter output")
 
         # Lock the pipeline to prevent further modifications
         self._pipeline_locked = True
@@ -432,12 +403,15 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
                 await self.export_processed(processed_item)
             else:
                 logger.debug("Skipping export of empty batch")
-        elif isinstance(processed_item, self.output_class):
+        elif self.validate_output_type(processed_item):
             await self.export_processed(processed_item)
         else:
             if raise_on_invalid:
-                raise ValueError(f"Processed item {processed_item} is not a valid output type. "
-                                 f"Expected {self.output_class} or list[{self.output_class}]")
+                logger.error("Invalid processed item type for export: %s (expected %s or list[%s])",
+                             type(processed_item),
+                             self.output_type,
+                             self.output_type)
+                raise ValueError("Invalid processed item type for export")
             logger.warning("Processed item %s is not a valid output type for export", processed_item)
 
     async def _continue_pipeline_after(self, source_processor: Processor, item: Any) -> None:
@@ -512,7 +486,7 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             event (IntermediateStep): The event to be exported.
         """
         # Convert IntermediateStep to PipelineInputT and create export task
-        if isinstance(event, self.input_class):
+        if self.validate_input_type(event):
             input_item: PipelineInputT = event  # type: ignore
             coro = self._export_with_processing(input_item)
             self._create_export_task(coro)

nat/observability/exporter/span_exporter.py

@@ -126,6 +126,7 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
 
         parent_span = None
         span_ctx = None
+        workflow_trace_id = self._context_state.workflow_trace_id.get()
 
         # Look up the parent span to establish hierarchy
         # event.parent_id is the UUID of the last START step with a different UUID from current step
@@ -141,6 +142,9 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
             parent_span = parent_span.model_copy() if isinstance(parent_span, Span) else None
             if parent_span and parent_span.context:
                 span_ctx = SpanContext(trace_id=parent_span.context.trace_id)
+        # No parent: adopt workflow trace id if available to keep all spans in the same trace
+        if span_ctx is None and workflow_trace_id:
+            span_ctx = SpanContext(trace_id=workflow_trace_id)
 
         # Extract start/end times from the step
         # By convention, `span_event_timestamp` is the time we started, `event_timestamp` is the time we ended.
@@ -154,28 +158,52 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
         else:
             sub_span_name = f"{event.payload.event_type}"
 
+        # Prefer parent/context trace id for attribute, else workflow trace id
+        _attr_trace_id = None
+        if span_ctx is not None:
+            _attr_trace_id = span_ctx.trace_id
+        elif parent_span and parent_span.context:
+            _attr_trace_id = parent_span.context.trace_id
+        elif workflow_trace_id:
+            _attr_trace_id = workflow_trace_id
+
+        attributes = {
+            f"{self._span_prefix}.event_type":
+                event.payload.event_type.value,
+            f"{self._span_prefix}.function.id":
+                event.function_ancestry.function_id if event.function_ancestry else "unknown",
+            f"{self._span_prefix}.function.name":
+                event.function_ancestry.function_name if event.function_ancestry else "unknown",
+            f"{self._span_prefix}.subspan.name":
+                event.payload.name or "",
+            f"{self._span_prefix}.event_timestamp":
+                event.event_timestamp,
+            f"{self._span_prefix}.framework":
+                event.payload.framework.value if event.payload.framework else "unknown",
+            f"{self._span_prefix}.conversation.id":
+                self._context_state.conversation_id.get() or "unknown",
+            f"{self._span_prefix}.workflow.run_id":
+                self._context_state.workflow_run_id.get() or "unknown",
+            f"{self._span_prefix}.workflow.trace_id": (f"{_attr_trace_id:032x}" if _attr_trace_id else "unknown"),
+        }
+
         sub_span = Span(name=sub_span_name,
                         parent=parent_span,
                         context=span_ctx,
-                        attributes={
-                            f"{self._span_prefix}.event_type":
-                                event.payload.event_type.value,
-                            f"{self._span_prefix}.function.id":
-                                event.function_ancestry.function_id if event.function_ancestry else "unknown",
-                            f"{self._span_prefix}.function.name":
-                                event.function_ancestry.function_name if event.function_ancestry else "unknown",
-                            f"{self._span_prefix}.subspan.name":
-                                event.payload.name or "",
-                            f"{self._span_prefix}.event_timestamp":
-                                event.event_timestamp,
-                            f"{self._span_prefix}.framework":
-                                event.payload.framework.value if event.payload.framework else "unknown",
-                        },
+                        attributes=attributes,
                         start_time=start_ns)
 
         span_kind = event_type_to_span_kind(event.event_type)
         sub_span.set_attribute(f"{self._span_prefix}.span.kind", span_kind.value)
 
+        # Enable session grouping by setting session.id from conversation_id
+        try:
+            conversation_id = self._context_state.conversation_id.get()
+            if conversation_id:
+                sub_span.set_attribute("session.id", conversation_id)
+        except (AttributeError, LookupError):
+            pass
+
         if event.payload.data and event.payload.data.input:
             match = re.search(r"Human:\s*Question:\s*(.*)", str(event.payload.data.input))
             if match:
@@ -252,7 +280,7 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
 
         end_metadata = event.payload.metadata or {}
 
-        if not isinstance(end_metadata, (dict, TraceMetadata)):
+        if not isinstance(end_metadata, dict | TraceMetadata):
             logger.warning("Invalid metadata type for step %s", event.UUID)
             return
 

nat/observability/exporter_manager.py

@@ -184,7 +184,7 @@ class ExporterManager:
             try:
                 await asyncio.wait_for(asyncio.gather(*cleanup_tasks, return_exceptions=True),
                                        timeout=self._shutdown_timeout)
-            except asyncio.TimeoutError:
+            except TimeoutError:
                 logger.warning("Some isolated exporters did not clean up within timeout")
 
         self._active_isolated_exporters.clear()
@@ -301,7 +301,7 @@ class ExporterManager:
             try:
                 task.cancel()
                 await asyncio.wait_for(task, timeout=self._shutdown_timeout)
-            except asyncio.TimeoutError:
+            except TimeoutError:
                 logger.warning("Exporter '%s' task did not shut down in time and may be stuck.", name)
                 stuck_tasks.append(name)
             except asyncio.CancelledError:

nat/observability/mixin/redaction_config_mixin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -25,9 +25,10 @@ class RedactionConfigMixin(BaseModel):
     """
     redaction_enabled: bool = Field(default=False, description="Whether to enable redaction processing.")
     redaction_value: str = Field(default="[REDACTED]", description="Value to replace redacted attributes with.")
-    redaction_attributes: list[str] = Field(default_factory=lambda: ["input.value", "output.value", "metadata"],
-                                            description="Span attributes to redact when redaction is triggered.")
+    redaction_attributes: list[str] = Field(default_factory=lambda: ["input.value", "output.value", "nat.metadata"],
+                                            description="Attributes to redact when redaction is triggered.")
     force_redaction: bool = Field(default=False, description="Always redact regardless of other conditions.")
+    redaction_tag: str | None = Field(default=None, description="Tag to add to spans when redaction is triggered.")
 
 
 class HeaderRedactionConfigMixin(RedactionConfigMixin):
@@ -38,4 +39,4 @@ class HeaderRedactionConfigMixin(RedactionConfigMixin):
 
     Note: The callback function must be provided directly to the processor at runtime.
     """
-    redaction_header: str = Field(default="x-redaction-key", description="Header to check for redaction decisions.")
+    redaction_headers: list[str] = Field(default_factory=list, description="Headers to check for redaction decisions.")

nat/observability/mixin/tagging_config_mixin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+from collections.abc import Mapping
 from enum import Enum
 from typing import Generic
 from typing import TypeVar
@@ -20,7 +22,17 @@ from typing import TypeVar
 from pydantic import BaseModel
 from pydantic import Field
 
-TagValueT = TypeVar("TagValueT")
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+TagMappingT = TypeVar("TagMappingT", bound=Mapping)
+
+
+class BaseTaggingConfigMixin(BaseModel, Generic[TagMappingT]):
+    """Base mixin for tagging spans."""
+    tags: TagMappingT | None = Field(default=None, description="Tags to add to the span.")
 
 
 class PrivacyLevel(str, Enum):
@@ -31,20 +43,20 @@ class PrivacyLevel(str, Enum):
     HIGH = "high"
 
 
-class TaggingConfigMixin(BaseModel, Generic[TagValueT]):
-    """Generic mixin for tagging spans with typed values.
+PrivacyTagSchema = TypedDict(
+    "PrivacyTagSchema",
+    {
+        "privacy.level": PrivacyLevel,
+    },
+    total=True,
+)
 
-    This mixin provides a flexible tagging system where both the tag key
-    and value type can be customized for different use cases.
-    """
-    tag_key: str | None = Field(default=None, description="Key to use when tagging traces.")
-    tag_value: TagValueT | None = Field(default=None, description="Value to tag the traces with.")
 
+class PrivacyTaggingConfigMixin(BaseTaggingConfigMixin[PrivacyTagSchema]):
+    """Mixin for privacy level tagging on spans."""
+    pass
 
-class PrivacyTaggingConfigMixin(TaggingConfigMixin[PrivacyLevel]):
-    """Mixin for privacy level tagging on spans.
 
-    Specializes TaggingConfigMixin to work with PrivacyLevel enum values,
-    providing a typed interface for privacy-related span tagging.
-    """
+class CustomTaggingConfigMixin(BaseTaggingConfigMixin[dict[str, str]]):
+    """Mixin for string key-value tagging on spans."""
     pass