nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/data_models/object_store.py

@@ -0,0 +1,44 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class ObjectStoreBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+ObjectStoreBaseConfigT = typing.TypeVar("ObjectStoreBaseConfigT", bound=ObjectStoreBaseConfig)
+
+
+class KeyAlreadyExistsError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"Key already exists: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))
+
+
+class NoSuchKeyError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"No object found with key: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))

aiq/data_models/profiler.py

@@ -0,0 +1,54 @@
+# 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 pydantic import BaseModel
+
+
+class PromptCachingConfig(BaseModel):
+    enable: bool = False
+    min_frequency: float = 0.5
+
+
+class BottleneckConfig(BaseModel):
+    enable_simple_stack: bool = False
+    enable_nested_stack: bool = False
+
+
+class ConcurrencySpikeConfig(BaseModel):
+    enable: bool = False
+    spike_threshold: int = 1
+
+
+class PrefixSpanConfig(BaseModel):
+    enable: bool = False
+    min_support: float = 2
+    min_coverage: float = 0
+    max_text_len: int = 1000
+    top_k: int = 10
+    chain_with_common_prefixes: bool = False
+
+
+class ProfilerConfig(BaseModel):
+
+    base_metrics: bool = False
+    token_usage_forecast: bool = False
+    token_uniqueness_forecast: bool = False
+    workflow_runtime_forecast: bool = False
+    compute_llm_metrics: bool = False
+    csv_exclude_io_text: bool = False
+    prompt_caching_prefixes: PromptCachingConfig = PromptCachingConfig()
+    bottleneck_analysis: BottleneckConfig = BottleneckConfig()
+    concurrency_spike_analysis: ConcurrencySpikeConfig = ConcurrencySpikeConfig()
+    prefix_span_analysis: PrefixSpanConfig = PrefixSpanConfig()

aiq/data_models/registry_handler.py

@@ -0,0 +1,26 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class RegistryHandlerBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+RegistryHandlerBaseConfigT = typing.TypeVar("RegistryHandlerBaseConfigT", bound=RegistryHandlerBaseConfig)

aiq/data_models/retriever.py

