nvidia-nat 1.2.0rc5__py3-none-any.whl → 1.2.0rc6__py3-none-any.whl

{aiq → nat}/front_ends/fastapi/message_handler.py

@@ -24,48 +24,58 @@ from pydantic import BaseModel
 from pydantic import ValidationError
 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
-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.interactive import HumanPromptNotification
-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
+from nat.authentication.interfaces import FlowHandlerBase
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import ChatResponseChunk
+from nat.data_models.api_server import Error
+from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ResponsePayloadOutput
+from nat.data_models.api_server import ResponseSerializable
+from nat.data_models.api_server import SystemResponseContent
+from nat.data_models.api_server import TextContent
+from nat.data_models.api_server import WebSocketMessageStatus
+from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketSystemInteractionMessage
+from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
+from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
+from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
+from nat.data_models.api_server import WebSocketUserMessage
+from nat.data_models.api_server import WorkflowSchemaType
+from nat.data_models.interactive import HumanPromptNotification
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import HumanResponseNotification
+from nat.data_models.interactive import InteractionPrompt
+from nat.front_ends.fastapi.message_validator import MessageValidator
+from nat.front_ends.fastapi.response_helpers import generate_streaming_response
+from nat.front_ends.fastapi.step_adaptor import StepAdaptor
+from nat.runtime.session import SessionManager
 
 logger = logging.getLogger(__name__)
 
 
 class WebSocketMessageHandler:
 
-    def __init__(self, socket: WebSocket, session_manager: AIQSessionManager, step_adaptor: StepAdaptor):
+    def __init__(self, socket: WebSocket, session_manager: SessionManager, step_adaptor: StepAdaptor):
         self._socket: WebSocket = socket
-        self._session_manager: AIQSessionManager = session_manager
+        self._session_manager: SessionManager = session_manager
         self._step_adaptor: StepAdaptor = step_adaptor
 
         self._message_validator: MessageValidator = MessageValidator()
         self._running_workflow_task: asyncio.Task | None = None
         self._message_parent_id: str = "default_id"
+        self._conversation_id: str | None = None
         self._workflow_schema_type: str = None
         self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
 
         self._flow_handler: FlowHandlerBase | None = None
 
+        self._schema_output_mapping: dict[str, type[BaseModel] | None] = {
+            WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
+            WorkflowSchemaType.CHAT: ChatResponse,
+            WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
+            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.workflow.streaming_output_schema,
+        }
+
     def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
         self._flow_handler = flow_handler
 
@@ -76,7 +86,7 @@ class WebSocketMessageHandler:
 
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
-        # TODO: Handle the exit
+        # TODO: Handle the exit  # pylint: disable=fixme
         pass
 
     async def run(self) -> None:
@@ -109,7 +119,7 @@ class WebSocketMessageHandler:
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
-                # TODO: Handle the disconnect
+                # TODO: Handle the disconnect  # pylint: disable=fixme
                 break
 
         return None
@@ -143,7 +153,7 @@ class WebSocketMessageHandler:
         try:
             self._message_parent_id = user_message_as_validated_type.id
             self._workflow_schema_type = user_message_as_validated_type.schema_type
-            conversation_id: str = user_message_as_validated_type.conversation_id
+            self._conversation_id = user_message_as_validated_type.conversation_id
 
             content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
 
@@ -152,13 +162,15 @@ class WebSocketMessageHandler:
 
             if isinstance(content, TextContent) and (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):
+                def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
                     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)
+                    self._run_workflow(content.text,
+                                       self._conversation_id,
+                                       result_type=self._schema_output_mapping[self._workflow_schema_type],
+                                       output_type=self._schema_output_mapping[
+                                           self._workflow_schema_type])).add_done_callback(_done_callback)
 
         except ValueError as e:
             logger.error("User message content not found: %s", str(e), exc_info=True)
@@ -196,18 +208,27 @@ class WebSocketMessageHandler:
 
             if issubclass(message_schema, WebSocketSystemResponseTokenMessage):
                 message = await self._message_validator.create_system_response_token_message(
-                    message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content,
+                    status=status)
 
             elif issubclass(message_schema, WebSocketSystemIntermediateStepMessage):
                 message = await self._message_validator.create_system_intermediate_step_message(
                     message_id=message_id,
                     parent_id=await self._message_validator.get_intermediate_step_parent_id(data_model),
+                    conversation_id=self._conversation_id,
                     content=content,
                     status=status)
 
             elif issubclass(message_schema, WebSocketSystemInteractionMessage):
                 message = await self._message_validator.create_system_interaction_message(
-                    message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content,
+                    status=status)
 
             elif isinstance(content, Error):
                 raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
