aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/builder/component_utils.py

@@ -1,316 +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
-import typing
-from collections.abc import Generator
-from collections.abc import Iterable
-
-import networkx as nx
-from pydantic import BaseModel
-
-from aiq.data_models.authentication import AuthProviderBaseConfig
-from aiq.data_models.common import TypedBaseModel
-from aiq.data_models.component import ComponentGroup
-from aiq.data_models.component_ref import ComponentRef
-from aiq.data_models.component_ref import ComponentRefNode
-from aiq.data_models.component_ref import generate_instance_id
-from aiq.data_models.config import AIQConfig
-from aiq.data_models.embedder import EmbedderBaseConfig
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.data_models.its_strategy import ITSStrategyBaseConfig
-from aiq.data_models.llm import LLMBaseConfig
-from aiq.data_models.memory import MemoryBaseConfig
-from aiq.data_models.object_store import ObjectStoreBaseConfig
-from aiq.data_models.retriever import RetrieverBaseConfig
-from aiq.utils.type_utils import DecomposedType
-
-logger = logging.getLogger(__name__)
-
-# Order in which we want to process the component groups
-_component_group_order = [
-    ComponentGroup.AUTHENTICATION,
-    ComponentGroup.EMBEDDERS,
-    ComponentGroup.LLMS,
-    ComponentGroup.MEMORY,
-    ComponentGroup.OBJECT_STORES,
-    ComponentGroup.RETRIEVERS,
-    ComponentGroup.ITS_STRATEGIES,
-    ComponentGroup.FUNCTIONS,
-]
-
-
-class ComponentInstanceData(BaseModel):
-    """A data model to hold component runtime instance metadata to support generating build sequences.
-
-    Args:
-        component_group (ComponentGroup): The component group in an AIQ Toolkit configuration object.
-        name (ComponentRef): The name of the component runtime instance.
-        config (TypedBaseModel): The runtime instance's configuration object.
-        instance_id (str): Unique identifier for each runtime instance.
-        is_root (bool): A flag to indicate if the runtime instance is the root of the workflow.
-    """
-
-    component_group: ComponentGroup
-    name: ComponentRef
-    config: TypedBaseModel
-    instance_id: str
-    is_root: bool = False
-
-
-def iterate_leaf_to_root(graph: nx.DiGraph) -> Generator[ComponentRefNode]:
-    """A recursive generator that yields leaf nodes from the bottom to the root of a directed graph.
-
-    Args:
-        graph (nx.DiGraph): A networkx directed graph object.
-
-    Yields:
-        ComponentRefNode: An object contain a ComponentRef and its component group.
-    """
-
-    leaf_nodes = [node for node, degree in graph.out_degree() if degree == 0]
-
-    if len(leaf_nodes) > 0:
-        for leaf_node in leaf_nodes:
-            yield leaf_node
-            graph.remove_node(leaf_node)
-
-        yield from iterate_leaf_to_root(graph)
-
-
-def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
-    """Determines the component group from a runtime instance configuration object.
-
-    Args:
-        component (TypedBaseModel): A runtime instance configuration object.
-
-    Returns:
-        ComponentGroup | None: The component group of the runtime instance configuration object. If the
-            component is not a valid runtime instance, None is returned.
-    """
-
-    if (isinstance(component, AuthProviderBaseConfig)):
-        return ComponentGroup.AUTHENTICATION
-    if (isinstance(component, EmbedderBaseConfig)):
-        return ComponentGroup.EMBEDDERS
-    if (isinstance(component, FunctionBaseConfig)):
-        return ComponentGroup.FUNCTIONS
-    if (isinstance(component, LLMBaseConfig)):
-        return ComponentGroup.LLMS
-    if (isinstance(component, MemoryBaseConfig)):
-        return ComponentGroup.MEMORY
-    if (isinstance(component, ObjectStoreBaseConfig)):
-        return ComponentGroup.OBJECT_STORES
-    if (isinstance(component, RetrieverBaseConfig)):
-        return ComponentGroup.RETRIEVERS
-    if (isinstance(component, ITSStrategyBaseConfig)):
-        return ComponentGroup.ITS_STRATEGIES
-
-    return None
-
-
-def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
-                                     type_hint: type[typing.Any]) -> Generator[tuple[str, ComponentRefNode]]:
-    """Discovers instances of ComponentRefs in a configuration object and updates the dependency graph.
-
-    Args:
-        cls (TypedBaseModel): A configuration object for a runtime instance.
-        value (typing.Any): The current traversed value from the configuration object.
-        type_hint (type[typing.Any]): The type of the current traversed value from the configuration object.
-    """
-
-    decomposed_type = DecomposedType(type_hint)
-
-    if (value is None):
-        return
-
-    if ((decomposed_type.origin is None) and (not issubclass(type(value), BaseModel))):
-        if issubclass(type(value), ComponentRef):
-            instance_id = generate_instance_id(cls)
-            value_node = ComponentRefNode(ref_name=value, component_group=value.component_group)
-            yield instance_id, value_node
-
-    elif ((decomposed_type.origin in (tuple, list, set)) and (isinstance(value, Iterable))):
-        for v in value:
-            yield from recursive_componentref_discovery(cls, v, decomposed_type.args[0])
-    elif ((decomposed_type.origin in (dict, type(typing.TypedDict))) and (isinstance(value, dict))):
-        for v in value.values():
-            yield from recursive_componentref_discovery(cls, v, decomposed_type.args[1])
-    elif (issubclass(type(value), BaseModel)):
-        for field, field_info in value.model_fields.items():
-            field_data = getattr(value, field)
-            yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
-    if (decomposed_type.is_union):
-        for arg in decomposed_type.args:
-            if arg is typing.Any or (isinstance(value, DecomposedType(arg).root)):
-                yield from recursive_componentref_discovery(cls, value, arg)
-    else:
-        for arg in decomposed_type.args:
-            yield from recursive_componentref_discovery(cls, value, arg)
-
-
-def update_dependency_graph(config: "AIQConfig", instance_config: TypedBaseModel,
-                            dependency_graph: nx.DiGraph) -> nx.DiGraph:
-    """Updates the hierarchical component instance dependency graph from a configuration runtime instance.
-
-    Args:
-        config (AIQConfig): An AIQ Toolkit configuration object with runtime instance details.
-        instance_config (TypedBaseModel): A component's runtime instance configuration object.
-        dependency_graph (nx.DiGraph): A graph tracking runtime instance component dependencies.
-
-    Returns:
-        nx.DiGraph: An dependency graph that has been updated with the provided runtime instance.
-    """
-
-    for field_name, field_info in instance_config.model_fields.items():
-
-        for instance_id, value_node in recursive_componentref_discovery(
-                instance_config,
-                getattr(instance_config, field_name),
-                field_info.annotation):  # type: ignore
-
-            # add immediate edge
-            dependency_graph.add_edge(instance_id, value_node)
-            # add dependency edge to ensure connections to leaf nodes exist
-            dependency_component_dict = getattr(config, value_node.component_group)
-            dependency_component_instance_config = dependency_component_dict.get(value_node.ref_name)
-            dependency_component_instance_id = generate_instance_id(dependency_component_instance_config)
-            dependency_graph.add_edge(value_node, dependency_component_instance_id)
-
-    return dependency_graph
-
-
-def config_to_dependency_objects(config: "AIQConfig") -> tuple[dict[str, ComponentInstanceData], nx.DiGraph]:
-    """Generates a map of component runtime instance IDs to use when generating a build sequence.
-
-    Args:
-        config (AIQConfig): The AIQ Toolkit workflow configuration object.
-
-    Returns:
-        tuple[dict[str, ComponentInstanceData], nx.DiGraph]: A tuple containing a map of component runtime instance
-            IDs to a component object containing its metadata and a dependency graph of nested components.
-    """
-
-    # Build map of every runtime instances
-    dependency_map: dict[str, ComponentInstanceData] = {}
-    dependency_graph: nx.DiGraph = nx.DiGraph()
-
-    # Create the dependency map preserving as much order as we can
-    for group in _component_group_order:
-
-        component_dict = getattr(config, group.value)
-
-        assert isinstance(component_dict, dict), "Config components must be a dictionary"
-
-        for component_instance_name, component_instance_config in component_dict.items():
-
-            instance_id = generate_instance_id(component_instance_config)
-            dependency_map[instance_id] = ComponentInstanceData(component_group=group,
-                                                                instance_id=instance_id,
-                                                                name=component_instance_name,
-                                                                config=component_instance_config)
-
-            dependency_graph = update_dependency_graph(config=config,
-                                                       instance_config=component_instance_config,
-                                                       dependency_graph=dependency_graph)
-
-    # Set the workflow flag on the workflow instance (must be last)
-    workflow_instance_id = generate_instance_id(config.workflow)
-
-    dependency_map[workflow_instance_id] = ComponentInstanceData(
-        component_group=ComponentGroup.FUNCTIONS,
-        instance_id=workflow_instance_id,
-        name="<workflow>",  # type: ignore
-        config=config.workflow,
-        is_root=True)
-
-    dependency_graph = update_dependency_graph(config=config,
-                                               instance_config=config.workflow,
-                                               dependency_graph=dependency_graph)
-
-    return dependency_map, dependency_graph
-
-
-def build_dependency_sequence(config: "AIQConfig") -> list[ComponentInstanceData]:
-    """Generates the depencency sequence from an AIQ Toolkit configuration object
-
-    Args:
-        config (AIQConfig): An AIQ Toolkit configuration object.
-
-    Returns:
-        list[ComponentInstanceData]: A list representing the instatiation sequence to ensure all valid
-            runtime instance references.
-    """
-
-    total_node_count = len(config.embedders) + len(config.functions) + len(config.llms) + len(config.memory) + len(
-        config.object_stores) + len(config.retrievers) + len(config.its_strategies) + len(
-            config.authentication) + 1  # +1 for the workflow
-
-    dependency_map: dict
-    dependency_graph: nx.DiGraph
-    dependency_map, dependency_graph = config_to_dependency_objects(config=config)
-
-    dependency_sequence: list[ComponentInstanceData] = []
-    instance_ids = set()
-    for node in iterate_leaf_to_root(dependency_graph.copy()):  # type: ignore
-
-        if (node not in dependency_sequence):
-
-            # Convert node to id
-            if (isinstance(node, ComponentRefNode) and issubclass(type(node.ref_name), ComponentRef)):
-
-                component_group_configs = getattr(config, node.component_group.value)
-                node_config = component_group_configs.get(node.ref_name, None)
-
-                # Only add nodes that are valid in the current instance configuration
-                if (node_config is None):
-                    continue
-
-                component_instance = ComponentInstanceData(
-                    name=node.ref_name,
-                    component_group=node.component_group.value,  # type: ignore
-                    config=node_config,
-                    instance_id=generate_instance_id(node_config))
-
-            else:
-
-                component_instance = dependency_map.get(node, None)
-
-                # Only add nodes that are valid in the current instance configuration
-                if (component_instance is None):
-                    continue
-
-            if (component_instance.instance_id not in instance_ids):
-
-                dependency_sequence.append(component_instance)
-                instance_ids.add(component_instance.instance_id)
-
-    remaining_dependency_sequence: list[ComponentInstanceData] = []
-
-    # Find the remaining nodes that are not in the sequence preserving order
-    for instance_id, instance in dependency_map.items():
-        if (instance_id not in instance_ids):
-            remaining_dependency_sequence.append(instance)
-
-    # Add the remaining at the front of the sequence
-    dependency_sequence = remaining_dependency_sequence + dependency_sequence
-
-    # Find the root node and make sure it is the last node in the sequence
-    dependency_sequence = [x for x in dependency_sequence if not x.is_root
-                           ] + [x for x in dependency_sequence if x.is_root]
-
-    assert len(dependency_sequence) == total_node_count, "Dependency sequence generation failed. Report as bug."
-
-    return dependency_sequence