@@ -0,0 +1,30 @@
+# 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 typing
+
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
+
+
+class RetrieverBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    The base level config object for a retriever object. Retrievers use different provider clients (e.g., Milvus) to
+    provide an interface for searching for and retrieving documents from the configured data store.
+    """
+    pass
+
+
+RetrieverBaseConfigT = typing.TypeVar("RetrieverBaseConfigT", bound=RetrieverBaseConfig)

aiq/data_models/retry_mixin.py

@@ -0,0 +1,35 @@
+# 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 pydantic import BaseModel
+from pydantic import Field
+
+
+class RetryMixin(BaseModel):
+    """Mixin class for retry configuration."""
+    do_auto_retry: bool = Field(default=True,
+                                description="Whether to automatically retry method calls"
+                                " that fail with a retryable error.",
+                                exclude=True)
+    num_retries: int = Field(default=5,
+                             description="Number of times to retry a method call that fails"
+                             " with a retryable error.",
+                             exclude=True)
+    retry_on_status_codes: list[int | str] = Field(default_factory=lambda: [429, 500, 502, 503, 504],
+                                                   description="List of HTTP status codes that should trigger a retry.",
+                                                   exclude=True)
+    retry_on_errors: list[str] | None = Field(default_factory=lambda: ["Too Many Requests"],
+                                              description="List of error substrings that should trigger a retry.",
+                                              exclude=True)

aiq/data_models/span.py

@@ -0,0 +1,187 @@
+# 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 time
+import uuid
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import field_validator
+
+logger = logging.getLogger(__name__)
+
+
+class SpanKind(Enum):
+    LLM = "LLM"
+    TOOL = "TOOL"
+    WORKFLOW = "WORKFLOW"
+    TASK = "TASK"
+    FUNCTION = "FUNCTION"
+    CUSTOM = "CUSTOM"
+    SPAN = "SPAN"
+    EMBEDDER = "EMBEDDER"
+    RETRIEVER = "RETRIEVER"
+    AGENT = "AGENT"
+    RERANKER = "RERANKER"
+    GUARDRAIL = "GUARDRAIL"
+    EVALUATOR = "EVALUATOR"
+    UNKNOWN = "UNKNOWN"
+
+
+EVENT_TYPE_TO_SPAN_KIND_MAP = {
+    "LLM_START": SpanKind.LLM,
+    "LLM_END": SpanKind.LLM,
+    "LLM_NEW_TOKEN": SpanKind.LLM,
+    "TOOL_START": SpanKind.TOOL,
+    "TOOL_END": SpanKind.TOOL,
+    "WORKFLOW_START": SpanKind.WORKFLOW,
+    "WORKFLOW_END": SpanKind.WORKFLOW,
+    "TASK_START": SpanKind.TASK,
+    "TASK_END": SpanKind.TASK,
+    "FUNCTION_START": SpanKind.FUNCTION,
+    "FUNCTION_END": SpanKind.FUNCTION,
+    "CUSTOM_START": SpanKind.CUSTOM,
+    "CUSTOM_END": SpanKind.CUSTOM,
+    "SPAN_START": SpanKind.SPAN,
+    "SPAN_END": SpanKind.SPAN,
+    "EMBEDDER_START": SpanKind.EMBEDDER,
+    "EMBEDDER_END": SpanKind.EMBEDDER,
+    "RETRIEVER_START": SpanKind.RETRIEVER,
+    "RETRIEVER_END": SpanKind.RETRIEVER,
+    "AGENT_START": SpanKind.AGENT,
+    "AGENT_END": SpanKind.AGENT,
+    "RERANKER_START": SpanKind.RERANKER,
+    "RERANKER_END": SpanKind.RERANKER,
+    "GUARDRAIL_START": SpanKind.GUARDRAIL,
+    "GUARDRAIL_END": SpanKind.GUARDRAIL,
+    "EVALUATOR_START": SpanKind.EVALUATOR,
+    "EVALUATOR_END": SpanKind.EVALUATOR,
+}
+
+
+def event_type_to_span_kind(event_type: str) -> SpanKind:
+    """Convert an event type to a span kind.
+
+    Args:
+        event_type (str): The event type to convert.
+
+    Returns:
+        SpanKind: The span kind.
+    """
+    return EVENT_TYPE_TO_SPAN_KIND_MAP.get(event_type, SpanKind.UNKNOWN)
+
+
+class SpanAttributes(Enum):
+    AIQ_SPAN_KIND = "aiq.span.kind"
+    INPUT_VALUE = "input.value"
+    INPUT_MIME_TYPE = "input.mime_type"
+    LLM_TOKEN_COUNT_PROMPT = "llm.token_count.prompt"
+    LLM_TOKEN_COUNT_COMPLETION = "llm.token_count.completion"
+    LLM_TOKEN_COUNT_TOTAL = "llm.token_count.total"
+    OUTPUT_VALUE = "output.value"
+    OUTPUT_MIME_TYPE = "output.mime_type"
+    AIQ_USAGE_NUM_LLM_CALLS = "aiq.usage.num_llm_calls"
+    AIQ_USAGE_SECONDS_BETWEEN_CALLS = "aiq.usage.seconds_between_calls"
+    AIQ_USAGE_TOKEN_COUNT_PROMPT = "aiq.usage.token_count.prompt"
+    AIQ_USAGE_TOKEN_COUNT_COMPLETION = "aiq.usage.token_count.completion"
+    AIQ_USAGE_TOKEN_COUNT_TOTAL = "aiq.usage.token_count.total"
+    AIQ_EVENT_TYPE = "aiq.event_type"
+
+
+class MimeTypes(Enum):
+    TEXT = "text/plain"
+    JSON = "application/json"
+
+
+class SpanStatusCode(Enum):
+    OK = "OK"
+    ERROR = "ERROR"
+    UNSET = "UNSET"
+
+
+class SpanEvent(BaseModel):
+    timestamp: float = Field(default_factory=lambda: int(time.time() * 1e9), description="The timestamp of the event.")
+    name: str = Field(description="The name of the event.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the event.")
+
+
+class SpanStatus(BaseModel):
+    code: SpanStatusCode = Field(default=SpanStatusCode.OK, description="The status code of the span.")
+    message: str | None = Field(default=None, description="The status message of the span.")
+
+
+class SpanContext(BaseModel):
+    trace_id: int = Field(default_factory=lambda: uuid.uuid4().int, description="The 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=lambda: uuid.uuid4().int & ((1 << 64) - 1),
+                         description="The 64-bit span ID of the span.")
+
+
+class Span(BaseModel):
+    name: str = Field(description="The name of the span.")
+    context: SpanContext | None = Field(default=None, description="The context of the span.")
+    parent: "Span | None" = Field(default=None, description="The parent span of the span.")
+    start_time: int = Field(default_factory=lambda: int(time.time() * 1e9), description="The start time of the span.")
+    end_time: int | None = Field(default=None, description="The end time of the span.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the span.")
+    events: list[SpanEvent] = Field(default_factory=list, description="The events of the span.")
+    status: SpanStatus = Field(default_factory=SpanStatus, description="The status of the span.")
+
+    @field_validator('context', mode='before')
+    @classmethod
+    def set_default_context(cls, v: SpanContext | None) -> SpanContext:
+        """Set the default context if the context is not provided.
+
+        Args:
+            v (SpanContext | None): The context to set.
+
+        Returns:
+            SpanContext: The context.
+        """
+        if v is None:
+            return SpanContext()
+        return v
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Set the attribute of the span.
+
+        Args:
+            key (str): The key of the attribute.
+            value (Any): The value of the attribute.
+        """
+        self.attributes[key] = value
+
+    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
+        """Add an event to the span.
+
+        Args:
+            name (str): The name of the event.
+            attributes (dict[str, Any] | None): The attributes of the event.
+        """
+        if attributes is None:
+            attributes = {}
+        self.events = self.events + [SpanEvent(name=name, attributes=attributes)]
+
+    def end(self, end_time: int | None = None) -> None:
+        """End the span.
+
+        Args:
+            end_time (int | None): The end time of the span.
+        """
+        if end_time is None:
+            end_time = int(time.time() * 1e9)
+        self.end_time = end_time

