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

aiq/builder/intermediate_step_manager.py

@@ -1,174 +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 dataclasses
-import logging
-import typing
-
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepState
-from aiq.utils.reactive.observable import OnComplete
-from aiq.utils.reactive.observable import OnError
-from aiq.utils.reactive.observable import OnNext
-from aiq.utils.reactive.subscription import Subscription
-
-if typing.TYPE_CHECKING:
-    from aiq.builder.context import AIQContextState
-
-logger = logging.getLogger(__name__)
-
-
-@dataclasses.dataclass
-class OpenStep:
-    step_id: str
-    step_name: str
-    step_type: str
-    step_parent_id: str
-    prev_stack: list[str]
-    active_stack: list[str]
-
-
-class IntermediateStepManager:
-    """
-    Manages updates to the AIQ Toolkit Event Stream for intermediate steps
-    """
-
-    def __init__(self, context_state: "AIQContextState"):  # noqa: F821
-        self._context_state = context_state
-
-        self._outstanding_start_steps: dict[str, OpenStep] = {}
-
-    def push_intermediate_step(self, payload: IntermediateStepPayload) -> None:
-        """
-        Pushes an intermediate step to the AIQ Toolkit Event Stream
-        """
-
-        if not isinstance(payload, IntermediateStepPayload):
-            raise TypeError(f"Payload must be of type IntermediateStepPayload, not {type(payload)}")
-
-        active_span_id_stack = self._context_state.active_span_id_stack.get()
-
-        if (payload.event_state == IntermediateStepState.START):
-
-            prev_stack = active_span_id_stack
-
-            parent_step_id = active_span_id_stack[-1]
-
-            # Note, this must not mutate the active_span_id_stack in place
-            active_span_id_stack = active_span_id_stack + [payload.UUID]
-            self._context_state.active_span_id_stack.set(active_span_id_stack)
-
-            self._outstanding_start_steps[payload.UUID] = OpenStep(step_id=payload.UUID,
-                                                                   step_name=payload.name or payload.UUID,
-                                                                   step_type=payload.event_type,
-                                                                   step_parent_id=parent_step_id,
-                                                                   prev_stack=prev_stack,
-                                                                   active_stack=active_span_id_stack)
-
-            logger.debug("Pushed start step %s, name %s, type %s, parent %s, stack id %s",
-                         payload.UUID,
-                         payload.name,
-                         payload.event_type,
-                         parent_step_id,
-                         id(active_span_id_stack))
-
-        elif (payload.event_state == IntermediateStepState.END):
-
-            # Remove the current step from the outstanding steps
-            open_step = self._outstanding_start_steps.pop(payload.UUID, None)
-
-            if (open_step is None):
-                logger.warning("Step id %s not found in outstanding start steps", payload.UUID)
-                return
-
-            parent_step_id = open_step.step_parent_id
-
-            # Get the current and previous active span id stack.
-            curr_stack = open_step.active_stack
-            prev_stack = open_step.prev_stack
-
-            # To restore the stack, we need to handle two scenarios:
-            # 1. This function is called from a coroutine. In this case, the context variable will be the same as the
-            #    one used in START. So we can just set the context variable to the previous stack.
-            # 2. This function is called from a task. In this case, the context variable will be separate from the one
-            #    used in START so calling set() will have no effect. However, we still have a reference to the list used
-            #    in START. So we update the reference to be equal to the old one.. So we need to update the current
-            #    reference stack to be equal to the previous stack.
-
-            # Scenario 1: Restore the previous active span id stack in case we are in a coroutine. Dont use reset here
-            # since we can be in different contexts
-            self._context_state.active_span_id_stack.set(prev_stack)
-
-            pop_count = 0
-
-            # Scenario 2: Remove all steps from the current stack until we reach the parent step id to make it equal to
-            # the previous stack. In the coroutine case, this will not have any effect.
-            while (curr_stack[-1] != parent_step_id):
-                curr_stack.pop()
-                pop_count += 1
-
-            if (pop_count != 1):
-                logger.warning(
-                    "Step id %s not the last step in the stack. "
-                    "Removing it from the stack but this is likely an error",
-                    payload.UUID)
-
-            # Verify that the stack is now equal to the previous stack
-            if (curr_stack != prev_stack):
-                logger.warning("Current span ID stack is not equal to the previous stack. "
-                               "This is likely an error. Report this to the AIQ team.")
-
-            logger.debug("Popped end step %s, name %s, type %s, parent %s, stack id %s",
-                         payload.UUID,
-                         payload.name,
-                         payload.event_type,
-                         parent_step_id,
-                         id(curr_stack))
-
-        elif (payload.event_state == IntermediateStepState.CHUNK):
-
-            # Get the current step from the outstanding steps
-            open_step = self._outstanding_start_steps.get(payload.UUID, None)
-
-            # Generate a warning if the parent step id is not set to the current step id
-            if (open_step is None):
-                logger.warning(
-                    "Created a chunk for step %s, but no matching start step was found. "
-                    "Chunks must be created with the same ID as the start step.",
-                    payload.UUID)
-                return
-
-            parent_step_id = open_step.step_parent_id
-        else:
-            assert False, "Invalid event state"
-
-        active_function = self._context_state.active_function.get()
-
-        intermediate_step = IntermediateStep(parent_id=parent_step_id,
-                                             function_ancestry=active_function,
-                                             payload=payload)
-
-        self._context_state.event_stream.get().on_next(intermediate_step)
-
-    def subscribe(self,
-                  on_next: OnNext[IntermediateStep],
-                  on_error: OnError = None,
-                  on_complete: OnComplete = None) -> Subscription:
-        """
-        Subscribes to the AIQ Toolkit Event Stream for intermediate steps
-        """
-
-        return self._context_state.event_stream.get().subscribe(on_next, on_error, on_complete)

