aiqtoolkit 1.1.0rc6__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiq/retriever/nemo_retriever/retriever.py

@@ -1,190 +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
-import logging
-import os
-import typing
-from functools import partial
-from urllib.parse import urljoin
-
-import httpx
-from langchain_core.retrievers import BaseRetriever
-from pydantic import BaseModel
-from pydantic import Field
-from pydantic import HttpUrl
-
-from aiq.retriever.interface import AIQRetriever
-from aiq.retriever.models import AIQDocument
-from aiq.retriever.models import RetrieverError
-from aiq.retriever.models import RetrieverOutput
-
-logger = logging.getLogger(__name__)
-
-
-class Collection(BaseModel):
-    id: str
-    name: str
-    meta: typing.Any
-    pipeline: str
-    created_at: str
-
-
-class RetrieverPayload(BaseModel):
-    query: str
-    top_k: int = Field(le=50, gt=0)
-
-
-class CollectionUnavailableError(RetrieverError):
-    pass
-
-
-class NemoRetriever(AIQRetriever):
-    """
-    Client for retrieving document chunks from a Nemo Retriever service.
-    """
-
-    def __init__(self, uri: str | HttpUrl, timeout: int = 60, nvidia_api_key: str = None, **kwargs):
-
-        self.base_url = str(uri)
-        self.timeout = timeout
-        self._search_func = self._search
-        self.api_key = nvidia_api_key if nvidia_api_key else os.getenv('NVIDIA_API_KEY')
-        self._bound_params = []
-        if not self.api_key:
-            logger.warning("No API key was specified as part of configuration or as an environment variable.")
-
-    def bind(self, **kwargs) -> None:
-        """
-        Bind default values to the search method. Cannot bind the 'query' parameter.
-
-        Args:
-          kwargs (dict): Key value pairs corresponding to the default values of search parameters.
-        """
-        if "query" in kwargs:
-            kwargs = {k: v for k, v in kwargs.items() if k != "query"}
-        self._search_func = partial(self._search_func, **kwargs)
-        self._bound_params = list(kwargs.keys())
-        logger.debug("Binding paramaters for search function: %s", kwargs)
-
-    def get_unbound_params(self) -> list[str]:
-        """
-        Returns a list of unbound parameters which will need to be passed to the search function.
-        """
-        return [param for param in ["query", "collection_name", "top_k"] if param not in self._bound_params]
-
-    async def get_collections(self, client) -> list[Collection]:
-        """
-        Get a list of all available collections as pydantic `Collection` objects
-        """
-        collection_response = await client.get(urljoin(self.base_url, "/v1/collections"))
-        collection_response.raise_for_status()
-        if not collection_response or len(collection_response.json().get('collections', [])) == 0:
-            raise CollectionUnavailableError(f"No collections available at {self.base_url}")
-
-        collections = [
-            Collection.model_validate(collection) for collection in collection_response.json()["collections"]
-        ]
-
-        return collections
-
-    async def get_collection_by_name(self, collection_name, client) -> Collection:
-        """
-        Retrieve a collection using it's name. Will return the first collection found if the name is ambiguous.
-        """
-        collections = await self.get_collections(client)
-        if (collection := next((c for c in collections if c.name == collection_name), None)) is None:
-            raise CollectionUnavailableError(f"Collection {collection_name} not found")
-        return collection
-
-    async def search(self, query: str, **kwargs):
-        return await self._search_func(query=query, **kwargs)
-
-    async def _search(
-        self,
-        query: str,
-        collection_name: str,
-        top_k: str,
-        output_fields: list[str] = None,
-    ):
-        """
-        Retrieve document chunks from the configured Nemo Retriever Service.
-        """
-        output = []
-        try:
-            async with httpx.AsyncClient(headers={"Authorization": f"Bearer {self.api_key}"},
-                                         timeout=self.timeout) as client:
-                collection = await self.get_collection_by_name(collection_name, client)
-                url = urljoin(self.base_url, f"/v1/collections/{collection.id}/search")
-
-                payload = RetrieverPayload(query=query, top_k=top_k)
-                response = await client.post(url, content=json.dumps(payload.model_dump(mode="python")))
-
-                logger.debug("response.status_code=%s", response.status_code)
-
-                response.raise_for_status()
-                output = response.json().get("chunks")
-
-                # Handle output fields
-                output = [_flatten(chunk, output_fields) for chunk in output]
-
-                return _wrap_nemo_results(output=output, content_field="content")
-
-        except Exception as e:
-            logger.exception("Encountered an error when retrieving results from Nemo Retriever: %s", e)
-            raise CollectionUnavailableError(
-                f"Error when retrieving documents from {collection_name} for query '{query}'") from e
-
-
-def _wrap_nemo_results(output: list[dict], content_field: str):
-    return RetrieverOutput(results=[_wrap_nemo_single_results(o, content_field=content_field) for o in output])
-
-
-def _wrap_nemo_single_results(output: dict, content_field: str):
-    return AIQDocument(page_content=output[content_field],
-                       metadata={
-                           k: v
-                           for k, v in output.items() if k != content_field
-                       })
-
-
-def _flatten(obj: dict, output_fields: list[str]) -> list[str]:
-    base_fields = [
-        "format",
-        "id",
-    ]
-    if not output_fields:
-        output_fields = [
-            "format",
-            "id",
-        ]
-        output_fields.extend(list(obj["metadata"].keys()))
-    data = {"content": obj.get("content")}
-    for field in base_fields:
-        if field in output_fields:
-            data.update({field: obj[field]})
-
-    data.update({k: v for k, v in obj['metadata'].items() if k in output_fields})
-    return data
-
-
-class NemoLangchainRetriever(BaseRetriever, BaseModel):
-    client: NemoRetriever
-
-    def _get_relevant_documents(self, query, *, run_manager, **kwargs):
-        raise NotImplementedError
-
-    async def _aget_relevant_documents(self, query, *, run_manager, **kwargs):
-        return await self.client.search(query, **kwargs)

