nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/front_ends/fastapi/message_validator.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -27,13 +27,16 @@ from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
 from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ObservabilityTraceContent
 from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponseObservabilityTrace
 from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import SystemIntermediateStepContent
 from nat.data_models.api_server import SystemResponseContent
 from nat.data_models.api_server import TextContent
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketObservabilityTraceMessage
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
 from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
@@ -67,11 +70,24 @@ class MessageValidator:
             WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE: WebSocketSystemIntermediateStepMessage,
             WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE: WebSocketSystemInteractionMessage,
             WebSocketMessageType.USER_INTERACTION_MESSAGE: WebSocketUserInteractionResponseMessage,
-            WebSocketMessageType.ERROR_MESSAGE: Error
+            WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE: WebSocketObservabilityTraceMessage,
+            WebSocketMessageType.ERROR_MESSAGE: Error,
         }
 
         self._message_parent_id: str = "default_id"
 
+    def _get_observability_trace_id_from_context(self) -> str | None:
+        """
+        Retrieves observability_trace_id from Context
+
+        :return: observability_trace_id if available, None otherwise.
+        """
+        try:
+            from nat.builder.context import Context
+            return Context.get().observability_trace_id
+        except (ImportError, AttributeError, KeyError):
+            return None
+
     async def validate_message(self, message: dict[str, Any]) -> BaseModel:
         """
         Validates an incoming WebSocket message against its expected schema.
@@ -147,6 +163,9 @@ class MessageValidator:
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_content = SystemIntermediateStepContent(name=data_model.name,
                                                                           payload=data_model.payload)
+            elif (isinstance(data_model, ResponseObservabilityTrace)):
+                validated_message_content = ObservabilityTraceContent(
+                    observability_trace_id=data_model.observability_trace_id)
             elif (isinstance(data_model, HumanPromptBase)):
                 validated_message_content = data_model
             elif (isinstance(data_model, SystemResponseContent)):
@@ -212,6 +231,9 @@ class MessageValidator:
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_type = WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE
 
+            elif (isinstance(data_model, ResponseObservabilityTrace)):
+                validated_message_type = WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE
+
             elif (isinstance(data_model, HumanPromptBase)):
                 validated_message_type = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
             else:
@@ -349,3 +371,33 @@ class MessageValidator:
         except Exception as e:
             logger.exception("Error creating system interaction message: %s", str(e))
             return None
+
+    async def create_observability_trace_message(
+        self,
+        *,
+        message_id: str | None = str(uuid.uuid4()),
+        parent_id: str = "default",
+        conversation_id: str | None = None,
+        content: ObservabilityTraceContent,
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
+    ) -> WebSocketObservabilityTraceMessage | None:
+        """
+        Creates an observability trace message.
+
+        :param message_id: Unique identifier for the message (default: generated UUID).
+        :param parent_id: ID of the user message that spawned child messages.
+        :param conversation_id: ID of the conversation this message belongs to (default: None).
+        :param content: Message content.
+        :param timestamp: Timestamp of the message (default: current UTC time).
+        :return: A WebSocketObservabilityTraceMessage instance.
+        """
+        try:
+            return WebSocketObservabilityTraceMessage(id=message_id,
+                                                      parent_id=parent_id,
+                                                      conversation_id=conversation_id,
+                                                      content=content,
+                                                      timestamp=timestamp)
+
+        except Exception as e:
+            logger.exception("Error creating observability trace message: %s", str(e))
+            return None

nat/front_ends/fastapi/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/response_helpers.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -23,20 +23,20 @@ from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.step_adaptor import StepAdaptorConfig
 from nat.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
 from nat.front_ends.fastapi.step_adaptor import StepAdaptor
-from nat.runtime.session import SessionManager
+from nat.runtime.session import Session
 from nat.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
 
 
 async def generate_streaming_response_as_str(payload: typing.Any,
                                              *,
-                                             session_manager: SessionManager,
+                                             session: Session,
                                              streaming: bool,
                                              step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                              result_type: type | None = None,
                                              output_type: type | None = None) -> AsyncGenerator[str]:
 
     async for item in generate_streaming_response(payload,
-                                                  session_manager=session_manager,
+                                                  session=session,
                                                   streaming=streaming,
                                                   step_adaptor=step_adaptor,
                                                   result_type=result_type,
@@ -51,13 +51,13 @@ async def generate_streaming_response_as_str(payload: typing.Any,
 
 async def generate_streaming_response(payload: typing.Any,
                                       *,
-                                      session_manager: SessionManager,
+                                      session: Session,
                                       streaming: bool,
                                       step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                       result_type: type | None = None,
                                       output_type: type | None = None) -> AsyncGenerator[ResponseSerializable]:
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
 
         q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
@@ -65,7 +65,7 @@ async def generate_streaming_response(payload: typing.Any,
         intermediate_complete = await pull_intermediate(q, step_adaptor)
 
         async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
+            if session.workflow.has_streaming_output and streaming:
                 async for chunk in runner.result_stream(to_type=output_type):
                     await q.put(chunk)
             else:
@@ -107,19 +107,20 @@ async def generate_streaming_response(payload: typing.Any,
 
 async def generate_single_response(
     payload: typing.Any,
-    session_manager: SessionManager,
+    session: Session,
     result_type: type | None = None,
 ) -> typing.Any:
-    if (not session_manager.workflow.has_single_output):
+
+    if not session.workflow.has_single_output:
         raise ValueError("Cannot get a single output value for streaming workflows")
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
         return await runner.result(to_type=result_type)
 
 
 async def generate_streaming_response_full(payload: typing.Any,
                                            *,
-                                           session_manager: SessionManager,
+                                           session: Session,
                                            streaming: bool,
                                            result_type: type | None = None,
                                            output_type: type | None = None,
@@ -137,14 +138,14 @@ async def generate_streaming_response_full(payload: typing.Any,
         else:
             allowed_types = set(filter_steps.split(','))
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
         q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
         # Start the intermediate stream without step adaptor
         intermediate_complete = await pull_intermediate(q, None)
 
         async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
+            if session.workflow.has_streaming_output and streaming:
                 async for chunk in runner.result_stream(to_type=output_type):
                     await q.put(chunk)
             else:
@@ -174,7 +175,7 @@ async def generate_streaming_response_full(payload: typing.Any,
 
 async def generate_streaming_response_full_as_str(payload: typing.Any,
                                                   *,
-                                                  session_manager: SessionManager,
+                                                  session: Session,
                                                   streaming: bool,
                                                   result_type: type | None = None,
                                                   output_type: type | None = None,
@@ -183,7 +184,7 @@ async def generate_streaming_response_full_as_str(payload: typing.Any,
     Similar to generate_streaming_response but converts the response to a string format.
     """
     async for item in generate_streaming_response_full(payload,
-                                                       session_manager=session_manager,
+                                                       session=session,
                                                        streaming=streaming,
                                                        result_type=result_type,
                                                        output_type=output_type,

nat/front_ends/fastapi/step_adaptor.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/utils.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,4 +18,3 @@
 
 from .console import register as console_register
 from .fastapi import register as fastapi_register
-from .mcp import register as mcp_register

nat/front_ends/simple_base/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -47,9 +47,11 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
                 click.echo(stream.getvalue())
 
-            workflow = await builder.build()
-            session_manager = SessionManager(workflow)
-            await self.run_workflow(session_manager)
+            session_manager = await SessionManager.create(config=self.full_config, shared_builder=builder)
+            try:
+                await self.run_workflow(session_manager)
+            finally:
+                await session_manager.shutdown()
 
     @abstractmethod
     async def run_workflow(self, session_manager: SessionManager):

nat/llm/aws_bedrock_llm.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/llm/azure_openai_llm.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,6 +16,7 @@
 from pydantic import AliasChoices
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import computed_field
 
 from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
@@ -59,6 +60,14 @@ class AzureOpenAIModelConfig(
                                            description="Top-p for distribution sampling.",
                                            space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
+    @computed_field
+    @property
+    def model_name(self) -> str:
+        """
+        Returns the model name for compatibility with other parts of the code base which expect a model_name attribute.
+        """
+        return self.azure_deployment
+
 
 @register_llm_provider(config_type=AzureOpenAIModelConfig)
 async def azure_openai_llm(config: AzureOpenAIModelConfig, _builder: Builder):

nat/llm/dynamo_llm.py

@@ -0,0 +1,363 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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.
+"""
+Dynamo LLM provider with automatic prefix header injection for KV cache optimization.
+
+This module provides a specialized OpenAI-compatible LLM that sends Dynamo prefix headers
+for optimal KV cache management and request routing. The prefix parameters are optimizable
+via the NAT optimizer.
+
+The implementation uses httpx event hooks to inject headers at the HTTP transport level,
+making it framework-agnostic (works with LangChain, LlamaIndex, etc.).
+
+Dynamo Prefix Parameters
+-------------------------
+
+prefix_osl (Output Sequence Length)
+    Hint for expected response length:
+
+    - LOW: decode_cost=1.0, short responses
+    - MEDIUM: decode_cost=2.0, typical responses
+    - HIGH: decode_cost=3.0, long responses
+
+prefix_iat (Inter-Arrival Time)
+    Hint for request pacing:
+
+    - LOW: iat_factor=1.5, rapid bursts -> high worker stickiness
+    - MEDIUM: iat_factor=1.0, normal pacing
+    - HIGH: iat_factor=0.6, slow requests -> more exploration
+
+prefix_total_requests
+    Expected requests per conversation:
+
+    - Higher values increase KV cache affinity and worker stickiness
+    - Lower values allow more load balancing
+"""
+
+import logging
+import uuid
+from collections.abc import Callable
+from collections.abc import Coroutine
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Literal
+
+if TYPE_CHECKING:
+    import httpx
+
+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.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
+from nat.llm.openai_llm import OpenAIModelConfig
+
+logger = logging.getLogger(__name__)
+
+# Define valid prefix hint values
+PrefixLevel = Literal["LOW", "MEDIUM", "HIGH"]
+
+# =============================================================================
+# CONTEXT MANAGEMENT FOR DYNAMO PREFIX ID
+# =============================================================================
+
+
+class DynamoPrefixContext:
+    """
+    Singleton class for managing Dynamo prefix IDs across LLM calls.
+
+    This allows evaluation code to set a prefix ID that persists across all LLM
+    calls for a single evaluation question (multi-turn conversation).
+
+    Usage::
+
+        from nat.llm.dynamo_llm import DynamoPrefixContext
+
+        # Set prefix ID at the start of each evaluation question
+        DynamoPrefixContext.set("eval-q001-abc123")
+
+        # ... perform LLM calls ...
+
+        # Clear when done
+        DynamoPrefixContext.clear()
+
+        # Or use as a context manager
+        with DynamoPrefixContext.scope("eval-q001-abc123"):
+            # ... perform LLM calls ...
+    """
+
+    _current_prefix_id: ContextVar[str | None] = ContextVar('dynamo_prefix_id', default=None)
+
+    @classmethod
+    def set(cls, prefix_id: str) -> None:
+        """
+        Set the Dynamo prefix ID for the current context.
+
+        Call this at the start of each evaluation question to ensure all LLM calls
+        for that question share the same prefix ID (enabling KV cache reuse).
+
+        Args:
+            prefix_id: The unique prefix ID (e.g., "eval-q001-abc123")
+        """
+        cls._current_prefix_id.set(prefix_id)
+        logger.debug("Set Dynamo prefix ID: %s", prefix_id)
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear the current Dynamo prefix ID context."""
+        cls._current_prefix_id.set(None)
+        logger.debug("Cleared Dynamo prefix ID")
+
+    @classmethod
+    def get(cls) -> str | None:
+        """Get the current Dynamo prefix ID from context, if any."""
+        return cls._current_prefix_id.get()
+
+    @classmethod
+    @contextmanager
+    def scope(cls, prefix_id: str) -> Iterator[None]:
+        """
+        Context manager for scoped prefix ID usage.
+
+        Automatically sets the prefix ID on entry and clears it on exit,
+        ensuring proper cleanup even if exceptions occur.
+
+        Args:
+            prefix_id: The unique prefix ID for this scope
+
+        Yields:
+            None
+
+        Usage:
+            with DynamoPrefixContext.scope("eval-q001"):
+                # All LLM calls here will use "eval-q001" prefix
+                await llm.ainvoke(...)
+        """
+        cls.set(prefix_id)
+        try:
+            yield
+        finally:
+            cls.clear()
+
+
+# =============================================================================
+# DYNAMO MODEL CONFIGURATION
+# =============================================================================
+
+
+class DynamoModelConfig(OpenAIModelConfig, name="dynamo"):
+    """
+    A Dynamo LLM provider with automatic prefix header injection for KV cache optimization.
+
+    This is a specialized OpenAI-compatible LLM that sends Dynamo prefix headers
+    for optimal KV cache management and request routing. Prefix headers are enabled
+    by default using the template "nat-dynamo-{uuid}". The prefix routing parameters
+    (prefix_total_requests, prefix_osl, prefix_iat) are optimizable via the NAT optimizer.
+
+    To disable prefix headers, set prefix_template to null/None in your config.
+    """
+
+    # =========================================================================
+    # DYNAMO PREFIX PARAMETERS
+    # =========================================================================
+
+    prefix_template: str | None = Field(
+        default="nat-dynamo-{uuid}",
+        description="Template for prefix ID. The {uuid} placeholder will be replaced with a unique ID. "
+        "Prefix headers are sent by default for KV cache optimization. "
+        "Set to null/None to disable prefix header injection.",
+    )
+
+    prefix_total_requests: int = OptimizableField(
+        default=10,
+        ge=1,
+        le=50,
+        description=("Expected number of requests for this conversation/prefix. "
+                     "Higher values increase worker stickiness and KV cache locality. "
+                     "Lower values allow more load balancing across workers."),
+        space=SearchSpace(low=1, high=20, step=5))
+
+    prefix_osl: PrefixLevel = OptimizableField(default="MEDIUM",
+                                               description=("Output Sequence Length hint for the Dynamo router. "
+                                                            "LOW=short responses (decode_cost=1.0), "
+                                                            "MEDIUM=typical (decode_cost=2.0), "
+                                                            "HIGH=long responses (decode_cost=3.0)."),
+                                               space=SearchSpace(values=["LOW", "MEDIUM", "HIGH"]))
+
+    prefix_iat: PrefixLevel = OptimizableField(default="MEDIUM",
+                                               description=("Inter-Arrival Time hint for the Dynamo router. "
+                                                            "LOW=rapid bursts (iat_factor=1.5, high stickiness), "
+                                                            "MEDIUM=normal (iat_factor=1.0), "
+                                                            "HIGH=slow requests (iat_factor=0.6, more exploration)."),
+                                               space=SearchSpace(values=["LOW", "MEDIUM", "HIGH"]))
+
+    request_timeout: float = Field(
+        default=600.0,
+        gt=0.0,
+        description="HTTP request timeout in seconds for LLM requests.",
+    )
+
+    # =========================================================================
+    # UTILITY METHODS
+    # =========================================================================
+
+    @staticmethod
+    def get_dynamo_field_names() -> frozenset[str]:
+        """
+        Get the set of Dynamo-specific field names for model_dump exclusion.
+
+        Use this when building config dicts for framework clients to exclude
+        Dynamo-specific parameters that should not be passed to the underlying client.
+
+        Returns:
+            A frozenset of Dynamo-specific field names.
+
+        Example::
+
+            config_dict = config.model_dump(
+                exclude={"type", "thinking", *DynamoModelConfig.get_dynamo_field_names()},
+                ...
+            )
+        """
+        return frozenset({
+            "prefix_template",
+            "prefix_total_requests",
+            "prefix_osl",
+            "prefix_iat",
+            "request_timeout",
+        })
+
+
+# =============================================================================
+# HTTPX EVENT HOOK FOR HEADER INJECTION
+# =============================================================================
+
+
+def _create_dynamo_request_hook(
+    prefix_template: str | None,
+    total_requests: int,
+    osl: str,
+    iat: str,
+) -> Callable[["httpx.Request"], Coroutine[Any, Any, None]]:
+    """
+    Create an httpx event hook that injects Dynamo prefix headers into requests.
+
+    This hook is called before each HTTP request is sent, allowing us to inject
+    headers dynamically. The prefix ID is generated ONCE when the hook is created,
+    ensuring all requests from the same client share the same prefix ID. This enables
+    Dynamo's KV cache optimization across multi-turn conversations.
+
+    The context variable can override this for scenarios where you need different
+    prefix IDs (e.g., per-question in batch evaluation).
+
+    Args:
+        prefix_template: Template string with {uuid} placeholder
+        total_requests: Expected number of requests for this prefix
+        osl: Output sequence length hint (LOW/MEDIUM/HIGH)
+        iat: Inter-arrival time hint (LOW/MEDIUM/HIGH)
+
+    Returns:
+        An async function suitable for use as an httpx event hook.
+    """
+    # Generate the default prefix ID ONCE when the hook is created
+    # This ensures all requests from this client share the same prefix ID
+    unique_id = uuid.uuid4().hex[:16]
+    if prefix_template:
+        default_prefix_id = prefix_template.format(uuid=unique_id)
+    else:
+        default_prefix_id = f"nat-dynamo-{unique_id}"
+
+    logger.debug("Created Dynamo request hook with default prefix ID: %s", default_prefix_id)
+
+    async def on_request(request):
+        """Inject Dynamo prefix headers before each request."""
+        # Check context variable first (allows per-question override in batch evaluation)
+        context_prefix_id = DynamoPrefixContext.get()
+
+        if context_prefix_id:
+            prefix_id = context_prefix_id
+            logger.debug("Using context prefix ID: %s", prefix_id)
+        else:
+            # Use the pre-generated prefix ID (same for all requests from this client)
+            prefix_id = default_prefix_id
+            logger.debug("Using default prefix ID: %s", prefix_id)
+
+        # Inject Dynamo headers
+        request.headers["x-prefix-id"] = prefix_id
+        request.headers["x-prefix-total-requests"] = str(total_requests)
+        request.headers["x-prefix-osl"] = osl.upper()
+        request.headers["x-prefix-iat"] = iat.upper()
+
+        logger.debug("Injected Dynamo headers: prefix_id=%s, total_requests=%d, osl=%s, iat=%s",
+                     prefix_id,
+                     total_requests,
+                     osl.upper(),
+                     iat.upper())
+
+    return on_request
+
+
+def create_httpx_client_with_dynamo_hooks(
+    prefix_template: str | None,
+    total_requests: int,
+    osl: str,
+    iat: str,
+    timeout: float = 600.0,
+) -> "httpx.AsyncClient":
+    """
+    Create an httpx.AsyncClient with Dynamo prefix header injection.
+
+    This client can be passed to the OpenAI SDK to inject headers at the HTTP level,
+    making it framework-agnostic.
+
+    Args:
+        prefix_template: Template string with {uuid} placeholder
+        total_requests: Expected number of requests for this prefix
+        osl: Output sequence length hint (LOW/MEDIUM/HIGH)
+        iat: Inter-arrival time hint (LOW/MEDIUM/HIGH)
+        timeout: HTTP request timeout in seconds
+
+    Returns:
+        An httpx.AsyncClient configured with Dynamo header injection.
+    """
+    import httpx
+
+    request_hook = _create_dynamo_request_hook(prefix_template, total_requests, osl, iat)
+
+    return httpx.AsyncClient(
+        event_hooks={"request": [request_hook]},
+        timeout=httpx.Timeout(timeout),
+    )
+
+
+# =============================================================================
+# PROVIDER REGISTRATION
+# =============================================================================
+# Note: Client registrations for each framework (LangChain, LlamaIndex, etc.)
+# are in the respective plugin packages under packages/nvidia_nat_<framework>/
+
+
+@register_llm_provider(config_type=DynamoModelConfig)
+async def dynamo_llm(config: DynamoModelConfig, _builder: Builder):
+    """Register the Dynamo LLM provider."""
+    yield LLMProviderInfo(
+        config=config,
+        description="A Dynamo-optimized model with automatic prefix headers for KV cache management.",
+    )