aiq/builder/llm.py

@@ -1,25 +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.data_models.llm import LLMBaseConfig
-
-
-class LLMProviderInfo:
-
-    def __init__(self, *, config: LLMBaseConfig, description: str):
-
-        self.config = config
-        self.provider_type = type(config).static_type()
-        self.description = description

aiq/builder/retriever.py

@@ -1,25 +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.retriever import RetrieverBaseConfig
-
-
-class RetrieverProviderInfo:
-
-    def __init__(self, *, config: RetrieverBaseConfig, description: str):
-
-        self.config = config
-        self.provider_type = type(config).static_type()
-        self.description = description

aiq/builder/user_interaction_manager.py

@@ -1,74 +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 time
-import uuid
-
-from aiq.data_models.interactive import HumanPrompt
-from aiq.data_models.interactive import HumanResponse
-from aiq.data_models.interactive import InteractionPrompt
-from aiq.data_models.interactive import InteractionResponse
-from aiq.data_models.interactive import InteractionStatus
-
-logger = logging.getLogger(__name__)
-
-
-class AIQUserInteractionManager:
-    """
-    AIQUserInteractionManager is responsible for requesting user input
-    at runtime. It delegates the actual prompting to a callback function
-    stored in AIQContextState.user_input_callback.
-
-    Type is not imported in __init__ to prevent partial import.
-    """
-
-    def __init__(self, context_state: "AIQContextState") -> None:  # noqa: F821
-        self._context_state = context_state
-
-    @staticmethod
-    async def default_callback_handler(prompt: InteractionPrompt) -> HumanResponse:
-        """
-        Default callback handler for user input. This is a no-op function
-        that simply returns the input text from the Interaction Content
-        object.
-
-        Args:
-            prompt (InteractionPrompt): The interaction to process.
-        """
-        raise NotImplementedError("No human prompt callback was registered. Unable to handle requested prompt.")
-
-    async def prompt_user_input(self, content: HumanPrompt) -> InteractionResponse:
-        """
-        Ask the user a question and wait for input. This calls out to
-        the callback from user_input_callback, which is typically
-        set by AIQSessionManager.
-
-        Returns the user's typed-in answer as a string.
-        """
-
-        uuid_req = str(uuid.uuid4())
-        status = InteractionStatus.IN_PROGRESS
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
-        sys_human_interaction = InteractionPrompt(id=uuid_req, status=status, timestamp=timestamp, content=content)
-
-        resp = await self._context_state.user_input_callback.get()(sys_human_interaction)
-
-        # Rebuild a InteractionResponse object with the response
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
-        status = InteractionStatus.COMPLETED
-        sys_human_interaction = InteractionResponse(id=uuid_req, status=status, timestamp=timestamp, content=resp)
-
-        return sys_human_interaction

aiq/builder/workflow.py