@@ -223,6 +244,7 @@ class WebSocketMessageHandler:
             logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
             message = await self._message_validator.create_system_response_token_message(
                 message_type=WebSocketMessageType.ERROR_MESSAGE,
+                conversation_id=self._conversation_id,
                 content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
 
         finally:
@@ -287,8 +309,8 @@ class WebSocketMessageHandler:
                                                                result_type=result_type,
                                                                output_type=output_type):
 
-                    if not isinstance(value, AIQResponseSerializable):
-                        value = AIQResponsePayloadOutput(payload=value)
+                    if not isinstance(value, ResponseSerializable):
+                        value = ResponsePayloadOutput(payload=value)
 
                     await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
 

{aiq → nat}/front_ends/fastapi/message_validator.py

@@ -23,38 +23,37 @@ 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
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import ChatResponseChunk
+from nat.data_models.api_server import Error
+from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponsePayloadOutput
+from nat.data_models.api_server import SystemIntermediateStepContent
+from nat.data_models.api_server import SystemResponseContent
+from nat.data_models.api_server import TextContent
+from nat.data_models.api_server import WebSocketMessageStatus
+from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketSystemInteractionMessage
+from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
+from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
+from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
+from nat.data_models.api_server import WebSocketUserMessage
+from nat.data_models.interactive import BinaryHumanPromptOption
+from nat.data_models.interactive import HumanPrompt
+from nat.data_models.interactive import HumanPromptBase
+from nat.data_models.interactive import HumanPromptBinary
+from nat.data_models.interactive import HumanPromptCheckbox
+from nat.data_models.interactive import HumanPromptDropdown
+from nat.data_models.interactive import HumanPromptRadio
+from nat.data_models.interactive import HumanPromptText
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import HumanResponseBinary
+from nat.data_models.interactive import HumanResponseCheckbox
+from nat.data_models.interactive import HumanResponseDropdown
+from nat.data_models.interactive import HumanResponseRadio
+from nat.data_models.interactive import HumanResponseText
+from nat.data_models.interactive import MultipleChoiceOption
 
 logger = logging.getLogger(__name__)
 
@@ -70,12 +69,7 @@ class MessageValidator:
             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:
@@ -138,13 +132,17 @@ class MessageValidator:
 
         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))):
+            if (isinstance(data_model, ResponsePayloadOutput)):
+                if hasattr(data_model.payload, 'model_dump_json'):
+                    text_content: str = data_model.payload.model_dump_json()
+                else:
+                    text_content: str = str(data_model.payload)
+                validated_message_content = SystemResponseContent(text=text_content)
+
+            elif (isinstance(data_model, (ChatResponse, ChatResponseChunk))):
                 validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
 
-            elif (isinstance(data_model, AIQResponseIntermediateStep)):
+            elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_content = SystemIntermediateStepContent(name=data_model.name,
                                                                           payload=data_model.payload)
             elif (isinstance(data_model, HumanPromptBase)):
@@ -206,10 +204,10 @@ class MessageValidator:
 
         validated_message_type: str = ""
         try:
-            if (isinstance(data_model, (AIQResponsePayloadOutput, AIQChatResponse, AIQChatResponseChunk))):
+            if (isinstance(data_model, (ResponsePayloadOutput, ChatResponse, ChatResponseChunk))):
                 validated_message_type = WebSocketMessageType.RESPONSE_MESSAGE
 
-            elif (isinstance(data_model, AIQResponseIntermediateStep)):
+            elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_type = WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE
 
             elif (isinstance(data_model, HumanPromptBase)):
@@ -225,11 +223,11 @@ class MessageValidator:
                          exc_info=True)
             return WebSocketMessageType.ERROR_MESSAGE
 
