aiqtoolkit 1.1.0a20250503__py3-none-any.whl

aiq/builder/intermediate_step_manager.py

@@ -0,0 +1,127 @@
+# 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 contextvars
+import dataclasses
+import logging
+
+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.data_models.invocation_node import InvocationNode
+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
+
+logger = logging.getLogger(__name__)
+
+_current_open_step_id = contextvars.ContextVar[str | None]("_current_open_step_id", default=None)
+
+
+@dataclasses.dataclass
+class OpenStep:
+    step_id: str
+    step_name: str
+    step_type: str
+    step_parent_id: str | None
+    context: contextvars.Context
+    token: contextvars.Token[str | None]
+
+
+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)}")
+
+        parent_step_id = _current_open_step_id.get()
+
+        if (payload.event_state == IntermediateStepState.START):
+
+            token = _current_open_step_id.set(payload.UUID)
+
+            self._outstanding_start_steps[payload.UUID] = OpenStep(step_id=payload.UUID,
+                                                                   step_name=payload.name,
+                                                                   step_type=payload.event_type,
+                                                                   step_parent_id=parent_step_id,
+                                                                   context=contextvars.copy_context(),
+                                                                   token=token)
+
+        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
+
+            # Restore the parent step ID directly instead of using a cross‑context token.
+            if parent_step_id == payload.UUID:
+                _current_open_step_id.set(open_step.step_parent_id)
+            else:
+                # Different context (e.g. thread‑pool); safely restore the parent ID **without**
+                # trying to use a token that belongs to another Context.
+                _current_open_step_id.set(open_step.step_parent_id)
+                parent_step_id = open_step.step_parent_id
+
+        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
+
+            if (parent_step_id != payload.UUID):
+                # Manually set the parent step ID. This happens when running on the thread pool
+                parent_step_id = open_step.step_parent_id
+
+        function_ancestry = InvocationNode(function_name=self._context_state.active_function.get().function_name,
+                                           function_id=self._context_state.active_function.get().function_id,
+                                           parent_id=parent_step_id,
+                                           parent_name=self._context_state.active_function.get().parent_name)
+
+        intermediate_step = IntermediateStep(function_ancestry=function_ancestry, 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,71 @@
+# 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 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
+
+
+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,143 @@
+# 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.memory.interfaces import MemoryEditor
+from aiq.runtime.runner import AIQRunner
+from aiq.utils.optional_imports import TelemetryOptionalImportError
+from aiq.utils.optional_imports import try_import_opentelemetry
+
+# Try to import OpenTelemetry modules
+# If the dependencies are not installed, use a dummy span exporter here
+try:
+    opentelemetry = try_import_opentelemetry()
+    from opentelemetry.sdk.trace.export import SpanExporter
+except TelemetryOptionalImportError:
+    from aiq.utils.optional_imports import DummySpanExporter  # pylint: disable=ungrouped-imports
+    SpanExporter = DummySpanExporter
+
+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,
+                 exporters: dict[str, SpanExporter] | None = None,
+                 retrievers: dict[str | None, RetrieverProviderInfo] | 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.retrievers = retrievers or {}
+
+        self._entry_fn = entry_fn
+
+        self._context_state = context_state
+
+        self._exporters = exporters or {}
+
+    @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) 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,
+                      exporters: dict[str, SpanExporter] | None = None,
+                      retrievers: dict[str | None, RetrieverProviderInfo] | 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,
+                            exporters=exporters,
+                            retrievers=retrievers,
+                            context_state=context_state)