@@ -1,148 +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 contextlib import asynccontextmanager
-from contextvars import ContextVar
-from typing import Any
-
-from aiq.builder.context import AIQContextState
-from aiq.builder.embedder import EmbedderProviderInfo
-from aiq.builder.function import Function
-from aiq.builder.function_base import FunctionBase
-from aiq.builder.function_base import InputT
-from aiq.builder.function_base import SingleOutputT
-from aiq.builder.function_base import StreamingOutputT
-from aiq.builder.llm import LLMProviderInfo
-from aiq.builder.retriever import RetrieverProviderInfo
-from aiq.data_models.config import AIQConfig
-from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
-from aiq.memory.interfaces import MemoryEditor
-from aiq.object_store.interfaces import ObjectStore
-from aiq.observability.exporter.base_exporter import BaseExporter
-from aiq.observability.exporter_manager import ExporterManager
-from aiq.runtime.runner import AIQRunner
-
-callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
-
-
-class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
-
-    def __init__(self,
-                 *,
-                 config: AIQConfig,
-                 entry_fn: Function[InputT, StreamingOutputT, SingleOutputT],
-                 functions: dict[str, Function] | None = None,
-                 llms: dict[str, LLMProviderInfo] | None = None,
-                 embeddings: dict[str, EmbedderProviderInfo] | None = None,
-                 memory: dict[str, MemoryEditor] | None = None,
-                 object_stores: dict[str, ObjectStore] | None = None,
-                 telemetry_exporters: dict[str, BaseExporter] | None = None,
-                 retrievers: dict[str | None, RetrieverProviderInfo] | None = None,
-                 its_strategies: dict[str, StrategyBase] | None = None,
-                 context_state: AIQContextState):
-
-        super().__init__(input_schema=entry_fn.input_schema,
-                         streaming_output_schema=entry_fn.streaming_output_schema,
-                         single_output_schema=entry_fn.single_output_schema)
-
-        self.config = config
-        self.functions = functions or {}
-        self.llms = llms or {}
-        self.embeddings = embeddings or {}
-        self.memory = memory or {}
-        self.telemetry_exporters = telemetry_exporters or {}
-        self.object_stores = object_stores or {}
-        self.retrievers = retrievers or {}
-
-        self._exporter_manager = ExporterManager.from_exporters(self.telemetry_exporters)
-        self.its_strategies = its_strategies or {}
-
-        self._entry_fn = entry_fn
-
-        self._context_state = context_state
-
-    @property
-    def has_streaming_output(self) -> bool:
-
-        return self._entry_fn.has_streaming_output
-
-    @property
-    def has_single_output(self) -> bool:
-
-        return self._entry_fn.has_single_output
-
-    @asynccontextmanager
-    async def run(self, message: InputT):
-        """
-        Called each time we start a new workflow run. We'll create
-        a new top-level workflow span here.
-        """
-
-        async with AIQRunner(input_message=message,
-                             entry_fn=self._entry_fn,
-                             context_state=self._context_state,
-                             exporter_manager=self._exporter_manager.get()) as runner:
-
-            # The caller can `yield runner` so they can do `runner.result()` or `runner.result_stream()`
-            yield runner
-
-    async def result_with_steps(self, message: InputT, to_type: type | None = None):
-
-        async with self.run(message) as runner:
-
-            from aiq.eval.runtime_event_subscriber import pull_intermediate
-
-            # Start the intermediate stream
-            pull_done, intermediate_steps = pull_intermediate()
-
-            # Wait on the result
-            result = await runner.result(to_type=to_type)
-
-            await pull_done.wait()
-
-            return result, intermediate_steps
-
-    @staticmethod
-    def from_entry_fn(*,
-                      config: AIQConfig,
-                      entry_fn: Function[InputT, StreamingOutputT, SingleOutputT],
-                      functions: dict[str, Function] | None = None,
-                      llms: dict[str, LLMProviderInfo] | None = None,
-                      embeddings: dict[str, EmbedderProviderInfo] | None = None,
-                      memory: dict[str, MemoryEditor] | None = None,
-                      object_stores: dict[str, ObjectStore] | None = None,
-                      telemetry_exporters: dict[str, BaseExporter] | None = None,
-                      retrievers: dict[str | None, RetrieverProviderInfo] | None = None,
-                      its_strategies: dict[str, StrategyBase] | None = None,
-                      context_state: AIQContextState) -> 'Workflow[InputT, StreamingOutputT, SingleOutputT]':
-
-        input_type: type = entry_fn.input_type
-        streaming_output_type = entry_fn.streaming_output_type
-        single_output_type = entry_fn.single_output_type
-
-        class WorkflowImpl(Workflow[input_type, streaming_output_type, single_output_type]):
-            pass
-
-        return WorkflowImpl(config=config,
-                            entry_fn=entry_fn,
-                            functions=functions,
-                            llms=llms,
-                            embeddings=embeddings,
-                            memory=memory,
-                            object_stores=object_stores,
-                            telemetry_exporters=telemetry_exporters,
-                            retrievers=retrievers,
-                            its_strategies=its_strategies,
-                            context_state=context_state)