aiq/data_models/step_adaptor.py

@@ -0,0 +1,64 @@
+# 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 enum import Enum
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
+
+from aiq.data_models.intermediate_step import IntermediateStepType
+
+logger = logging.getLogger(__name__)
+
+
+class StepAdaptorMode(str, Enum):
+    DEFAULT = "default"
+    CUSTOM = "custom"
+    OFF = "off"
+
+
+class StepAdaptorConfig(BaseModel):
+    """
+    Configures how intermediate steps are filtered and normalized by the StepAdaptor.
+
+    Args:
+        mode (StepAdaptorMode): One of:
+            - 'current' => pass only LLM (all LLM_* events) + TOOL_END
+            - 'end_events_only' => pass only LLM_END and TOOL_END
+            - 'custom' => pass only the events in custom_event_types
+        custom_event_types (list[IntermediateStepType]):
+            If mode == 'custom', we only pass events whose event_type is in this list.
+            Otherwise, this field is ignored.
+    """
+    mode: StepAdaptorMode = StepAdaptorMode.DEFAULT
+    custom_event_types: list[IntermediateStepType] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def check_custom_event_types(self) -> "StepAdaptorConfig":
+        """
+        Validates custom configurations
+        """
+        if self.mode != StepAdaptorMode.CUSTOM and self.custom_event_types:
+            logger.warning("Ignoring custom_event_types because mode is not 'custom'")
+            self.custom_event_types = []
+        elif self.mode == StepAdaptorMode.CUSTOM and not self.custom_event_types:
+            logger.warning("No custom_event_types provided for custom mode. Defaulting to CUSTOM_START and CUSTOM_END")
+            self.custom_event_types = [IntermediateStepType.CUSTOM_START, IntermediateStepType.CUSTOM_END]
+        elif self.mode == StepAdaptorMode.OFF:
+            logger.warning("StepAdaptor is disabled. Ignoring all intermediate event types")
+            self.custom_event_types = []
+        return self

