aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/data_models/memory.py

@@ -1,27 +0,0 @@
-# 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 MemoryBaseConfig(TypedBaseModel, BaseModelRegistryTag):
-    """ The base level config object for a memory object. Memories provide an interface for storing and retrieving. """
-    pass
-
-
-MemoryBaseConfigT = typing.TypeVar("MemoryBaseConfigT", bound=MemoryBaseConfig)

aiq/data_models/object_store.py

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

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

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

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

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

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

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

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

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

@@ -1,26 +0,0 @@
-# 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/embedder/langchain_client.py

@@ -1,41 +0,0 @@
-# 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)