-    async def get_intermediate_step_parent_id(self, data_model: AIQResponseIntermediateStep) -> str:
+    async def get_intermediate_step_parent_id(self, data_model: ResponseIntermediateStep) -> str:
         """
-        Retrieves intermediate step parent_id from AIQResponseIntermediateStep instance.
+        Retrieves intermediate step parent_id from ResponseIntermediateStep instance.
 
-        :param data_model: AIQResponseIntermediateStep Data Model instance.
+        :param data_model: ResponseIntermediateStep Data Model instance.
         :return: Intermediate step parent_id or "default".
         """
         return data_model.parent_id or "root"
@@ -241,6 +239,7 @@ class MessageValidator:
         message_id: str | None = str(uuid.uuid4()),
         thread_id: str = "default",
         parent_id: str = "default",
+        conversation_id: str | None = None,
         content: SystemResponseContent
         | Error = SystemResponseContent(),
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
@@ -253,6 +252,7 @@ class MessageValidator:
         :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 conversation_id: ID of the conversation this message belongs to (default: None).
         :param content: Message content.
         :param status: Status of the message (default: IN_PROGRESS).
         :param timestamp: Timestamp of the message (default: current UTC time).
@@ -263,6 +263,7 @@ class MessageValidator:
                                                        id=message_id,
                                                        thread_id=thread_id,
                                                        parent_id=parent_id,
+                                                       conversation_id=conversation_id,
                                                        content=content,
                                                        status=status,
                                                        timestamp=timestamp)
@@ -278,6 +279,7 @@ class MessageValidator:
         message_id: str = str(uuid.uuid4()),
         thread_id: str = "default",
         parent_id: str = "default",
+        conversation_id: str | None = None,
         content: SystemIntermediateStepContent = SystemIntermediateStepContent(name="default", payload="default"),
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
         timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
@@ -289,6 +291,7 @@ class MessageValidator:
         :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 conversation_id: ID of the conversation this message belongs to (default: None).
         :param content: Message content
         :param status: Status of the message (default: IN_PROGRESS).
         :param timestamp: Timestamp of the message (default: current UTC time).
@@ -299,6 +302,7 @@ class MessageValidator:
                                                           id=message_id,
                                                           thread_id=thread_id,
                                                           parent_id=parent_id,
+                                                          conversation_id=conversation_id,
                                                           content=content,
                                                           status=status,
                                                           timestamp=timestamp)
@@ -315,6 +319,7 @@ class MessageValidator:
         message_id: str | None = str(uuid.uuid4()),
         thread_id: str = "default",
         parent_id: str = "default",
+        conversation_id: str | None = None,
         content: HumanPrompt,
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
         timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
@@ -326,6 +331,7 @@ class MessageValidator:
         :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 conversation_id: ID of the conversation this message belongs to (default: None).
         :param content: Message content
         :param status: Status of the message (default: IN_PROGRESS).
         :param timestamp: Timestamp of the message (default: current UTC time).
@@ -336,6 +342,7 @@ class MessageValidator:
                                                      id=message_id,
                                                      thread_id=thread_id,
                                                      parent_id=parent_id,
+                                                     conversation_id=conversation_id,
                                                      content=content,
                                                      status=status,
                                                      timestamp=timestamp)

{aiq → nat}/front_ends/fastapi/register.py

@@ -13,13 +13,13 @@
 # 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
+from nat.cli.register_workflow import register_front_end
+from nat.data_models.config import Config
+from nat.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
+async def register_fastapi_front_end(config: FastApiFrontEndConfig, full_config: Config):
+    from nat.front_ends.fastapi.fastapi_front_end_plugin import FastApiFrontEndPlugin
 
     yield FastApiFrontEndPlugin(full_config=full_config)

{aiq → nat}/front_ends/fastapi/response_helpers.py

@@ -17,19 +17,19 @@ 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
+from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponsePayloadOutput
+from nat.data_models.api_server import ResponseSerializable
+from nat.data_models.step_adaptor import StepAdaptorConfig
+from nat.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
+from nat.front_ends.fastapi.step_adaptor import StepAdaptor
+from nat.runtime.session import SessionManager
+from nat.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
 
 
 async def generate_streaming_response_as_str(payload: typing.Any,
                                              *,
-                                             session_manager: AIQSessionManager,
+                                             session_manager: SessionManager,
                                              streaming: bool,
                                              step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                              result_type: type | None = None,
@@ -42,24 +42,24 @@ async def generate_streaming_response_as_str(payload: typing.Any,
                                                   result_type=result_type,
                                                   output_type=output_type):
 
-        if (isinstance(item, AIQResponseSerializable)):
+        if (isinstance(item, ResponseSerializable)):
             yield item.get_stream_data()
         else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
+            raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
                              str(type(item)))
 
 
 async def generate_streaming_response(payload: typing.Any,
                                       *,
-                                      session_manager: AIQSessionManager,
+                                      session_manager: SessionManager,
                                       streaming: bool,
                                       step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                       result_type: type | None = None,
-                                      output_type: type | None = None) -> AsyncGenerator[AIQResponseSerializable]:
+                                      output_type: type | None = None) -> AsyncGenerator[ResponseSerializable]:
 
     async with session_manager.run(payload) as runner:
 