aiq/retriever/register.py

@@ -1,22 +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.
-
-# pylint: disable=unused-import
-# flake8: noqa
-# isort:skip_file
-
-# Import any providers which need to be automatically registered here
-import aiq.retriever.milvus.register
-import aiq.retriever.nemo_retriever.register

aiq/runtime/__init__.py

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

aiq/runtime/loader.py

@@ -1,188 +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 __future__ import annotations
-
-import importlib.metadata
-import logging
-import time
-from contextlib import asynccontextmanager
-from enum import IntFlag
-from enum import auto
-from functools import reduce
-
-from aiq.builder.workflow_builder import WorkflowBuilder
-from aiq.cli.type_registry import GlobalTypeRegistry
-from aiq.data_models.config import AIQConfig
-from aiq.runtime.session import AIQSessionManager
-from aiq.utils.data_models.schema_validator import validate_schema
-from aiq.utils.debugging_utils import is_debugger_attached
-from aiq.utils.io.yaml_tools import yaml_load
-from aiq.utils.type_utils import StrPath
-
-logger = logging.getLogger(__name__)
-
-
-class PluginTypes(IntFlag):
-    COMPONENT = auto()
-    """
-    A plugin that is a component of the workflow. This includes tools, LLMs, retrievers, etc.
-    """
-    FRONT_END = auto()
-    """
-    A plugin that is a front end for the workflow. This includes FastAPI, Gradio, etc.
-    """
-    EVALUATOR = auto()
-    """
-    A plugin that is an evaluator for the workflow. This includes evaluators like RAGAS, SWE-bench, etc.
-    """
-    REGISTRY_HANDLER = auto()
-
-    # Convenience flag for groups of plugin types
-    CONFIG_OBJECT = COMPONENT | FRONT_END | EVALUATOR
-    """
-    Any plugin that can be specified in the AIQ Toolkit configuration file.
-    """
-    ALL = COMPONENT | FRONT_END | EVALUATOR | REGISTRY_HANDLER
-    """
-    All plugin types
-    """
-
-
-def load_config(config_file: StrPath) -> AIQConfig:
-    """
-    This is the primary entry point for loading an AIQ Toolkit configuration file. It ensures that all plugins are
-    loaded and then validates the configuration file against the AIQConfig schema.
-
-    Parameters
-    ----------
-    config_file : StrPath
-        The path to the configuration file
-
-    Returns
-    -------
-    AIQConfig
-        The validated AIQConfig object
-    """
-
-    # Ensure all of the plugins are loaded
-    discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
-
-    config_yaml = yaml_load(config_file)
-
-    # Validate configuration adheres to AIQ Toolkit schemas
-    validated_aiq_config = validate_schema(config_yaml, AIQConfig)
-
-    return validated_aiq_config
-
-
-@asynccontextmanager
-async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
-    """
-    Load the AIQ Toolkit configuration file and create an AIQRunner object. This is the primary entry point for running
-    AIQ Toolkit workflows.
-
-    Parameters
-    ----------
-    config_file : StrPath
-        The path to the configuration file
-    max_concurrency : int, optional
-        The maximum number of parallel workflow invocations to support. Specifying 0 or -1 will allow an unlimited
-        count, by default -1
-    """
-
-    # Load the config object
-    config = load_config(config_file)
-
-    # Must yield the workflow function otherwise it cleans up
-    async with WorkflowBuilder.from_config(config=config) as workflow:
-
-        yield AIQSessionManager(workflow.build(), max_concurrency=max_concurrency)
-
-
-def discover_entrypoints(plugin_type: PluginTypes):
-    """
-    Discover all the requested plugin types which were registered via an entry point group and return them.
-    """
-
-    entry_points = importlib.metadata.entry_points()
-
-    plugin_groups = []
-
-    # Add the specified plugin type to the list of groups to load
-    if (plugin_type & PluginTypes.COMPONENT):
-        plugin_groups.extend(["aiq.plugins", "aiq.components"])
-    if (plugin_type & PluginTypes.FRONT_END):
-        plugin_groups.append("aiq.front_ends")
-    if (plugin_type & PluginTypes.REGISTRY_HANDLER):
-        plugin_groups.append("aiq.registry_handlers")
-    if (plugin_type & PluginTypes.EVALUATOR):
-        plugin_groups.append("aiq.evaluators")
-
-    # Get the entry points for the specified groups
-    aiq_plugins = reduce(lambda x, y: x + y, [entry_points.select(group=y) for y in plugin_groups])
-
-    return aiq_plugins
-
-
-def discover_and_register_plugins(plugin_type: PluginTypes):
-    """
-    Discover all the requested plugin types which were registered via an entry point group and register them into the
-    GlobalTypeRegistry.
-    """
-
-    # Get the entry points for the specified groups
-    aiq_plugins = discover_entrypoints(plugin_type)
-
-    count = 0
-
-    # Pause registration hooks for performance. This is useful when loading a large number of plugins.
-    with GlobalTypeRegistry.get().pause_registration_changed_hooks():
-
-        for entry_point in aiq_plugins:
-            try:
-                logger.debug("Loading module '%s' from entry point '%s'...", entry_point.module, entry_point.name)
-
-                start_time = time.time()
-
-                entry_point.load()
-
-                elapsed_time = (time.time() - start_time) * 1000
-
-                logger.debug("Loading module '%s' from entry point '%s'...Complete (%f ms)",
-                             entry_point.module,
-                             entry_point.name,
-                             elapsed_time)
-
-                # Log a warning if the plugin took a long time to load. This can be useful for debugging slow imports.
-                # The threshold is 300 ms if no plugins have been loaded yet, and 100 ms otherwise. Triple the threshold
-                # if a debugger is attached.
-                if (elapsed_time > (300.0 if count == 0 else 100.0) * (3 if is_debugger_attached() else 1)):
-                    logger.warning(
-                        "Loading module '%s' from entry point '%s' took a long time (%f ms). "
-                        "Ensure all imports are inside your registered functions.",
-                        entry_point.module,
-                        entry_point.name,
-                        elapsed_time)
-
-            except ImportError:
-                logger.warning("Failed to import plugin '%s'", entry_point.name, exc_info=True)
-                # Optionally, you can mark the plugin as unavailable or take other actions
-
-            except Exception:
-                logger.exception("An error occurred while loading plugin '%s': {e}", entry_point.name, exc_info=True)
-
-            finally:
-                count += 1

