aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/runtime/loader.py

@@ -48,14 +48,18 @@ class PluginTypes(IntFlag):
     """
     A plugin that is an evaluator for the workflow. This includes evaluators like RAGAS, SWE-bench, etc.
     """
+    AUTHENTICATION = auto()
+    """
+    A plugin that is an API authentication provider for the workflow. This includes Oauth2, API Key, etc.
+    """
     REGISTRY_HANDLER = auto()
 
     # Convenience flag for groups of plugin types
-    CONFIG_OBJECT = COMPONENT | FRONT_END | EVALUATOR
+    CONFIG_OBJECT = COMPONENT | FRONT_END | EVALUATOR | AUTHENTICATION
     """
     Any plugin that can be specified in the AIQ Toolkit configuration file.
     """
-    ALL = COMPONENT | FRONT_END | EVALUATOR | REGISTRY_HANDLER
+    ALL = COMPONENT | FRONT_END | EVALUATOR | REGISTRY_HANDLER | AUTHENTICATION
     """
     All plugin types
     """
@@ -130,9 +134,11 @@ def discover_entrypoints(plugin_type: PluginTypes):
         plugin_groups.append("aiq.registry_handlers")
     if (plugin_type & PluginTypes.EVALUATOR):
         plugin_groups.append("aiq.evaluators")
+    if (plugin_type & PluginTypes.AUTHENTICATION):
+        plugin_groups.append("aiq.authentication_providers")
 
     # Get the entry points for the specified groups
-    aiq_plugins = reduce(lambda x, y: x + y, [entry_points.select(group=y) for y in plugin_groups])
+    aiq_plugins = reduce(lambda x, y: list(x) + list(y), [entry_points.select(group=y) for y in plugin_groups])
 
     return aiq_plugins
 

aiq/runtime/runner.py

@@ -21,7 +21,7 @@ from aiq.builder.context import AIQContext
 from aiq.builder.context import AIQContextState
 from aiq.builder.function import Function
 from aiq.data_models.invocation_node import InvocationNode
-from aiq.observability.async_otel_listener import AsyncOtelSpanListener
+from aiq.observability.exporter_manager import ExporterManager
 from aiq.utils.reactive.subject import Subject
 
 logger = logging.getLogger(__name__)
@@ -44,7 +44,11 @@ _T = typing.TypeVar("_T")
 
 class AIQRunner:
 
-    def __init__(self, input_message: typing.Any, entry_fn: Function, context_state: AIQContextState):
+    def __init__(self,
+                 input_message: typing.Any,
+                 entry_fn: Function,
+                 context_state: AIQContextState,
+                 exporter_manager: ExporterManager):
         """
         The AIQRunner class is used to run a workflow. It handles converting input and output data types and running the
         workflow with the specified concurrency.
@@ -57,6 +61,8 @@ class AIQRunner:
             The entry function to the workflow
         context_state : AIQContextState
             The context state to use
+        exporter_manager : ExporterManager
+            The exporter manager to use
         """
 
         if (entry_fn is None):
@@ -73,7 +79,7 @@ class AIQRunner:
         # Before we start, we need to convert the input message to the workflow input type
         self._input_message = input_message
 
-        self._span_manager = AsyncOtelSpanListener(context_state=context_state)
+        self._exporter_manager = exporter_manager
 
     @property
     def context(self) -> AIQContext:
@@ -130,19 +136,23 @@ class AIQRunner:
             if (not self._entry_fn.has_single_output):
                 raise ValueError("Workflow does not support single output")
 
-            async with self._span_manager.start():
+            async with self._exporter_manager.start(context_state=self._context_state):
                 # Run the workflow
                 result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)
 
                 # Close the intermediate stream
-                self._context_state.event_stream.get().on_complete()
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
 
             self._state = AIQRunnerState.COMPLETED
 
             return result
         except Exception as e:
             logger.exception("Error running workflow: %s", e)
-            self._context_state.event_stream.get().on_complete()
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
             self._state = AIQRunnerState.FAILED
 
             raise
@@ -159,18 +169,22 @@ class AIQRunner:
                 raise ValueError("Workflow does not support streaming output")
 
             # Run the workflow
-            async with self._span_manager.start():
+            async with self._exporter_manager.start(context_state=self._context_state):
                 async for m in self._entry_fn.astream(self._input_message, to_type=to_type):
                     yield m
 
                 self._state = AIQRunnerState.COMPLETED
 
                 # Close the intermediate stream
-                self._context_state.event_stream.get().on_complete()
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
 
         except Exception as e:
             logger.exception("Error running workflow: %s", e)