aiq/data_models/streaming.py

@@ -0,0 +1,33 @@
+# 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 inspect
+import typing
+from collections.abc import Callable
+
+_T = typing.TypeVar("_T")
+
+
+class Streaming(typing.Generic[_T]):
+
+    def __init__(self, *, convert: Callable[[typing.Any], _T] = None):
+        self.convert = convert
+
+        sig = inspect.signature(self.convert)
+
+        if (sig.return_annotation == inspect.Parameter.empty):
+            raise ValueError("convert must have a return annotation")
+
+        self.single_output_type = sig.return_annotation

aiq/data_models/swe_bench_model.py

@@ -0,0 +1,54 @@
+# 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
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import field_validator
+
+
+class SWEBenchInput(BaseModel):
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    repo: str
+    instance_id: str
+    base_commit: str
+    patch: str
+    test_patch: str
+    problem_statement: str
+    hints_text: str
+    created_at: str | int
+    version: float
+    FAIL_TO_PASS: list[str]
+    PASS_TO_PASS: list[str]
+    environment_setup_commit: str
+
+    # Handle improperly formatted JSON strings for list fields
+    @field_validator("FAIL_TO_PASS", "PASS_TO_PASS", mode="before")
+    def parse_list_fields(cls, value):  # pylint: disable=no-self-argument
+        if isinstance(value, str):
+            # Attempt to parse the string as a list
+            return json.loads(value)
+        return value
+
+
+class SWEBenchOutput(BaseModel):
+    model_config = ConfigDict(protected_namespaces=(), populate_by_name=True)
+
+    instance_id: str
+    model_name_or_path: str
+    model_patch: str

aiq/data_models/telemetry_exporter.py

@@ -0,0 +1,26 @@
+# 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 typing
+
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
+
+
+class TelemetryExporterBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+TelemetryExporterConfigT = typing.TypeVar("TelemetryExporterConfigT", bound=TelemetryExporterBaseConfig)

aiq/data_models/ttc_strategy.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.
+
+import typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class TTCStrategyBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration class for Test Time Compute (TTC) strategy.
+    This class is used to define the structure of TTC strategy configurations.
+    """
+    pass
+
+
+TTCStrategyBaseConfigT = typing.TypeVar("TTCStrategyBaseConfigT", bound=TTCStrategyBaseConfig)

aiq/embedder/langchain_client.py

@@ -0,0 +1,41 @@
+# 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 aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_embedder_client
+from aiq.embedder.nim_embedder import NIMEmbedderModelConfig
+
+
+@register_embedder_client(config_type=NIMEmbedderModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def nim_langchain(embedder_config: NIMEmbedderModelConfig, builder: Builder):
+
+    from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
+
+    yield NVIDIAEmbeddings(**embedder_config.model_dump(exclude={"type"}, by_alias=True))
+
+
+@register_embedder_client(config_type=NIMEmbedderModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def nim_llamaindex(embedder_config: NIMEmbedderModelConfig, builder: Builder):
+
+    from llama_index.embeddings.nvidia import NVIDIAEmbedding  # pylint: disable=no-name-in-module
+
+    config_obj = {
+        **embedder_config.model_dump(exclude={"type", "model_name"}, by_alias=True),
+        "model":
+            embedder_config.model_name,
+    }
+
+    yield NVIDIAEmbedding(**config_obj)