-        q: AsyncIOProducerConsumerQueue[AIQResponseSerializable] = AsyncIOProducerConsumerQueue()
+        q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
         # Start the intermediate stream
         intermediate_complete = await pull_intermediate(q, step_adaptor)
@@ -94,10 +94,10 @@ async def generate_streaming_response(payload: typing.Any,
 
             async for item in q:
 
-                if (isinstance(item, AIQResponseSerializable)):
+                if (isinstance(item, ResponseSerializable)):
                     yield item
                 else:
-                    yield AIQResponsePayloadOutput(payload=item)
+                    yield ResponsePayloadOutput(payload=item)
         except Exception as e:
             # Handle exceptions here
             raise e
@@ -107,7 +107,7 @@ async def generate_streaming_response(payload: typing.Any,
 
 async def generate_single_response(
     payload: typing.Any,
-    session_manager: AIQSessionManager,
+    session_manager: SessionManager,
     result_type: type | None = None,
 ) -> typing.Any:
     if (not session_manager.workflow.has_single_output):
@@ -119,13 +119,13 @@ async def generate_single_response(
 
 async def generate_streaming_response_full(payload: typing.Any,
                                            *,
-                                           session_manager: AIQSessionManager,
+                                           session_manager: SessionManager,
                                            streaming: bool,
                                            result_type: type | None = None,
                                            output_type: type | None = None,
-                                           filter_steps: str | None = None) -> AsyncGenerator[AIQResponseSerializable]:
+                                           filter_steps: str | None = None) -> AsyncGenerator[ResponseSerializable]:
     """
-    Similar to generate_streaming_response but provides raw AIQResponseIntermediateStep objects
+    Similar to generate_streaming_response but provides raw ResponseIntermediateStep objects
     without any step adaptor translations.
     """
     # Parse filter_steps into a set of allowed types if provided
@@ -138,7 +138,7 @@ async def generate_streaming_response_full(payload: typing.Any,
             allowed_types = set(filter_steps.split(','))
 
     async with session_manager.run(payload) as runner:
-        q: AsyncIOProducerConsumerQueue[AIQResponseSerializable] = AsyncIOProducerConsumerQueue()
+        q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
         # Start the intermediate stream without step adaptor
         intermediate_complete = await pull_intermediate(q, None)
@@ -159,12 +159,12 @@ async def generate_streaming_response_full(payload: typing.Any,
             asyncio.create_task(pull_result())
 
             async for item in q:
-                if (isinstance(item, AIQResponseIntermediateStep)):
+                if (isinstance(item, ResponseIntermediateStep)):
                     # 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)
+                    yield ResponsePayloadOutput(payload=item)
         except Exception as e:
             # Handle exceptions here
             raise e
@@ -174,7 +174,7 @@ async def generate_streaming_response_full(payload: typing.Any,
 
 async def generate_streaming_response_full_as_str(payload: typing.Any,
                                                   *,
-                                                  session_manager: AIQSessionManager,
+                                                  session_manager: SessionManager,
                                                   streaming: bool,
                                                   result_type: type | None = None,
                                                   output_type: type | None = None,
@@ -188,8 +188,8 @@ async def generate_streaming_response_full_as_str(payload: typing.Any,
                                                        result_type=result_type,
                                                        output_type=output_type,
                                                        filter_steps=filter_steps):
-        if (isinstance(item, AIQResponseIntermediateStep) or isinstance(item, AIQResponsePayloadOutput)):
+        if (isinstance(item, ResponseIntermediateStep) or isinstance(item, ResponsePayloadOutput)):
             yield item.get_stream_data()
         else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
+            raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
                              str(type(item)))