-            self._context_state.event_stream.get().on_complete()
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
             self._state = AIQRunnerState.FAILED
 
             raise

aiq/runtime/session.py

@@ -21,11 +21,14 @@ from collections.abc import Callable
 from contextlib import asynccontextmanager
 from contextlib import nullcontext
 
-from fastapi import Request
+from starlette.requests import HTTPConnection
 
 from aiq.builder.context import AIQContext
 from aiq.builder.context import AIQContextState
 from aiq.builder.workflow import Workflow
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthProviderBaseConfig
 from aiq.data_models.config import AIQConfig
 from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import InteractionPrompt
@@ -86,8 +89,11 @@ class AIQSessionManager:
     @asynccontextmanager
     async def session(self,
                       user_manager=None,
-                      request: Request = None,
-                      user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None):
+                      request: HTTPConnection | None = None,
+                      conversation_id: str | None = None,
+                      user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None,
+                      user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType],
+                                                             Awaitable[AuthenticatedContext | None]] = None):
 
         token_user_input = None
         if user_input_callback is not None:
@@ -97,7 +103,14 @@ class AIQSessionManager:
         if user_manager is not None:
             token_user_manager = self._context_state.user_manager.set(user_manager)
 
-        self.set_request_attributes(request)
+        token_user_authentication = None
+        if user_authentication_callback is not None:
+            token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
+
+        if conversation_id is not None and request is None:
+            self._context_state.conversation_id.set(conversation_id)
+
+        self.set_metadata_from_http_request(request)
 
         try:
             yield self
@@ -106,6 +119,8 @@ class AIQSessionManager:
                 self._context_state.user_manager.reset(token_user_manager)
             if token_user_input is not None:
                 self._context_state.user_input_callback.reset(token_user_input)
+            if token_user_authentication is not None:
+                self._context_state.user_auth_callback.reset(token_user_authentication)
 
     @asynccontextmanager
     async def run(self, message):
@@ -120,15 +135,15 @@ class AIQSessionManager:
             async with self._workflow.run(message) as runner:
                 yield runner
 
-    def set_request_attributes(self, request: Request) -> None:
+    def set_metadata_from_http_request(self, request: HTTPConnection | None) -> None:
         """
-        Extracts and sets request attributes from an HTTP request.
+        Extracts and sets user metadata request attributes from a HTTP request.
         If request is None, no attributes are set.
         """
         if request is None:
             return
 
-        self._context.metadata._request.method = request.method
+        self._context.metadata._request.method = getattr(request, "method", None)
         self._context.metadata._request.url_path = request.url.path
         self._context.metadata._request.url_port = request.url.port
         self._context.metadata._request.url_scheme = request.url.scheme
@@ -138,3 +153,6 @@ class AIQSessionManager:
         self._context.metadata._request.client_host = request.client.host
         self._context.metadata._request.client_port = request.client.port
         self._context.metadata._request.cookies = request.cookies
+
+        if request.headers.get("conversation-id"):
+            self._context_state.conversation_id.set(request.headers["conversation-id"])

aiq/runtime/user_metadata.py

@@ -21,9 +21,8 @@ from aiq.data_models.api_server import Request
 
 class RequestAttributes:
     """
-    The RequestAttributes class is responsible for managing user-defined
-    metadata and attributes. It provides a way to store and
-    expose user-defined attributes to workflow tools.
+    The RequestAttributes class is responsible for managing user http and webscoket session
+    metadata. It provides a way to store and expose session attributes to workflow tools.
     """
 
     def __init__(self) -> None:

aiq/tool/chat_completion.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.
+"""
+Simple Completion Function for AIQ Toolkit
+
+This module provides a simple completion function that can handle
+natural language queries and perform basic text completion tasks.
+"""
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import LLMRef
+from aiq.data_models.function import FunctionBaseConfig
+
+
+class ChatCompletionConfig(FunctionBaseConfig, name="chat_completion"):
+    """Configuration for the Chat Completion Function."""
+
+    system_prompt: str = Field(("You are a helpful AI assistant. Provide clear, accurate, and helpful "
+                                "responses to user queries. You can give general advice, recommendations, "
+                                "tips, and engage in conversation. Be helpful and informative."),
+                               description="The system prompt to use for chat completion.")
+
+    llm_name: LLMRef = Field(description="The LLM to use for generating responses.")
+
+
+@register_function(config_type=ChatCompletionConfig)
+async def register_chat_completion(config: ChatCompletionConfig, builder: Builder):
+    """Registers a chat completion function that can handle natural language queries."""
+
+    # Get the LLM from the builder context using the configured LLM reference
+    # Use LangChain framework wrapper since we're using LangChain-based LLM
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    async def _chat_completion(query: str) -> str:
+        """A simple chat completion function that responds to natural language queries.
+
+        Args:
+            query: The user's natural language query
+
+        Returns:
+            A helpful response to the query
+        """
+        try:
+            # Create a simple prompt with the system message and user query
+            prompt = f"{config.system_prompt}\n\nUser: {query}\n\nAssistant:"
+
+            # Generate response using the LLM
+            response = await llm.ainvoke(prompt)
+
+            return response
+
+        except Exception as e:
+            # Fallback response if LLM call fails
+            return (f"I apologize, but I encountered an error while processing your "
+                    f"query: '{query}'. Please try rephrasing your question or try "
+                    f"again later. Error: {str(e)}")
+
+    yield _chat_completion