aiq/builder/context.py

@@ -1,264 +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
-import uuid
-from collections.abc import Awaitable
-from collections.abc import Callable
-from contextlib import contextmanager
-from contextvars import ContextVar
-
-from aiq.builder.intermediate_step_manager import IntermediateStepManager
-from aiq.builder.user_interaction_manager import AIQUserInteractionManager
-from aiq.data_models.authentication import AuthenticatedContext
-from aiq.data_models.authentication import AuthFlowType
-from aiq.data_models.authentication import AuthProviderBaseConfig
-from aiq.data_models.interactive import HumanResponse
-from aiq.data_models.interactive import InteractionPrompt
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.invocation_node import InvocationNode
-from aiq.runtime.user_metadata import RequestAttributes
-from aiq.utils.reactive.subject import Subject
-
-
-class Singleton(type):
-
-    def __init__(cls, name, bases, dict):  # pylint: disable=W0622
-        super(Singleton, cls).__init__(name, bases, dict)
-        cls.instance = None
-
-    def __call__(cls, *args, **kw):
-        if cls.instance is None:
-            cls.instance = super(Singleton, cls).__call__(*args, **kw)
-        return cls.instance
-
-
-class ActiveFunctionContextManager:
-
-    def __init__(self):
-        self._output: typing.Any | None = None
-
-    @property
-    def output(self) -> typing.Any | None:
-        return self._output
-
-    def set_output(self, output: typing.Any):
-        self._output = output
-
-
-class AIQContextState(metaclass=Singleton):
-
-    def __init__(self):
-        self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
-        self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
-        self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
-        self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
-        self.event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=Subject())
-        self.active_function: ContextVar[InvocationNode] = ContextVar("active_function",
-                                                                      default=InvocationNode(function_id="root",
-                                                                                             function_name="root"))
-        self.active_span_id_stack: ContextVar[list[str]] = ContextVar("active_span_id_stack", default=["root"])
-
-        # Default is a lambda no-op which returns NoneType
-        self.user_input_callback: ContextVar[Callable[[InteractionPrompt], Awaitable[HumanResponse | None]]
-                                             | None] = ContextVar(
-                                                 "user_input_callback",
-                                                 default=AIQUserInteractionManager.default_callback_handler)
-        self.user_auth_callback: ContextVar[Callable[[AuthProviderBaseConfig, AuthFlowType],
-                                                     Awaitable[AuthenticatedContext]]
-                                            | None] = ContextVar("user_auth_callback", default=None)
-
-    @staticmethod
-    def get() -> "AIQContextState":
-        return AIQContextState()
-
-
-class AIQContext:
-
-    def __init__(self, context: AIQContextState):
-        self._context_state = context
-
-    @property
-    def input_message(self):
-        """
-            Retrieves the input message from the context state.
-
-            The input_message property is used to access the message stored in the
-            context state. This property returns the message as it is currently
-            maintained in the context.
-
-            Returns:
-                str: The input message retrieved from the context state.
-        """
-        return self._context_state.input_message.get()
-
-    @property
-    def user_manager(self):
-        """
-        Retrieves the user manager instance from the current context state.
-
-        This property provides access to the user manager through the context
-        state, allowing interaction with user management functionalities.
-
-        Returns:
-            UserManager: The instance of the user manager retrieved from the
-                context state.
-        """
-        return self._context_state.user_manager.get()
-
-    @property
-    def metadata(self):
-        """
-        Retrieves the request attributes instance from the current context state
-        providing access to user-defined metadata.
-
-        Returns:
-            RequestAttributes: The instance of the request attributes
-                retrieved from the context state.
-        """
-        return self._context_state.metadata.get()
-
-    @property
-    def user_interaction_manager(self) -> AIQUserInteractionManager:
-        """
-        Return an instance of AIQUserInteractionManager that uses
-        the current context's user_input_callback.
-        """
-        return AIQUserInteractionManager(self._context_state)
-
-    @property
-    def intermediate_step_manager(self) -> IntermediateStepManager:
-        """
-        Retrieves the intermediate step manager instance from the current context state.
-
-        This property provides access to the intermediate step manager through the context
-        state, allowing interaction with intermediate step management functionalities.
-
-        Returns:
-            IntermediateStepManager: The instance of the intermediate step manager retrieved
-                from the context state.
-        """
-        return IntermediateStepManager(self._context_state)
-
-    @property
-    def conversation_id(self) -> str | None:
-        """
-        This property retrieves the conversation ID which is the unique identifier for the current chat conversation.
-
-        Returns:
-            str | None
-        """
-        return self._context_state.conversation_id.get()
-
-    @contextmanager
-    def push_active_function(self, function_name: str, input_data: typing.Any | None):
-        """
-        Set the 'active_function' in context, push an invocation node,
-        AND create an OTel child span for that function call.
-        """
-        parent_function_node = self._context_state.active_function.get()
-        current_function_id = str(uuid.uuid4())
-        current_function_node = InvocationNode(function_id=current_function_id,
-                                               function_name=function_name,
-                                               parent_id=parent_function_node.function_id,
-                                               parent_name=parent_function_node.function_name)
-
-        # 1) Set the active function in the contextvar
-        fn_token = self._context_state.active_function.set(current_function_node)
-
-        # 2) Optionally record function start as an intermediate step
-        step_manager = self.intermediate_step_manager
-        step_manager.push_intermediate_step(
-            IntermediateStepPayload(UUID=current_function_id,
-                                    event_type=IntermediateStepType.FUNCTION_START,
-                                    name=function_name,
-                                    data=StreamEventData(input=input_data)))
-
-        manager = ActiveFunctionContextManager()
-
-        try:
-            yield manager  # run the function body
-        finally:
-            # 3) Record function end
-
-            data = StreamEventData(input=input_data, output=manager.output)
-
-            step_manager.push_intermediate_step(
-                IntermediateStepPayload(UUID=current_function_id,
-                                        event_type=IntermediateStepType.FUNCTION_END,
-                                        name=function_name,
-                                        data=data))
-
-            # 4) Unset the function contextvar
-            self._context_state.active_function.reset(fn_token)
-
-    @property
-    def active_function(self) -> InvocationNode:
-        """
-        Retrieves the active function from the context state.
-
-        This property is used to access the active function stored in the context
-        state. The active function is the function that is currently being executed.
-        """
-        return self._context_state.active_function.get()
-
-    @property
-    def active_span_id(self) -> str:
-        """
-        Retrieves the active span ID from the context state.
-
-        This property provides access to the active span ID stored in the context state. The active span ID represents
-        the currently running function/tool/llm/agent/etc and can be used to group telemetry data together.
-
-        Returns:
-            str: The active span ID.
-        """
-        return self._context_state.active_span_id_stack.get()[-1]
-
-    @property
-    def user_auth_callback(self) -> Callable[[AuthProviderBaseConfig, AuthFlowType], Awaitable[AuthenticatedContext]]:
-        """
-        Retrieves the user authentication callback function from the context state.
-
-        This property provides access to the user authentication callback function stored in the context state.
-        The callback function is responsible for handling user authentication based on the provided configuration.
-
-        Returns:
-            Callable[[AuthenticationBaseConfig], Awaitable[AuthenticatedContext]]: The user authentication
-            callback function.
-
-        Raises:
-            RuntimeError: If the user authentication callback is not set in the context.
-        """
-        callback = self._context_state.user_auth_callback.get()
-        if callback is None:
-            raise RuntimeError("User authentication callback is not set in the context.")
-        return callback
-
-    @staticmethod
-    def get() -> "AIQContext":
-        """
-        Static method to retrieve the current AIQContext instance.
-
-        This method creates and returns an instance of the AIQContext class
-        by obtaining the current state from the AIQContextState.
-
-        Returns:
-            AIQContext: The created AIQContext instance.
-        """
-        return AIQContext(AIQContextState.get())

aiq/builder/embedder.py

@@ -1,24 +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 aiq.data_models.embedder import EmbedderBaseConfig
-
-
-class EmbedderProviderInfo:
-
-    def __init__(self, *, config: EmbedderBaseConfig, description: str):
-        self.config = config
-        self.provider_type = type(config).static_type()
-        self.description = description