aiq/runtime/runner.py

@@ -1,176 +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 typing
-from enum import Enum
-
-from aiq.builder.context import AIQContext
-from aiq.builder.context import AIQContextState
-from aiq.builder.function import Function
-from aiq.data_models.invocation_node import InvocationNode
-from aiq.observability.async_otel_listener import AsyncOtelSpanListener
-from aiq.utils.reactive.subject import Subject
-
-logger = logging.getLogger(__name__)
-
-
-class UserManagerBase:
-    pass
-
-
-class AIQRunnerState(Enum):
-    UNINITIALIZED = 0
-    INITIALIZED = 1
-    RUNNING = 2
-    COMPLETED = 3
-    FAILED = 4
-
-
-_T = typing.TypeVar("_T")
-
-
-class AIQRunner:
-
-    def __init__(self, input_message: typing.Any, entry_fn: Function, context_state: AIQContextState):
-        """
-        The AIQRunner class is used to run a workflow. It handles converting input and output data types and running the
-        workflow with the specified concurrency.
-
-        Parameters
-        ----------
-        input_message : typing.Any
-            The input message to the workflow
-        entry_fn : Function
-            The entry function to the workflow
-        context_state : AIQContextState
-            The context state to use
-        """
-
-        if (entry_fn is None):
-            raise ValueError("entry_fn cannot be None")
-
-        self._entry_fn = entry_fn
-        self._context_state = context_state
-        self._context = AIQContext(self._context_state)
-
-        self._state = AIQRunnerState.UNINITIALIZED
-
-        self._input_message_token = None
-
-        # Before we start, we need to convert the input message to the workflow input type
-        self._input_message = input_message
-
-        self._span_manager = AsyncOtelSpanListener(context_state=context_state)
-
-    @property
-    def context(self) -> AIQContext:
-        return self._context
-
-    def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
-        return self._entry_fn.convert(value, to_type)
-
-    async def __aenter__(self):
-
-        # Set the input message on the context
-        self._input_message_token = self._context_state.input_message.set(self._input_message)
-
-        # Create reactive event stream
-        self._context_state.event_stream.set(Subject())
-        self._context_state.active_function.set(InvocationNode(
-            function_name="root",
-            function_id="root",
-        ))
-
-        if (self._state == AIQRunnerState.UNINITIALIZED):
-            self._state = AIQRunnerState.INITIALIZED
-        else:
-            raise ValueError("Cannot enter the context more than once")
-
-        return self
-
-    async def __aexit__(self, exc_type, exc_value, traceback):
-
-        if (self._input_message_token is None):
-            raise ValueError("Cannot exit the context without entering it")
-
-        self._context_state.input_message.reset(self._input_message_token)
-
-        if (self._state not in (AIQRunnerState.COMPLETED, AIQRunnerState.FAILED)):
-            raise ValueError("Cannot exit the context without completing the workflow")
-
-    @typing.overload
-    async def result(self) -> typing.Any:
-        ...
-
-    @typing.overload
-    async def result(self, to_type: type[_T]) -> _T:
-        ...
-
-    async def result(self, to_type: type | None = None):
-
-        if (self._state != AIQRunnerState.INITIALIZED):
-            raise ValueError("Cannot run the workflow without entering the context")
-
-        try:
-            self._state = AIQRunnerState.RUNNING
-
-            if (not self._entry_fn.has_single_output):
-                raise ValueError("Workflow does not support single output")
-
-            async with self._span_manager.start():
-                # Run the workflow
-                result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)
-
-                # Close the intermediate stream
-                self._context_state.event_stream.get().on_complete()
-
-            self._state = AIQRunnerState.COMPLETED
-
-            return result
-        except Exception as e:
-            logger.exception("Error running workflow: %s", e)
-            self._context_state.event_stream.get().on_complete()
-            self._state = AIQRunnerState.FAILED
-
-            raise
-
-    async def result_stream(self, to_type: type | None = None):
-
-        if (self._state != AIQRunnerState.INITIALIZED):
-            raise ValueError("Cannot run the workflow without entering the context")
-
-        try:
-            self._state = AIQRunnerState.RUNNING
-
-            if (not self._entry_fn.has_streaming_output):
-                raise ValueError("Workflow does not support streaming output")
-
-            # Run the workflow
-            async with self._span_manager.start():
-                async for m in self._entry_fn.astream(self._input_message, to_type=to_type):
-                    yield m
-
-                self._state = AIQRunnerState.COMPLETED
-
-                # Close the intermediate stream
-                self._context_state.event_stream.get().on_complete()
-
-        except Exception as e:
-            logger.exception("Error running workflow: %s", e)
-            self._context_state.event_stream.get().on_complete()
-            self._state = AIQRunnerState.FAILED
-
-            raise