nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/builder/intermediate_step_manager.py

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

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

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

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

@@ -0,0 +1,148 @@
+# 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.test_time_compute.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,
+                 ttc_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.ttc_strategies = ttc_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,
+                      ttc_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,
+                            ttc_strategies=ttc_strategies,
+                            context_state=context_state)