aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__py3-none-any.whl

aiq/front_ends/fastapi/message_handler.py

@@ -15,14 +15,19 @@
 
 import asyncio
 import logging
+import typing
 import uuid
 from typing import Any
 
 from fastapi import WebSocket
 from pydantic import BaseModel
 from pydantic import ValidationError
-from starlette.endpoints import WebSocketEndpoint
+from starlette.websockets import WebSocketDisconnect
 
+from aiq.authentication.interfaces import FlowHandlerBase
+from aiq.data_models.api_server import AIQChatResponse
+from aiq.data_models.api_server import AIQResponsePayloadOutput
+from aiq.data_models.api_server import AIQResponseSerializable
 from aiq.data_models.api_server import Error
 from aiq.data_models.api_server import ErrorTypes
 from aiq.data_models.api_server import SystemResponseContent
@@ -39,74 +44,72 @@ from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import HumanResponseNotification
 from aiq.data_models.interactive import InteractionPrompt
 from aiq.front_ends.fastapi.message_validator import MessageValidator
+from aiq.front_ends.fastapi.response_helpers import generate_streaming_response
+from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
+from aiq.runtime.session import AIQSessionManager
 
 logger = logging.getLogger(__name__)
 
 
-class MessageHandler:
+class WebSocketMessageHandler:
+
+    def __init__(self, socket: WebSocket, session_manager: AIQSessionManager, step_adaptor: StepAdaptor):
+        self._socket: WebSocket = socket
+        self._session_manager: AIQSessionManager = session_manager
+        self._step_adaptor: StepAdaptor = step_adaptor
 
-    def __init__(self, websocket_reference: WebSocketEndpoint):
-        self._websocket_reference: WebSocketEndpoint = websocket_reference
         self._message_validator: MessageValidator = MessageValidator()
-        self._messages_queue: asyncio.Queue[dict[str, str]] = asyncio.Queue()
-        self._out_going_messages_queue: asyncio.Queue[dict] = asyncio.Queue()
-        self._process_messages_task: asyncio.Task | None = None
-        self._process_out_going_messages_task: asyncio.Task = None
-        self._background_task: asyncio.Task = None
+        self._running_workflow_task: asyncio.Task | None = None
         self._message_parent_id: str = "default_id"
         self._workflow_schema_type: str = None
-        self._user_interaction_response: asyncio.Future[TextContent] = asyncio.Future()
+        self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
 
-    @property
-    def messages_queue(self) -> asyncio.Queue[dict[str, str]]:
-        return self._messages_queue
+        self._flow_handler: FlowHandlerBase | None = None
 
-    @property
-    def background_task(self) -> asyncio.Task:
-        return self._background_task
+    def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
+        self._flow_handler = flow_handler
 
-    @property
-    def process_messages_task(self) -> asyncio.Task | None:
-        return self._process_messages_task
+    async def __aenter__(self) -> "WebSocketMessageHandler":
+        await self._socket.accept()
 
-    @process_messages_task.setter
-    def process_messages_task(self, process_messages_task) -> None:
-        self._process_messages_task = process_messages_task
+        return self
 
-    @property
-    def process_out_going_messages_task(self) -> asyncio.Task:
-        return self._process_out_going_messages_task
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
-    @process_out_going_messages_task.setter
-    def process_out_going_messages_task(self, process_out_going_messages_task) -> None:
-        self._process_out_going_messages_task = process_out_going_messages_task
+        # TODO: Handle the exit
+        pass
 
-    async def process_messages(self) -> None:
+    async def run(self) -> None:
         """
         Processes received messages from websocket and routes them appropriately.
         """
         while True:
 
             try:
-                message: dict[str, Any] = await self._messages_queue.get()
+
+                message: dict[str, Any] = await self._socket.receive_json()
 
                 validated_message: BaseModel = await self._message_validator.validate_message(message)
 
+                # Received a request to start a workflow
                 if (isinstance(validated_message, WebSocketUserMessage)):