aiq/tool/code_execution/README.md

@@ -0,0 +1,152 @@
+<!--
+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.
+-->
+
+# Code Execution Sandbox
+
+A secure, containerized Python code execution environment that allows safe execution of Python code with comprehensive error handling and debugging capabilities.
+
+## Overview
+
+The Code Execution Sandbox provides:
+- **Secure code execution** in isolated Docker containers
+- **Comprehensive error handling** with detailed stdout/stderr capture
+- **Multiple input formats** including raw code, dictionary format, and markdown
+- **Dependency management** with pre-installed libraries
+- **Flexible configuration** with customizable timeouts and output limits
+- **Robust debugging** with extensive logging and error reporting
+
+## Quick Start
+
+### Step 1: Start the Sandbox Server
+
+Navigate to the local sandbox directory and start the server:
+
+```bash
+cd src/aiq/tool/code_execution/local_sandbox
+./start_local_sandbox.sh
+```
+
+The script will:
+- Build the Docker image if it doesn't exist
+- Start the sandbox server on port 6000
+- Mount your working directory for file operations
+
+#### Advanced Usage:
+```bash
+# Custom container name
+./start_local_sandbox.sh my-sandbox
+
+# Custom output directory
+./start_local_sandbox.sh my-sandbox /path/to/output
+
+# Using environment variable
+export OUTPUT_DATA_PATH=/path/to/output
+./start_local_sandbox.sh
+```
+
+### Step 2: Test the Installation
+
+Run the comprehensive test suite to verify everything is working:
+
+```bash
+cd src/aiq/tool/code_execution
+pytest test_code_execution_sandbox.py
+```
+
+Note: a running instance of a local sandbox is required.
+
+## Using the Code Execution Tool
+
+### Basic Usage
+
+The sandbox accepts HTTP POST requests to `http://localhost:6000/execute` with JSON payloads:
+
+```bash
+curl -X POST \
+  -H "Content-Type: application/json" \
+  -d '{
+    "generated_code": "print(\"Hello, World!\")",
+    "timeout": 30,
+    "language": "python"
+  }' \
+  http://localhost:6000/execute
+```
+
+### Supported Input Formats
+
+#### 1. Raw Python Code
+```json
+{
+  "generated_code": "import numpy as np\nprint(np.array([1, 2, 3]))",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+#### 2. Dictionary Format
+```json
+{
+  "generated_code": "{'generated_code': 'print(\"Hello from dict format\")'}",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+#### 3. Markdown Code Blocks
+```json
+{
+  "generated_code": "```python\nprint('Hello from markdown')\n```",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+### Response Format
+
+The sandbox returns JSON responses with the following structure:
+
+```json
+{
+  "process_status": "completed|error|timeout",
+  "stdout": "Standard output content",
+  "stderr": "Standard error content"
+}
+```
+
+## Configuration Options
+
+### Sandbox Configuration
+
+- **URI**: Default `http://127.0.0.1:6000`
+- **Timeout**: Default 10 seconds (configurable)
+- **Max Output Characters**: Default 1000 characters
+- **Memory Limit**: 10GB (configurable in Docker)
+- **Working Directory**: Mounted volume for file operations
+
+### Environment Variables
+
+- `OUTPUT_DATA_PATH`: Custom path for file operations
+- `SANDBOX_HOST`: Custom sandbox host
+- `SANDBOX_PORT`: Custom sandbox port
+
+## Security Considerations
+
+- **Isolated execution**: All code runs in Docker containers
+- **Resource limits**: Memory and CPU limits prevent resource exhaustion
+- **Network isolation**: Containers have limited network access
+- **File system isolation**: Mounted volumes provide controlled file access
+- **Process isolation**: Each execution runs in a separate process