aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/front_ends/fastapi/message_validator.py

@@ -1,345 +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 datetime
-import json
-import logging
-import uuid
-from typing import Any
-from typing import Literal
-
-from pydantic import BaseModel
-from pydantic import ValidationError
-
-from aiq.data_models.api_server import AIQChatResponse
-from aiq.data_models.api_server import AIQChatResponseChunk
-from aiq.data_models.api_server import AIQResponseIntermediateStep
-from aiq.data_models.api_server import AIQResponsePayloadOutput
-from aiq.data_models.api_server import Error
-from aiq.data_models.api_server import ErrorTypes
-from aiq.data_models.api_server import SystemIntermediateStepContent
-from aiq.data_models.api_server import SystemResponseContent
-from aiq.data_models.api_server import TextContent
-from aiq.data_models.api_server import WebSocketMessageStatus
-from aiq.data_models.api_server import WebSocketMessageType
-from aiq.data_models.api_server import WebSocketSystemInteractionMessage
-from aiq.data_models.api_server import WebSocketSystemIntermediateStepMessage
-from aiq.data_models.api_server import WebSocketSystemResponseTokenMessage
-from aiq.data_models.api_server import WebSocketUserInteractionResponseMessage
-from aiq.data_models.api_server import WebSocketUserMessage
-from aiq.data_models.api_server import WorkflowSchemaType
-from aiq.data_models.interactive import BinaryHumanPromptOption
-from aiq.data_models.interactive import HumanPrompt
-from aiq.data_models.interactive import HumanPromptBase
-from aiq.data_models.interactive import HumanPromptBinary
-from aiq.data_models.interactive import HumanPromptCheckbox
-from aiq.data_models.interactive import HumanPromptDropdown
-from aiq.data_models.interactive import HumanPromptRadio
-from aiq.data_models.interactive import HumanPromptText
-from aiq.data_models.interactive import HumanResponse
-from aiq.data_models.interactive import HumanResponseBinary
-from aiq.data_models.interactive import HumanResponseCheckbox
-from aiq.data_models.interactive import HumanResponseDropdown
-from aiq.data_models.interactive import HumanResponseRadio
-from aiq.data_models.interactive import HumanResponseText
-from aiq.data_models.interactive import MultipleChoiceOption
-
-logger = logging.getLogger(__name__)
-
-
-class MessageValidator:
-
-    def __init__(self):
-        self._message_type_schema_mapping: dict[str, type[BaseModel]] = {
-            WebSocketMessageType.USER_MESSAGE: WebSocketUserMessage,
-            WebSocketMessageType.RESPONSE_MESSAGE: WebSocketSystemResponseTokenMessage,
-            WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE: WebSocketSystemIntermediateStepMessage,
-            WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE: WebSocketSystemInteractionMessage,
-            WebSocketMessageType.USER_INTERACTION_MESSAGE: WebSocketUserInteractionResponseMessage,
-            WebSocketMessageType.ERROR_MESSAGE: Error
-        }
-        self._data_type_schema_mapping: dict[str, type[BaseModel]] = {
-            WorkflowSchemaType.GENERATE: AIQResponsePayloadOutput,
-            WorkflowSchemaType.CHAT: AIQChatResponse,
-            WorkflowSchemaType.CHAT_STREAM: AIQChatResponseChunk,
-            WorkflowSchemaType.GENERATE_STREAM: AIQResponseIntermediateStep,
-        }
-        self._message_parent_id: str = "default_id"
-
-    async def validate_message(self, message: dict[str, Any]) -> BaseModel:
-        """
-        Validates an incoming WebSocket message against its expected schema.
-        If validation fails, returns a system response error message.
-
-        :param message: Incoming WebSocket message as a dictionary.
-        :return: A validated Pydantic model.
-        """
-        validated_message: BaseModel
-
-        try:
-            message_type = message.get("type")
-            if not message_type:
-                raise ValueError(f"Missing message type: {json.dumps(message)}")
-
-            schema: type[BaseModel] = await self.get_message_schema_by_type(message_type)
-
-            if issubclass(schema, Error):
-                raise TypeError(
-                    f"An error was encountered processing an incoming WebSocket message of type: {message_type}")
-
-            validated_message = schema(**message)
-            return validated_message
-
-        except (ValidationError, TypeError, ValueError) as e:
-            logger.error("A data validation error %s occurred for message: %s", str(e), str(message), exc_info=True)
-            return await self.create_system_response_token_message(message_type=WebSocketMessageType.ERROR_MESSAGE,
-                                                                   content=Error(code=ErrorTypes.INVALID_MESSAGE,
-                                                                                 message="Error validating message.",
-                                                                                 details=str(e)))
-
-    async def get_message_schema_by_type(self, message_type: str) -> type[BaseModel]:
-        """
-        Retrieves the corresponding Pydantic model schema based on the message type.
-
-        :param message_type: The type of message as a string.
-        :return: A Pydantic schema class if found, otherwise None.
-        """
-        try:
-            schema: type[BaseModel] | None = self._message_type_schema_mapping.get(message_type)
-
-            if schema is None:
-                raise ValueError(f"Unknown message type: {message_type}")
-
-            return schema
-
-        except (TypeError, ValueError) as e:
-            logger.error("Error retrieving schema for message type '%s': %s", message_type, str(e), exc_info=True)
-            return Error
-
-    async def convert_data_to_message_content(self, data_model: BaseModel) -> BaseModel:
-        """
-        Converts a Pydantic data model to a WebSocket message content instance.
-
-        :param data_model: Pydantic Data Model instance.
-        :return: A WebSocket Message Content Data Model instance.
-        """
-
-        validated_message_content: BaseModel = None
-        try:
-            if (isinstance(data_model, AIQResponsePayloadOutput)):
-                validated_message_content = SystemResponseContent(text=data_model.payload)
-
-            elif (isinstance(data_model, (AIQChatResponse, AIQChatResponseChunk))):
-                validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
-
-            elif (isinstance(data_model, AIQResponseIntermediateStep)):
-                validated_message_content = SystemIntermediateStepContent(name=data_model.name,
-                                                                          payload=data_model.payload)
-            elif (isinstance(data_model, HumanPromptBase)):
-                validated_message_content = data_model
-            elif (isinstance(data_model, SystemResponseContent)):
-                return data_model
-            else:
-                raise ValueError(
-                    f"Input data could not be converted to validated message content: {data_model.model_dump_json()}")
-
-            return validated_message_content
-
-        except ValueError as e:
-            logger.error("Input data could not be converted to validated message content: %s", str(e), exc_info=True)
-            return Error(code=ErrorTypes.INVALID_DATA_CONTENT, message="Input data not supported.", details=str(e))
-
-    async def convert_text_content_to_human_response(self, text_content: TextContent,
-                                                     human_prompt: HumanPromptBase) -> HumanResponse:
-        """
-        Converts Message Text Content data model to a Human Response Base data model instance.
-
-        :param text_content: Pydantic TextContent Data Model instance.
-        :param human_prompt: Pydantic HumanPrompt Data Model instance.
-        :return: A Human Response Data Model instance.
-        """
-
-        human_response: HumanResponse = None
-        try:
-            if (isinstance(human_prompt, HumanPromptText)):
-                human_response = HumanResponseText(text=text_content.text)
-
-            elif (isinstance(human_prompt, HumanPromptBinary)):
-                human_response = HumanResponseBinary(selected_option=BinaryHumanPromptOption(value=text_content.text))
-
-            elif (isinstance(human_prompt, HumanPromptRadio)):
-                human_response = HumanResponseRadio(selected_option=MultipleChoiceOption(value=text_content.text))
-
-            elif (isinstance(human_prompt, HumanPromptCheckbox)):
-                human_response = HumanResponseCheckbox(selected_option=MultipleChoiceOption(value=text_content.text))
-
-            elif (isinstance(human_prompt, HumanPromptDropdown)):
-                human_response = HumanResponseDropdown(selected_option=MultipleChoiceOption(value=text_content.text))
-            else:
-                raise ValueError("Message content type not found")
-
-            return human_response
-
-        except ValueError as e:
-            logger.error("Error human response content not found: %s", str(e), exc_info=True)
-            return HumanResponseText(text=str(e))
-
-    async def resolve_message_type_by_data(self, data_model: BaseModel) -> str:
-        """
-        Resolve message type from a validated model
-
-        :param data_model: Pydantic Data Model instance.
-        :return: A WebSocket Message Content Data Model instance.
-        """
-
-        validated_message_type: str = ""
-        try:
-            if (isinstance(data_model, (AIQResponsePayloadOutput, AIQChatResponse, AIQChatResponseChunk))):
-                validated_message_type = WebSocketMessageType.RESPONSE_MESSAGE
-
-            elif (isinstance(data_model, AIQResponseIntermediateStep)):
-                validated_message_type = WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE
-
-            elif (isinstance(data_model, HumanPromptBase)):
-                validated_message_type = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
-            else:
-                raise ValueError("Data type not found")
-
-            return validated_message_type
-
-        except ValueError as e:
-            logger.error("Error type not found converting data to validated websocket message content: %s",
-                         str(e),
-                         exc_info=True)
-            return WebSocketMessageType.ERROR_MESSAGE
-
-    async def get_intermediate_step_parent_id(self, data_model: AIQResponseIntermediateStep) -> str:
-        """
-        Retrieves intermediate step parent_id from AIQResponseIntermediateStep instance.
-
-        :param data_model: AIQResponseIntermediateStep Data Model instance.
-        :return: Intermediate step parent_id or "default".
-        """
-        return data_model.parent_id or "root"
-
-    async def create_system_response_token_message(  # pylint: disable=R0917:too-many-positional-arguments
-        self,
-        message_type: Literal[WebSocketMessageType.RESPONSE_MESSAGE,
-                              WebSocketMessageType.ERROR_MESSAGE] = WebSocketMessageType.RESPONSE_MESSAGE,
-        message_id: str | None = str(uuid.uuid4()),
-        thread_id: str = "default",
-        parent_id: str = "default",
-        content: SystemResponseContent
-        | Error = SystemResponseContent(),
-        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
-    ) -> WebSocketSystemResponseTokenMessage | None:
-        """
-        Creates a system response token message with default values.
-
-        :param message_type: Type of WebSocket message.
-        :param message_id: Unique identifier for the message (default: generated UUID).
-        :param thread_id: ID of the thread the message belongs to (default: "default").
-        :param parent_id: ID of the user message that spawned child messages.
-        :param content: Message content.
-        :param status: Status of the message (default: IN_PROGRESS).
-        :param timestamp: Timestamp of the message (default: current UTC time).
-        :return: A WebSocketSystemResponseTokenMessage instance.
-        """
-        try:
-            return WebSocketSystemResponseTokenMessage(type=message_type,
-                                                       id=message_id,
-                                                       thread_id=thread_id,
-                                                       parent_id=parent_id,
-                                                       content=content,
-                                                       status=status,
-                                                       timestamp=timestamp)
-
-        except Exception as e:
-            logger.error("Error creating system response token message: %s", str(e), exc_info=True)
-            return None
-
-    async def create_system_intermediate_step_message(  # pylint: disable=R0917:too-many-positional-arguments
-        self,
-        message_type: Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE] = (
-            WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE),
-        message_id: str = str(uuid.uuid4()),
-        thread_id: str = "default",
-        parent_id: str = "default",
-        content: SystemIntermediateStepContent = SystemIntermediateStepContent(name="default", payload="default"),
-        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
-    ) -> WebSocketSystemIntermediateStepMessage | None:
-        """
-        Creates a system intermediate step message with default values.
-
-        :param message_type: Type of WebSocket message.
-        :param message_id: Unique identifier for the message (default: generated UUID).
-        :param thread_id: ID of the thread the message belongs to (default: "default").
-        :param parent_id: ID of the user message that spawned child messages.
-        :param content: Message content
-        :param status: Status of the message (default: IN_PROGRESS).
-        :param timestamp: Timestamp of the message (default: current UTC time).
-        :return: A WebSocketSystemIntermediateStepMessage instance.
-        """
-        try:
-            return WebSocketSystemIntermediateStepMessage(type=message_type,
-                                                          id=message_id,
-                                                          thread_id=thread_id,
-                                                          parent_id=parent_id,
-                                                          content=content,
-                                                          status=status,
-                                                          timestamp=timestamp)
-
-        except Exception as e:
-            logger.error("Error creating system intermediate step message: %s", str(e), exc_info=True)
-            return None
-
-    async def create_system_interaction_message(  # pylint: disable=R0917:too-many-positional-arguments
-        self,
-        *,
-        message_type: Literal[WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = (
-            WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE),
-        message_id: str | None = str(uuid.uuid4()),
-        thread_id: str = "default",
-        parent_id: str = "default",
-        content: HumanPrompt,
-        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
-    ) -> WebSocketSystemInteractionMessage | None:  # noqa: E125 continuation line with same indent as next logical line
-        """
-        Creates a system interaction message with default values.
-
-        :param message_type: Type of WebSocket message.
-        :param message_id: Unique identifier for the message (default: generated UUID).
-        :param thread_id: ID of the thread the message belongs to (default: "default").
-        :param parent_id: ID of the user message that spawned child messages.
-        :param content: Message content
-        :param status: Status of the message (default: IN_PROGRESS).
-        :param timestamp: Timestamp of the message (default: current UTC time).
-        :return: A WebSocketSystemInteractionMessage instance.
-        """
-        try:
-            return WebSocketSystemInteractionMessage(type=message_type,
-                                                     id=message_id,
-                                                     thread_id=thread_id,
-                                                     parent_id=parent_id,
-                                                     content=content,
-                                                     status=status,
-                                                     timestamp=timestamp)
-
-        except Exception as e:
-            logger.error("Error creating system interaction message: %s", str(e), exc_info=True)
-            return None

aiq/front_ends/fastapi/register.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.cli.register_workflow import register_front_end
-from aiq.data_models.config import AIQConfig
-from aiq.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
-
-
-@register_front_end(config_type=FastApiFrontEndConfig)
-async def register_fastapi_front_end(config: FastApiFrontEndConfig, full_config: AIQConfig):
-    from aiq.front_ends.fastapi.fastapi_front_end_plugin import FastApiFrontEndPlugin
-
-    yield FastApiFrontEndPlugin(full_config=full_config)

aiq/front_ends/fastapi/response_helpers.py

@@ -1,195 +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 asyncio
-import typing
-from collections.abc import AsyncGenerator
-
-from aiq.data_models.api_server import AIQResponseIntermediateStep
-from aiq.data_models.api_server import AIQResponsePayloadOutput
-from aiq.data_models.api_server import AIQResponseSerializable
-from aiq.data_models.step_adaptor import StepAdaptorConfig
-from aiq.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
-from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
-from aiq.runtime.session import AIQSessionManager
-from aiq.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
-
-
-async def generate_streaming_response_as_str(payload: typing.Any,
-                                             *,
-                                             session_manager: AIQSessionManager,
-                                             streaming: bool,
-                                             step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
-                                             result_type: type | None = None,
-                                             output_type: type | None = None) -> AsyncGenerator[str]:
-
-    async for item in generate_streaming_response(payload,
-                                                  session_manager=session_manager,
-                                                  streaming=streaming,
-                                                  step_adaptor=step_adaptor,
-                                                  result_type=result_type,
-                                                  output_type=output_type):
-
-        if (isinstance(item, AIQResponseSerializable)):
-            yield item.get_stream_data()
-        else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
-                             str(type(item)))
-
-
-async def generate_streaming_response(payload: typing.Any,
-                                      *,
-                                      session_manager: AIQSessionManager,
-                                      streaming: bool,
-                                      step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
-                                      result_type: type | None = None,
-                                      output_type: type | None = None) -> AsyncGenerator[AIQResponseSerializable]:
-
-    async with session_manager.run(payload) as runner:
-
-        q: AsyncIOProducerConsumerQueue[AIQResponseSerializable] = AsyncIOProducerConsumerQueue()
-
-        # Start the intermediate stream
-        intermediate_complete = await pull_intermediate(q, step_adaptor)
-
-        async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
-                async for chunk in runner.result_stream(to_type=output_type):
-                    await q.put(chunk)
-            else:
-                result = await runner.result(to_type=result_type)
-                await q.put(runner.convert(result, output_type))
-
-            # Wait until the intermediate subscription is done before closing q
-            # But we have no direct "intermediate_done" reference here
-            # because it's encapsulated in pull_intermediate. So we can do:
-            #    await some_event.wait()
-            # If needed. Alternatively, you can skip that if the intermediate
-            # subscriber won't block the main flow.
-            #
-            # For example, if you *need* to guarantee the subscriber is done before
-            # closing the queue, you can structure the code to store or return
-            # the 'intermediate_done' event from pull_intermediate.
-            #
-
-            await intermediate_complete.wait()
-
-            await q.close()
-
-        try:
-            # Start the result stream
-            asyncio.create_task(pull_result())
-
-            async for item in q:
-
-                if (isinstance(item, AIQResponseSerializable)):
-                    yield item
-                else:
-                    yield AIQResponsePayloadOutput(payload=item)
-        except Exception as e:
-            # Handle exceptions here
-            raise e
-        finally:
-            await q.close()
-
-
-async def generate_single_response(
-    payload: typing.Any,
-    session_manager: AIQSessionManager,
-    result_type: type | None = None,
-) -> typing.Any:
-    if (not session_manager.workflow.has_single_output):
-        raise ValueError("Cannot get a single output value for streaming workflows")
-
-    async with session_manager.run(payload) as runner:
-        return await runner.result(to_type=result_type)
-
-
-async def generate_streaming_response_full(payload: typing.Any,
-                                           *,
-                                           session_manager: AIQSessionManager,
-                                           streaming: bool,
-                                           result_type: type | None = None,
-                                           output_type: type | None = None,
-                                           filter_steps: str | None = None) -> AsyncGenerator[AIQResponseSerializable]:
-    """
-    Similar to generate_streaming_response but provides raw AIQResponseIntermediateStep objects
-    without any step adaptor translations.
-    """
-    # Parse filter_steps into a set of allowed types if provided
-    # Special case: if filter_steps is "none", suppress all steps
-    allowed_types = None
-    if filter_steps:
-        if filter_steps.lower() == "none":
-            allowed_types = set()  # Empty set means no steps allowed
-        else:
-            allowed_types = set(filter_steps.split(','))
-
-    async with session_manager.run(payload) as runner:
-        q: AsyncIOProducerConsumerQueue[AIQResponseSerializable] = AsyncIOProducerConsumerQueue()
-
-        # Start the intermediate stream without step adaptor
-        intermediate_complete = await pull_intermediate(q, None)
-
-        async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
-                async for chunk in runner.result_stream(to_type=output_type):
-                    await q.put(chunk)
-            else:
-                result = await runner.result(to_type=result_type)
-                await q.put(runner.convert(result, output_type))
-
-            await intermediate_complete.wait()
-            await q.close()
-
-        try:
-            # Start the result stream
-            asyncio.create_task(pull_result())
-
-            async for item in q:
-                if (isinstance(item, AIQResponseIntermediateStep)):
-                    # Filter intermediate steps if filter_steps is provided
-                    if allowed_types is None or item.type in allowed_types:
-                        yield item
-                else:
-                    yield AIQResponsePayloadOutput(payload=item)
-        except Exception as e:
-            # Handle exceptions here
-            raise e
-        finally:
-            await q.close()
-
-
-async def generate_streaming_response_full_as_str(payload: typing.Any,
-                                                  *,
-                                                  session_manager: AIQSessionManager,
-                                                  streaming: bool,
-                                                  result_type: type | None = None,
-                                                  output_type: type | None = None,
-                                                  filter_steps: str | None = None) -> AsyncGenerator[str]:
-    """
-    Similar to generate_streaming_response but converts the response to a string format.
-    """
-    async for item in generate_streaming_response_full(payload,
-                                                       session_manager=session_manager,
-                                                       streaming=streaming,
-                                                       result_type=result_type,
-                                                       output_type=output_type,
-                                                       filter_steps=filter_steps):
-        if (isinstance(item, AIQResponseIntermediateStep) or isinstance(item, AIQResponsePayloadOutput)):
-            yield item.get_stream_data()
-        else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
-                             str(type(item)))