aiqtoolkit 1.1.0a20250503__py3-none-any.whl

aiq/front_ends/fastapi/step_adaptor.py

@@ -0,0 +1,315 @@
+# 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 html
+import logging
+from functools import reduce
+from textwrap import dedent
+
+from aiq.data_models.api_server import AIQResponseIntermediateStep
+from aiq.data_models.api_server import AIQResponseSerializable
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepCategory
+from aiq.data_models.intermediate_step import IntermediateStepPayload
+from aiq.data_models.intermediate_step import IntermediateStepType
+from aiq.data_models.invocation_node import InvocationNode
+from aiq.data_models.step_adaptor import StepAdaptorConfig
+from aiq.data_models.step_adaptor import StepAdaptorMode
+from aiq.utils.type_utils import is_valid_json
+
+logger = logging.getLogger(__name__)
+
+
+class StepAdaptor:
+
+    def __init__(self, config: StepAdaptorConfig):
+
+        self._history: list[IntermediateStep] = []
+        self.config = config
+
+    def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool:
+        """
+        Returns True if this intermediate step should be included (based on the config.mode).
+        """
+
+        if config.mode == StepAdaptorMode.OFF:
+            return False
+
+        if config.mode == StepAdaptorMode.DEFAULT:
+            # default existing behavior: show LLM events + TOOL_END + FUNCTION events
+            if step.event_category == IntermediateStepCategory.LLM:
+                return True
+            if step.event_type == IntermediateStepType.TOOL_END:
+                return True
+            if step.event_type in [IntermediateStepType.FUNCTION_START, IntermediateStepType.FUNCTION_END]:
+                return True
+            return False
+
+        if config.mode == StepAdaptorMode.CUSTOM:
+            # pass only what the user explicitly listed
+            return step.event_type in config.custom_event_types
+
+        return False
+
+    def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> AIQResponseSerializable | None:
+        input_str: str | None = None
+        output_str: str | None = None
+
+        # Find the start in the history with matching run_id
+        start_step = next(
+            (x for x in self._history if x.event_type == IntermediateStepType.LLM_START and x.UUID == step.UUID), None)
+
+        if not start_step:
+            # If we don't have a start step, we can't do anything
+            return None
+
+        input_str = str(start_step.data.input)
+
+        if step.event_type == IntermediateStepType.LLM_NEW_TOKEN:
+
+            # Find all of the previous LLM chunks and concatenate them
+            output_str = reduce(
+                lambda x, y: x + y,
+                (str(x.data.chunk)
+                 for x in self._history if x.event_type == IntermediateStepType.LLM_NEW_TOKEN and x.UUID == step.UUID),
+                "")
+
+        elif step.event_type == IntermediateStepType.LLM_END:
+            output_str = str(step.data.output)
+
+        if not input_str and not output_str:
+            return None
+
+        escaped_input = html.escape(input_str, quote=False)
+
+        # Dont use f-strings here because the payload is markdown and screws up the dedent
+        payload = dedent("""
+        **Input:**
+        ```python
+        {input_value}
+        ```
+        """).strip("\n").format(input_value=escaped_input)
+
+        if (output_str):
+            escaped_output = html.escape(output_str, quote=False) if output_str else ""
+
+            # Dont use f-strings here because the payload is markdown and screws up the dedent
+            payload = dedent("""
+            {payload}
+
+            **Output:**
+            {output_value}
+            """).strip("\n").format(payload=payload, output_value=escaped_output)
+
+        event = AIQResponseIntermediateStep(id=step.UUID,
+                                            name=step.name or "",
+                                            payload=payload,
+                                            parent_id=ancestry.function_id)
+
+        return event
+
+    def _handle_tool_end(self, payload: IntermediateStepPayload,
+                         ancestry: InvocationNode) -> AIQResponseSerializable | None:
+        """
+        Handles the TOOL_END event
+        """
+        escaped_tool_input = html.escape(str(payload.data.input), quote=False)
+        escaped_tool_output = html.escape(str(payload.data.output), quote=False)
+
+        escaped_tool_input = escaped_tool_input.replace("\n", "")
+        escaped_tool_output = escaped_tool_output.replace("\n", "")
+
+        # Determine the format
+        format_input_type = "json" if is_valid_json(escaped_tool_input) else "python"
+        format_output_type = "json" if is_valid_json(escaped_tool_output) else "python"
+
+        # Dont use f-strings here because the payload is markdown and screws up the dedent
+        payload_str = dedent("""
+        **Input:**
+        ```{format_input_type}
+        {input_value}
+        ```
+        **Output:**
+        ```{format_output_type}
+        {output_value}
+        ```
+        """).strip("\n").format(input_value=escaped_tool_input,
+                                output_value=escaped_tool_output,
+                                format_input_type=format_input_type,
+                                format_output_type=format_output_type)
+
+        event = AIQResponseIntermediateStep(id=payload.UUID,
+                                            name=f"Tool: {payload.name}",
+                                            payload=payload_str,
+                                            parent_id=ancestry.function_id)
+
+        return event
+
+    def _handle_function(self, step: IntermediateStepPayload,
+                         ancestry: InvocationNode) -> AIQResponseSerializable | None:
+        """
+        Handles the FUNCTION_START and FUNCTION_END events
+        """
+        input_str: str | None = None
+        output_str: str | None = None
+
+        if step.event_type == IntermediateStepType.FUNCTION_START:
+            # For function start events, display input data
+            if step.data and hasattr(step.data, 'input'):
+                input_str = str(step.data.input)
+            elif step.data:
+                input_str = str(step.data)
+
+            if not input_str:
+                return None
+
+            escaped_input = html.escape(input_str, quote=False)
+            format_input_type = "json" if is_valid_json(escaped_input) else "python"
+
+            # Create payload for function start
+            payload_str = dedent("""
+            **Function Input:**
+            ```{format_input_type}
+            {input_value}
+            ```
+            """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
+
+            event = AIQResponseIntermediateStep(id=step.UUID,
+                                                name=f"Function Start: {step.name}",
+                                                payload=payload_str,
+                                                parent_id=ancestry.parent_id)
+            return event
+
+        if step.event_type == IntermediateStepType.FUNCTION_END:
+            # Find the start event with matching UUID
+            start_step = next(
+                (x
+                 for x in self._history if x.event_type == IntermediateStepType.FUNCTION_START and x.UUID == step.UUID),
+                None)
+
+            # For function end events, display output data
+            if step.data and hasattr(step.data, 'output'):
+                output_str = str(step.data.output)
+            elif step.data:
+                output_str = str(step.data)
+
+            if not output_str:
+                return None
+
+            escaped_output = html.escape(output_str, quote=False)
+            format_output_type = "json" if is_valid_json(escaped_output) else "python"
+
+            # Get input from start step if available
+            input_payload = ""
+            if start_step and start_step.data:
+                if hasattr(start_step.data, 'input'):
+                    input_str = str(start_step.data.input)
+                else:
+                    input_str = str(start_step.data)
+
+                if input_str:
+                    escaped_input = html.escape(input_str, quote=False)
+                    format_input_type = "json" if is_valid_json(escaped_input) else "python"
+                    input_payload = dedent("""
+                    **Function Input:**
+                    ```{format_input_type}
+                    {input_value}
+                    ```
+                    """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
+
+            # Create payload for function end
+            payload_str = dedent("""
+            {input_payload}**Function Output:**
+            ```{format_output_type}
+            {output_value}
+            ```
+            """).strip("\n").format(input_payload=input_payload,
+                                    output_value=escaped_output,
+                                    format_output_type=format_output_type)
+
+            event = AIQResponseIntermediateStep(id=step.UUID,
+                                                name=f"Function Complete: {step.name}",
+                                                payload=payload_str,
+                                                parent_id=ancestry.parent_id)
+            return event
+
+        return None
+
+    def _handle_custom(self, payload: IntermediateStepPayload,
+                       ancestry: InvocationNode) -> AIQResponseSerializable | None:
+        """
+        Handles the CUSTOM event
+        """
+        escaped_payload = html.escape(str(payload), quote=False)
+        escaped_payload = escaped_payload.replace("\n", "")
+
+        # Attempt to determine type
+        format_type = "json" if is_valid_json(escaped_payload) else "python"
+
+        # Don't use f-strings here because the payload is markdown and screws up the dedent
+        payload_str = dedent("""
+        ```{format_type}
+        {payload}
+        ```
+        """).strip("\n").format(payload=escaped_payload, format_type=format_type)
+
+        # Return the event
+        event = AIQResponseIntermediateStep(id=payload.UUID,
+                                            name=f"{payload.event_type}",
+                                            payload=payload_str,
+                                            parent_id=ancestry.function_id)
+
+        return event
+
+    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:
+        # Track the chunk
+        self._history.append(step)
+        payload = step.payload
+        ancestry = step.function_ancestry
+
+        if not self._step_matches_filter(step, self.config):
+            return None
+
+        try:
+
+            if self.config.mode == StepAdaptorMode.DEFAULT:
+
+                if step.event_category == IntermediateStepCategory.LLM:
+                    return self._handle_llm(payload, ancestry)
+
+                if step.event_type == IntermediateStepType.TOOL_END:
+                    return self._handle_tool_end(payload, ancestry)
+
+                if step.event_type in [IntermediateStepType.FUNCTION_START, IntermediateStepType.FUNCTION_END]:
+                    return self._handle_function(payload, ancestry)
+
+            if self.config.mode == StepAdaptorMode.CUSTOM:
+                # Now we're processing user defined custom types
+
+                if step.event_category == IntermediateStepCategory.LLM:
+                    return self._handle_llm(payload, ancestry)
+
+                if step.event_type == IntermediateStepType.TOOL_END:
+                    return self._handle_tool_end(payload, ancestry)
+
+                if step.event_type in [IntermediateStepType.FUNCTION_START, IntermediateStepType.FUNCTION_END]:
+                    return self._handle_function(payload, ancestry)
+
+                return self._handle_custom(payload, ancestry)
+
+        except Exception as e:
+            logger.error("Error processing intermediate step: %s", e, exc_info=True)
+
+        return None

aiq/front_ends/fastapi/websocket.py

@@ -0,0 +1,148 @@
+# 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 logging
+import typing
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+
+from fastapi import WebSocket
+from fastapi import WebSocketException
+from starlette.endpoints import WebSocketEndpoint
+from starlette.websockets import WebSocketDisconnect
+
+from aiq.data_models.api_server import AIQChatRequest
+from aiq.data_models.api_server import AIQChatResponse
+from aiq.data_models.api_server import AIQChatResponseChunk
+from aiq.data_models.api_server import AIQResponsePayloadOutput
+from aiq.data_models.api_server import AIQResponseSerializable
+from aiq.data_models.api_server import WebSocketMessageStatus
+from aiq.data_models.api_server import WorkflowSchemaType
+from aiq.front_ends.fastapi.message_handler import MessageHandler
+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 AIQWebSocket(WebSocketEndpoint):
+    encoding = "json"
+
+    def __init__(self, session_manager: AIQSessionManager, step_adaptor: StepAdaptor, *args, **kwargs):
+        self._session_manager: AIQSessionManager = session_manager
+        self._message_handler: MessageHandler = MessageHandler(self)
+        self._process_response_event: asyncio.Event = asyncio.Event()
+        self._workflow_schema_type: dict[str, Callable[..., Awaitable[Any]]] = {
+            WorkflowSchemaType.GENERATE_STREAM: self.process_generate_stream,
+            WorkflowSchemaType.CHAT_STREAM: self.process_chat_stream,
+            WorkflowSchemaType.GENERATE: self.process_generate,
+            WorkflowSchemaType.CHAT: self.process_chat
+        }
+        self._step_adaptor = step_adaptor
+        super().__init__(*args, **kwargs)
+
+    @property
+    def workflow_schema_type(self) -> dict[str, Callable[..., Awaitable[Any]]]:
+        return self._workflow_schema_type
+
+    @property
+    def process_response_event(self) -> asyncio.Event:
+        return self._process_response_event
+
+    async def on_connect(self, websocket: WebSocket):
+        try:
+            # Accept the websocket connection
+            await websocket.accept()
+            try:
+                # Start background message processors
+                self._message_handler.process_messages_task = asyncio.create_task(
+                    self._message_handler.process_messages())
+
+                self._message_handler.process_out_going_messages_task = asyncio.create_task(
+                    self._message_handler.process_out_going_messages(websocket))
+
+            except asyncio.CancelledError:
+                pass
+
+        except (WebSocketDisconnect, WebSocketException):
+            logger.error("A WebSocket error occured during `on_connect`. Ignoring the connection.", exc_info=True)
+
+    async def on_send(self, websocket: WebSocket, data: dict[str, str]):
+        try:
+            await websocket.send_json(data)
+        except (WebSocketDisconnect, WebSocketException, Exception):
+            logger.error("A WebSocket error occurred during `on_send`. Ignoring the connection.", exc_info=True)
+
+    async def on_receive(self, websocket: WebSocket, data: dict[str, Any]):
+        try:
+            await self._message_handler.messages_queue.put(data)
+        except (Exception):
+            logger.error("An unxpected error occurred during `on_receive`. Ignoring the exception", exc_info=True)
+
+    async def on_disconnect(self, websocket: WebSocket, close_code: Any):
+        try:
+            if self._message_handler.process_messages_task:
+                self._message_handler.process_messages_task.cancel()
+
+            if self._message_handler.process_out_going_messages_task:
+                self._message_handler.process_out_going_messages_task.cancel()
+
+            if self._message_handler.background_task:
+                self._message_handler.background_task.cancel()
+
+        except (WebSocketDisconnect, asyncio.CancelledError):
+            pass
+
+    async def _process_message(self,
+                               payload: typing.Any,
+                               result_type: type | None = None,
+                               output_type: type | None = None) -> None:
+
+        async with self._session_manager.session(
+                user_input_callback=self._message_handler.human_interaction) as session:
+
+            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):
+
+                await self._process_response_event.wait()
+
+                if not isinstance(value, AIQResponseSerializable):
+                    value = AIQResponsePayloadOutput(payload=value)
+
+                await self._message_handler.create_websocket_message(data_model=value,
+                                                                     status=WebSocketMessageStatus.IN_PROGRESS)
+
+    async def process_generate_stream(self, payload: str):
+
+        return await self._process_message(payload, result_type=None, output_type=None)
+
+    async def process_chat_stream(self, payload: AIQChatRequest):
+
+        return await self._process_message(payload, result_type=AIQChatResponse, output_type=AIQChatResponseChunk)
+
+    async def process_generate(self, payload: typing.Any):
+
+        return await self._process_message(payload)
+
+    async def process_chat(self, payload: AIQChatRequest):
+
+        return await self._process_message(payload, result_type=AIQChatResponse)

aiq/front_ends/mcp/__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/front_ends/mcp/mcp_front_end_config.py

@@ -0,0 +1,32 @@
+# 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 pydantic import Field
+
+from aiq.data_models.front_end import FrontEndBaseConfig
+
+
+class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
+    """MCP front end configuration.
+
+    A simple MCP (Modular Communication Protocol) front end for AIQ.
+    """
+
+    name: str = Field(default="AIQ MCP", description="Name of the MCP server")
+    host: str = Field(default="localhost", description="Host to bind the server to")
+    port: int = Field(default=9901, description="Port to bind the server to", ge=0, le=65535)
+    debug: bool = Field(default=False, description="Enable debug mode")
+    log_level: str = Field(default="INFO", description="Log level for the MCP server")
+    tool_names: list[str] = Field(default_factory=list, description="The list of tools MCP server will expose.")

aiq/front_ends/mcp/mcp_front_end_plugin.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from aiq.builder.front_end import FrontEndBase
+from aiq.builder.function import Function
+from aiq.builder.workflow import Workflow
+from aiq.builder.workflow_builder import WorkflowBuilder
+from aiq.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
+    """MCP front end plugin implementation."""
+
+    async def run(self) -> None:
+        """Run the MCP server."""
+        # Import FastMCP
+        from mcp.server.fastmcp import FastMCP
+
+        from aiq.front_ends.mcp.tool_converter import register_function_with_mcp
+
+        # Create an MCP server with the configured parameters
+        mcp = FastMCP(
+            self.front_end_config.name,
+            host=self.front_end_config.host,
+            port=self.front_end_config.port,
+            debug=self.front_end_config.debug,
+            log_level=self.front_end_config.log_level,
+        )
+
+        # Build the workflow and register all functions with MCP
+        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+            # Build the workflow
+            workflow = builder.build()
+
+            # Get all functions from the workflow
+            functions = self._get_all_functions(workflow)
+
+            # Filter functions based on tool_names if provided
+            if self.front_end_config.tool_names:
+                logger.info(f"Filtering functions based on tool_names: {self.front_end_config.tool_names}")
+                filtered_functions: dict[str, Function] = {}
+                for function_name, function in functions.items():
+                    if function_name in self.front_end_config.tool_names:
+                        filtered_functions[function_name] = function
+                    else:
+                        logger.debug(f"Skipping function {function_name} as it's not in tool_names")
+                functions = filtered_functions
+
+            # Register each function with MCP
+            for function_name, function in functions.items():
+                register_function_with_mcp(mcp, function_name, function)
+
+            # Add a simple fallback function if no functions were found
+            if not functions:
+                raise RuntimeError("No functions found in workflow. Please check your configuration.")
+
+            # Start the MCP server
+            await mcp.run_sse_async()
+
+    def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
+        """Get all functions from the workflow.
+
+        Args:
+            workflow: The AIQ workflow.
+
+        Returns:
+            Dict mapping function names to Function objects.
+        """
+        functions: dict[str, Function] = {}
+
+        # Extract all functions from the workflow
+        for function_name, function in workflow.functions.items():
+            functions[function_name] = function
+
+        functions[workflow.config.workflow.type] = workflow
+
+        return functions

aiq/front_ends/mcp/register.py

@@ -0,0 +1,27 @@
+# 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 collections.abc import AsyncIterator
+
+from aiq.cli.register_workflow import register_front_end
+from aiq.data_models.config import AIQConfig
+from aiq.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+
+
+@register_front_end(config_type=MCPFrontEndConfig)
+async def register_mcp_front_end(config: MCPFrontEndConfig, full_config: AIQConfig) -> AsyncIterator:
+    from aiq.front_ends.mcp.mcp_front_end_plugin import MCPFrontEndPlugin
+
+    yield MCPFrontEndPlugin(full_config=full_config)