-                    await self.process_user_message(validated_message)
+                    await self.process_workflow_request(validated_message)
 
-                if isinstance(
+                elif isinstance(
                         validated_message,
                     (  # noqa: E131
                         WebSocketSystemResponseTokenMessage,
                         WebSocketSystemIntermediateStepMessage,
                         WebSocketSystemInteractionMessage)):
-                    await self._out_going_messages_queue.put(validated_message.model_dump())
+                    # These messages are already handled by self.create_websocket_message(data_model=value, …)
+                    # No further processing is needed here.
+                    pass
 
-                if (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
+                elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
-            except (asyncio.CancelledError):
+            except (asyncio.CancelledError, WebSocketDisconnect):
+                # TODO: Handle the disconnect
                 break
 
         return None
@@ -130,7 +133,7 @@ class MessageHandler:
 
         return None
 
-    async def process_user_message(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
+    async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
         """
         Process user messages and routes them appropriately.
 
@@ -147,13 +150,15 @@ class MessageHandler:
             if content is None:
                 raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
 
-            if isinstance(content, TextContent) and (self._background_task is None):
+            if isinstance(content, TextContent) and (self._running_workflow_task is None):
 
-                await self._process_response()
-                self._background_task = asyncio.create_task(
-                    self._websocket_reference.workflow_schema_type.get(self._workflow_schema_type)(
-                        content.text, conversation_id)).add_done_callback(
-                            lambda task: asyncio.create_task(self._on_process_stream_task_done(task)))
+                def _done_callback(task: asyncio.Task):
+                    self._running_workflow_task = None
+
+                # await self._process_response()
+                self._running_workflow_task = asyncio.create_task(
+                    self._run_workflow(content.text, conversation_id,
+                                       result_type=AIQChatResponse)).add_done_callback(_done_callback)
 
         except ValueError as e:
             logger.error("User message content not found: %s", str(e), exc_info=True)
@@ -221,60 +226,73 @@ class MessageHandler:
                 content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
 
         finally:
-            await self._messages_queue.put(message.model_dump())
-
-    async def _on_process_stream_task_done(self, task: asyncio.Task) -> None:
-        await self.create_websocket_message(data_model=SystemResponseContent(),
-                                            message_type=WebSocketMessageType.RESPONSE_MESSAGE,
-                                            status=WebSocketMessageStatus.COMPLETE)
+            if (message is not None):
+                await self._socket.send_json(message.model_dump())
 
-        return None
-
-    async def process_out_going_messages(self, websocket: WebSocket) -> None:
+    async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse:
         """
-        Spawns out going message processing task.
+        Registered human interaction callback that processes human interactions and returns
+        responses from websocket connection.
 
-        :param websocket: Websocket instance.
+        :param prompt: Incoming interaction content data model.
+        :return: A Text Content Base Pydantic model.
         """
-        while True:
-            try:
-                out_going_message = await self._out_going_messages_queue.get()
-                await self._websocket_reference.on_send(websocket, out_going_message)
 
-            except (asyncio.CancelledError, ValidationError):
-                break
+        # First create a future from the loop for the human response
+        human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
 
-        return None
+        # Then add the future to the outstanding human prompts dictionary
+        self._user_interaction_response = human_response_future
+
+        try:
 
-    async def _process_response(self):
-        self._websocket_reference.process_response_event.set()
+            await self.create_websocket_message(data_model=prompt.content,
+                                                message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
+                                                status=WebSocketMessageStatus.IN_PROGRESS)
 
-    async def _pause_response(self):
-        self._websocket_reference.process_response_event.clear()
+            if (isinstance(prompt.content, HumanPromptNotification)):
 
-    async def __reset_user_interaction_response(self):
-        self._user_interaction_response = asyncio.Future()
+                return HumanResponseNotification()
 
-    async def human_interaction(self, prompt: InteractionPrompt) -> HumanResponse:
-        """
-        Registered human interaction callback that processes human interactions and returns
-        responses from websocket connection.
+            # Wait for the human response future to complete
+            interaction_response: HumanResponse = await human_response_future
 
-        :param prompt: Incoming interaction content data model.
-        :return: A Text Content Base Pydantic model.
-        """
-        await self.create_websocket_message(data_model=prompt.content,
-                                            message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
-                                            status=WebSocketMessageStatus.IN_PROGRESS)
+            interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
+                interaction_response, prompt.content)
 
-        if (isinstance(prompt.content, HumanPromptNotification)):
-            return HumanResponseNotification()
+            return interaction_response
 
-        user_message_repsonse_content: TextContent = await self._user_interaction_response
-        interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
-            user_message_repsonse_content, prompt.content)
+        finally:
+            # Delete the future from the outstanding human prompts dictionary
+            self._user_interaction_response = None
+
+    async def _run_workflow(self,
+                            payload: typing.Any,
+                            conversation_id: str | None = None,
+                            result_type: type | None = None,
+                            output_type: type | None = None) -> None:
+
+        try:
+            async with self._session_manager.session(
+                    conversation_id=conversation_id,
+                    request=self._socket,
+                    user_input_callback=self.human_interaction_callback,
+                    user_authentication_callback=(self._flow_handler.authenticate
+                                                  if self._flow_handler else None)) as session:
 
-        await self.__reset_user_interaction_response()
-        await self._process_response()
+                async for value in generate_streaming_response(payload,
+                                                               session_manager=session,
+                                                               streaming=True,
+                                                               step_adaptor=self._step_adaptor,
+                                                               result_type=result_type,
+                                                               output_type=output_type):
 
-        return interaction_response
+                    if not isinstance(value, AIQResponseSerializable):
+                        value = AIQResponsePayloadOutput(payload=value)
+
+                    await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
+
+        finally:
+            await self.create_websocket_message(data_model=SystemResponseContent(),
+                                                message_type=WebSocketMessageType.RESPONSE_MESSAGE,
+                                                status=WebSocketMessageStatus.COMPLETE)

aiq/front_ends/fastapi/step_adaptor.py

@@ -291,7 +291,8 @@ class StepAdaptor:
 
         return event
 
-    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:
+    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:  # pylint: disable=R1710
+
         # Track the chunk
         self._history.append(step)
         payload = step.payload

aiq/llm/aws_bedrock_llm.py

@@ -21,9 +21,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.llm import LLMProviderInfo
 from aiq.cli.register_workflow import register_llm_provider
 from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

aiq/llm/nim_llm.py

@@ -22,9 +22,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.llm import LLMProviderInfo
 from aiq.cli.register_workflow import register_llm_provider
 from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 
-class NIMModelConfig(LLMBaseConfig, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

aiq/llm/openai_llm.py

@@ -21,9 +21,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.llm import LLMProviderInfo
 from aiq.cli.register_workflow import register_llm_provider
 from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 
-class OpenAIModelConfig(LLMBaseConfig, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")

aiq/object_store/__init__.py

@@ -0,0 +1,20 @@
+# 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 Toolkit Object Store Module
+
+This package provides foundational classes and interfaces
+for managing object storage in AIQ Toolkit's LLM-based agents.
+"""

aiq/object_store/in_memory_object_store.py

@@ -0,0 +1,74 @@
+# 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.builder.builder import Builder
+from aiq.cli.register_workflow import register_object_store
+from aiq.data_models.object_store import KeyAlreadyExistsError
+from aiq.data_models.object_store import NoSuchKeyError
+from aiq.data_models.object_store import ObjectStoreBaseConfig
+from aiq.utils.type_utils import override
+
+from .interfaces import ObjectStore
+from .models import ObjectStoreItem
+
+
+class InMemoryObjectStoreConfig(ObjectStoreBaseConfig, name="in_memory"):
+    """
+    Object store that stores objects in memory. Objects are not persisted when the process shuts down.
+    """
+    pass
+
+
+class InMemoryObjectStore(ObjectStore):
+    """
+    Implementation of ObjectStore that stores objects in memory. Objects are not persisted when the process shuts down.
+    """
+
+    def __init__(self) -> None:
+        self._store: dict[str, ObjectStoreItem] = {}
+
+    @override
+    async def put_object(self, key: str, item: ObjectStoreItem) -> None:
+        if key in self._store:
+            raise KeyAlreadyExistsError(key)
+
+        self._store[key] = item
+        return
+
+    @override
+    async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
+        self._store[key] = item
+        return
+
+    @override
+    async def get_object(self, key: str) -> ObjectStoreItem:
+
+        if key not in self._store:
+            raise NoSuchKeyError(key)
+
+        return self._store[key]
+
+    @override
+    async def delete_object(self, key: str) -> None:
+        if key not in self._store:
+            raise NoSuchKeyError(key)
+
+        self._store.pop(key)
+        return
+
+
+@register_object_store(config_type=InMemoryObjectStoreConfig)
+async def in_memory_object_store(config: InMemoryObjectStoreConfig, builder: Builder):
+    yield InMemoryObjectStore()

aiq/object_store/interfaces.py

@@ -0,0 +1,84 @@
+# 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 abc import ABC
+from abc import abstractmethod
+
+from .models import ObjectStoreItem
+
+
+class ObjectStore(ABC):
+    """
+    Abstract interface for an object store.
+
+    Implementations may integrate with various object stores,
+    such as S3, MySQL, etc.
+    """
+
+    @abstractmethod
+    async def put_object(self, key: str, item: ObjectStoreItem) -> None:
+        """
+        Save an ObjectStoreItem in the object store with the given key.
+        If the key already exists, raise an error.
+
+        Args:
+            key (str): The key to save the item under.
+            item (ObjectStoreItem): The item to save.
+
+        Raises:
+            KeyAlreadyExistsError: If the key already exists.
+        """
+        pass
+
+    @abstractmethod
+    async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
+        """
+        Save an ObjectStoreItem in the object store with the given key.
+        If the key already exists, update the item.
+
+        Args:
+            key (str): The key to save the item under.
+            item (ObjectStoreItem): The item to save.
+        """
+        pass
+
+    @abstractmethod
+    async def get_object(self, key: str) -> ObjectStoreItem:
+        """
+        Get an ObjectStoreItem from the object store by key.
+
+        Args:
+            key (str): The key to get the item from.
+
+        Returns:
+            ObjectStoreItem: The item retrieved from the object store.
+
+        Raises:
+            NoSuchKeyError: If the item does not exist.
+        """
+        pass
+
+    @abstractmethod
+    async def delete_object(self, key: str) -> None:
+        """
+        Delete an ObjectStoreItem from the object store by key.
+
+        Args:
+            key (str): The key to delete the item from.
+
+        Raises:
+            NoSuchKeyError: If the item does not exist.
+        """
+        pass

aiq/object_store/models.py

@@ -0,0 +1,36 @@
+# 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 pydantic import BaseModel
+from pydantic import Field
+
+
+class ObjectStoreItem(BaseModel):
+    """
+    Represents an object store item consisting of bytes and associated metadata.
+
+    Attributes
+    ----------
+    data : bytes
+        The data to store in the object store.
+    content_type : str | None
+        The content type of the data.
+    metadata : dict[str, str] | None
+        Metadata providing context and utility for management operations.
+    """
+
+    data: bytes = Field(description="The data to store in the object store.")
+    content_type: str | None = Field(description="The content type of the data.", default=None)
+    metadata: dict[str, str] | None = Field(description="The metadata of the data.", default=None)

aiq/object_store/register.py

@@ -0,0 +1,20 @@
+# 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
+
+from . import in_memory_object_store

aiq/observability/__init__.py

@@ -0,0 +1,14 @@
+# 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/observability/exporter/__init__.py

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